Update visualization.py

visualization.py

@@ -24,12 +24,19 @@ def visualize_prediction(
     fig, ax = plt.subplots(figsize=(12, 12))
     ax.imshow(pil_img)
     if display_mask and mask is not None:
+        # Convert the mask image to a numpy array
+        mask_arr = np.asarray(mask)
+        
         # Create a new mask with white objects and black background
-        new_mask = np.zeros_like(mask)
-        new_mask[mask > 0] = 255
+        new_mask = np.zeros_like(mask_arr)
+        new_mask[mask_arr > 0] = 255
+        
+        # Convert the numpy array back to a PIL Image
+        new_mask = Image.fromarray(new_mask)
         
         # 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