Update app.py

app.py

@@ -132,37 +132,55 @@ def video_vision(video_input_path, prompt, video_interval):
         _seg_idx = 0
         pred_masks = result['prediction_masks'][_seg_idx]
         seg_frames = []
+        masked_only_frames = []  # New list for masked-only frames
+
         for frame_idx in range(len(vid_frames)):
             pred_mask = pred_masks[frame_idx]
             temp_dir = tempfile.mkdtemp()
             os.makedirs(temp_dir, exist_ok=True)
+
+            # Create visualized frame with segmentation overlay
             seg_frame = visualize(pred_mask, image_paths[frame_idx], temp_dir)
             seg_frames.append(seg_frame)
 
+            # Create a binary mask image (white mask on black background)
+            binary_mask = (pred_mask.astype('uint8') * 255)  # Convert mask to 0/255
+            binary_mask_path = os.path.join(temp_dir, f"binary_mask_{frame_idx}.png")
+            cv2.imwrite(binary_mask_path, binary_mask)
+            masked_only_frames.append(binary_mask_path)
+
         output_video = "output_video.mp4"
+        masked_video = "masked_only_video.mp4"  # New video file for masked areas only
 
         # Read the first image to get the size (resolution)
         frame = cv2.imread(seg_frames[0])
         height, width, layers = frame.shape
 
-        # Define the video codec and create VideoWriter object
+        # Define the video codec and create VideoWriter objects
         fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # Codec for MP4
         video = cv2.VideoWriter(output_video, fourcc, new_fps, (width, height))
+        masked_video_writer = cv2.VideoWriter(masked_video, fourcc, new_fps, (width, height), isColor=False)
 
-        # Iterate over the image paths and write to the video
-        for img_path in seg_frames:
-            frame = cv2.imread(img_path)
-            video.write(frame)
+        # Write frames to the videos
+        for idx, (seg_frame_path, mask_frame_path) in enumerate(zip(seg_frames, masked_only_frames)):
+            seg_frame = cv2.imread(seg_frame_path)
+            mask_frame = cv2.imread(mask_frame_path, cv2.IMREAD_GRAYSCALE)  # Read the binary mask in grayscale
 
-        # Release the video writer
+            video.write(seg_frame)
+            masked_video_writer.write(mask_frame)
+
+        # Release the video writers
         video.release()
+        masked_video_writer.release()
 
         print(f"Video created successfully at {output_video}")
+        print(f"Masked-only video created successfully at {masked_video}")
+
+        return result['prediction'], output_video, masked_video
 
-        return result['prediction'], output_video
             
     else:
-        return result['prediction'], None
+        return result['prediction'], None, None
     
 
 
@@ -214,11 +232,12 @@ with gr.Blocks(analytics_enabled=False) as demo:
                 with gr.Column():
                     vid_output_res = gr.Textbox(label="Response")
                     output_video = gr.Video(label="Segmentation")
+                    masked_output = gr.Video(label="Masked video")
             
             submit_video_btn.click(
                 fn = video_vision,
                 inputs = [video_input, vid_instruction, frame_interval],
-                outputs = [vid_output_res, output_video]
+                outputs = [vid_output_res, output_video, masked_output]
             )
 
 demo.queue().launch(show_api=False, show_error=True)