Update app.py

app.py

@@ -1,130 +1,79 @@
-import spaces
 import gradio as gr
 import cv2
 import numpy as np
+import tempfile
 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()
-
+def process_video(video, mode, progress=gr.Progress()):
+    remover = Remover(mode='fast' if mode == 'Fast' else 'base')
+    
     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
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    
+    output_filename = f"{random.randint(111111111, 999999999)}.mp4"
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    out = cv2.VideoWriter(output_filename, fourcc, fps, (width, height))
+
+    prev_frame = None
 
-    while cap.isOpened():
+    for frame_num in range(total_frames):
         ret, frame = cap.read()
-        if ret is False:
+        if not ret:
             break
 
-        if time.time() - start_time >= 20 * 60 - 5:
-            print("GPU Timeout is coming")
-            cap.release()
-            if writer:
-                writer.release()
-            return str(tmpname) + '.mp4'
+        progress(frame_num / total_frames, desc=f"Processing frame {frame_num+1}/{total_frames}")
 
-        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')
-        
-        # 마스크 생성 (RGB to grayscale)
-        mask = cv2.cvtColor(np.array(out), cv2.COLOR_RGB2GRAY)
+        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        pil_image = Image.fromarray(rgb_frame)
         
-        # 마스크를 float32로 변환하고 0-1 범위로 정규화
-        mask = mask.astype(np.float32) / 255.0
-        
-        # 마스크 후처리
-        mask = post_process_mask(mask)
-        
-        # 시간적 평활화 적용
-        mask = apply_temporal_smoothing(mask, previous_mask)
-        previous_mask = mask
+        # 배경 제거
+        output = remover.process(pil_image)
+        output = np.array(output)
 
-        # 마스크를 다시 0-255 범위의 uint8로 변환
-        mask = (mask * 255).astype(np.uint8)
-        
-        # 마스크 적용 및 색상 보정
-        mask_3d = cv2.cvtColor(mask, cv2.COLOR_GRAY2RGB)
-        
-        # frame을 float32로 변환하고 0-1 범위로 정규화
-        frame = frame.astype(np.float32) / 255.0
-        mask_3d = mask_3d.astype(np.float32) / 255.0
-        
-        result = cv2.multiply(frame, mask_3d)
-        result = cv2.addWeighted(result, 1.1, frame, 0, 0)  # 색상 보정
-        
-        # 0-255 범위로 변환
-        result = (result * 255).astype(np.uint8)
+        # 알파 채널이 있다면 마스크로 사용, 없다면 그레이스케일로 변환
+        if output.shape[2] == 4:
+            mask = output[:,:,3]
+        else:
+            mask = cv2.cvtColor(output, cv2.COLOR_RGB2GRAY)
 
-        writer.write(cv2.cvtColor(result, cv2.COLOR_RGB2BGR))
+        # 마스크 임계값 처리
+        _, mask = cv2.threshold(mask, 128, 255, cv2.THRESH_BINARY)
 
-    cap.release()
-    if writer:
-        writer.release()
-    return str(tmpname) + '.mp4'
+        # 움직임 검출
+        if prev_frame is not None:
+            diff = cv2.absdiff(frame, prev_frame)
+            motion_mask = cv2.threshold(cv2.cvtColor(diff, cv2.COLOR_BGR2GRAY), 30, 255, cv2.THRESH_BINARY)[1]
+            motion_mask = cv2.dilate(motion_mask, np.ones((5,5), np.uint8), iterations=2)
+            mask = cv2.bitwise_or(mask, motion_mask)
 
-def post_process_mask(mask, kernel_size=5):
-    kernel = np.ones((kernel_size, kernel_size), np.float32)
-    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
-    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
-    return cv2.GaussianBlur(mask, (kernel_size, kernel_size), 0)
+        prev_frame = frame.copy()
 
-def apply_temporal_smoothing(current_mask, previous_mask, alpha=0.9):
-    if previous_mask is None:
-        return current_mask
-    return cv2.addWeighted(previous_mask, alpha, current_mask, 1-alpha, 0)
+        # 마스크 적용
+        mask = cv2.GaussianBlur(mask, (5, 5), 0)
+        mask = mask.astype(float) / 255.0
+        mask = np.stack([mask] * 3, axis=2)
+        result = frame.astype(float) * mask
 
-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.*
-"""
+        # 결과 저장
+        out.write(result.astype(np.uint8))
+
+    cap.release()
+    out.release()
 
-examples = [['./input.mp4']]
+    return output_filename
 
 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.')],
+    fn=process_video,
+    inputs=["video", gr.Radio(["Normal", "Fast"], label="Processing mode")],
     outputs="video",
-    examples=examples,
-    title=title,
-    description=description
+    title="Video Background Removal",
+    description="Upload a video to remove its background."
 )
 
 iface.launch()