Update visualization.py

visualization.py

@@ -7,15 +7,6 @@ from constants import COLORS
 from utils import fig2img
 
 
-def visualize_mask(mask, img, alpha=0.5):
-    mask = mask.cpu().squeeze().numpy()
-    img = img.cpu().squeeze().permute(1, 2, 0).numpy()
-    plt.imshow(img)
-    plt.imshow(mask, alpha=alpha)
-    plt.axis("off")
-    return fig2img(plt.gcf())
-
-
 def visualize_prediction(
     pil_img, output_dict, threshold=0.7, id2label=None, display_mask=False, mask=None
 ):
@@ -32,8 +23,13 @@ def visualize_prediction(
 
     fig, ax = plt.subplots(figsize=(12, 12))
     ax.imshow(pil_img)
-    if display_mask:
-        ax.imshow(mask, alpha=0.5)
+    if display_mask and mask is not None:
+        # Create a new mask with white objects and black background
+        new_mask = np.zeros_like(mask)
+        new_mask[mask > 0] = 255
+        
+        # Display the new mask as a semi-transparent overlay
+        ax.imshow(new_mask, alpha=0.5, cmap='gray')
     colors = COLORS * 100
     for score, (xmin, ymin, xmax, ymax), label, color in zip(
         scores, boxes, labels, colors