Update app.py

app.py

@@ -3,69 +3,25 @@ import numpy as np
 import gradio as gr
 
 def remove_green_screen(video_file, bg_image):
-    # Load the background image
-    bg = cv2.imread(bg_image)
-
-    # Create a video capture object to read from the video file
     cap = cv2.VideoCapture(video_file)
-
-    # Get the video dimensions
-    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-
-    # Create a video writer object to write the output to a new video file
-    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
-    out = cv2.VideoWriter("output.mp4", fourcc, 30, (width, height))
-
-    # Loop over each frame in the video
+    background = cv2.imread(bg_image)
     while True:
-        # Read a frame from the video
         ret, frame = cap.read()
-
-        # If we have reached the end of the video, break out of the loop
         if not ret:
             break
-
-        # Convert the frame to HSV color space
         hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
-
-        # Define the range of green color in HSV
-        lower_green = np.array([40, 50, 50])
+        lower_green = np.array([35, 0, 0])
         upper_green = np.array([90, 255, 255])
-
-        # Create a mask for the green color range
         mask = cv2.inRange(hsv, lower_green, upper_green)
-
-        # Invert the mask
-        mask = cv2.bitwise_not(mask)
-
-        # Apply the mask to the frame to remove the green screen
-        fg = cv2.bitwise_and(frame, frame, mask=mask)
-
-        # Apply the mask to the background image
-        bg_mask = cv2.bitwise_not(mask)
-        bg_mask = cv2.cvtColor(bg_mask, cv2.COLOR_GRAY2BGR)
-        bg_fg = cv2.bitwise_and(bg, bg_mask)
-
-        # Combine the foreground and background images
-        output = cv2.add(fg, bg_fg)
-
-        # Write the output frame to the video file
-        out.write(output)
-
-    # Release the video capture and video writer objects
+        mask_inv = cv2.bitwise_not(mask)
+        fg = cv2.bitwise_and(frame, frame, mask=mask_inv)
+        bg = cv2.bitwise_and(background, background, mask=mask)
+        final = cv2.add(fg, bg)
+        cv2.imshow("output", final)
+        if cv2.waitKey(1) & 0xFF == ord("q"):
+            break
     cap.release()
-    out.release()
-
-    return "output.mp4"
-
-# Define the input and output components for the app
-video_file = gr.inputs.Video(label="Input Video")
-bg_image = gr.inputs.Image(label="Background Image")
-output_file = gr.outputs.Video(label="Output Video")
-
-# Create the app interface
-iface = gr.Interface(fn=remove_green_screen, inputs=[video_file, bg_image], outputs=output_file)
+    cv2.destroyAllWindows()
 
-# Launch the app
+iface = gr.Interface(fn=remove_green_screen, inputs=["file", "image"], outputs=None, title="Green Screen Remover")
 iface.launch()