app.py

"""
SAM3 Floor Plan Detection — Lightweight proxy
Calls the SAM3 demo space for segmentation, converts masks to JSON coordinates.
"""

import gradio as gr
import numpy as np
from PIL import Image
from gradio_client import Client, handle_file
import cv2
import json
import time
import tempfile
import os

SAM3_DEMO = "prithivMLmods/SAM3-Demo"


def mask_to_lines(mask_img: np.ndarray, min_length: int = 20) -> list:
    """Convert a segmentation mask image to line segments."""
    # Convert to grayscale if needed
    if len(mask_img.shape) == 3:
        gray = cv2.cvtColor(mask_img, cv2.COLOR_RGB2GRAY)
    else:
        gray = mask_img

    # Threshold to binary — segmented regions are colored, background is not
    _, binary = cv2.threshold(gray, 10, 255, cv2.THRESH_BINARY)

    # Skeletonize to get thin lines
    try:
        skeleton = cv2.ximgproc.thinning(binary)
    except AttributeError:
        # Fallback: use Canny edge detection
        skeleton = cv2.Canny(binary, 50, 150)

    # Detect line segments
    lines = cv2.HoughLinesP(
        skeleton, rho=1, theta=np.pi / 180,
        threshold=25, minLineLength=min_length, maxLineGap=15,
    )

    if lines is None:
        return []

    result = []
    for line in lines:
        x1, y1, x2, y2 = line[0]
        result.append({"position": [[int(x1), int(y1)], [int(x2), int(y2)]]})

    return merge_close_lines(result)


def mask_to_bboxes(mask_img: np.ndarray, min_area: int = 80) -> list:
    """Convert a segmentation mask image to bounding boxes."""
    if len(mask_img.shape) == 3:
        gray = cv2.cvtColor(mask_img, cv2.COLOR_RGB2GRAY)
    else:
        gray = mask_img

    _, binary = cv2.threshold(gray, 10, 255, cv2.THRESH_BINARY)
    contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    bboxes = []
    for contour in contours:
        if cv2.contourArea(contour) < min_area:
            continue
        x, y, w, h = cv2.boundingRect(contour)
        bboxes.append({"bbox": [int(x), int(y), int(x + w), int(y + h)]})

    return bboxes


def merge_close_lines(lines: list, threshold: int = 10) -> list:
    """Merge line segments that are close and roughly parallel."""
    if not lines:
        return lines

    merged = []
    used = set()

    for i, la in enumerate(lines):
        if i in used:
            continue
        p = la["position"]
        x1, y1, x2, y2 = p[0][0], p[0][1], p[1][0], p[1][1]
        dx, dy = abs(x2 - x1), abs(y2 - y1)
        horiz = dx > dy

        if horiz and dy < 5:
            avg_y = (y1 + y2) // 2
            y1 = y2 = avg_y
        elif not horiz and dx < 5:
            avg_x = (x1 + x2) // 2
            x1 = x2 = avg_x

        for j, lb in enumerate(lines):
            if j <= i or j in used:
                continue
            q = lb["position"]
            bx1, by1, bx2, by2 = q[0][0], q[0][1], q[1][0], q[1][1]
            bdx, bdy = abs(bx2 - bx1), abs(by2 - by1)
            b_horiz = bdx > bdy

            if horiz != b_horiz:
                continue
            if horiz and abs(y1 - (by1 + by2) // 2) < threshold:
                x1, x2 = min(x1, bx1, bx2), max(x2, bx1, bx2)
                used.add(j)
            elif not horiz and abs(x1 - (bx1 + bx2) // 2) < threshold:
                y1, y2 = min(y1, by1, by2), max(y2, by1, by2)
                used.add(j)

        merged.append({"position": [[x1, y1], [x2, y2]]})
    return merged


def detect_floor_plan(image: Image.Image) -> dict:
    """Run SAM3 on a floor plan image via the demo space."""
    if image is None:
        return {"error": "No image provided"}

    start = time.time()
    image = image.convert("RGB")
    w, h = image.size
    print(f"[SAM3] Processing {w}x{h} image...")

    # Save image to temp file for gradio_client
    tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
    image.save(tmp.name)
    tmp.close()

    results = {"walls": [], "doors": [], "windows": [], "rooms": [],
               "_imgWidth": w, "_imgHeight": h}

    try:
        client = Client(SAM3_DEMO)

        # Detect walls ("black line")
        print("[SAM3] Detecting walls...")
        try:
            wall_result = client.predict(
                source_img=handle_file(tmp.name),
                text_query="black line",
                api_name="/run_image_segmentation",
            )
            # Result is a tuple: (output_image_path, ...)
            if wall_result and isinstance(wall_result, (tuple, list)):
                mask_path = wall_result[0] if isinstance(wall_result[0], str) else wall_result[0].get("path", "")
                if mask_path and os.path.exists(mask_path):
                    mask_img = cv2.imread(mask_path)
                    if mask_img is not None:
                        results["walls"] = mask_to_lines(mask_img)
                        print(f"[SAM3] Found {len(results['walls'])} walls")
        except Exception as e:
            print(f"[SAM3] Wall detection error: {e}")

        # Detect doors ("curved line")
        print("[SAM3] Detecting doors...")
        try:
            door_result = client.predict(
                source_img=handle_file(tmp.name),
                text_query="curved line",
                api_name="/run_image_segmentation",
            )
            if door_result and isinstance(door_result, (tuple, list)):
                mask_path = door_result[0] if isinstance(door_result[0], str) else door_result[0].get("path", "")
                if mask_path and os.path.exists(mask_path):
                    mask_img = cv2.imread(mask_path)
                    if mask_img is not None:
                        results["doors"] = mask_to_bboxes(mask_img)
                        print(f"[SAM3] Found {len(results['doors'])} doors")
        except Exception as e:
            print(f"[SAM3] Door detection error: {e}")

    finally:
        os.unlink(tmp.name)

    elapsed = time.time() - start
    results["_elapsed"] = round(elapsed, 2)
    results["_source"] = "sam3"
    print(f"[SAM3] Done in {elapsed:.1f}s")

    return results


demo = gr.Interface(
    fn=detect_floor_plan,
    inputs=gr.Image(type="pil", label="Floor Plan Image"),
    outputs=gr.JSON(label="Detected Elements"),
    title="SAM3 Floor Plan Detection",
    description="Detects walls and doors in floor plan images using Meta SAM3.",
)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)