Update app.py

app.py

@@ -28,48 +28,78 @@ transform_image = transforms.Compose(
 
 
 @spaces.GPU
-def fn(vid, bg_type="Color", bg_image=None, color="#00FF00", fps=0):
-    # Load the video using moviepy
-    video = mp.VideoFileClip(vid)
-
-    # Load original fps if fps value is equal to 0
-    if fps == 0:
-        fps = video.fps
-
-    # Extract audio from the video
-    audio = video.audio
-
-    # Extract frames at the specified FPS
-    frames = video.iter_frames(fps=fps)
-
-    # Process each frame for background removal
-    processed_frames = []
-    yield gr.update(visible=True), gr.update(visible=False)
-    for frame in frames:
-        pil_image = Image.fromarray(frame)
-        if bg_type == "Color":
-            processed_image = process(pil_image, color)
+def fn(vid, bg_type="Color", bg_image=None, bg_video=None, color="#00FF00", fps=0, video_handling="slow_down"):
+    try:
+        # Load the video using moviepy
+        video = mp.VideoFileClip(vid)
+
+        # Load original fps if fps value is equal to 0
+        if fps == 0:
+            fps = video.fps
+
+        # Extract audio from the video
+        audio = video.audio
+
+        # Extract frames at the specified FPS
+        frames = video.iter_frames(fps=fps)
+
+        # Process each frame for background removal
+        processed_frames = []
+        yield gr.update(visible=True), gr.update(visible=False)
+
+        if bg_type == "Video":
+            background_video = mp.VideoFileClip(bg_video)
+            if background_video.duration < video.duration:
+                if video_handling == "slow_down":
+                    background_video = background_video.fx(mp.vfx.speedx, factor=video.duration / background_video.duration)
+                else:  # video_handling == "loop"
+                    background_video = mp.concatenate_videoclips([background_video] * int(video.duration / background_video.duration + 1))
+            background_frames = background_video.iter_frames(fps=fps)
         else:
-            processed_image = process(pil_image, bg_image)
-        processed_frames.append(np.array(processed_image))
-        yield processed_image, None
+            background_frames = None
+
+        for i, frame in enumerate(frames):
+            pil_image = Image.fromarray(frame)
+            if bg_type == "Color":
+                processed_image = process(pil_image, color)
+            elif bg_type == "Image":
+                processed_image = process(pil_image, bg_image)
+            elif bg_type == "Video":
+                try:
+                    background_frame = next(background_frames)
+                    background_image = Image.fromarray(background_frame)
+                    processed_image = process(pil_image, background_image)
+                except StopIteration:
+                    # Handle case where background video is shorter than input video
+                    processed_image = process(pil_image, "#000000")  # Default to black background
+            else:
+                processed_image = pil_image  # Default to original image if no background is selected
+
+            processed_frames.append(np.array(processed_image))
+            yield processed_image, None
+
+        # Create a new video from the processed frames
+        processed_video = mp.ImageSequenceClip(processed_frames, fps=fps)
+
+        # Add the original audio back to the processed video
+        processed_video = processed_video.set_audio(audio)
+
+        # Save the processed video to a temporary file
+        temp_dir = "temp"
+        os.makedirs(temp_dir, exist_ok=True)
+        unique_filename = str(uuid.uuid4()) + ".mp4"
+        temp_filepath = os.path.join(temp_dir, unique_filename)
+        processed_video.write_videofile(temp_filepath, codec="libx264")
+
+        yield gr.update(visible=False), gr.update(visible=True)
+        # Return the path to the temporary file
+        yield processed_image, temp_filepath
+
+    except Exception as e:
+        print(f"Error: {e}")
+        yield gr.update(visible=False), gr.update(visible=True)
+        yield None, f"Error processing video: {e}"
 
-    # Create a new video from the processed frames
-    processed_video = mp.ImageSequenceClip(processed_frames, fps=fps)
-
-    # Add the original audio back to the processed video
-    processed_video = processed_video.set_audio(audio)
-
-    # Save the processed video to a temporary file
-    temp_dir = "temp"
-    os.makedirs(temp_dir, exist_ok=True)
-    unique_filename = str(uuid.uuid4()) + ".mp4"
-    temp_filepath = os.path.join(temp_dir, unique_filename)
-    processed_video.write_videofile(temp_filepath, codec="libx264")
-
-    yield gr.update(visible=False), gr.update(visible=True)
-    # Return the path to the temporary file
-    yield processed_image, temp_filepath
 
 
 def process(image, bg):
@@ -82,11 +112,13 @@ def process(image, bg):
     pred_pil = transforms.ToPILImage()(pred)
     mask = pred_pil.resize(image_size)
 
-    if bg.startswith("#"):
-        color_rgb = tuple(int(bg[i : i + 2], 16) for i in (1, 3, 5))
+    if isinstance(bg, str) and bg.startswith("#"):
+        color_rgb = tuple(int(bg[i:i+2], 16) for i in (1, 3, 5))
         background = Image.new("RGBA", image_size, color_rgb + (255,))
+    elif isinstance(bg, Image.Image):
+        background = bg.convert("RGBA").resize(image_size)
     else:
-        background = Image.open(bg).convert("RGBA").resize(image_size)
+        background = Image.new("RGBA", image_size, (0, 0, 0, 255))  # Default to black background
 
     # Composite the image onto the background using the mask
     image = Image.composite(image, background, mask)
@@ -108,22 +140,34 @@ with gr.Blocks(theme=gr.themes.Ocean()) as demo:
             value=0,
             label="Output FPS (0 will inherit the original fps value)",
         )
-        bg_type = gr.Radio(["Color", "Image"], label="Background Type", value="Color")
+        bg_type = gr.Radio(["Color", "Image", "Video"], label="Background Type", value="Color")
         color_picker = gr.ColorPicker(label="Background Color", value="#00FF00", visible=True)
         bg_image = gr.Image(label="Background Image", type="filepath", visible=False)
+        bg_video = gr.Video(label="Background Video", visible=False)
+        with gr.Column(visible=False) as video_handling_options:
+            video_handling_radio = gr.Radio(["slow_down", "loop"], label="Video Handling", value="slow_down")
 
     def update_visibility(bg_type):
         if bg_type == "Color":
-            return gr.update(visible=True), gr.update(visible=False)
+            return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+        elif bg_type == "Image":
+            return gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)
+        elif bg_type == "Video":
+            return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True), gr.update(visible=True)
         else:
-            return gr.update(visible=False), gr.update(visible=True)
+            return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+
 
-    bg_type.change(update_visibility, inputs=bg_type, outputs=[color_picker, bg_image])
+    bg_type.change(update_visibility, inputs=bg_type, outputs=[color_picker, bg_image, bg_video, video_handling_options])
 
 
     examples = gr.Examples(
-        [["rickroll-2sec.mp4", "Image", "images.webp"], ["rickroll-2sec.mp4", "Color", None]],
-        inputs=[in_video, bg_type, bg_image],
+        [
+            ["rickroll-2sec.mp4", "Image", "images.webp", None],
+            ["rickroll-2sec.mp4", "Color", None, None],
+            ["rickroll-2sec.mp4", "Video", None, "background.mp4"] 
+        ],
+        inputs=[in_video, bg_type, bg_image, bg_video],
         outputs=[stream_image, out_video],
         fn=fn,
         cache_examples=True,
@@ -133,7 +177,7 @@ with gr.Blocks(theme=gr.themes.Ocean()) as demo:
 
     submit_button.click(
         fn,
-        inputs=[in_video, bg_type, bg_image, color_picker, fps_slider],
+        inputs=[in_video, bg_type, bg_image, bg_video, color_picker, fps_slider, video_handling_radio],
         outputs=[stream_image, out_video],
     )