Update video_mask_app.py

video_mask_app.py

@@ -93,6 +93,7 @@ def process_video(input_video_path):
 
     # Process each frame
     frames = []
+    mask_frames = []
     for frame in tqdm(video_clip.iter_frames()):
         frame_pil = Image.fromarray(frame)
         frame_no_bg, mask_resized = remove_background(frame_pil)
@@ -109,12 +110,23 @@ def process_video(input_video_path):
         output_np = np.array(output)
         frames.append(output_np)
 
+        # Create a mask frame with white foreground and black background
+        mask_frame = np.array(mask_resized)
+        mask_frame = np.stack([mask_frame, mask_frame, mask_frame], axis=-1)  # Convert to 3 channels
+        mask_frame[mask_frame > 0] = 255  # Set foreground to white
+        mask_frames.append(mask_frame)
+
     # Save the processed frames as a new video
     output_video_path = os.path.join(output_folder, "no_bg_video.mp4")
     processed_clip = ImageSequenceClip(frames, fps=video_clip.fps)
     processed_clip.write_videofile(output_video_path, codec='libx264', ffmpeg_params=['-pix_fmt', 'yuva420p'])
 
-    return output_video_path
+    # Save the mask frames as a new video
+    mask_video_path = os.path.join(output_folder, "mask_video.mp4")
+    mask_clip = ImageSequenceClip(mask_frames, fps=video_clip.fps)
+    mask_clip.write_videofile(mask_video_path, codec='libx264')
+
+    return output_video_path, mask_video_path
 
 # Gradio components
 slider1 = ImageSlider(label="RMBG-2.0", type="pil")
@@ -125,6 +137,7 @@ text = gr.Textbox(label="Paste an image URL")
 png_file = gr.File(label="output png file")
 video_input = gr.Video(label="Upload a video")
 video_output = gr.Video(label="Processed video")
+mask_video_output = gr.Video(label="Mask video")
 
 # Example videos
 example_videos = [
@@ -140,7 +153,7 @@ tab1 = gr.Interface(
 
 tab2 = gr.Interface(fn, inputs=text, outputs=[slider2, gr.File(label="output png file")], examples=["http://farm9.staticflickr.com/8488/8228323072_76eeddfea3_z.jpg"], api_name="text")
 #tab3 = gr.Interface(process_file, inputs=image2, outputs=png_file, examples=["giraffe.jpg"], api_name="png")
-tab4 = gr.Interface(process_video, inputs=video_input, outputs=video_output, examples=example_videos, api_name="video", cache_examples = False)
+tab4 = gr.Interface(process_video, inputs=video_input, outputs=[video_output, mask_video_output], examples=example_videos, api_name="video", cache_examples = False)
 
 # Gradio tabbed interface
 demo = gr.TabbedInterface(