Update app.py

app.py

@@ -1,31 +1,81 @@
-# app.py  (key lines only)
+# app.py  – Hugging‑Face Space backend
+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 fastapi import FastAPI, UploadFile, File
-from backend_src import archs                    # ← path fixed
-import torch, albumentations as A, cv2, numpy as np
+# ---------- model code lives inside backend_src -------------
+from backend_src import archs                                    # <- you copied this
+# ------------------------------------------------------------
 
-# ------------------------------- CFG ---------------------------------
-CKPT_PATH = "models/best_model.pth"
-ARCH_NAME  = "NestedUNet"
-NUM_CLASSES = 1
-# ---------------------------------------------------------------------
+### CONFIG #### -------------------------------------------------
+CKPT          = "models/best_model.pth"   # uploaded via Git‑LFS
+ARCH_NAME     = "NestedUNet"
+INPUT_H = INPUT_W = 512                   # what you trained with
+THRESH        = 0.5                       # binarisation
+DEEP_SUP      = False
+ALPHA         = 0.40                      # overlay transparency
+ALLOWED_FRONT = "https://amitabhm1.github.io"  # GitHub‑Pages origin
+################ -------------------------------------------------
 
-app = FastAPI()
+# ---------------- FastAPI (plus simple CORS) ------------------
+app = FastAPI(title="Polyp Segmentation API")
 
-# ----- CORS (GitHub Pages) -----
-from fastapi.middleware.cors import CORSMiddleware
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["https://amitabhm.github.io",  # or "*" while testing
-                   "http://localhost:4173"],                   # vite / local dev
+    allow_origins=[ALLOWED_FRONT, "http://localhost:4173", "http://127.0.0.1:5500"],
     allow_methods=["POST"],
     allow_headers=["*"],
 )
-
-# ------------------------------ LOAD MODEL ---------------------------
-model = archs.__dict__[ARCH_NAME](NUM_CLASSES, 3, False)
-state = torch.load(CKPT_PATH, map_location="cpu")
-# strip "module." keys if you trained with DDP
-state = {k.replace("module.", ""): v for k, v in state.items()}
+# --------------- load model once at startup -------------------
+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)
+
+    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
+    buf = io.BytesIO()
+    Image.fromarray(blended).save(buf, format="PNG")
+    buf.seek(0)
+    return buf.getvalue()
+# ---------------------------------------------------------------
+
+@app.post("/segment", summary="Upload an image, get overlay PNG")
+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)}