not-lain KingNish commited on
Commit
2189235
1 Parent(s): 3d705e5

Update app.py (#3)

Browse files

- Update app.py (3730e1073a7bdd6a026ab547730fafafef1075f4)


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

Files changed (1) hide show
  1. app.py +25 -19
app.py CHANGED
@@ -26,26 +26,27 @@ transform_image = transforms.Compose(
26
  ]
27
  )
28
 
 
29
  @spaces.GPU
30
- def fn(vid):
31
  # Load the video using moviepy
32
  video = mp.VideoFileClip(vid)
33
 
34
  # Extract audio from the video
35
  audio = video.audio
36
 
37
- # Extract frames at 12 fps
38
- frames = video.iter_frames(fps=12)
39
 
40
  # Process each frame for background removal
41
  processed_frames = []
42
  for frame in frames:
43
  pil_image = Image.fromarray(frame)
44
- processed_image = process(pil_image)
45
  processed_frames.append(np.array(processed_image))
46
 
47
  # Create a new video from the processed frames
48
- processed_video = mp.ImageSequenceClip(processed_frames, fps=12)
49
 
50
  # Add the original audio back to the processed video
51
  processed_video = processed_video.set_audio(audio)
@@ -61,7 +62,7 @@ def fn(vid):
61
  return temp_filepath
62
 
63
 
64
- def process(image):
65
  image_size = image.size
66
  input_images = transform_image(image).unsqueeze(0).to("cuda")
67
  # Prediction
@@ -71,32 +72,37 @@ def process(image):
71
  pred_pil = transforms.ToPILImage()(pred)
72
  mask = pred_pil.resize(image_size)
73
 
74
- # Create a green screen image
75
- green_screen = Image.new("RGBA", image_size, (0, 255, 0, 255))
 
 
 
76
 
77
- # Composite the image onto the green screen using the mask
78
- image = Image.composite(image, green_screen, mask)
79
 
80
  return image
81
 
82
 
83
- def process_file(f):
84
  name_path = f.rsplit(".", 1)[0] + ".png"
85
  im = load_img(f, output_type="pil")
86
  im = im.convert("RGB")
87
- transparent = process(im)
88
  transparent.save(name_path)
89
  return name_path
90
 
91
 
92
- in_video = gr.Video(label="birefnet")
93
- out_video = gr.Video()
94
-
95
-
96
- demo = gr.Interface(
97
- fn, inputs=in_video, outputs=out_video, api_name="video"
98
- )
99
 
 
 
 
100
 
101
  if __name__ == "__main__":
102
  demo.launch(show_error=True)
 
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)
 
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
 
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__":
108
  demo.launch(show_error=True)