not-lain KingNish commited on
Commit
4902bd9
1 Parent(s): 361c20d

Update app.py (#4)

Browse files

- Update app.py (f1cdcfc2c95766fedcfeaeff0c03c649569c1a30)
- Update app.py (ff98b478fbc03e108b49fc3f293940cec52ea952)


Co-authored-by: Nishith Jain <KingNish@users.noreply.huggingface.co>

Files changed (1) hide show
  1. app.py +64 -53
app.py CHANGED
@@ -26,82 +26,93 @@ transform_image = transforms.Compose(
26
  ]
27
  )
28
 
29
-
30
  @spaces.GPU
31
- def fn(vid, fps, color):
32
  # Load the video using moviepy
33
  video = mp.VideoFileClip(vid)
34
-
35
- # Extract audio from the video
36
  audio = video.audio
37
-
38
- # Extract frames at the specified FPS
39
- frames = video.iter_frames(fps=fps)
40
-
41
- # Process each frame for background removal
42
- processed_frames = []
43
- for frame in frames:
 
 
 
 
44
  pil_image = Image.fromarray(frame)
45
- processed_image = process(pil_image, color)
46
- processed_frames.append(np.array(processed_image))
47
-
48
- # Create a new video from the processed frames
49
- processed_video = mp.ImageSequenceClip(processed_frames, fps=fps)
50
-
51
- # Add the original audio back to the processed video
52
- processed_video = processed_video.set_audio(audio)
53
-
54
- # Save the processed video to a temporary file
 
 
 
 
 
55
  temp_dir = "temp"
56
  os.makedirs(temp_dir, exist_ok=True)
57
- unique_filename = str(uuid.uuid4()) + ".mp4"
58
- temp_filepath = os.path.join(temp_dir, unique_filename)
 
 
59
  processed_video.write_videofile(temp_filepath, codec="libx264")
60
-
61
- # Return the path to the temporary file
62
- return temp_filepath
63
-
 
 
 
 
64
 
65
  def process(image, color_hex):
66
  image_size = image.size
67
  input_images = transform_image(image).unsqueeze(0).to("cuda")
68
- # Prediction
69
  with torch.no_grad():
70
  preds = birefnet(input_images)[-1].sigmoid().cpu()
71
  pred = preds[0].squeeze()
72
  pred_pil = transforms.ToPILImage()(pred)
73
  mask = pred_pil.resize(image_size)
74
-
75
- # Convert hex color to RGB tuple
76
- color_rgb = tuple(int(color_hex[i : i + 2], 16) for i in (1, 3, 5))
77
-
78
- # Create a background image with the chosen color
79
  background = Image.new("RGBA", image_size, color_rgb + (255,))
80
-
81
- # Composite the image onto the background using the mask
82
  image = Image.composite(image, background, mask)
 
 
83
 
84
- return image
85
-
86
-
87
- def process_file(f, color="#00FF00"):
88
- name_path = f.rsplit(".", 1)[0] + ".png"
89
- im = load_img(f, output_type="pil")
90
- im = im.convert("RGB")
91
- transparent = process(im, color)
92
- transparent.save(name_path)
93
- return name_path
94
-
95
 
96
  with gr.Blocks() as demo:
97
- in_video = gr.Video(label="birefnet")
98
- out_video = gr.Video()
99
- fps_slider = gr.Slider(minimum=1, maximum=60, step=1, value=12, label="Output FPS")
100
- color_picker = gr.ColorPicker(label="Background Color", value="#00FF00")
101
- submit_button = gr.Button("Process Video")
102
-
 
 
 
 
 
 
 
 
 
 
103
  submit_button.click(
104
- fn, inputs=[in_video, fps_slider, color_picker], outputs=out_video
 
 
105
  )
106
 
107
  if __name__ == "__main__":
 
26
  ]
27
  )
28
 
 
29
  @spaces.GPU
30
+ def fn(vid, fps, color, progress=gr.Progress()):
31
  # Load the video using moviepy
32
  video = mp.VideoFileClip(vid)
 
 
33
  audio = video.audio
34
+ frames = list(video.iter_frames(fps=fps))
35
+ total_frames = len(frames)
36
+
37
+ processed_frames_no_bg = []
38
+ processed_frames_changed_bg = []
39
+
40
+ # Create a live preview state
41
+ preview_no_bg = None
42
+ preview_with_bg = None
43
+
44
+ for idx, frame in enumerate(progress.tqdm(frames)):
45
  pil_image = Image.fromarray(frame)
46
+ processed_image, mask = process(pil_image, color)
47
+
48
+ processed_frames_no_bg.append(np.array(processed_image))
49
+
50
+ background = Image.new("RGBA", pil_image.size, color + (255,))
51
+ composed_image = Image.composite(pil_image, background, mask)
52
+ processed_frames_changed_bg.append(np.array(composed_image))
53
+
54
+ # Update preview every 10 frames or on the last frame
55
+ if idx % 10 == 0 or idx == total_frames - 1:
56
+ preview_no_bg = np.array(processed_image)
57
+ preview_with_bg = np.array(composed_image)
58
+ yield preview_no_bg, preview_with_bg, None, None
59
+
60
+ # Create videos from processed frames
61
  temp_dir = "temp"
62
  os.makedirs(temp_dir, exist_ok=True)
63
+
64
+ processed_video = mp.ImageSequenceClip(processed_frames_changed_bg, fps=fps)
65
+ processed_video = processed_video.set_audio(audio)
66
+ temp_filepath = os.path.join(temp_dir, f"{uuid.uuid4()}.mp4")
67
  processed_video.write_videofile(temp_filepath, codec="libx264")
68
+
69
+ processed_video_no_bg = mp.ImageSequenceClip(processed_frames_no_bg, fps=fps)
70
+ processed_video_no_bg = processed_video_no_bg.set_audio(audio)
71
+ temp_filepath_no_bg = os.path.join(temp_dir, f"{uuid.uuid4()}.webm")
72
+ processed_video_no_bg.write_videofile(temp_filepath_no_bg, codec="libvpx")
73
+
74
+ # Final yield with completed videos
75
+ yield None, None, temp_filepath_no_bg, temp_filepath
76
 
77
  def process(image, color_hex):
78
  image_size = image.size
79
  input_images = transform_image(image).unsqueeze(0).to("cuda")
 
80
  with torch.no_grad():
81
  preds = birefnet(input_images)[-1].sigmoid().cpu()
82
  pred = preds[0].squeeze()
83
  pred_pil = transforms.ToPILImage()(pred)
84
  mask = pred_pil.resize(image_size)
85
+
86
+ color_rgb = tuple(int(color_hex[i:i + 2], 16) for i in (1, 3, 5))
 
 
 
87
  background = Image.new("RGBA", image_size, color_rgb + (255,))
 
 
88
  image = Image.composite(image, background, mask)
89
+
90
+ return image, mask
91
 
92
+ def change_color(in_video, fps_slider, color_picker):
93
+ return gr.update(visible=True), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)
 
 
 
 
 
 
 
 
 
94
 
95
  with gr.Blocks() as demo:
96
+ with gr.Row():
97
+ in_video = gr.Video(label="Input Video")
98
+ no_bg_video = gr.Video(label="No BG Video", visible=True)
99
+ out_video = gr.Video(label="Output Video", visible=True)
100
+
101
+ with gr.Row(visible=False) as preview_row:
102
+ preview_no_bg = gr.Image(label="Live Preview (No Background)", visible=True)
103
+ preview_with_bg = gr.Image(label="Live Preview (With Background)", visible=True)
104
+
105
+ with gr.Row():
106
+ fps_slider = gr.Slider(minimum=1, maximum=60, step=1, value=12, label="Output FPS")
107
+ color_picker = gr.ColorPicker(label="Background Color", value="#00FF00")
108
+
109
+ submit_button = gr.Button("Change Background")
110
+
111
+ # Handle visibility changes and processing
112
  submit_button.click(
113
+ fn=fn,
114
+ inputs=[in_video, fps_slider, color_picker],
115
+ outputs=[preview_no_bg, preview_with_bg, no_bg_video, out_video]
116
  )
117
 
118
  if __name__ == "__main__":