Add app.py

app.py

@@ -0,0 +1,57 @@
+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
+
+# ----- CORS: allow your GitHub Pages origin -----
+app = FastAPI()
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["https://<your-github-username>.github.io"],
+    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")