Update app.py

app.py

@@ -3,6 +3,7 @@ import requests
 import cv2
 import numpy as np
 import torch
+import base64
 from fastapi import FastAPI, UploadFile, File
 from fastapi.responses import JSONResponse
 from detectron2.config import get_cfg
@@ -33,12 +34,22 @@ cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rc
 cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
 cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
 cfg.MODEL.WEIGHTS = MODEL_PATH
-cfg.MODEL.DEVICE = "cpu"   # Hugging Face Spaces often has no GPU
+cfg.MODEL.DEVICE = "cpu"   # Hugging Face Spaces default (no GPU)
 
 predictor = DefaultPredictor(cfg)
 
 # -----------------------------
-# 3. API Endpoint
+# 3. Helper: Encode mask to Base64
+# -----------------------------
+def encode_mask(mask: np.ndarray) -> str:
+    """Convert mask numpy array to base64 PNG string."""
+    mask_img = Image.fromarray(mask.astype(np.uint8))
+    buf = io.BytesIO()
+    mask_img.save(buf, format="PNG")
+    return base64.b64encode(buf.getvalue()).decode("utf-8")
+
+# -----------------------------
+# 4. API Endpoint
 # -----------------------------
 @app.post("/predict")
 async def predict(file: UploadFile = File(...)):
@@ -52,15 +63,24 @@ async def predict(file: UploadFile = File(...)):
     instances = outputs["instances"].to("cpu")
 
     results = []
+    mask_b64 = None
+
     if instances.has("pred_masks"):
         masks = instances.pred_masks.numpy()
         boxes = instances.pred_boxes.tensor.numpy()
         scores = instances.scores.numpy()
 
+        # Combine masks into one
+        combined_mask = np.any(masks, axis=0).astype(np.uint8) * 255
+        mask_b64 = encode_mask(combined_mask)
+
         for i in range(len(masks)):
             results.append({
                 "box": boxes[i].tolist(),
                 "score": float(scores[i])
             })
 
-    return JSONResponse({"predictions": results})
+    return JSONResponse({
+        "predictions": results,
+        "mask": mask_b64  # base64 string (PNG)
+    })