Create app.py

app.py

@@ -0,0 +1,66 @@
+import gradio as gr
+import torch
+import numpy as np
+from PIL import Image
+from detectron2 import model_zoo
+from detectron2.config import get_cfg
+from detectron2.engine import DefaultPredictor
+from detectron2.utils.visualizer import Visualizer
+from detectron2.data import MetadataCatalog
+from huggingface_hub import hf_hub_download
+import os
+
+# Only 1 class
+ROOM_CLASSES = ["room"]
+
+# Hugging Face repo info
+HF_REPO = "TallManager267/SG_Room_Segmentation"
+WEIGHTS_FILE = "sg_room_segmentation_726_4000itr_0.004lr.pth"
+
+# Download weights from Hugging Face
+os.makedirs("model", exist_ok=True)
+weights_path = hf_hub_download(repo_id=HF_REPO, filename=WEIGHTS_FILE)
+
+def load_model(weights_path):
+    cfg = get_cfg()
+    cfg.merge_from_file(model_zoo.get_config_file(
+        "COCO-InstanceSegmentation/mask_rcnn_R_101_FPN_3x.yaml"
+    ))
+    cfg.MODEL.WEIGHTS = weights_path
+    cfg.MODEL.DEVICE = "cpu"
+    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
+    cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1 
+
+    # Metadata for visualization
+    metadata = MetadataCatalog.get("room_metadata")
+    metadata.set(thing_classes=ROOM_CLASSES)
+
+    predictor = DefaultPredictor(cfg)
+    return predictor, metadata
+
+predictor, metadata = load_model(weights_path)
+
+def predict(pil_img):
+    # Convert to RGB and ensure dtype is uint8
+    img = np.array(pil_img.convert("RGB"), dtype=np.uint8)
+
+    outputs = predictor(img)
+    instances = outputs["instances"].to("cpu")
+
+    v = Visualizer(
+        img[:, :, ::-1],
+        metadata=metadata,
+        scale=1.0
+    )
+    out = v.draw_instance_predictions(instances)
+    out_img = out.get_image()[:, :, ::-1]
+
+    return Image.fromarray(out_img)
+
+gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil", label="Upload floor plan"),
+    outputs=gr.Image(type="pil", label="Room segmentation"),
+    title="Room Segmentation (Detectron2)",
+    description="Upload a floor plan image to segment the room using Detectron2."
+).launch(server_name="0.0.0.0", server_port=7860)