Update app.py

app.py

@@ -1,84 +1,123 @@
-# ---------- app.py  (backend root of your HF Space) -----------------
-import io, cv2, torch, albumentations as A, numpy as np
+# app.py  ─── Hugging‑Face Space backend  (Python 3.10 + FastAPI + Torch 2.x)
+# ──────────────────────────────────────────────────────────────────────────
+import io, cv2, numpy as np, torch
 from PIL import Image
-from fastapi import FastAPI, File, UploadFile
+from fastapi import FastAPI, UploadFile, File
 from fastapi.middleware.cors import CORSMiddleware
-from starlette.responses import StreamingResponse, JSONResponse
-
-# -------------------------------------------------------------------
-CKPT          = "models/best_model.pth"
-ARCH_NAME     = "NestedUNet"
-INPUT_H = INPUT_W = 512
-THRESH        = 0.50
-ALPHA         = 0.40
-DEEP_SUP      = False
-FRONT_ORIGIN  = "https://amitabhm1.github.io"
-# -------------------------------------------------------------------
-
-# ---------- import model definition from backend_src ----------------
-from backend_src import archs
-# --------------------------------------------------------------------
-
-app = FastAPI(title="Polyp‑Seg‑API")
+from starlette.responses import StreamingResponse, PlainTextResponse
+
+# ------------------------------------------------------------------------
+# 1️⃣  MODEL + TRAINING PARAMS
+# ------------------------------------------------------------------------
+CKPT_PATH      = "models/best_model.pth"        # tracked via Git‑LFS
+ARCH_NAME      = "NestedUNet"
+INPUT_H = INPUT_W = 512                         # resize used during training
+DEEP_SUP       = False
+THRESH_PRIMARY = 0.50                           # first try
+THRESH_FALLBACK= 0.25                           # if mask is empty
+ALPHA          = 0.60                           # overlay transparency (0‑1)
+
+# ------------------------------------------------------------------------
+# 2️⃣  ALLOWED FRONT‑ENDS  (GitHub Pages, localhost for dev, etc.)
+# ------------------------------------------------------------------------
+FRONTENDS = [
+    "https://amitabhm1.github.io",      # ← your GitHub‑Pages root
+    "http://localhost:4173",            # vite / React dev server
+    "http://127.0.0.1:5500"             # plain HTML preview
+]
+
+# ------------------------------------------------------------------------
+# 3️⃣  FASTAPI APP + CORS
+# ------------------------------------------------------------------------
+app = FastAPI(title="Polyp Segmentation API", version="1.0")
 
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=[FRONT_ORIGIN, "http://localhost:4173"],
-    allow_methods=["POST"],
+    allow_origins=FRONTENDS,
+    allow_methods=["POST", "GET", "OPTIONS"],
     allow_headers=["*"],
 )
 
-# ---------------- load checkpoint (fail fast) -----------------------
-if not torch.cuda.is_available():
-    print("⚠ CUDA not available – running on CPU")
-
+# ------------------------------------------------------------------------
+# 4️⃣  LOAD NETWORK ONCE AT STARTUP
+# ------------------------------------------------------------------------
+from backend_src import archs
 device = "cuda" if torch.cuda.is_available() else "cpu"
-net = archs.__dict__[ARCH_NAME](1, 3, DEEP_SUP).to(device)
-
-try:
-    state = torch.load(CKPT, map_location=device)
-except FileNotFoundError as e:
-    raise SystemExit(f"❌ Checkpoint {CKPT} not found in the Space") from e
 
-state = {k.replace("module.", ""): v for k, v in state.items()}
-missing, unexpected = net.load_state_dict(state, strict=False)
-if missing:
-    raise SystemExit(f"❌ Checkpoint is missing keys: {missing[:5]} …")
+model = archs.__dict__[ARCH_NAME](1, 3, DEEP_SUP).to(device)
+state = torch.load(CKPT_PATH, map_location=device)
+state = {k.replace("module.", ""): v for k, v in state.items()}  # DDP → single‑GPU
+model.load_state_dict(state, strict=False)
+model.eval()
 
-net.eval()
-print("✓ model loaded on", device)
+print(f"✓ model loaded ({device})  –  params: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")
 
-# ---------------- preprocessing -------------------------------------
-xform = A.Compose([A.Resize(INPUT_H, INPUT_W)])
+# ------------------------------------------------------------------------
+# 5️⃣  PRE‑ / POST‑PROCESS HELPERS
+# ------------------------------------------------------------------------
+import albumentations as A
+_resize = A.Compose([A.Resize(INPUT_H, INPUT_W, interpolation=cv2.INTER_LINEAR)])
 
-def predict_overlay(pil: Image.Image) -> bytes:
+def preprocess(pil: Image.Image) -> np.ndarray:
+    """PIL → float32 tensor 1×3×H×W on device"""
     img = np.array(pil.convert("RGB"))
-    img_rz = xform(image=img)["image"]
-    ten = torch.from_numpy(img_rz.astype("float32")/255).permute(2,0,1)
-    ten = ten.unsqueeze(0).to(device)
-
-    with torch.no_grad():
-        logits = net(ten)
-        if DEEP_SUP and isinstance(logits,(list,tuple)): logits = logits[-1]
-    mask = (torch.sigmoid(logits)[0,0].cpu().numpy() > THRESH).astype("uint8")
-
-    # colourise mask in red
-    red = np.zeros_like(img_rz); red[:] = (0,0,255)
-    blend = (img_rz*(1-ALPHA) + red*ALPHA).astype("uint8")
-    out = img_rz.copy()
-    out[mask==1] = blend[mask==1]
-
-    buf = io.BytesIO()
-    Image.fromarray(out).save(buf, format="PNG")
-    buf.seek(0);  return buf
-
-# ---------------- endpoint -----------------------------------------
-@app.post("/segment")
+    img = _resize(image=img)["image"].astype("float32") / 255.0
+    ten = torch.from_numpy(img.transpose(2,0,1)).unsqueeze(0).to(device)
+    return ten
+
+def make_mask(prob: np.ndarray) -> np.ndarray:
+    """adaptive mask – fallback if THRESH_PRIMARY gives an empty mask"""
+    mask = (prob >= THRESH_PRIMARY).astype("uint8")
+    if mask.sum() == 0:
+        mask = (prob >= THRESH_FALLBACK).astype("uint8")
+    return mask
+
+def overlay_contour(img: np.ndarray, mask: np.ndarray,
+                    color=(0,0,255), thickness=2) -> np.ndarray:
+    """Draw red contours over the original image"""
+    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    out = img.copy()
+    cv2.drawContours(out, contours, -1, color, thickness)
+    return out
+
+# ------------------------------------------------------------------------
+# 6️⃣  HEALTH‑CHECK
+# ------------------------------------------------------------------------
+@app.get("/ping", summary="Liveness/ready probe")
+def ping():
+    return PlainTextResponse("pong")
+
+# ------------------------------------------------------------------------
+# 7️⃣  MAIN INFERENCE ENDPOINT
+# ------------------------------------------------------------------------
+@app.post("/segment", summary="Upload an image, get PNG with mask overlay")
 async def segment(file: UploadFile = File(...)):
     try:
-        pil = Image.open(io.BytesIO(await file.read()))
-    except Exception:
-        return JSONResponse({"error":"invalid image"}, status_code=400)
+        raw_bytes = await file.read()
+        pil_img   = Image.open(io.BytesIO(raw_bytes)).convert("RGB")
+        inp       = preprocess(pil_img)
+
+        with torch.inference_mode():
+            out = model(inp)
+            if DEEP_SUP and isinstance(out, (list, tuple)):
+                out = out[-1]
+
+        prob = torch.sigmoid(out)[0,0].cpu().numpy()
+        mask = make_mask(prob)
+
+        # ---- overlay  ----
+        resized = pil_img.resize((INPUT_W, INPUT_H), Image.BILINEAR)
+        overlay = overlay_contour(np.array(resized), mask)
+
+        buf = io.BytesIO()
+        Image.fromarray(overlay).save(buf, format="PNG")
+        buf.seek(0)
+
+        # ● optional debugging in Space logs
+        print(f"foreground px: {mask.sum():,}  –  file: {file.filename}")
+
+        return StreamingResponse(buf, media_type="image/png")
 
-    png = predict_overlay(pil)
-    return StreamingResponse(png, media_type="image/png")
+    except Exception as e:
+        # keep same structure so JS can always parse JSON on error
+        return {"error": str(e)}