File size: 8,959 Bytes
e392e21
a72119e
 
 
 
 
496112d
a2b9299
496112d
 
8365126
 
 
0791cf5
a8fd4c9
a72119e
0af2d38
22b8c91
 
9248f9f
2f818fd
22b8c91
2f818fd
9248f9f
 
9e41d90
 
 
 
 
 
2f818fd
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9248f9f
de54836
2f818fd
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
262a1a2
2f818fd
 
 
 
 
 
 
 
 
 
 
 
 
 
 
7b27191
2f818fd
 
9e41d90
2f818fd
 
 
 
 
6764406
 
 
9e41d90
6764406
9e41d90
6764406
 
 
 
 
9e41d90
6764406
 
 
 
 
 
 
9e41d90
6764406
 
 
5d2dafa
a8fd4c9
9e41d90
4902bd9
70e42a3
b1d6fce
d3daa33
9e41d90
402afc5
9e41d90
4902bd9
d3daa33
 
 
 
 
 
402afc5
d3daa33
402afc5
 
 
 
9e41d90
55e1949
402afc5
9e41d90
e5e4f17
9e41d90
 
 
d3daa33
 
 
55e1949
 
 
 
 
d3daa33
55e1949
 
 
d3daa33
 
55e1949
1b81f82
55e1949
 
 
 
7b27191
0f83d78
d3daa33
 
 
26a50b2
2189235
0f83d78
e87311d
7b27191
2189235
d3daa33
a72119e
9e41d90
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
import gradio as gr
from loadimg import load_img
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
import os
import tempfile
import uuid
import time
from concurrent.futures import ThreadPoolExecutor

torch.set_float32_matmul_precision("medium")
device = "cuda" if torch.cuda.is_available() else "cpu"

# Load both BiRefNet models
birefnet = AutoModelForImageSegmentation.from_pretrained("ZhengPeng7/BiRefNet", trust_remote_code=True)
birefnet.to(device)
birefnet_lite = AutoModelForImageSegmentation.from_pretrained("ZhengPeng7/BiRefNet_lite", trust_remote_code=True)
birefnet_lite.to(device)

transform_image = transforms.Compose([
    transforms.Resize((1024, 1024)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])

# Function to process a single frame
def process_frame(frame, bg_type, bg, fast_mode, bg_frame_index, background_frames, color):
    try:
        pil_image = Image.fromarray(frame)
        if bg_type == "Color":
            processed_image = process(pil_image, color, fast_mode)
        elif bg_type == "Image":
            processed_image = process(pil_image, bg, fast_mode)
        elif bg_type == "Video":
            background_frame = background_frames[bg_frame_index % len(background_frames)]
            bg_frame_index += 1
            background_image = Image.fromarray(background_frame)
            processed_image = process(pil_image, background_image, fast_mode)
        else:
            processed_image = pil_image  # Default to original image if no background is selected
        return np.array(processed_image), bg_frame_index
    except Exception as e:
        print(f"Error processing frame: {e}")
        return frame, bg_frame_index

@spaces.GPU
def fn(vid, bg_type="Color", bg_image=None, bg_video=None, color="#00FF00", fps=0, video_handling="slow_down", fast_mode=True, max_workers=6):
    try:
        start_time = time.time()  # Start the timer
        video = mp.VideoFileClip(vid)
        if fps == 0:
            fps = video.fps
        
        audio = video.audio
        frames = list(video.iter_frames(fps=fps))
        
        processed_frames = []
        yield gr.update(visible=True), gr.update(visible=False), f"Processing started... Elapsed time: 0 seconds"
        
        if bg_type == "Video":
            background_video = mp.VideoFileClip(bg_video)
            if background_video.duration < video.duration:
                if video_handling == "slow_down":
                    background_video = background_video.fx(mp.vfx.speedx, factor=video.duration / background_video.duration)
                else:  # video_handling == "loop"
                    background_video = mp.concatenate_videoclips([background_video] * int(video.duration / background_video.duration + 1))
            background_frames = list(background_video.iter_frames(fps=fps))
        else:
            background_frames = None
        
        bg_frame_index = 0  # Initialize background frame index

        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            futures = [executor.submit(process_frame, frames[i], bg_type, bg_image, fast_mode, bg_frame_index, background_frames, color) for i in range(len(frames))]
            for future in futures:
                result, bg_frame_index = future.result()
                processed_frames.append(result)
                elapsed_time = time.time() - start_time
                yield result, None, f"Processing frame {len(processed_frames)}... Elapsed time: {elapsed_time:.2f} seconds"
        
        processed_video = mp.ImageSequenceClip(processed_frames, fps=fps)
        processed_video = processed_video.set_audio(audio)
        
        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as temp_file:
            temp_filepath = temp_file.name
            processed_video.write_videofile(temp_filepath, codec="libx264")
        
        elapsed_time = time.time() - start_time
        yield gr.update(visible=False), gr.update(visible=True), f"Processing complete! Elapsed time: {elapsed_time:.2f} seconds"
        yield processed_frames[-1], temp_filepath, f"Processing complete! Elapsed time: {elapsed_time:.2f} seconds"
    
    except Exception as e:
        print(f"Error: {e}")
        elapsed_time = time.time() - start_time
        yield gr.update(visible=False), gr.update(visible=True), f"Error processing video: {e}. Elapsed time: {elapsed_time:.2f} seconds"
        yield None, f"Error processing video: {e}", f"Error processing video: {e}. Elapsed time: {elapsed_time:.2f} seconds"

def process(image, bg, fast_mode=False):
    image_size = image.size
    input_images = transform_image(image).unsqueeze(0).to(device)
    model = birefnet_lite if fast_mode else birefnet
    
    with torch.no_grad():
        preds = model(input_images)[-1].sigmoid().cpu()
    pred = preds[0].squeeze()
    pred_pil = transforms.ToPILImage()(pred)
    mask = pred_pil.resize(image_size)
    
    if isinstance(bg, str) and bg.startswith("#"):
        color_rgb = tuple(int(bg[i:i+2], 16) for i in (1, 3, 5))
        background = Image.new("RGBA", image_size, color_rgb + (255,))
    elif isinstance(bg, Image.Image):
        background = bg.convert("RGBA").resize(image_size)
    else:
        background = Image.open(bg).convert("RGBA").resize(image_size)
    
    image = Image.composite(image, background, mask)
    return image

with gr.Blocks(theme=gr.themes.Ocean()) as demo:
    gr.Markdown("# Video Background Remover & Changer\n### You can replace image background with any color, image or video.\nNOTE: As this Space is running on ZERO GPU it has limit. It can handle approx 200 frames at once. So, if you have a big video than use small chunks or Duplicate this space.")
    
    with gr.Row():
        in_video = gr.Video(label="Input Video", interactive=True)
        stream_image = gr.Image(label="Streaming Output", visible=False)
        out_video = gr.Video(label="Final Output Video")
    
    submit_button = gr.Button("Change Background", interactive=True)
    
    with gr.Row():
        fps_slider = gr.Slider(
            minimum=0,
            maximum=60,
            step=1,
            value=0,
            label="Output FPS (0 will inherit the original fps value)",
            interactive=True
        )
        bg_type = gr.Radio(["Color", "Image", "Video"], label="Background Type", value="Color", interactive=True)
        color_picker = gr.ColorPicker(label="Background Color", value="#00FF00", visible=True, interactive=True)
        bg_image = gr.Image(label="Background Image", type="filepath", visible=False, interactive=True)
        bg_video = gr.Video(label="Background Video", visible=False, interactive=True)
        
        with gr.Column(visible=False) as video_handling_options:
            video_handling_radio = gr.Radio(["slow_down", "loop"], label="Video Handling", value="slow_down", interactive=True)
        
        fast_mode_checkbox = gr.Checkbox(label="Fast Mode (Use BiRefNet_lite)", value=True, interactive=True)
        max_workers_slider = gr.Slider( minimum=1, maximum=32, step=1, value=6, label="Max Workers", info="Determines how many frames to process in parallel", interactive=True )

    time_textbox = gr.Textbox(label="Time Elapsed", interactive=False)

    def update_visibility(bg_type):
        if bg_type == "Color":
            return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
        elif bg_type == "Image":
            return gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)
        elif bg_type == "Video":
            return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True), gr.update(visible=True)
        else:
            return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)

    bg_type.change(update_visibility, inputs=bg_type, outputs=[color_picker, bg_image, bg_video, video_handling_options])

    examples = gr.Examples(
        [
            ["rickroll-2sec.mp4", "Video", None, "background.mp4"],
            ["rickroll-2sec.mp4", "Image", "images.webp", None],
            ["rickroll-2sec.mp4", "Color", None, None],
        ],
        inputs=[in_video, bg_type, bg_image, bg_video],
        outputs=[stream_image, out_video, time_textbox],
        fn=fn,
        cache_examples=True,
        cache_mode="eager",
    )

    submit_button.click(
        fn,
        inputs=[in_video, bg_type, bg_image, bg_video, color_picker, fps_slider, video_handling_radio, fast_mode_checkbox, max_workers_slider],
        outputs=[stream_image, out_video, time_textbox],
    )

if __name__ == "__main__":
    demo.launch(show_error=True)