Update app.py

app.py

@@ -1,20 +1,15 @@
-# app.py — Gradio app for panoramic radiograph segmentation (Detectron2, CPU/GPU)
-
 import os, json, time
 import numpy as np, cv2, 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
 
-# --------- 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 ---------
 cfg = get_cfg()
 cfg.merge_from_file(CFG_PATH)
 cfg.MODEL.DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
@@ -31,16 +26,14 @@ 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) ---------
 MAX_SIDE = 1600
 
 def segment(rgb: np.ndarray):
     t0 = time.time()
     h0, w0 = rgb.shape[:2]
-    long_side = max(h0, w0)
     scale = 1.0
-    if long_side > MAX_SIDE:
-        scale = MAX_SIDE / long_side
+    if max(h0, w0) > MAX_SIDE:
+        scale = MAX_SIDE / max(h0, w0)
         rgb_small = cv2.resize(rgb, (int(w0*scale), int(h0*scale)), interpolation=cv2.INTER_AREA)
     else:
         rgb_small = rgb
@@ -52,14 +45,15 @@ def segment(rgb: np.ndarray):
     overlay_rgb = vis.draw_instance_predictions(inst).get_image()
 
     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 []
-    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]  # 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)})
+    if inst.has("pred_boxes"):
+        boxes = inst.pred_boxes.tensor.numpy().tolist()
+        scores = inst.scores.numpy().tolist() if inst.has("scores") else [None]*len(boxes)
+        classes_idx = inst.pred_classes.numpy().tolist() if inst.has("pred_classes") else [0]*len(boxes)
+        inv = (1.0/scale) if scale != 1.0 else 1.0
+        for b, s, c in zip(boxes, scores, classes_idx):
+            b = [float(x*inv) for x in b]
+            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,
@@ -76,4 +70,4 @@ with gr.Blocks(title="Panoramic Radiograph Segmentation") as demo:
     gr.Button("Run").click(segment, inputs=img_in, outputs=[img_out, json_out], api_name="/predict")
 
 if __name__ == "__main__":
-    demo.launch()  # no share=True on Spaces
+    demo.launch()  # Spaces provides the permanent URL