Update app.py

app.py

@@ -5,7 +5,6 @@ from transformers import AutoModelForImageSegmentation
 import torch
 from torchvision import transforms
 import moviepy.editor as mp
-from pydub import AudioSegment
 from PIL import Image
 import numpy as np
 import os
@@ -28,7 +27,19 @@ transform_image = transforms.Compose(
 )
 
 BATCH_SIZE = 3
-executor = ThreadPoolExecutor(max_workers=4)
+executor = ThreadPoolExecutor(max_workers=4)  # Adjust as needed
+
+def get_background_image(bg_type, bg_image, background_frames, current_frame_index, video_handling, slow_down_factor):
+    if bg_type == "Video":
+        if video_handling == "slow_down":
+            frame_index = int(current_frame_index / slow_down_factor)
+        else:
+            frame_index = current_frame_index
+        return Image.fromarray(background_frames[frame_index % len(background_frames)])
+    elif bg_type == "Image":
+        return bg_image  # Directly returns the image path
+    else:  # bg_type == "Color"
+        return bg_image  # bg_image here is the color string
 
 @spaces.GPU
 def fn(vid, bg_type="Color", bg_image=None, bg_video=None, color="#00FF00", fps=0, video_handling="slow_down"):
@@ -38,66 +49,62 @@ def fn(vid, bg_type="Color", bg_image=None, bg_video=None, color="#00FF00", fps=
             audio = video.audio
         except AttributeError:
             audio = None
+
         if fps == 0:
             fps = video.fps
-
         frames = video.iter_frames(fps=fps)
         processed_frames = []
-        yield gr.update(visible=True), gr.update(visible=False)
+        yield gr.update(visible=True), gr.update(visible=False)  # Update Gradio display
 
         if bg_type == "Video":
             background_video = mp.VideoFileClip(bg_video)
-
             if background_video.duration < video.duration and video_handling == "slow_down":
                 slow_down_factor = video.duration / background_video.duration
             else:
                 slow_down_factor = 1
             background_frames = list(background_video.iter_frames(fps=fps))
-
         else:
             background_frames = None
-            slow_down_factor = None
-
+            slow_down_factor = None  # Not needed for image or color backgrounds
 
         bg_frame_index = 0
         frame_batch = []
 
-
         for i, frame in enumerate(frames):
             frame_batch.append(frame)
-
             if len(frame_batch) == BATCH_SIZE or i == int(video.fps * video.duration) - 1:
-
                 pil_images = [Image.fromarray(f) for f in frame_batch]
 
                 if bg_type == "Video":
                     processed_images = list(executor.map(process, pil_images, [get_background_image(bg_type, bg_image, background_frames, bg_frame_index + j, video_handling, slow_down_factor) for j in range(len(pil_images))]))
                     bg_frame_index += len(frame_batch)
                 elif bg_type == "Color":
-                    processed_images = list(executor.map(process, pil_images, [color] * len(pil_images)))
+                    processed_images = list(executor.map(process, pil_images, [color] * len(pil_images)))  # Use color directly
                 elif bg_type == "Image":
-                    processed_images = list(executor.map(process, pil_images, [bg_image] * len(pil_images)))
+                    processed_images = list(executor.map(process, pil_images, [bg_image] * len(pil_images))) # Use image path directly
                 else:
-                    processed_images = pil_images
+                    processed_images = pil_images  # No processing needed
 
                 for processed_image in processed_images:
                     processed_frames.append(np.array(processed_image))
-                    yield processed_image, None
+                    yield processed_image, None  # Update Gradio with processed images
                 frame_batch = []
 
+
         processed_video = mp.ImageSequenceClip(processed_frames, fps=fps)
         if audio:
             processed_video = processed_video.set_audio(audio)
 
+        # Save 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", logger=None)
 
-        yield gr.update(visible=False), gr.update(visible=True)
-        yield processed_image, temp_filepath
+
+        yield gr.update(visible=False), gr.update(visible=True) # Update Gradio display
+        yield processed_image, temp_filepath # Return final output
 
     except Exception as e:
         print(f"Error: {e}")
@@ -105,30 +112,30 @@ def fn(vid, bg_type="Color", bg_image=None, bg_video=None, color="#00FF00", fps=
         yield None, f"Error processing video: {e}"
 
 
+
+
 def process(image, bg):
     image_size = image.size
     input_images = transform_image(image).unsqueeze(0).to("cuda")
-    # Prediction
     with torch.no_grad():
         preds = birefnet(input_images)[-1].sigmoid().cpu()
     pred = preds[0].squeeze()
     pred_pil = transforms.ToPILImage()(pred)
     mask = pred_pil.resize(image_size)
 
-    if isinstance(bg, str) and bg.startswith("#"):
+    if isinstance(bg, str) and bg.startswith("#"):  # If bg is a color
         color_rgb = tuple(int(bg[i:i+2], 16) for i in (1, 3, 5))
-        background = Image.new("RGBA", image_size, color_rgb + (255,))
+        background = Image.new("RGBA", image_size, color_rgb + (255,)) # Create image with color
     elif isinstance(bg, Image.Image):
-        background = bg.convert("RGBA").resize(image_size)
-    else:
-        background = Image.open(bg).convert("RGBA").resize(image_size)
+        background = bg.convert("RGBA").resize(image_size) #Resize if bg is an image
+    else: #If bg is an image path
+        background = Image.open(bg).convert("RGBA").resize(image_size) # Open and resize image
 
-    # Composite the image onto the background using the mask
     image = Image.composite(image, background, mask)
-
     return image
 
 
+
 with gr.Blocks(theme=gr.themes.Ocean()) as demo:
     with gr.Row():
         in_video = gr.Video(label="Input Video", interactive=True)