Update app.py

app.py

@@ -1,18 +1,20 @@
-# app.py — Gradio app for panoramic radiograph segmentation (Detectron2)
+# app.py — Gradio app for panoramic radiograph segmentation (Detectron2, CPU/GPU)
+
+import os, json, time
+import numpy as np, cv2, torch, gradio as gr
 
-import os, json, time, cv2, numpy as np, torch, gradio as gr
 from detectron2.config import get_cfg
 from detectron2.engine import DefaultPredictor
 from detectron2.data import MetadataCatalog
 from detectron2.utils.visualizer import Visualizer, ColorMode
 
-# --- Local artifacts (uploaded to this Space) ---
-LOAD_DIR   = "./artifacts"
-WEIGHTS    = f"{LOAD_DIR}/model_final.pth"
-CFG_PATH   = f"{LOAD_DIR}/config.yaml"
-CLASSES_PATH = f"{LOAD_DIR}/classes.json"
+# --------- Artifacts already in the repo ---------
+LOAD_DIR = "./artifacts"
+WEIGHTS = os.path.join(LOAD_DIR, "model_final.pth")
+CFG_PATH = os.path.join(LOAD_DIR, "config.yaml")
+CLASSES_PATH = os.path.join(LOAD_DIR, "classes.json")
 
-# --- Build cfg & load model ---
+# --------- Build cfg & load model ---------
 cfg = get_cfg()
 cfg.merge_from_file(CFG_PATH)
 cfg.MODEL.DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
@@ -29,7 +31,7 @@ predictor = DefaultPredictor(cfg)
 meta = MetadataCatalog.get("inference_only")
 meta.thing_classes = classes if classes else [f"class_{i}" for i in range(cfg.MODEL.ROI_HEADS.NUM_CLASSES)]
 
-# --- Inference (auto-downscale very wide panoramics) ---
+# --------- Inference (auto-downscale very wide panoramics) ---------
 MAX_SIDE = 1600
 
 def segment(rgb: np.ndarray):
@@ -43,8 +45,7 @@ def segment(rgb: np.ndarray):
     else:
         rgb_small = rgb
 
-    # Detectron2 expects BGR; our array is RGB
-    outputs = predictor(rgb_small[:, :, ::-1])
+    outputs = predictor(rgb_small[:, :, ::-1])  # predictor expects BGR
     inst = outputs["instances"].to("cpu")
 
     vis = Visualizer(rgb_small, metadata=meta, scale=1.0, instance_mode=ColorMode.IMAGE_BW)
@@ -52,18 +53,18 @@ def segment(rgb: np.ndarray):
 
     dets = []
     boxes_small = inst.pred_boxes.tensor.numpy().tolist() if inst.has("pred_boxes") else []
-    scores      = inst.scores.numpy().tolist() if inst.has("scores") else []
+    scores = inst.scores.numpy().tolist() if inst.has("scores") else []
     classes_idx = inst.pred_classes.numpy().tolist() if inst.has("pred_classes") else []
     inv = (1.0/scale) if scale != 1.0 else 1.0
     for b, s, c in zip(boxes_small, scores, classes_idx):
-        b = [float(x*inv) for x in b]  # map back to original size
+        b = [float(x*inv) for x in b]  # back to original image coords
         label = meta.thing_classes[c] if 0 <= c < len(meta.thing_classes) else str(c)
         dets.append({"box": b, "class": label, "score": float(s)})
 
     return overlay_rgb, {
         "instances": dets,
         "original_size": [int(h0), int(w0)],
-        "latency_ms": int((time.time() - t0)*1000),
+        "latency_ms": int((time.time()-t0)*1000),
     }
 
 with gr.Blocks(title="Panoramic Radiograph Segmentation") as demo:
@@ -71,8 +72,7 @@ with gr.Blocks(title="Panoramic Radiograph Segmentation") as demo:
     img_in = gr.Image(type="numpy", label="Panoramic Radiograph")
     with gr.Row():
         img_out = gr.Image(label="Overlay")
-        json_out = gr.JSON(label="Detections (boxes in original image coords)")
-    # Keep a stable API route for future programmatic calls
+        json_out = gr.JSON(label="Detections (boxes in original coords)")
     gr.Button("Run").click(segment, inputs=img_in, outputs=[img_out, json_out], api_name="/predict")
 
 if __name__ == "__main__":