embed / app.py
om-app's picture
Update app.py
1874dee
raw
history blame
2.41 kB
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()