Update app.py

app.py

@@ -1,46 +1,62 @@
+import gradio as gr
 import cv2
 import numpy as np
-import gradio as gr
 from tqdm import tqdm
 
 
-def remove_green_screen(video_file, bg_image):
-
-    def process_frame(frame):
-        # remove green screen
-        hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
-        mask = cv2.inRange(hsv, (36, 25, 25), (86, 255, 255))
-        mask = cv2.merge((mask, mask, mask))
-        masked = cv2.bitwise_and(frame, mask)
+def process_frame(frame):
+    hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
+    lower_green = np.array([36, 25, 25])
+    upper_green = np.array([86, 255, 255])
+    mask = cv2.inRange(hsv, lower_green, upper_green)
+    mask_inv = cv2.bitwise_not(mask)
+    bg = cv2.imread(bg_image)
+    bg = cv2.resize(bg, (frame.shape[1], frame.shape[0]))
+    bg_masked = cv2.bitwise_and(bg, bg, mask=mask)
+    fg_masked = cv2.bitwise_and(frame, frame, mask=mask_inv)
+    return cv2.add(bg_masked, fg_masked)
 
-        # add background
-        bg = cv2.imread(bg_image)
-        bg = cv2.resize(bg, (frame.shape[1], frame.shape[0]))
-        masked[np.where((masked == [0, 0, 0]).all(axis=2))] = bg[np.where((masked == [0, 0, 0]).all(axis=2))]
 
-        return masked
-
-    cap = cv2.VideoCapture(video_file)
-    frames = []
+def remove_green_screen(video, bg_image):
+    cap = cv2.VideoCapture(video.name)
+    fps = cap.get(cv2.CAP_PROP_FPS)
     frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    out = cv2.VideoWriter('out.mp4', cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
+    frames = []
     for i in tqdm(range(frame_count)):
-        success, frame = cap.read()
-        if not success:
-            continue
+        ret, frame = cap.read()
+        if not ret:
+            break
         frames.append(process_frame(frame))
-
     cap.release()
+    for frame in tqdm(frames):
+        out.write(frame)
+    out.release()
+    return 'out.mp4'
+
+
+def predict_video(video, bg_image):
+    result_file = remove_green_screen(video, bg_image)
+    with open(result_file, 'rb') as f:
+        result = f.read()
+    return result
+
 
-    return np.array(frames)
+inputs = [
+    gr.inputs.Video(label="Video"),
+    gr.inputs.Image(label="Background Image")
+]
 
+outputs = gr.outputs.Video(label="Processed Video", type="mp4", source="file")
 
-iface = gr.Interface(
-    fn=remove_green_screen,
-    inputs=[
-        gr.inputs.Video(type="mp4", label="Input video"),
-        gr.inputs.Image(label="Background image")
-    ],
-    outputs=gr.outputs.Video(label="Processed video")
-)
+title = "Green Screen Video Remover"
+description = "Remove the green background from a video and replace it with another image."
+article = "<p style='text-align: center'><a href='https://huggingface.co/blog/how-to-build-a-web-app-for-a-transformer-in-pytorch'>How to Build a Web App for a Transformer in PyTorch</a></p>"
+examples = [
+    ["./input.mp4", "./bg.png"],
+    ["./input2.mp4", "./bg2.png"]
+]
 
-iface.launch()
+gr.Interface(predict_video, inputs, outputs, title=title, description=description, article=article, examples=examples, live=True).launch()