Spaces:
Sleeping
Sleeping
| import io | |
| import json | |
| import os | |
| import tempfile | |
| import zipfile | |
| from dataclasses import dataclass | |
| from pathlib import Path | |
| import cairosvg | |
| import cv2 | |
| import gradio as gr | |
| import numpy as np | |
| from PIL import Image | |
| MODEL_AVAILABLE = False | |
| MODEL_ERROR = None | |
| processor = None | |
| seg_pipeline = None | |
| try: | |
| from transformers import pipeline | |
| MODEL_AVAILABLE = True | |
| except Exception as e: | |
| MODEL_ERROR = str(e) | |
| class PlanConfig: | |
| 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: | |
| arr = np.array(img.convert("RGB")) | |
| return cv2.cvtColor(arr, cv2.COLOR_RGB2BGR) | |
| def bgr_to_pil(arr: np.ndarray) -> Image.Image: | |
| rgb = cv2.cvtColor(arr, cv2.COLOR_BGR2RGB) | |
| return Image.fromarray(rgb) | |
| 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": | |
| return [(x1, y1, pos - gap, y2), (pos + gap, y1, x2, y2)] | |
| return [(x1, y1, x2, pos - gap), (x1, pos + gap, x2, y2)] | |
| 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 | |
| roi = mask[max(0, y1):max(0, y2), max(0, x1):max(0, x2)] | |
| if roi.size == 0: | |
| return False | |
| return (roi > 0).mean() >= fill_ratio | |
| 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() | |
| 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 * 4): | |
| if len(rooms) >= target: | |
| break | |
| areas = [max(0, (r[2] - r[0]) * (r[3] - r[1])) for r in rooms] | |
| idx = int(np.argmax(areas)) | |
| r = rooms.pop(idx) | |
| rx1, ry1, rx2, ry2 = r | |
| w, h = rx2 - rx1, ry2 - ry1 | |
| 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.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" | |
| 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 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 | |
| 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(boundary_mask) | |
| cv2.rectangle(room_mask, (x1, y1), (x2, y2), 255, thickness=-1) | |
| room_mask = cv2.bitwise_and(room_mask, room_seed_mask) | |
| c = polygon_from_mask_component(room_mask) | |
| if c is None: | |
| continue | |
| 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) | |
| if M["m00"]: | |
| cx = int(M["m10"] / M["m00"]) | |
| cy = int(M["m01"] / M["m00"]) | |
| cv2.putText(canvas, f"Room {i+1}", (cx - 28, cy), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (20, 20, 20), 1, cv2.LINE_AA) | |
| 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.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": {"image_size": [int(w), int(h)], "vision_note": vision_note}, | |
| } | |
| return canvas, result | |
| def export_svg(result, width, height): | |
| def pts_to_path(pts): | |
| if not pts: | |
| return "" | |
| 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"]): | |
| parts.append(f'<path d="{pts_to_path(room["polygon"])}" fill="{colors[i % len(colors)]}" stroke="#111" stroke-width="2" />') | |
| parts.append(f'<path d="{pts_to_path(result["boundary"])}" fill="none" stroke="#000" stroke-width="6" />') | |
| parts.append('</svg>') | |
| return "\n".join(parts) | |
| 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, use_vision_model=use_vision_model) | |
| img_bgr = pil_to_bgr(image) | |
| 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(room_seed_mask) | |
| tmpdir = Path(tempfile.mkdtemp(prefix="floorplan_space_")) | |
| zip_path = tmpdir / "floorplan_bundle.zip" | |
| with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf: | |
| zf.writestr("floorplan.json", json_str) | |
| zf.writestr("floorplan.svg", svg) | |
| return preview, binary_preview, json_str, str(zip_path) | |
| 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. 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): | |
| inp = gr.File(file_types=[".png", ".jpg", ".jpeg", ".svg"], label="Closed shape file (PNG/JPG/SVG)") | |
| 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="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, use_vision_model], outputs=[out, binary, structured, archive]) | |
| if __name__ == "__main__": | |
| demo.launch(theme=gr.themes.Soft(), ssr_mode=False) |