Update app.py

app.py

@@ -1,62 +1,67 @@
 import gradio as gr
 import cv2
 import numpy as np
-from tqdm import tqdm
-
-
-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)
-
-
-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))
+
+def remove_green_screen(video):
+    def process_frame(frame):
+        # Convert to HSV color space
+        hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
+
+        # Define the range of green color
+        lower_green = np.array([25, 52, 72])
+        upper_green = np.array([102, 255, 255])
+
+        # Threshold the HSV image to get only green colors
+        mask = cv2.inRange(hsv, lower_green, upper_green)
+
+        # Bitwise-AND mask and original image
+        res = cv2.bitwise_and(frame, frame, mask=mask)
+
+        return res
+
+    # Load background image
+    bg_image = cv2.imread("background.jpg")
+
+    # Open video file
+    cap = cv2.VideoCapture(video)
+
+    # Get video codec and fps
+    fourcc = int(cap.get(cv2.CAP_PROP_FOURCC))
+    fps = int(cap.get(cv2.CAP_PROP_FPS))
+
+    # Get video dimensions
     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))
+
+    # Create output video writer
+    out = cv2.VideoWriter("output.mp4", fourcc, fps, (width, height))
+
     frames = []
-    for i in tqdm(range(frame_count)):
+
+    # Loop through video frames
+    while True:
         ret, frame = cap.read()
         if not ret:
             break
+
+        # Process frame
+        bg = cv2.imread(bg_image)
         frames.append(process_frame(frame))
+
+    # Release video file and output video writer
     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
 
+    # Write output video
+    out = cv2.VideoWriter("output.mp4", fourcc, fps, (width, height))
+    for frame in frames:
+        out.write(frame)
+    out.release()
 
-inputs = [
-    gr.inputs.Video(label="Video"),
-    gr.inputs.Image(label="Background Image")
-]
+    return "output.mp4"
 
-outputs = gr.outputs.Video(label="Processed Video", type="mp4", source="file")
 
-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"]
-]
+input_video = gr.inputs.Video(label="Input Video")
+outputs = gr.outputs.Video(label="Processed Video", type="mp4")
 
-gr.Interface(predict_video, inputs, outputs, title=title, description=description, article=article, examples=examples, live=True).launch() 
+gr.Interface(remove_green_screen, inputs=input_video, outputs=outputs).launch()