app.py

"""
Floor Plan Segmentation API — Hugging Face Space
Segments rooms, walls, doors, windows from floor plan images.

Uses Mask2Former for instance segmentation with fallback to
color-based contour detection for robustness.

API endpoint: POST /api/predict
"""

import gradio as gr
import numpy as np
import cv2
import json
from PIL import Image
import io
import base64

# ─── Model loading ───────────────────────────────────────────

MODEL = None
PROCESSOR = None
USE_MASK2FORMER = False

def load_model():
    """Try loading Mask2Former; fall back to OpenCV contour detection."""
    global MODEL, PROCESSOR, USE_MASK2FORMER
    try:
        from transformers import AutoImageProcessor, Mask2FormerForInstanceSegmentation
        PROCESSOR = AutoImageProcessor.from_pretrained(
            "Hyunwoo1605/mask2former-floorplan-instance-segmentation"
        )
        MODEL = Mask2FormerForInstanceSegmentation.from_pretrained(
            "Hyunwoo1605/mask2former-floorplan-instance-segmentation"
        )
        MODEL.eval()
        USE_MASK2FORMER = True
        print("[INFO] Mask2Former model loaded successfully")
    except Exception as e:
        print(f"[WARN] Could not load Mask2Former: {e}")
        print("[INFO] Using OpenCV contour-based fallback")
        USE_MASK2FORMER = False

load_model()


# ─── Mask2Former inference ───────────────────────────────────

def segment_mask2former(image: np.ndarray) -> dict:
    """Run Mask2Former instance segmentation on floor plan image."""
    import torch

    pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    inputs = PROCESSOR(images=pil_image, return_tensors="pt")

    with torch.no_grad():
        outputs = MODEL(**inputs)

    # Post-process: get instance masks and labels
    result = PROCESSOR.post_process_instance_segmentation(
        outputs, target_sizes=[pil_image.size[::-1]]
    )[0]

    rooms = []
    walls = []
    h, w = image.shape[:2]

    for seg_info in result["segments_info"]:
        mask = (result["segmentation"] == seg_info["id"]).numpy().astype(np.uint8)
        label_id = seg_info["label_id"]
        score = float(seg_info["score"])

        # Extract contour from mask
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        if not contours:
            continue

        # Largest contour = room boundary
        contour = max(contours, key=cv2.contourArea)
        area_px = cv2.contourArea(contour)
        if area_px < 100:  # skip tiny segments
            continue

        # Simplify contour to polygon
        epsilon = 0.02 * cv2.arcLength(contour, True)
        approx = cv2.approxPolyDP(contour, epsilon, True)
        boundary = [{"x": float(p[0][0]) / w, "y": float(p[0][1]) / h} for p in approx]

        # Get label name from model config
        label_name = MODEL.config.id2label.get(label_id, f"class_{label_id}")

        # Map to room type
        room_type = classify_label(label_name)

        if room_type == "wall":
            # Extract wall segments from contour
            for i in range(len(approx)):
                j = (i + 1) % len(approx)
                walls.append({
                    "start": {"x": float(approx[i][0][0]) / w, "y": float(approx[i][0][1]) / h},
                    "end": {"x": float(approx[j][0][0]) / w, "y": float(approx[j][0][1]) / h},
                    "is_exterior": False,
                })
        else:
            rooms.append({
                "name": label_name,
                "type": room_type,
                "boundary": boundary,
                "area_estimate_m2": 0,  # needs scale info
                "has_door": False,
                "has_window": False,
                "confidence": score,
                "floor_type": "parkett",
            })

    return {"rooms": rooms, "walls": walls, "doors": [], "windows": [], "method": "mask2former"}


# ─── OpenCV contour-based fallback ──────────────────────────

def segment_opencv(image: np.ndarray) -> dict:
    """
    OpenCV-based floor plan segmentation using adaptive thresholding
    and contour detection. Works without GPU or ML models.
    """
    h, w = image.shape[:2]
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # Adaptive threshold to detect walls (dark lines on light/white background)
    thresh = cv2.adaptiveThreshold(
        gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 11, 5
    )

    # Morphological operations to clean up wall detection
    kernel_close = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    walls_mask = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel_close, iterations=2)

    # Dilate walls slightly to close small gaps
    kernel_dilate = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
    walls_dilated = cv2.dilate(walls_mask, kernel_dilate, iterations=1)

    # Invert to get room regions (white = room interior)
    rooms_mask = cv2.bitwise_not(walls_dilated)

    # Find room contours
    contours, _ = cv2.findContours(rooms_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    rooms = []
    min_room_area = (w * h) * 0.005  # min 0.5% of image area
    max_room_area = (w * h) * 0.5    # max 50% of image area

    for i, contour in enumerate(contours):
        area = cv2.contourArea(contour)
        if area < min_room_area or area > max_room_area:
            continue

        # Simplify contour
        epsilon = 0.015 * cv2.arcLength(contour, True)
        approx = cv2.approxPolyDP(contour, epsilon, True)

        if len(approx) < 3:
            continue

        boundary = [{"x": float(p[0][0]) / w, "y": float(p[0][1]) / h} for p in approx]

        # Compute bounding rect for aspect ratio
        x_r, y_r, w_r, h_r = cv2.boundingRect(approx)
        aspect = max(w_r, h_r) / max(min(w_r, h_r), 1)

        # Simple room classification by size and shape
        relative_area = area / (w * h)
        if aspect > 5:
            room_type = "hallway"
            name = f"Flur {i+1}"
        elif relative_area > 0.08:
            room_type = "living_room"
            name = f"Raum {i+1}"
        elif relative_area < 0.02:
            room_type = "wc"
            name = f"WC/Bad {i+1}"
        else:
            room_type = "custom"
            name = f"Raum {i+1}"

        rooms.append({
            "name": name,
            "type": room_type,
            "boundary": boundary,
            "area_estimate_m2": 0,
            "has_door": False,
            "has_window": False,
            "confidence": 0.6,
            "floor_type": "parkett",
        })

    # Extract wall segments using Hough Line Transform
    walls = []
    lines = cv2.HoughLinesP(walls_mask, 1, np.pi / 180, 80, minLineLength=30, maxLineGap=10)
    if lines is not None:
        for line in lines[:200]:  # cap at 200 wall segments
            x1, y1, x2, y2 = line[0]
            walls.append({
                "start": {"x": float(x1) / w, "y": float(y1) / h},
                "end": {"x": float(x2) / w, "y": float(y2) / h},
                "is_exterior": False,
            })

    # Detect doors (arcs / small circular segments)
    doors = []
    circles = cv2.HoughCircles(
        gray, cv2.HOUGH_GRADIENT, 1, 50,
        param1=100, param2=30, minRadius=15, maxRadius=80
    )
    if circles is not None:
        for circle in circles[0][:20]:
            cx, cy, r = circle
            doors.append({
                "position": {"x": float(cx) / w, "y": float(cy) / h},
                "width_mm": int(r * 2 * 10),  # rough estimate
                "type": "standard",
            })

    return {"rooms": rooms, "walls": walls, "doors": doors, "windows": [], "method": "opencv"}


# ─── Label mapping ──────────────────────────────────────────

LABEL_MAP = {
    "wall": "wall",
    "room": "custom",
    "living": "living_room",
    "living_room": "living_room",
    "bedroom": "bedroom",
    "bathroom": "bathroom",
    "kitchen": "kitchen",
    "hallway": "hallway",
    "corridor": "hallway",
    "closet": "storage",
    "storage": "storage",
    "balcony": "terrace",
    "door": "door",
    "window": "window",
    "dining": "dining_room",
    "office": "study",
    "garage": "garage",
    "stairs": "staircase",
    "toilet": "wc",
    "wc": "wc",
    "utility": "utility_room",
    "laundry": "utility_room",
    "entrance": "entrance",
}

def classify_label(label: str) -> str:
    """Map model output label to standard room type."""
    label_lower = label.lower().strip()
    for key, value in LABEL_MAP.items():
        if key in label_lower:
            return value
    return "custom"


# ─── Main API function ──────────────────────────────────────

def analyze_floor_plan(image: np.ndarray) -> dict:
    """Analyze a floor plan image and return segmentation results."""
    if image is None:
        return {"error": "No image provided", "rooms": [], "walls": [], "doors": [], "windows": []}

    # Ensure BGR format
    if len(image.shape) == 2:
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
    elif image.shape[2] == 4:
        image = cv2.cvtColor(image, cv2.COLOR_RGBA2BGR)

    # Run segmentation
    if USE_MASK2FORMER:
        result = segment_mask2former(image)
    else:
        result = segment_opencv(image)

    result["image_width"] = image.shape[1]
    result["image_height"] = image.shape[0]
    result["_version"] = "hf-space-v1"
    result["_coordSystem"] = "normalized"  # all coords 0..1
    result["notes"] = f"Analyzed using {result.get('method', 'unknown')} method. {len(result.get('rooms', []))} rooms detected."

    return result


# ─── Gradio Interface ────────────────────────────────────────

def gradio_predict(image):
    """Gradio wrapper that returns JSON string + annotated image."""
    result = analyze_floor_plan(image)

    # Draw annotations on image for visualization
    annotated = image.copy()
    h, w = annotated.shape[:2]

    colors = [
        (66, 133, 244), (234, 67, 53), (251, 188, 4), (52, 168, 83),
        (171, 71, 188), (255, 112, 67), (0, 172, 193), (124, 179, 66),
    ]

    for i, room in enumerate(result.get("rooms", [])):
        color = colors[i % len(colors)]
        boundary = room.get("boundary", [])
        if len(boundary) < 3:
            continue

        pts = np.array([[int(p["x"] * w), int(p["y"] * h)] for p in boundary], dtype=np.int32)

        # Semi-transparent fill
        overlay = annotated.copy()
        cv2.fillPoly(overlay, [pts], color)
        cv2.addWeighted(overlay, 0.3, annotated, 0.7, 0, annotated)

        # Boundary outline
        cv2.polylines(annotated, [pts], True, color, 2)

        # Label
        M = cv2.moments(pts)
        if M["m00"] > 0:
            cx = int(M["m10"] / M["m00"])
            cy = int(M["m01"] / M["m00"])
            label = room.get("name", f"Room {i+1}")
            cv2.putText(annotated, label, (cx - 30, cy), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
            cv2.putText(annotated, label, (cx - 30, cy), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)

    return annotated, json.dumps(result, indent=2, ensure_ascii=False)


with gr.Blocks(title="Floor Plan Segmentation API") as demo:
    gr.Markdown("""
    # Floor Plan Segmentation
    Upload a floor plan image to detect rooms, walls, doors, and windows.

    **API Usage:** `POST /api/predict` with `{"data": [<base64_image>]}`
    """)

    with gr.Row():
        with gr.Column():
            input_image = gr.Image(label="Floor Plan", type="numpy")
            analyze_btn = gr.Button("Analyze", variant="primary")
        with gr.Column():
            output_image = gr.Image(label="Segmentation Result")
            output_json = gr.Textbox(label="JSON Result", lines=15, max_lines=30)

    analyze_btn.click(
        fn=gradio_predict,
        inputs=[input_image],
        outputs=[output_image, output_json],
        api_name="predict",
    )

demo.launch(show_api=True)