Update app.py

app.py

@@ -1,32 +1,65 @@
+ in Python with Transformers. I used a SegFormer-based segmentation pipeline as the first demoable model because Hugging Face supports semantic segmentation directly through pipeline(task="image-segmentation").
+
+Why this helps
+Your earlier version only split the detected outer mask geometrically, so if the outer region was clean and undivided, the engine naturally returned one room. A vision model introduces an extra prior by estimating interior regions from the uploaded image before the partition logic runs, which gives the downstream room generator more structure to work with.
+
+Important limitation
+This is still a demo-grade improvement, not a true architectural floor-plan model. Generic segmentation models like SegFormer or SAM are not trained specifically on floor-plan semantics, so they can improve region proposals but will not reliably infer architecturally correct rooms, walls, doors, and circulation from arbitrary closed shapes.
+
+Better next model
+If you want proper floor-plan generation, the stronger path is:
+
+use SAM/SAM2 for boundary and mask proposals,
+
+or plug in a floor-plan-specific segmentation model such as a U-Net/ResNet room-segmentation model trained on floor-plan datasets.
+
+The files are updated so you can redeploy the Space directly.
+
+Follow-ups
+
+How to use SAM2 for video floor plan generation
+
+Add door and window detection to multi-room floor plan
+
+Fine-tune floor plan segmentation model on custom dataset
+
+Convert SAM masks to SVG floor plan with measurements
+
+Compare SAM vs U-Net ResNet for room segmentation accuracy
+app
 import io
 import json
-import math
 import os
 import tempfile
 import zipfile
-from pathlib import Path
 from dataclasses import dataclass
+from pathlib import Path
 
+import cairosvg
 import cv2
 import gradio as gr
 import numpy as np
-from PIL import Image, ImageDraw
-import cairosvg
+from PIL import Image
+
+MODEL_AVAILABLE = False
+MODEL_ERROR = None
+processor = None
+seg_pipeline = None
 
 try:
-    from shapely.geometry import Polygon, LineString, mapping
-    from shapely.ops import unary_union
-    SHAPELY_OK = True
-except Exception:
-    SHAPELY_OK = False
+    from transformers import pipeline
+    MODEL_AVAILABLE = True
+except Exception as e:
+    MODEL_ERROR = str(e)
 
 
 @dataclass
 class PlanConfig:
-    min_room_area_ratio: float = 0.035
+    min_room_area_ratio: float = 0.02
     wall_thickness: int = 6
     variation: str = "balanced"
     max_rooms: int = 6
+    use_vision_model: bool = False
 
 
 def pil_to_bgr(img: Image.Image) -> np.ndarray:
@@ -39,41 +72,84 @@ def bgr_to_pil(arr: np.ndarray) -> Image.Image:
     return Image.fromarray(rgb)
 
 
-def order_quad_points(pts):
-    pts = np.array(pts, dtype=np.float32)
-    s = pts.sum(axis=1)
-    d = np.diff(pts, axis=1)
-    return np.array([
-        pts[np.argmin(s)],
-        pts[np.argmin(d)],
-        pts[np.argmax(s)],
-        pts[np.argmax(d)],
-    ], dtype=np.float32)
+def load_any_image(image_file):
+    if image_file is None:
+        raise gr.Error("Upload a PNG, JPG, or SVG file first.")
+    path = image_file if isinstance(image_file, str) else getattr(image_file, "name", None)
+    if not path:
+        raise gr.Error("Could not read uploaded file.")
+    ext = os.path.splitext(path)[1].lower()
+    if ext == ".svg":
+        png_bytes = cairosvg.svg2png(url=path)
+        return Image.open(io.BytesIO(png_bytes)).convert("RGB"), ext
+    return Image.open(path).convert("RGB"), ext
+
+
+def init_segmentation_pipeline():
+    global seg_pipeline, MODEL_AVAILABLE, MODEL_ERROR
+    if seg_pipeline is not None:
+        return seg_pipeline
+    if not MODEL_AVAILABLE:
+        return None
+    try:
+        seg_pipeline = pipeline(
+            task="image-segmentation",
+            model="nvidia/segformer-b0-finetuned-ade-512-512",
+            device=-1,
+        )
+        return seg_pipeline
+    except Exception as e:
+        MODEL_ERROR = str(e)
+        return None
 
 
 def preprocess_and_extract_boundary(img_bgr: np.ndarray):
     gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
     blur = cv2.GaussianBlur(gray, (5, 5), 0)
     _, th = cv2.threshold(blur, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
-
     kernel = np.ones((5, 5), np.uint8)
     th = cv2.morphologyEx(th, cv2.MORPH_CLOSE, kernel, iterations=2)
     th = cv2.morphologyEx(th, cv2.MORPH_OPEN, np.ones((3, 3), np.uint8), iterations=1)
-
     contours, _ = cv2.findContours(th, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
     if not contours:
         raise ValueError("No closed shape detected. Use a darker outline on light background.")
-
     contour = max(contours, key=cv2.contourArea)
     peri = cv2.arcLength(contour, True)
     approx = cv2.approxPolyDP(contour, 0.01 * peri, True)
-
     mask = np.zeros_like(gray)
     cv2.drawContours(mask, [contour], -1, 255, thickness=cv2.FILLED)
-
     return th, contour[:, 0, :], approx[:, 0, :], mask
 
 
+def mask_from_vision_model(image: Image.Image, boundary_mask: np.ndarray):
+    pipe = init_segmentation_pipeline()
+    if pipe is None:
+        return None, f"Vision model unavailable: {MODEL_ERROR}" if MODEL_ERROR else "Vision model unavailable"
+    try:
+        preds = pipe(image)
+        if not preds:
+            return None, "Vision model returned no masks"
+        room_mask = np.zeros(boundary_mask.shape, dtype=np.uint8)
+        for pred in preds:
+            label = str(pred.get("label", "")).lower()
+            if any(k in label for k in ["floor", "rug", "carpet", "room", "interior"]):
+                seg = pred.get("mask")
+                if isinstance(seg, Image.Image):
+                    seg = np.array(seg.convert("L"))
+                else:
+                    seg = np.array(seg)
+                if seg.ndim == 3:
+                    seg = seg[..., 0]
+                seg = cv2.resize(seg.astype(np.uint8), (boundary_mask.shape[1], boundary_mask.shape[0]))
+                room_mask = np.maximum(room_mask, (seg > 0).astype(np.uint8) * 255)
+        room_mask = cv2.bitwise_and(room_mask, boundary_mask)
+        if np.count_nonzero(room_mask) == 0:
+            return None, "Vision model found no interior-like regions"
+        return room_mask, "Vision model used for interior proposal"
+    except Exception as e:
+        return None, f"Vision inference failed: {e}"
+
+
 def axis_aligned_split(bounds, orientation, pos, gap=8):
     x1, y1, x2, y2 = bounds
     if orientation == "v":
@@ -81,7 +157,7 @@ def axis_aligned_split(bounds, orientation, pos, gap=8):
     return [(x1, y1, x2, pos - gap), (x1, pos + gap, x2, y2)]
 
 
-def rect_inside_mask(rect, mask, fill_ratio=0.84):
+def rect_inside_mask(rect, mask, fill_ratio=0.72):
     x1, y1, x2, y2 = [int(v) for v in rect]
     if x2 <= x1 or y2 <= y1:
         return False
@@ -94,12 +170,10 @@ def rect_inside_mask(rect, mask, fill_ratio=0.84):
 def generate_rooms(mask: np.ndarray, cfg: PlanConfig):
     ys, xs = np.where(mask > 0)
     x1, y1, x2, y2 = xs.min(), ys.min(), xs.max(), ys.max()
-    root = (x1, y1, x2, y2)
-    rooms = [root]
-    target = {"compact": 4, "balanced": 5, "open": 3}.get(cfg.variation, 5)
+    rooms = [(x1, y1, x2, y2)]
+    target = {"compact": 6, "balanced": 5, "open": 3}.get(cfg.variation, 5)
     target = min(target, cfg.max_rooms)
-
-    for _ in range(target * 2):
+    for _ in range(target * 4):
         if len(rooms) >= target:
             break
         areas = [max(0, (r[2] - r[0]) * (r[3] - r[1])) for r in rooms]
@@ -110,43 +184,51 @@ def generate_rooms(mask: np.ndarray, cfg: PlanConfig):
         if w * h < mask.shape[0] * mask.shape[1] * cfg.min_room_area_ratio:
             rooms.append(r)
             continue
-
         if cfg.variation == "open":
-            orientation = "v" if w > h * 1.35 else "h"
+            orientation = "v" if w > h * 1.25 else "h"
+            cuts = [0.5]
         elif cfg.variation == "compact":
             orientation = "h" if h > w else "v"
+            cuts = [0.45, 0.55]
         else:
             orientation = "v" if w >= h else "h"
-
-        pos = int((rx1 + rx2) / 2) if orientation == "v" else int((ry1 + ry2) / 2)
-        kids = axis_aligned_split(r, orientation, pos, gap=cfg.wall_thickness)
-        valid = [k for k in kids if rect_inside_mask(k, mask)]
-        if len(valid) == 2:
-            rooms.extend(valid)
-        else:
+            cuts = [0.5, 0.4, 0.6]
+        split_done = False
+        for c in cuts:
+            pos = int(rx1 + w * c) if orientation == "v" else int(ry1 + h * c)
+            kids = axis_aligned_split(r, orientation, pos, gap=cfg.wall_thickness)
+            valid = [k for k in kids if rect_inside_mask(k, mask)]
+            if len(valid) == 2:
+                rooms.extend(valid)
+                split_done = True
+                break
+        if not split_done:
             rooms.append(r)
-
     return rooms
 
 
-def draw_floorplan(mask, rooms, contour_points, approx_points, cfg: PlanConfig):
-    h, w = mask.shape
-    canvas = np.full((h, w, 3), 255, np.uint8)
+def polygon_from_mask_component(component_mask):
+    contours, _ = cv2.findContours(component_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    if not contours:
+        return None
+    c = max(contours, key=cv2.contourArea)
+    return c
 
-    cv2.drawContours(canvas, [contour_points.reshape(-1, 1, 2)], -1, (0, 0, 0), thickness=max(3, cfg.wall_thickness + 2))
 
+def draw_floorplan(boundary_mask, room_seed_mask, rooms, contour_points, approx_points, cfg: PlanConfig, vision_note: str):
+    h, w = boundary_mask.shape
+    canvas = np.full((h, w, 3), 255, np.uint8)
     palette = [(239, 246, 255), (236, 253, 245), (255, 247, 237), (250, 245, 255), (254, 242, 242), (240, 253, 250)]
     room_json = []
 
     for i, r in enumerate(rooms):
         x1, y1, x2, y2 = [int(v) for v in r]
-        room_mask = np.zeros_like(mask)
+        room_mask = np.zeros_like(boundary_mask)
         cv2.rectangle(room_mask, (x1, y1), (x2, y2), 255, thickness=-1)
-        room_mask = cv2.bitwise_and(room_mask, mask)
-        contours, _ = cv2.findContours(room_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-        if not contours:
+        room_mask = cv2.bitwise_and(room_mask, room_seed_mask)
+        c = polygon_from_mask_component(room_mask)
+        if c is None:
             continue
-        c = max(contours, key=cv2.contourArea)
         cv2.drawContours(canvas, [c], -1, palette[i % len(palette)], thickness=cv2.FILLED)
         cv2.drawContours(canvas, [c], -1, (0, 0, 0), thickness=2)
         M = cv2.moments(c)
@@ -157,14 +239,14 @@ def draw_floorplan(mask, rooms, contour_points, approx_points, cfg: PlanConfig):
         room_json.append({"id": f"room_{i+1}", "bbox": [x1, y1, x2, y2], "area_px": int(cv2.contourArea(c)), "polygon": c[:, 0, :].tolist()})
 
     cv2.drawContours(canvas, [contour_points.reshape(-1, 1, 2)], -1, (0, 0, 0), thickness=max(3, cfg.wall_thickness + 2))
-    cv2.polylines(canvas, [approx_points.reshape(-1, 1, 2)], True, (90, 90, 90), 1)
+    cv2.putText(canvas, vision_note, (16, 24), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (60, 60, 60), 1, cv2.LINE_AA)
 
     result = {
         "boundary": contour_points.tolist(),
         "boundary_simplified": approx_points.tolist(),
         "rooms": room_json,
         "config": cfg.__dict__,
-        "meta": {"shapely_available": SHAPELY_OK, "image_size": [int(w), int(h)]},
+        "meta": {"image_size": [int(w), int(h)], "vision_note": vision_note},
     }
     return canvas, result
 
@@ -176,7 +258,6 @@ def export_svg(result, width, height):
         start = f"M {pts[0][0]} {pts[0][1]}"
         rest = " ".join([f"L {p[0]} {p[1]}" for p in pts[1:]])
         return f"{start} {rest} Z"
-
     colors = ["#dbeafe", "#dcfce7", "#ffedd5", "#f3e8ff", "#fee2e2", "#ccfbf1"]
     parts = [f'<svg xmlns="http://www.w3.org/2000/svg" width="{width}" height="{height}" viewBox="0 0 {width} {height}">', '<rect width="100%" height="100%" fill="white"/>']
     for i, room in enumerate(result["rooms"]):
@@ -186,36 +267,28 @@ def export_svg(result, width, height):
     return "\n".join(parts)
 
 
-def load_any_image(image_file):
-    if image_file is None:
-        raise gr.Error("Upload a PNG, JPG, or SVG file first.")
-    path = image_file if isinstance(image_file, str) else getattr(image_file, "name", None)
-    if not path:
-        raise gr.Error("Could not read uploaded file.")
-    ext = os.path.splitext(path)[1].lower()
-    if ext == ".svg":
-        png_bytes = cairosvg.svg2png(url=path)
-        return Image.open(io.BytesIO(png_bytes)).convert("RGB"), ext
-    return Image.open(path).convert("RGB"), ext
-
-
-def process(image_file, variation, wall_thickness, max_rooms):
+def process(image_file, variation, wall_thickness, max_rooms, use_vision_model):
     image, source_ext = load_any_image(image_file)
-
-    cfg = PlanConfig(variation=variation, wall_thickness=wall_thickness, max_rooms=max_rooms)
+    cfg = PlanConfig(variation=variation, wall_thickness=wall_thickness, max_rooms=max_rooms, use_vision_model=use_vision_model)
     img_bgr = pil_to_bgr(image)
-    binary, contour, approx, mask = preprocess_and_extract_boundary(img_bgr)
-    rooms = generate_rooms(mask, cfg)
-    rendered, result = draw_floorplan(mask, rooms, contour, approx, cfg)
+    binary, contour, approx, boundary_mask = preprocess_and_extract_boundary(img_bgr)
+
+    room_seed_mask = boundary_mask.copy()
+    vision_note = "Heuristic partition mode"
+    if use_vision_model:
+        proposed_mask, note = mask_from_vision_model(image, boundary_mask)
+        if proposed_mask is not None:
+            room_seed_mask = proposed_mask
+        vision_note = note
+
+    rooms = generate_rooms(room_seed_mask, cfg)
+    rendered, result = draw_floorplan(boundary_mask, room_seed_mask, rooms, contour, approx, cfg, vision_note)
+    result["meta"]["source_ext"] = source_ext
 
     svg = export_svg(result, rendered.shape[1], rendered.shape[0])
     json_str = json.dumps(result, indent=2)
-
     preview = bgr_to_pil(rendered)
-    binary_preview = Image.fromarray(binary)
-
-    result["meta"]["source_ext"] = source_ext
-    json_str = json.dumps(result, indent=2)
+    binary_preview = Image.fromarray(room_seed_mask)
 
     tmpdir = Path(tempfile.mkdtemp(prefix="floorplan_space_"))
     zip_path = tmpdir / "floorplan_bundle.zip"
@@ -229,7 +302,7 @@ def process(image_file, variation, wall_thickness, max_rooms):
 with gr.Blocks(title="Closed Shape to Floor Plan") as demo:
     gr.Markdown(
         "# Closed Shape to Floor Plan\n"
-        "Upload a closed-outline PNG, JPG, or SVG. This Space rasterizes SVG input when needed, extracts the outer boundary, partitions the space into rooms, and exports structured JSON + SVG."
+        "Upload a closed-outline PNG, JPG, or SVG. Enable the vision-model option to let a Hugging Face segmentation model propose interior regions before room partitioning."
     )
     with gr.Row():
         with gr.Column(scale=1):
@@ -237,14 +310,15 @@ with gr.Blocks(title="Closed Shape to Floor Plan") as demo:
             variation = gr.Radio(["balanced", "compact", "open"], value="balanced", label="Variation")
             wall_thickness = gr.Slider(2, 16, value=6, step=1, label="Wall thickness")
             max_rooms = gr.Slider(2, 8, value=5, step=1, label="Max rooms")
+            use_vision_model = gr.Checkbox(value=True, label="Use vision model for interior proposal")
             btn = gr.Button("Generate floor plan", variant="primary")
         with gr.Column(scale=1):
             out = gr.Image(label="Generated floor plan")
-            binary = gr.Image(label="Detected binary mask")
+            binary = gr.Image(label="Interior / room seed mask")
     structured = gr.Code(label="Structured JSON", language="json")
     archive = gr.File(label="Download JSON + SVG")
 
-    btn.click(process, inputs=[inp, variation, wall_thickness, max_rooms], outputs=[out, binary, structured, archive])
+    btn.click(process, inputs=[inp, variation, wall_thickness, max_rooms, use_vision_model], outputs=[out, binary, structured, archive])
 
 if __name__ == "__main__":
     demo.launch(theme=gr.themes.Soft(), ssr_mode=False)