Update app.py

app.py

@@ -1,5 +1,6 @@
 from fastapi import FastAPI, UploadFile, File
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import StreamingResponse
 from PIL import Image
 from io import BytesIO
 import numpy as np
@@ -7,7 +8,6 @@ import tensorflow as tf
 
 app = FastAPI()
 
-# CORS config (optional)
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -15,13 +15,26 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+# Load your trained segmentation model here
+# model = tf.keras.models.load_model("seg_model_path")
+
 @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))  # or whatever your model expects
+    img = img.resize((256, 256))
     arr = np.array(img) / 255.0
     arr = np.expand_dims(arr, 0)
-    # prediction = model.predict(arr)
-    # result = do_something_with_prediction(prediction)
-    return {"success": True}
+
+    # Prediction
+    prediction = model.predict(arr)   # (1, 256, 256, num_classes)
+    mask = np.argmax(prediction[0], axis=-1).astype(np.uint8)  # (256, 256)
+
+    # Convert to image (you can colorize or just multiply for visualization)
+    mask_img = Image.fromarray(mask * 50)  # Optional scaling for visibility
+
+    buf = BytesIO()
+    mask_img.save(buf, format='PNG')
+    buf.seek(0)
+
+    return StreamingResponse(buf, media_type="image/png")