Update app.py (#1)

- Update app.py (7e8bbe825e852025ee36e5c1593bc9ac9ed126ab)
- Update requirements.txt (6f54a09c84f667ea2d3e8d8e42fae3b5a1f6e8d4)


Co-authored-by: Nishith Jain <[email protected]>

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"
 )
 
 

requirements.txt

@@ -13,4 +13,6 @@ huggingface_hub
 transformers>=4.39.1
 gradio
 gradio_imageslider
-loadimg>=0.1.1
+loadimg>=0.1.1
+moviepy 
+pydub