Update app.py

app.py

@@ -313,6 +313,9 @@ def propagate_to_all(video_in, checkpoint, stored_inference_state, stored_frame_
     # Define a directory to save the JPEG images
     frames_output_dir = "frames_output_images"
     os.makedirs(frames_output_dir, exist_ok=True)
+
+    mask_frames_output_dir = "mask_frames_output_images"
+    os.makedirs(mask_frames_output_dir, exist_ok=True)
     
     # Initialize a list to store file paths of saved images
     jpeg_images = []
@@ -354,8 +357,13 @@ def propagate_to_all(video_in, checkpoint, stored_inference_state, stored_frame_
             available_frames_to_check.append(f"frame_{out_frame_idx}.jpg")
 
         # Save the raw binary mask as a separate image
-        mask_filename = os.path.join(frames_output_dir, f"mask_{out_frame_idx}.jpg")
-        binary_mask = (out_mask * 255).astype(np.uint8)  # Scale mask to 0-255
+        mask_filename = os.path.join(mask_frames_output_dir, f"mask_{out_frame_idx}.jpg")
+        binary_mask = np.squeeze(out_mask)  # Ensure the mask is 2D
+        binary_mask = (binary_mask * 255).astype(np.uint8)  # Scale mask to 0-255
+
+        if binary_mask.ndim != 2:  # Ensure it's 2D for PIL
+            raise ValueError(f"Mask has invalid dimensions: {binary_mask.shape}")
+
         mask_image = Image.fromarray(binary_mask)
         mask_image.save(mask_filename)  # Save the mask as a JPEG
         masks_frames.append(mask_filename)  # Append to the list of masks