Update app.py

app.py

@@ -104,18 +104,42 @@ def analyze_image(image):
 
 
 def show_mask(mask, ax, random_color=False):
-    if random_color:
-        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
-    else:
-        color = np.array([1.0, 0.0, 0.0, 0.5])
-
-    if len(mask.shape) == 4:
-        mask = mask[0, 0]
+    try:
+        # Debug print to understand mask type
+        print(f"show_mask input type: {type(mask)}")
+        
+        # Convert mask if it's a tuple
+        if isinstance(mask, tuple):
+            if len(mask) > 0 and mask[0] is not None:
+                mask = mask[0]
+            else:
+                raise ValueError("Invalid mask tuple")
+
+        # Convert torch tensor to numpy if needed
+        if torch.is_tensor(mask):
+            mask = mask.cpu().numpy()
+            
+        # Handle 4D tensor/array case
+        if len(mask.shape) == 4:
+            mask = mask[0, 0]
+        # Handle 3D tensor/array case
+        elif len(mask.shape) == 3:
+            mask = mask[0]
+
+        if random_color:
+            color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
+        else:
+            color = np.array([1.0, 0.0, 0.0, 0.5])
 
-    mask_image = np.zeros((*mask.shape, 4), dtype=np.float32)
-    mask_image[mask > 0] = color
+        mask_image = np.zeros((*mask.shape, 4), dtype=np.float32)
+        mask_image[mask > 0] = color
 
-    ax.imshow(mask_image)
+        ax.imshow(mask_image)
+        
+    except Exception as e:
+        print(f"show_mask error: {str(e)}")
+        print(f"mask shape: {getattr(mask, 'shape', 'no shape')}")
+        raise
 
 
 def process_image_detection(image, target_label, surprise_rating):