LuhanM/reefscope-inference

1import base642import io3import logging4 5import numpy as np6import torch7from fastapi import FastAPI, Request, HTTPException8from fastapi.responses import JSONResponse9from PIL import Image10from transformers import AutoImageProcessor, SegformerForSemanticSegmentation11 12logging.basicConfig(level=logging.INFO)13logger = logging.getLogger(__name__)14 15MODEL_ID = "EPFL-ECEO/segformer-b2-finetuned-coralscapes-1024-1024"16 17ID_TO_LABEL = {18    0: "unlabeled", 1: "seagrass", 2: "trash", 3: "other coral dead",19    4: "other coral bleached", 5: "sand", 6: "other coral alive", 7: "human",20    8: "transect tools", 9: "fish", 10: "algae covered substrate",21    11: "other animal", 12: "unknown hard substrate", 13: "background water",22    14: "dark", 15: "transect line", 16: "massive/meandering bleached",23    17: "massive/meandering alive", 18: "rubble", 19: "branching bleached",24    20: "branching dead", 21: "millepora", 22: "branching alive",25    23: "massive/meandering dead", 24: "clam", 25: "acropora alive",26    26: "sea cucumber", 27: "turbinaria", 28: "table acropora alive",27    29: "sponge", 30: "anemone", 31: "pocillopora alive",28    32: "table acropora dead", 33: "meandering bleached", 34: "stylophora alive",29    35: "sea urchin", 36: "meandering alive", 37: "meandering dead",30    38: "crown of thorn", 39: "dead clam",31}32 33app = FastAPI()34 35logger.info("Loading model %s ...", MODEL_ID)36processor = AutoImageProcessor.from_pretrained(MODEL_ID, use_fast=True)37model = SegformerForSemanticSegmentation.from_pretrained(MODEL_ID)38model.eval()39logger.info("Model ready.")40 41 42@app.get("/health")43def health():44    return {"status": "ok"}45 46 47@app.post("/predict")48async def predict(request: Request):49    image_bytes = await request.body()50    if not image_bytes:51        raise HTTPException(status_code=400, detail="Empty request body")52 53    try:54        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")55    except Exception:56        raise HTTPException(status_code=400, detail="Could not decode image")57 58    # Cap longest side at 512px for CPU inference speed59    w, h = image.size60    if max(w, h) > 512:61        scale = 512 / max(w, h)62        image = image.resize((int(w * scale), int(h * scale)), Image.BILINEAR)63 64    inputs = processor(images=image, return_tensors="pt")65    with torch.no_grad():66        outputs = model(**inputs)67 68    upsampled = torch.nn.functional.interpolate(69        outputs.logits,70        size=(image.height, image.width),71        mode="bilinear",72        align_corners=False,73    )74    label_map = upsampled.argmax(dim=1).squeeze().numpy().astype(np.uint8)75 76    segments = []77    for class_id in np.unique(label_map):78        mask = ((label_map == class_id) * 255).astype(np.uint8)79        buf = io.BytesIO()80        Image.fromarray(mask).save(buf, format="PNG")81        segments.append({82            "label": ID_TO_LABEL.get(int(class_id), str(class_id)),83            "mask": base64.b64encode(buf.getvalue()).decode(),84            "score": 1.0,85        })86 87    return JSONResponse(segments)88