Update app.py (#3)

- Update app.py (3730e1073a7bdd6a026ab547730fafafef1075f4)


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

app.py

@@ -26,26 +26,27 @@ transform_image = transforms.Compose(
     ]
 )
 
+
 @spaces.GPU
-def fn(vid):
+def fn(vid, fps, color):
     # 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)
+    # Extract frames at the specified FPS
+    frames = video.iter_frames(fps=fps)
 
     # Process each frame for background removal
     processed_frames = []
     for frame in frames:
         pil_image = Image.fromarray(frame)
-        processed_image = process(pil_image)
+        processed_image = process(pil_image, color)
         processed_frames.append(np.array(processed_image))
 
     # Create a new video from the processed frames
-    processed_video = mp.ImageSequenceClip(processed_frames, fps=12)
+    processed_video = mp.ImageSequenceClip(processed_frames, fps=fps)
 
     # Add the original audio back to the processed video
     processed_video = processed_video.set_audio(audio)
@@ -61,7 +62,7 @@ def fn(vid):
     return temp_filepath
 
 
-def process(image):
+def process(image, color_hex):
     image_size = image.size
     input_images = transform_image(image).unsqueeze(0).to("cuda")
     # Prediction
@@ -71,32 +72,37 @@ def process(image):
     pred_pil = transforms.ToPILImage()(pred)
     mask = pred_pil.resize(image_size)
 
-    # Create a green screen image
-    green_screen = Image.new("RGBA", image_size, (0, 255, 0, 255))
+    # Convert hex color to RGB tuple
+    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,))
 
-    # Composite the image onto the green screen using the mask
-    image = Image.composite(image, green_screen, mask)
+    # Composite the image onto the background using the mask
+    image = Image.composite(image, background, mask)
 
     return image
 
 
-def process_file(f):
+def process_file(f, color="#00FF00"):
     name_path = f.rsplit(".", 1)[0] + ".png"
     im = load_img(f, output_type="pil")
     im = im.convert("RGB")
-    transparent = process(im)
+    transparent = process(im, color)
     transparent.save(name_path)
     return name_path
 
 
-in_video = gr.Video(label="birefnet")
-out_video = gr.Video()
-
-
-demo = gr.Interface(
-    fn, inputs=in_video, outputs=out_video, api_name="video"
-)
+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")
 
+    submit_button.click(
+        fn, inputs=[in_video, fps_slider, color_picker], outputs=out_video
+    )
 
 if __name__ == "__main__":
     demo.launch(show_error=True)