Update app.py

app.py

@@ -39,16 +39,14 @@ def fn(vid, fps, color):
  frames = video.iter_frames(fps=fps)
 
  # Process each frame for background removal
- processed_frames_no_bg = []
- processed_frames_changed_bg = []
+ processed_frames = []
  for frame in frames:
  pil_image = Image.fromarray(frame)
- processed_image, mask = process(pil_image, color) 
- processed_frames_no_bg.append(np.array(mask)) 
- processed_frames_changed_bg.append(np.array(processed_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_changed_bg, fps=fps)
+ processed_video = mp.ImageSequenceClip(processed_frames, fps=fps)
 
  # Add the original audio back to the processed video
  processed_video = processed_video.set_audio(audio)
@@ -60,13 +58,8 @@ def fn(vid, fps, color):
  temp_filepath = os.path.join(temp_dir, unique_filename)
  processed_video.write_videofile(temp_filepath, codec="libx264")
 
- # Create and save no-background video
- processed_video_no_bg = mp.ImageSequenceClip(processed_frames_no_bg, fps=fps)
- processed_video_no_bg = processed_video_no_bg.set_audio(audio)
- temp_filepath_no_bg = os.path.join(temp_dir, str(uuid.uuid4()) + ".webm")
- processed_video_no_bg.write_videofile(temp_filepath_no_bg, codec="libvpx")
-
- return temp_filepath_no_bg, temp_filepath
+ # Return the path to the temporary file
+ return temp_filepath
 
 
 def process(image, color_hex):
@@ -80,7 +73,7 @@ def process(image, color_hex):
  mask = pred_pil.resize(image_size)
 
  # Convert hex color to RGB tuple
- color_rgb = tuple(int(color_hex[i:i + 2], 16) for i in (1, 3, 5))
+ 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,))
@@ -88,22 +81,29 @@ def process(image, color_hex):
  # Composite the image onto the background using the mask
  image = Image.composite(image, background, mask)
 
- return image, mask # Return both the processed image and the mask
+ return image
+
+
+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, color)
+ transparent.save(name_path)
+ return name_path
 
 
 with gr.Blocks() as demo:
  with gr.Row():
  in_video = gr.Video(label="Input Video")
- no_bg_video = gr.Video(label="No BG Video") # Added for no-background video
- out_video = gr.Video(label="Output Video") # This will be the changed-background video
+ out_video = gr.Video(label="Output Video") 
  submit_button = gr.Button("Change Background")
  with gr.Row():
  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.click(
- fn, inputs=[in_video, fps_slider, color_picker], outputs=[no_bg_video, out_video] 
+ fn, inputs=[in_video, fps_slider, color_picker], outputs=out_video
  )
 
 if __name__ == "__main__":