Update app.py

app.py

@@ -1,4 +1,3 @@
-import os
 import io
 import base64
 from typing import Optional
@@ -12,7 +11,6 @@ from fastapi.middleware.cors import CORSMiddleware
 
 from detectron2.engine import DefaultPredictor
 from detectron2.config import get_cfg
-from detectron2 import model_zoo
 from detectron2.data import MetadataCatalog
 
 # -------------------
@@ -26,9 +24,9 @@ cfg.merge_from_file(
 )
 cfg.MODEL.ROI_HEADS.NUM_CLASSES = len(classes)
 cfg.MODEL.POINT_HEAD.NUM_CLASSES = len(classes)
-cfg.MODEL.WEIGHTS = "model_final.pth"  # path to your trained PointRend model
+cfg.MODEL.WEIGHTS = "model_final.pth"  # make sure this file is in repo
 cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
-cfg.MODEL.DEVICE = "cpu"  # change to "cuda" if GPU available
+cfg.MODEL.DEVICE = "cpu"
 MetadataCatalog.get("__unused__").thing_classes = classes
 
 predictor = DefaultPredictor(cfg)
@@ -43,13 +41,10 @@ app.add_middleware(
     allow_methods=["*"], allow_headers=["*"],
 )
 
-
 # -------------------
-# Post-processing helpers
+# Helpers
 # -------------------
 def postprocess_simplified(mask: np.ndarray, epsilon_factor: float = 0.01) -> np.ndarray:
-    if mask is None:
-        return None
     mask_uint8 = (mask * 255).astype(np.uint8) if mask.max() <= 1 else mask.astype(np.uint8)
     bw = (mask_uint8 > 127).astype(np.uint8) * 255
     contours, _ = cv2.findContours(bw, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
@@ -62,10 +57,7 @@ def postprocess_simplified(mask: np.ndarray, epsilon_factor: float = 0.01) -> np
     cv2.fillPoly(simp, [approx], 255)
     return simp
 
-
 def mask_to_polygon(mask: np.ndarray, epsilon_factor: float = 0.01, min_area: int = 150) -> Optional[np.ndarray]:
-    if mask is None:
-        return None
     mask_uint8 = (mask * 255).astype(np.uint8) if mask.max() <= 1 else mask.astype(np.uint8)
     bw = (mask_uint8 > 127).astype(np.uint8) * 255
     contours, _ = cv2.findContours(bw, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
@@ -78,16 +70,14 @@ def mask_to_polygon(mask: np.ndarray, epsilon_factor: float = 0.01, min_area: in
     approx = cv2.approxPolyDP(contour, epsilon, True)
     return approx.reshape(-1, 2)
 
-
 def im_to_b64_png(im: np.ndarray) -> str:
     ok, buf = cv2.imencode(".png", im)
     if not ok:
         raise RuntimeError("Failed to encode image")
     return base64.b64encode(buf).decode("utf-8")
 
-
 # -------------------
-# Prediction endpoint: return only one G-Zone + overlay + polygon
+# Prediction endpoint
 # -------------------
 @app.post("/polygon")
 async def polygon_endpoint(file: UploadFile = File(...)):
@@ -104,17 +94,14 @@ async def polygon_endpoint(file: UploadFile = File(...)):
     if len(instances) == 0:
         return JSONResponse(content={"chosen": None, "polygon": None, "image": None})
 
-    # --- Pick one G-Zone: first instance (can change to largest if needed) ---
     idx = 0
     cls_id = int(instances.pred_classes[idx])
     cls_name = classes[cls_id]
     raw_mask = instances.pred_masks[idx].numpy().astype(np.uint8)
 
-    # --- Simplify mask and convert to polygon ---
     simp_mask = postprocess_simplified(raw_mask)
     poly = mask_to_polygon(simp_mask)
 
-    # --- Overlay polygon on original image ---
     overlay = im.copy()
     poly_list = None
     if poly is not None:
@@ -123,8 +110,4 @@ async def polygon_endpoint(file: UploadFile = File(...)):
 
     img_b64 = im_to_b64_png(overlay)
 
-    return {
-        "chosen": cls_name,
-        "polygon": poly_list,
-        "image": img_b64
-    }
+    return {"chosen": cls_name, "polygon": poly_list, "image": img_b64}