Update app.py

app.py

@@ -1,57 +1,31 @@
-import io, cv2, numpy as np, torch
-from fastapi import FastAPI, File, UploadFile
-from fastapi.middleware.cors import CORSMiddleware
-from starlette.responses import StreamingResponse
-from PIL import Image
-import archs  # copy your model definition file into this repo
+# app.py  (key lines only)
+
+from fastapi import FastAPI, UploadFile, File
+from backend_src import archs                    # ← path fixed
+import torch, albumentations as A, cv2, numpy as np
+
+# ------------------------------- CFG ---------------------------------
+CKPT_PATH = "models/best_model.pth"
+ARCH_NAME  = "NestedUNet"
+NUM_CLASSES = 1
+# ---------------------------------------------------------------------
 
-# ----- CORS: allow your GitHub Pages origin -----
 app = FastAPI()
+
+# ----- CORS (GitHub Pages) -----
+from fastapi.middleware.cors import CORSMiddleware
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["https://<your-github-username>.github.io"],
+    allow_origins=["https://amitabhm.github.io",  # or "*" while testing
+                   "http://localhost:4173"],                   # vite / local dev
     allow_methods=["POST"],
     allow_headers=["*"],
 )
 
-# ----- load model once -----
-MODEL_PATH = "model.pth"
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-model = archs.NestedUNet(num_classes=1, input_channels=3, deep_supervision=False)
-model.load_state_dict(torch.load(MODEL_PATH, map_location=DEVICE))
-model.to(DEVICE).eval()
-
-# ----- utils -----
-def preprocess(pil):
-    im = pil.resize((512, 512))               # same as training
-    arr = np.asarray(im).astype("float32")/255
-    ten = torch.from_numpy(arr.transpose(2,0,1)).unsqueeze(0)
-    return ten.to(DEVICE)
-
-def postprocess(pred, alpha=0.4):
-    mask = (torch.sigmoid(pred)[0,0].cpu().numpy() > .5).astype("uint8")
-    mask_rgb = np.zeros((*mask.shape,3), np.uint8)
-    mask_rgb[mask==1] = (255,0,0)
-    return mask_rgb, mask
-
-def overlay(img, mask_rgb, alpha=0.4):
-    blend = (img*(1-alpha) + mask_rgb*alpha).astype("uint8")
-    out = img.copy()
-    out[mask_rgb[:,:,0]>0] = blend[mask_rgb[:,:,0]>0]
-    return out
-
-# ----- endpoint -----
-@app.post("/segment")
-async def segment(file: UploadFile = File(...)):
-    raw = await file.read()
-    pil = Image.open(io.BytesIO(raw)).convert("RGB")
-    input_t = preprocess(pil)
-    with torch.no_grad():
-        pred = model(input_t)
-    if isinstance(pred,(list,tuple)): pred = pred[-1]
-    mask_rgb,_ = postprocess(pred)
-    result = overlay(np.array(pil.resize((512,512))), mask_rgb)
-    buf = io.BytesIO()
-    Image.fromarray(result).save(buf, format="PNG")
-    buf.seek(0)
-    return StreamingResponse(buf, media_type="image/png")
+# ------------------------------ 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()}
+model.load_state_dict(state, strict=False)
+model.eval()