|
import gradio as gr |
|
import cv2 |
|
import numpy as np |
|
|
|
def process_frame(frame, bg_image): |
|
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) |
|
hsv = cv2.cvtColor(frame, cv2.COLOR_RGB2HSV) |
|
lower_green = np.array([45, 100, 50]) |
|
upper_green = np.array([75, 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])) |
|
fg = cv2.bitwise_and(frame, frame, mask=mask_inv) |
|
bg = cv2.bitwise_and(bg, bg, mask=mask) |
|
result = cv2.add(bg, fg) |
|
return result |
|
|
|
def remove_green_screen(input_video, bg_image): |
|
cap = cv2.VideoCapture(input_video.name) |
|
codec = cv2.VideoWriter_fourcc(*"mp4v") |
|
fps = int(cap.get(cv2.CAP_PROP_FPS)) |
|
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) |
|
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) |
|
out = cv2.VideoWriter("output.mp4", codec, fps, (frame_width, frame_height)) |
|
while True: |
|
ret, frame = cap.read() |
|
if not ret: |
|
break |
|
result = process_frame(frame, bg_image) |
|
out.write(result) |
|
cap.release() |
|
out.release() |
|
|
|
def predict(input_video, bg_image): |
|
remove_green_screen(input_video, bg_image) |
|
return "output.mp4" |
|
|
|
inputs = [ |
|
gr.inputs.Video(label="Input Video"), |
|
gr.inputs.Image(label="Background Image") |
|
] |
|
|
|
outputs = gr.outputs.Video(label="Processed Video", type="auto") |
|
|
|
title = "Green Screen Remover" |
|
description = "Upload a video and an image to use as the background to remove the green screen." |
|
article = "<p style='text-align: center'><a href='https://github.com/gradio-app/examples/blob/master/green_screen_removal.py'>This code</a> was made into an interactive interface using Gradio. See the full tutorial at <a href='https://blog.gradio.app/green-screen-removal-with-opencv/'>this blog post</a> on the Gradio blog.</p>" |
|
examples = [ |
|
[ |
|
"https://www.youtube.com/watch?v=clD6_yXKo2I", |
|
"https://i.imgur.com/lxIhsG6.jpg" |
|
], |
|
[ |
|
"https://www.youtube.com/watch?v=6DfZ6UOZi0A", |
|
"https://i.imgur.com/6UaTvfo.jpg" |
|
], |
|
] |
|
|
|
iface = gr.Interface( |
|
fn=predict, |
|
inputs=inputs, |
|
outputs=outputs, |
|
title=title, |
|
description=description, |
|
article=article, |
|
examples=examples, |
|
analytics_enabled=False, |
|
server_port=8000 |
|
) |
|
|
|
if __name__ == '__main__': |
|
iface.launch() |
|
|