k

app.py

@@ -97,10 +97,11 @@ def segment_endpoint():
     overlay = request.args.get("overlay", "false").lower() == "true"
     if overlay:
         colored = np.array(image_pil).copy()
-        colored[mask] = [255, 0, 0]  # red overlay
+        mask_np = mask.cpu().numpy() if hasattr(mask, "cpu") else np.array(mask)
+        colored[mask_np] = [255, 0, 0]  # red overlay
         out_img = Image.fromarray(colored)
     else:
-        out_img = Image.fromarray((mask * 255).astype(np.uint8))
+        out_img = Image.fromarray((mask.astype(np.uint8) * 255) if hasattr(mask, "astype") else (np.array(mask, dtype=np.uint8) * 255))
 
     buf = io.BytesIO()
     out_img.save(buf, format="PNG")