Update app.py

app.py

@@ -1,16 +1,14 @@
 from fastapi import FastAPI, UploadFile, File
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import StreamingResponse
 from PIL import Image
+from transformers import MobileNetV2ForSemanticSegmentation, AutoImageProcessor
+import torch
 from io import BytesIO
+import base64
 import numpy as np
-import tensorflow as tf
-
-# --------- LOAD YOUR SEGMENTATION MODEL HERE ---------
-model = tf.keras.models.load_model("seg_model")   # <<<<=== THIS LINE!
-# -----------------------------------------------------
 
 app = FastAPI()
+
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -18,20 +16,25 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+# Load processor and model
+processor = AutoImageProcessor.from_pretrained("seg_model")
+model = MobileNetV2ForSemanticSegmentation.from_pretrained("seg_model")
+
 @app.post("/predict")
 async def predict(file: UploadFile = File(...)):
     contents = await file.read()
     img = Image.open(BytesIO(contents)).convert("RGB")
-    img = img.resize((256, 256))
-    arr = np.array(img) / 255.0
-    arr = np.expand_dims(arr, 0)
 
-    prediction = model.predict(arr)
-    mask = np.argmax(prediction[0], axis=-1).astype(np.uint8)
-    mask_img = Image.fromarray(mask * 50)  # For visualization
+    inputs = processor(images=img, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model(**inputs)
+    logits = outputs.logits  # (batch, num_labels, H, W)
+    mask = torch.argmax(logits, dim=1)[0].numpy().astype(np.uint8)
 
+    # Optionally, you can convert mask to RGB with a color map for visualization
+    mask_img = Image.fromarray(mask)
     buf = BytesIO()
-    mask_img.save(buf, format='PNG')
+    mask_img.save(buf, format="PNG")
     buf.seek(0)
-
-    return StreamingResponse(buf, media_type="image/png")
+    b64 = base64.b64encode(buf.read()).decode()
+    return {"success": True, "mask": "data:image/png;base64," + b64}