Update app.py

app.py

@@ -5,6 +5,11 @@ import spaces
 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
+
 
 torch.set_float32_matmul_precision(["high", "highest"][0])
 
@@ -20,13 +25,33 @@ transform_image = transforms.Compose(
     ]
 )
 
+
 def fn(vid):
-    # TODO 
-    # loop over video and extract images and process each one
-    im = load_img(vid, output_type="pil")
-    im = im.convert("RGB")
-    image = process(im)
-    return image
+    # Load the video using moviepy
+    video = mp.VideoFileClip(vid)
+
+    # Extract audio from the video
+    audio = video.audio
+
+    # Extract frames at 12 fps
+    frames = video.iter_frames(fps=12)
+
+    # Process each frame for background removal
+    processed_frames = []
+    for frame in frames:
+        pil_image = Image.fromarray(frame)
+        processed_image = process(pil_image)
+        processed_frames.append(np.array(processed_image))
+
+    # Create a new video from the processed frames
+    processed_video = mp.ImageSequenceClip(processed_frames, fps=12)
+
+    # Add the original audio back to the processed video
+    processed_video = processed_video.set_audio(audio)
+
+    # Return the processed video
+    return processed_video
+
 
 @spaces.GPU
 def process(image):
@@ -38,24 +63,32 @@ def process(image):
     pred = preds[0].squeeze()
     pred_pil = transforms.ToPILImage()(pred)
     mask = pred_pil.resize(image_size)
-    image.putalpha(mask)
+
+    # Create a green screen image
+    green_screen = Image.new("RGBA", image_size, (0, 255, 0, 255))
+
+    # Composite the image onto the green screen using the mask
+    image = Image.composite(image, green_screen, mask)
+
     return image
-  
+
+
 def process_file(f):
-    name_path = f.rsplit(".",1)[0]+".png"
+    name_path = f.rsplit(".", 1)[0] + ".png"
     im = load_img(f, output_type="pil")
     im = im.convert("RGB")
     transparent = process(im)
     transparent.save(name_path)
     return name_path
 
+
 in_video = gr.Video(label="birefnet")
 out_video = gr.Video()
 
 
 url = "https://hips.hearstapps.com/hmg-prod/images/gettyimages-1229892983-square.jpg"
 demo = gr.Interface(
-    fn, inputs=in_video, outputs=out_video, api_name="image"
+    fn, inputs=in_video, outputs=out_video, api_name="video"
 )