Update app.py

app.py

@@ -1,81 +1,84 @@
-# app.py  – Hugging‑Face Space backend
+# ---------- app.py  (backend root of your HF Space) -----------------
 import io, cv2, torch, albumentations as A, numpy as np
 from PIL import Image
 from fastapi import FastAPI, File, UploadFile
 from fastapi.middleware.cors import CORSMiddleware
-from starlette.responses import StreamingResponse
+from starlette.responses import StreamingResponse, JSONResponse
 
-# ---------- model code lives inside backend_src -------------
-from backend_src import archs                                    # <- you copied this
-# ------------------------------------------------------------
-
-### CONFIG #### -------------------------------------------------
-CKPT          = "models/best_model.pth"   # uploaded via Git‑LFS
+# -------------------------------------------------------------------
+CKPT          = "models/best_model.pth"
 ARCH_NAME     = "NestedUNet"
-INPUT_H = INPUT_W = 512                   # what you trained with
-THRESH        = 0.5                       # binarisation
+INPUT_H = INPUT_W = 512
+THRESH        = 0.50
+ALPHA         = 0.40
 DEEP_SUP      = False
-ALPHA         = 0.40                      # overlay transparency
-ALLOWED_FRONT = "https://amitabhm1.github.io"  # GitHub‑Pages origin
-################ -------------------------------------------------
+FRONT_ORIGIN  = "https://amitabhm1.github.io"
+# -------------------------------------------------------------------
+
+# ---------- import model definition from backend_src ----------------
+from backend_src import archs
+# --------------------------------------------------------------------
 
-# ---------------- FastAPI (plus simple CORS) ------------------
-app = FastAPI(title="Polyp Segmentation API")
+app = FastAPI(title="Polyp‑Seg‑API")
 
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=[ALLOWED_FRONT, "http://localhost:4173", "http://127.0.0.1:5500"],
-    allow_methods=["POST", "OPTIONS"],
+    allow_origins=[FRONT_ORIGIN, "http://localhost:4173"],
+    allow_methods=["POST"],
     allow_headers=["*"],
 )
-# --------------- load model once at startup -------------------
+
+# ---------------- load checkpoint (fail fast) -----------------------
+if not torch.cuda.is_available():
+    print("⚠ CUDA not available – running on CPU")
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
-model  = archs.__dict__[ARCH_NAME](num_classes=1,
-                                   input_channels=3,
-                                   deep_supervision=DEEP_SUP).to(device)
-state  = torch.load(CKPT, map_location=device)
-# strip “module.” if checkpoint came from DDP
-state  = {k.replace("module.", ""): v for k, v in state.items()}
-model.load_state_dict(state, strict=False)
-model.eval()
-print("✓ model loaded on", device)
-
-# ---------- transforms = pad/resize exactly as during train ----
-transform = A.Compose([
-    A.Resize(height=INPUT_H, width=INPUT_W, interpolation=cv2.INTER_LINEAR),
-])
-
-# ------------------------- helpers -----------------------------
-def infer_overlay(pil_img: Image.Image) -> bytes:
-    orig = np.array(pil_img.convert("RGB"))
-    trans = transform(image=orig)["image"]
-    ten   = torch.from_numpy(trans.astype("float32")/255.0).permute(2,0,1).unsqueeze(0).to(device)
+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]} …")
+
+net.eval()
+print("✓ model loaded on", device)
+
+# ---------------- preprocessing -------------------------------------
+xform = A.Compose([A.Resize(INPUT_H, INPUT_W)])
+
+def predict_overlay(pil: Image.Image) -> bytes:
+    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():
-        out = model(ten)
-        if DEEP_SUP and isinstance(out, (list,tuple)):
-            out = out[-1]
-
-    mask = (torch.sigmoid(out)[0,0].cpu().numpy() > THRESH).astype("uint8")
-    # colour mask
-    mask_rgb        = np.zeros((*mask.shape,3), dtype="uint8")
-    mask_rgb[mask==1] = (255,   0,   0)          # red overlay
-    # blend
-    blended = (trans*(1-ALPHA) + mask_rgb*ALPHA).astype("uint8")
-
-    # encode PNG
+        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(blended).save(buf, format="PNG")
-    buf.seek(0)
-    return buf.getvalue()
-# ---------------------------------------------------------------
+    Image.fromarray(out).save(buf, format="PNG")
+    buf.seek(0);  return buf
 
-@app.post("/segment", summary="Upload an image, get overlay PNG")
+# ---------------- endpoint -----------------------------------------
+@app.post("/segment")
 async def segment(file: UploadFile = File(...)):
     try:
-        img_bytes = await file.read()
-        pil       = Image.open(io.BytesIO(img_bytes))
-        png_bytes = infer_overlay(pil)
-        return StreamingResponse(io.BytesIO(png_bytes), media_type="image/png")
-    except Exception as e:
-        return {"error": str(e)}
+        pil = Image.open(io.BytesIO(await file.read()))
+    except Exception:
+        return JSONResponse({"error":"invalid image"}, status_code=400)
+
+    png = predict_overlay(pil)
+    return StreamingResponse(png, media_type="image/png")