Update app.py

app.py

@@ -6,8 +6,11 @@ from io import BytesIO
 import numpy as np
 import tensorflow as tf
 
-app = FastAPI()
+# --------- LOAD YOUR SEGMENTATION MODEL HERE ---------
+model = tf.keras.models.load_model("seg_model")   # <<<<=== THIS LINE!
+# -----------------------------------------------------
 
+app = FastAPI()
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -15,9 +18,6 @@ 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()
@@ -26,12 +26,9 @@ async def predict(file: UploadFile = File(...)):
     arr = np.array(img) / 255.0
     arr = np.expand_dims(arr, 0)
 
-    # 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
+    prediction = model.predict(arr)
+    mask = np.argmax(prediction[0], axis=-1).astype(np.uint8)
+    mask_img = Image.fromarray(mask * 50)  # For visualization
 
     buf = BytesIO()
     mask_img.save(buf, format='PNG')