Update app.py

app.py

@@ -40,8 +40,10 @@ def load_model(weights_path):
 
 predictor, metadata = load_model(weights_path)
 
+import random
+
 def predict(pil_img):
-    # Convert to RGB and ensure dtype is uint8
+    # Convert to RGB uint8
     img = np.array(pil_img.convert("RGB"), dtype=np.uint8)
 
     outputs = predictor(img)
@@ -52,9 +54,23 @@ def predict(pil_img):
         metadata=metadata,
         scale=1.0
     )
-    out = v.draw_instance_predictions(instances)
-    out_img = out.get_image()[:, :, ::-1]
 
+    # Draw ONLY masks with random colors
+    if instances.has("pred_masks"):
+        for mask in instances.pred_masks:
+            random_color = (
+                random.random(),  # R
+                random.random(),  # G
+                random.random()   # B
+            )
+
+            v.draw_binary_mask(
+                mask.numpy(),
+                color=random_color,
+                alpha=0.6
+            )
+
+    out_img = v.get_image()[:, :, ::-1]
     return Image.fromarray(out_img)
 
 gr.Interface(