Update app.py

app.py

@@ -1,123 +1,73 @@
-# app.py  ─── Hugging‑Face Space backend  (Python 3.10 + FastAPI + Torch 2.x)
-# ──────────────────────────────────────────────────────────────────────────
-import io, cv2, numpy as np, torch
+# app.py  (debug‑friendly overlay)
+import io, cv2, numpy as np, torch, albumentations as A
 from PIL import Image
 from fastapi import FastAPI, UploadFile, File
 from fastapi.middleware.cors import CORSMiddleware
 from starlette.responses import StreamingResponse, PlainTextResponse
+from backend_src import archs
 
-# ------------------------------------------------------------------------
-# 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=FRONTENDS,
-    allow_methods=["POST", "GET", "OPTIONS"],
-    allow_headers=["*"],
-)
+# ─────────── config ────────────
+CKPT          = "models/best_model.pth"
+ARCH          = "NestedUNet"
+INP_H = INP_W = 512
+DEEP_SUP      = False
+TH_MAIN, TH_FALL = .50, .25
+ALPHA         = .60
+FRONTENDS     = ["https://amitabhm1.github.io","http://localhost:4173"]
+# ───────────────────────────────
 
-# ------------------------------------------------------------------------
-# 4️⃣  LOAD NETWORK ONCE AT STARTUP
-# ------------------------------------------------------------------------
-from backend_src import archs
 device = "cuda" if torch.cuda.is_available() else "cpu"
-
-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  = archs.__dict__[ARCH](1,3,DEEP_SUP).to(device)
+state  = torch.load(CKPT, map_location=device)
+state  = {k.removeprefix("module."):v for k,v in state.items()}
 model.load_state_dict(state, strict=False)
 model.eval()
+print("✓ model loaded on", device)
 
-print(f"✓ model loaded ({device})  –  params: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")
-
-# ------------------------------------------------------------------------
-# 5️⃣  PRE‑ / POST‑PROCESS HELPERS
-# ------------------------------------------------------------------------
-import albumentations as A
-_resize = A.Compose([A.Resize(INPUT_H, INPUT_W, interpolation=cv2.INTER_LINEAR)])
+resize = A.Compose([A.Resize(INP_H, INP_W)])
 
-def preprocess(pil: Image.Image) -> np.ndarray:
-    """PIL → float32 tensor 1×3×H×W on device"""
-    img = np.array(pil.convert("RGB"))
-    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 to_tensor(pil:Image.Image)->torch.Tensor:
+    arr = np.array(pil.convert("RGB")).astype("float32")/255.0  # ★ match training!
+    arr = resize(image=arr)["image"]
+    return torch.from_numpy(arr.transpose(2,0,1)).unsqueeze(0).to(device)
 
-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 mask_from_prob(prob:np.ndarray)->np.ndarray:
+    m = (prob>=TH_MAIN).astype("uint8")
+    return m if m.sum() else (prob>=TH_FALL).astype("uint8")
 
-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)
+def overlay(img,np_mask):                               # red contour
+    cnts,_ = cv2.findContours(np_mask,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE)
     out = img.copy()
-    cv2.drawContours(out, contours, -1, color, thickness)
+    cv2.drawContours(out,cnts,-1,(0,0,255),2)
     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:
-        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]
+def heatmap(img,prob):
+    hm = cv2.applyColorMap((prob*255).astype("uint8"),cv2.COLORMAP_JET)
+    return cv2.addWeighted(img,1-ALPHA,hm,ALPHA,0)
 
-        prob = torch.sigmoid(out)[0,0].cpu().numpy()
-        mask = make_mask(prob)
+app = FastAPI()
+app.add_middleware(CORSMiddleware,
+    allow_origins=FRONTENDS, allow_methods=["POST","GET","OPTIONS"], allow_headers=["*"])
 
-        # ---- overlay  ----
-        resized = pil_img.resize((INPUT_W, INPUT_H), Image.BILINEAR)
-        overlay = overlay_contour(np.array(resized), mask)
+@app.get("/ping")
+def ping(): return PlainTextResponse("pong")
 
-        buf = io.BytesIO()
-        Image.fromarray(overlay).save(buf, format="PNG")
-        buf.seek(0)
+@app.post("/segment")
+async def segment(file:UploadFile=File(...)):
+    raw = await file.read()
+    pil = Image.open(io.BytesIO(raw))
+    inp = to_tensor(pil)
 
-        # ● optional debugging in Space logs
-        print(f"foreground px: {mask.sum():,}  –  file: {file.filename}")
+    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()
+    mmax, mmin = prob.max(), prob.min()
+    mask = mask_from_prob(prob)
+    print(f"{file.filename}: min={mmin:.3f}  max={mmax:.3f}  fg={mask.sum()}")
 
-        return StreamingResponse(buf, media_type="image/png")
+    frame = np.array(pil.resize((INP_W,INP_H),Image.BILINEAR))
+    vis   = overlay(frame,mask) if mask.sum() else heatmap(frame,prob)
 
-    except Exception as e:
-        # keep same structure so JS can always parse JSON on error
-        return {"error": str(e)}
+    buf = io.BytesIO(); Image.fromarray(vis).save(buf,format="PNG"); buf.seek(0)
+    return StreamingResponse(buf, media_type="image/png")