Create app.py

app.py

@@ -0,0 +1,100 @@
+import spaces
+import gradio as gr
+import cv2
+import numpy as np
+import time
+import random
+from PIL import Image
+import torch
+from torchvision import transforms
+from transparent_background import Remover
+
+torch.jit.script = lambda f: f
+
+def apply_temporal_smoothing(current_mask, previous_mask, alpha=0.9):
+    if previous_mask is None:
+        return current_mask
+    return alpha * previous_mask + (1 - alpha) * current_mask
+
+def post_process_mask(mask, kernel_size=5):
+    kernel = np.ones((kernel_size, kernel_size), np.uint8)
+    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
+    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
+    mask = cv2.GaussianBlur(mask, (kernel_size, kernel_size), 0)
+    return mask
+
+@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))
+    writer = None
+    tmpname = random.randint(111111111, 999999999)
+    processed_frames = 0
+    start_time = time.time()
+    previous_mask = None
+
+    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')
+        
+        # 마스크 추출 및 후처리
+        mask = np.array(out)[:,:,3]
+        mask = post_process_mask(mask)
+        
+        # 시간적 평활화 적용
+        mask = apply_temporal_smoothing(mask, previous_mask)
+        previous_mask = mask
+
+        # 마스크 적용 및 색상 보정
+        result = cv2.multiply(frame, cv2.cvtColor(mask, cv2.COLOR_GRAY2RGB).astype(float) / 255.0)
+        result = cv2.addWeighted(result, 1.1, frame, 0, 0)  # 색상 보정
+
+        writer.write(cv2.cvtColor(result.astype(np.uint8), cv2.COLOR_RGB2BGR))
+
+    cap.release()
+    writer.release()
+    return str(tmpname) + '.mp4'
+
+title = "🎞️ Enhanced 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.
+This enhanced version includes improved mask processing, temporal smoothing, and color correction for better results.*
+"""
+
+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()