Update app.py

app.py

@@ -39,14 +39,20 @@ def fn(vid, fps, color):
     frames = video.iter_frames(fps=fps)
 
     # Process each frame for background removal
-    processed_frames = []
+    processed_frames_no_bg = []
+    processed_frames_changed_bg = []
     for frame in frames:
         pil_image = Image.fromarray(frame)
-        processed_image = process(pil_image, color)
-        processed_frames.append(np.array(processed_image))
+        processed_image, mask = process(pil_image, color)  # Get both processed image and mask
+        processed_frames_no_bg.append(np.array(processed_image))  # Save no-background frame
+        
+        # Compose with background for changed background video
+        background = Image.new("RGBA", pil_image.size, color + (255,))
+        composed_image = Image.composite(pil_image, background, mask)
+        processed_frames_changed_bg.append(np.array(composed_image))
 
     # Create a new video from the processed frames
-    processed_video = mp.ImageSequenceClip(processed_frames, fps=fps)
+    processed_video = mp.ImageSequenceClip(processed_frames_changed_bg, fps=fps)
 
     # Add the original audio back to the processed video
     processed_video = processed_video.set_audio(audio)
@@ -58,8 +64,13 @@ def fn(vid, fps, color):
     temp_filepath = os.path.join(temp_dir, unique_filename)
     processed_video.write_videofile(temp_filepath, codec="libx264")
 
-    # Return the path to the temporary file
-    return temp_filepath
+    # Create and save no-background video
+    processed_video_no_bg = mp.ImageSequenceClip(processed_frames_no_bg, fps=fps)
+    processed_video_no_bg = processed_video_no_bg.set_audio(audio)
+    temp_filepath_no_bg = os.path.join(temp_dir, str(uuid.uuid4()) + ".webm")
+    processed_video_no_bg.write_videofile(temp_filepath_no_bg, codec="libvpx")
+
+    return temp_filepath_no_bg, temp_filepath
 
 
 def process(image, color_hex):
@@ -73,7 +84,7 @@ def process(image, color_hex):
     mask = pred_pil.resize(image_size)
 
     # Convert hex color to RGB tuple
-    color_rgb = tuple(int(color_hex[i : i + 2], 16) for i in (1, 3, 5))
+    color_rgb = tuple(int(color_hex[i:i + 2], 16) for i in (1, 3, 5))
 
     # Create a background image with the chosen color
     background = Image.new("RGBA", image_size, color_rgb + (255,))
@@ -81,7 +92,7 @@ def process(image, color_hex):
     # Composite the image onto the background using the mask
     image = Image.composite(image, background, mask)
 
-    return image
+    return image, mask  # Return both the processed image and the mask
 
 
 def process_file(f, color="#00FF00"):
@@ -93,15 +104,24 @@ def process_file(f, color="#00FF00"):
     return name_path
 
 
+def change_color(in_video, fps_slider, color_picker):
+    no_bg_video_path, changed_bg_video_path = fn(in_video, fps_slider, color_picker) 
+    return no_bg_video_path, changed_bg_video_path
+
+
 with gr.Blocks() as demo:
-    in_video = gr.Video(label="birefnet")
-    out_video = gr.Video()
-    fps_slider = gr.Slider(minimum=1, maximum=60, step=1, value=12, label="Output FPS")
-    color_picker = gr.ColorPicker(label="Background Color", value="#00FF00")
-    submit_button = gr.Button("Process Video")
+    with gr.Row():
+        in_video = gr.Video(label="Input Video")
+        no_bg_video = gr.Video(label="No BG Video")  # Added for no-background video
+        out_video = gr.Video(label="Output Video")  # This will be the changed-background video
+    submit_button = gr.Button("Change Background")
+    with gr.Row():
+        fps_slider = gr.Slider(minimum=1, maximum=60, step=1, value=12, label="Output FPS")
+        color_picker = gr.ColorPicker(label="Background Color", value="#00FF00")
+
 
     submit_button.click(
-        fn, inputs=[in_video, fps_slider, color_picker], outputs=out_video
+        change_color, inputs=[in_video, fps_slider, color_picker], outputs=[no_bg_video, out_video] 
     )
 
 if __name__ == "__main__":