Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| import torch | |
| import numpy as np | |
| from PIL import Image | |
| from transformers import Sam3Processor, Sam3Model | |
| # Load model on startup | |
| print("Loading SAM 3 model...") | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| model = Sam3Model.from_pretrained("facebook/sam3", token=True).to(device) | |
| processor = Sam3Processor.from_pretrained("facebook/sam3", token=True) | |
| print(f"Model loaded on {device}") | |
| # Reference dimensions (inches) | |
| REFERENCE_DIMENSIONS = { | |
| "light switch wall plate": {"height": 4.5, "width": 2.75}, | |
| "electrical wall outlet": {"height": 4.5, "width": 2.75}, | |
| "room door": {"height": 80, "width": 36}, | |
| "window": {"height": 48, "width": 36}, | |
| } | |
| def get_mask_dimensions(mask): | |
| mask_np = mask.cpu().numpy() if hasattr(mask, 'cpu') else mask | |
| rows = np.any(mask_np, axis=1) | |
| cols = np.any(mask_np, axis=0) | |
| return np.sum(rows), np.sum(cols) | |
| def calculate_mask_area(mask): | |
| if hasattr(mask, 'sum'): | |
| return mask.sum().item() | |
| return np.sum(mask) | |
| def pixel_area_to_sqft(pixel_area, pixels_per_inch): | |
| if pixels_per_inch is None or pixels_per_inch == 0: | |
| return 0 | |
| square_inches = pixel_area / (pixels_per_inch ** 2) | |
| return square_inches / 144 | |
| def overlay_masks_colored(image, masks, alpha=0.5, color=(255, 0, 0)): | |
| if isinstance(image, np.ndarray): | |
| image = Image.fromarray(image.astype('uint8')) | |
| image = image.convert("RGBA") | |
| for mask in masks: | |
| mask_np = mask.cpu().numpy() if hasattr(mask, 'cpu') else mask | |
| if mask_np.max() <= 1.0: | |
| mask_np = (mask_np * 255).astype(np.uint8) | |
| else: | |
| mask_np = mask_np.astype(np.uint8) | |
| mask_img = Image.fromarray(mask_np) | |
| overlay = Image.new("RGBA", image.size, color + (0,)) | |
| alpha_mask = mask_img.point(lambda v: int(v * alpha)) | |
| overlay.putalpha(alpha_mask) | |
| image = Image.alpha_composite(image, overlay) | |
| return image | |
| def analyze_room(image): | |
| if image is None: | |
| return None, {"error": "No image provided"} | |
| if isinstance(image, np.ndarray): | |
| image = Image.fromarray(image).convert("RGB") | |
| else: | |
| image = image.convert("RGB") | |
| room_concepts = [ | |
| "wall", "floor", "rug", "carpet", "ceiling", | |
| "room door", "window", "electrical wall outlet", | |
| "light switch wall plate" | |
| ] | |
| room_segments = {} | |
| for concept in room_concepts: | |
| inputs = processor(images=image, text=concept, return_tensors="pt").to(device) | |
| with torch.no_grad(): | |
| outputs = model(**inputs) | |
| results = processor.post_process_instance_segmentation( | |
| outputs, | |
| threshold=0.3, | |
| mask_threshold=0.3, | |
| target_sizes=inputs.get("original_sizes").tolist() | |
| )[0] | |
| room_segments[concept] = results | |
| floor_masks = [] | |
| for key in ["floor", "rug", "carpet"]: | |
| if key in room_segments and len(room_segments[key]['masks']) > 0: | |
| for mask in room_segments[key]['masks']: | |
| floor_masks.append(mask) | |
| if floor_masks: | |
| room_segments["total_floor"] = {'masks': floor_masks} | |
| pixels_per_inch = None | |
| for ref_type in ["light switch wall plate", "electrical wall outlet", "room door", "window"]: | |
| if ref_type in room_segments and len(room_segments[ref_type]['masks']) > 0: | |
| mask = room_segments[ref_type]['masks'][0] | |
| pixel_height, _ = get_mask_dimensions(mask) | |
| pixels_per_inch = pixel_height / REFERENCE_DIMENSIONS[ref_type]["height"] | |
| break | |
| if pixels_per_inch is None: | |
| pixels_per_inch = 5.0 | |
| measurements = {} | |
| for concept in ["wall", "ceiling", "window", "room door"]: | |
| if concept in room_segments and len(room_segments[concept]['masks']) > 0: | |
| total_pixels = sum([calculate_mask_area(m) for m in room_segments[concept]['masks']]) | |
| sqft = pixel_area_to_sqft(total_pixels, pixels_per_inch) | |
| measurements[concept] = round(sqft, 1) | |
| if "total_floor" in room_segments and len(room_segments["total_floor"]['masks']) > 0: | |
| total_pixels = sum([calculate_mask_area(m) for m in room_segments["total_floor"]['masks']]) | |
| sqft = pixel_area_to_sqft(total_pixels, pixels_per_inch) | |
| measurements["floor"] = round(sqft, 1) | |
| floor_sqft = measurements.get("floor", 0) | |
| ceiling_sqft = measurements.get("ceiling", 0) | |
| if ceiling_sqft < (floor_sqft * 0.5) and floor_sqft > 0: | |
| measurements["ceiling"] = floor_sqft | |
| wall_sqft = measurements.get("wall", 0) | |
| window_sqft = measurements.get("window", 0) | |
| paint_wall = max(0, wall_sqft - window_sqft) | |
| materials = { | |
| "paint_gallons": max(1, round((paint_wall * 2) / 350, 1)), | |
| "flooring_boxes": max(1, round((floor_sqft * 1.1) / 20)) if floor_sqft else 0 | |
| } | |
| c_paint = (materials["paint_gallons"] * 35) + (paint_wall * 2) | |
| c_floor = (floor_sqft * 4) + (floor_sqft * 3) if floor_sqft else 0 | |
| total_cost = round(c_paint + c_floor) | |
| overlay_img = image.copy() | |
| if "wall" in room_segments and len(room_segments["wall"]['masks']) > 0: | |
| overlay_img = overlay_masks_colored(overlay_img, room_segments["wall"]['masks'], alpha=0.4, color=(239, 68, 68)) | |
| if "total_floor" in room_segments and len(room_segments["total_floor"]['masks']) > 0: | |
| overlay_img = overlay_masks_colored(overlay_img, room_segments["total_floor"]['masks'], alpha=0.4, color=(16, 185, 129)) | |
| if "window" in room_segments and len(room_segments["window"]['masks']) > 0: | |
| overlay_img = overlay_masks_colored(overlay_img, room_segments["window"]['masks'], alpha=0.4, color=(99, 102, 241)) | |
| if overlay_img.mode == 'RGBA': | |
| background = Image.new('RGB', overlay_img.size, (255, 255, 255)) | |
| background.paste(overlay_img, mask=overlay_img.split()[3]) | |
| overlay_img = background | |
| results = { | |
| "measurements": measurements, | |
| "materials": materials, | |
| "costs": { | |
| "paint_and_labor": round(c_paint), | |
| "flooring_and_labor": round(c_floor), | |
| "total": total_cost | |
| } | |
| } | |
| return overlay_img, results | |
| # Create and launch Gradio interface | |
| demo = gr.Interface( | |
| fn=analyze_room, | |
| inputs=gr.Image(type="pil", label="Upload Room Photo"), | |
| outputs=[ | |
| gr.Image(label="Segmentation Overlay"), | |
| gr.JSON(label="Room Analysis Results") | |
| ], | |
| title="Room Estimator", | |
| description="Upload a room photo to get measurements and material estimates" | |
| ) | |
| demo.launch(ssr_mode=False) |