|
|
|
|
|
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 |
|
|
|
|
|
|
|
|
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"] |
|
|
|
|
|
|
|
|
device = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
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) |
|
|
|
|
|
resize = A.Compose([A.Resize(INP_H, INP_W)]) |
|
|
|
|
|
def to_tensor(pil:Image.Image)->torch.Tensor: |
|
|
arr = np.array(pil.convert("RGB")).astype("float32")/255.0 |
|
|
arr = resize(image=arr)["image"] |
|
|
return torch.from_numpy(arr.transpose(2,0,1)).unsqueeze(0).to(device) |
|
|
|
|
|
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(img,np_mask): |
|
|
cnts,_ = cv2.findContours(np_mask,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE) |
|
|
out = img.copy() |
|
|
cv2.drawContours(out,cnts,-1,(0,0,255),2) |
|
|
return out |
|
|
|
|
|
def heatmap(img,prob): |
|
|
hm = cv2.applyColorMap((prob*255).astype("uint8"),cv2.COLORMAP_JET) |
|
|
return cv2.addWeighted(img,1-ALPHA,hm,ALPHA,0) |
|
|
|
|
|
app = FastAPI() |
|
|
app.add_middleware(CORSMiddleware, |
|
|
allow_origins=FRONTENDS, allow_methods=["POST","GET","OPTIONS"], allow_headers=["*"]) |
|
|
|
|
|
@app.get("/ping") |
|
|
def ping(): return PlainTextResponse("pong") |
|
|
|
|
|
@app.post("/segment") |
|
|
async def segment(file:UploadFile=File(...)): |
|
|
raw = await file.read() |
|
|
pil = Image.open(io.BytesIO(raw)) |
|
|
inp = to_tensor(pil) |
|
|
|
|
|
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()}") |
|
|
|
|
|
frame = np.array(pil.resize((INP_W,INP_H),Image.BILINEAR)) |
|
|
vis = overlay(frame,mask) if mask.sum() else heatmap(frame,prob) |
|
|
|
|
|
buf = io.BytesIO(); Image.fromarray(vis).save(buf,format="PNG"); buf.seek(0) |
|
|
return StreamingResponse(buf, media_type="image/png") |
