File size: 2,174 Bytes
7de892f
de2aa9b
3cd966d
6ce3893
caa1229
85d769e
6ce3893
b2c8d9d
 
 
 
6ce3893
ba4d1a9
27202f7
a63ebbb
 
a3504d6
5e5187b
 
a63ebbb
6ce3893
a63ebbb
6ce3893
85d769e
6ce3893
dded11c
a63ebbb
6ce3893
 
a63ebbb
6ce3893
 
dded11c
15a2bfc
 
 
 
 
a63ebbb
6ce3893
a63ebbb
6ce3893
a63ebbb
d4ae210
6ce3893
a63ebbb
 
6ce3893
 
a63ebbb
6ce3893
 
dc913e2
ce6434e
a63ebbb
cabaac2
927e8af
a63ebbb
927e8af
a3504d6
a63ebbb
 
 
 
 
 
a3504d6
a63ebbb
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
import spaces
import gradio as gr
import cv2
import numpy as np
import time
import random
from PIL import Image
import torch

torch.jit.script = lambda f: f

from transparent_background import Remover

@spaces.GPU()
def doo(video, mode, progress=gr.Progress()):
    if mode == 'Fast':
        remover = Remover(mode='fast')
    else:
        remover = Remover()

    cap = cv2.VideoCapture(video)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))  # Get total frames
    writer = None
    tmpname = random.randint(111111111, 999999999)
    processed_frames = 0
    start_time = time.time()

    while cap.isOpened():
        ret, frame = cap.read()

        if ret is False:
            break

        if time.time() - start_time >= 20 * 60 - 5:
            print("GPU Timeout is coming")
            cap.release()
            writer.release()
            return str(tmpname) + '.mp4'
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        img = Image.fromarray(frame).convert('RGB')

        if writer is None:
            writer = cv2.VideoWriter(str(tmpname) + '.mp4', cv2.VideoWriter_fourcc(*'mp4v'), cap.get(cv2.CAP_PROP_FPS), img.size)

        processed_frames += 1
        print(f"Processing frame {processed_frames}")
        progress(processed_frames / total_frames, desc=f"Processing frame {processed_frames}/{total_frames}")
        out = remover.process(img, type='green')
        writer.write(cv2.cvtColor(np.array(out), cv2.COLOR_BGR2RGB))

    cap.release()
    writer.release()
    return str(tmpname) + '.mp4'

title = "🎞️ Video Background Removal Tool 🎥"
description = """*Please note that if your video file is long (has a high number of frames), there is a chance that processing break due to GPU timeout. In this case, consider trying Fast mode."""

examples = [['./input.mp4']]

iface = gr.Interface(
    fn=doo,
    inputs=["video", gr.components.Radio(['Normal', 'Fast'], label='Select mode', value='Normal', info='Normal is more accurate, but takes longer. | Fast has lower accuracy so the process will be faster.')],
    outputs="video",
    examples=examples,
    title=title,
    description=description
)
iface.launch()