Create app.py

app.py

@@ -0,0 +1,31 @@
+from fastapi import FastAPI, UploadFile, File
+from fastapi.responses import JSONResponse
+from PIL import Image
+import torch, torchvision.transforms as T
+from transformers import MobileNetV2ForSemanticSegmentation
+import io
+
+# Load the model
+model = MobileNetV2ForSemanticSegmentation.from_pretrained("seg_model")
+model.eval()
+
+preprocess = T.Compose([
+    T.Resize(513),
+    T.ToTensor(),
+    T.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225])
+])
+
+app = FastAPI()
+
+@app.get("/")
+def root():
+    return {"status": "API up for segmentation"}
+
+@app.post("/predict")
+async def predict(file: UploadFile = File(...)):
+    img = Image.open(await file.read()).convert("RGB")
+    x = preprocess(img).unsqueeze(0)
+    with torch.no_grad():
+        outputs = model(x).logits
+        seg = outputs.argmax(1)[0].tolist()
+    return JSONResponse(content={"segmentation_mask": seg})