Update app.py

app.py

@@ -343,6 +343,18 @@ def propagate_to_all(video_in, checkpoint, stored_inference_state, stored_frame_
         for out_obj_id, out_mask in video_segments[out_frame_idx].items():
             show_mask(out_mask, plt.gca(), obj_id=out_obj_id)
 
+            # Save the raw binary mask as a separate image
+            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
+
         # Define the output filename and save the figure as a JPEG file
         output_filename = os.path.join(frames_output_dir, f"frame_{out_frame_idx}.jpg")
         plt.savefig(output_filename, format='jpg')
@@ -356,17 +368,7 @@ def propagate_to_all(video_in, checkpoint, stored_inference_state, stored_frame_
         if f"frame_{out_frame_idx}.jpg" not in available_frames_to_check:
             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(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
+        
 
     torch.cuda.empty_cache()
     print(f"JPEG_IMAGES: {jpeg_images}")