|
|
import cv2 |
|
|
import mediapipe as mp |
|
|
import numpy as np |
|
|
import gradio as gr |
|
|
import os |
|
|
import time |
|
|
|
|
|
|
|
|
os.environ["MPLCONFIGDIR"] = "/tmp/matplotlib-cache" |
|
|
|
|
|
|
|
|
mp_selfie_segmentation = mp.solutions.selfie_segmentation |
|
|
segmentation = mp_selfie_segmentation.SelfieSegmentation(model_selection=1) |
|
|
|
|
|
|
|
|
settings = { |
|
|
"seg_enabled": True, |
|
|
"blur_bg": False, |
|
|
"set_bg": False, |
|
|
"set_color": False, |
|
|
"bg_color": (0, 0, 0), |
|
|
"blur_intensity": 15 |
|
|
} |
|
|
bg_image = None |
|
|
|
|
|
def process_frame(frame, seg_enabled, blur_bg, set_bg, set_color, bg_color, blur_intensity, custom_image=None): |
|
|
global bg_image |
|
|
settings.update({ |
|
|
"seg_enabled": seg_enabled, |
|
|
"blur_bg": blur_bg, |
|
|
"set_bg": set_bg, |
|
|
"set_color": set_color, |
|
|
"bg_color": tuple(map(int, bg_color.split(","))) if set_color and bg_color else (0, 0, 0), |
|
|
"blur_intensity": blur_intensity |
|
|
}) |
|
|
|
|
|
if custom_image is not None and set_bg: |
|
|
bg_image = custom_image |
|
|
|
|
|
if frame is None: |
|
|
return None, "No video frame received" |
|
|
|
|
|
process_start = time.time() |
|
|
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) |
|
|
result = segmentation.process(frame_rgb) |
|
|
mask = result.segmentation_mask |
|
|
|
|
|
|
|
|
alpha = mask > 0.5 |
|
|
alpha = alpha.astype(np.uint8) * 255 |
|
|
alpha = cv2.merge([alpha, alpha, alpha]) |
|
|
|
|
|
output_frame = frame.copy() |
|
|
|
|
|
if settings["seg_enabled"]: |
|
|
if settings["blur_bg"]: |
|
|
bg = cv2.resize(frame, None, fx=0.1, fy=0.1, interpolation=cv2.INTER_LINEAR) |
|
|
ksize = settings["blur_intensity"] |
|
|
if ksize % 2 == 0: |
|
|
ksize -= 1 |
|
|
bg = cv2.GaussianBlur(bg, (ksize, ksize), 0) |
|
|
bg = cv2.resize(bg, (frame.shape[1], frame.shape[0]), interpolation=cv2.INTER_LINEAR) |
|
|
output_frame = np.where(alpha == 255, frame, bg) |
|
|
elif settings["set_bg"] and bg_image is not None: |
|
|
if bg_image.shape[:2] != frame.shape[:2]: |
|
|
bg_image = cv2.resize(bg_image, (frame.shape[1], frame.shape[0])) |
|
|
output_frame = np.where(alpha == 255, frame, bg_image) |
|
|
elif settings["set_color"]: |
|
|
bg = np.full_like(frame, settings["bg_color"]) |
|
|
output_frame = np.where(alpha == 255, frame, bg) |
|
|
else: |
|
|
bg = np.zeros_like(frame) |
|
|
output_frame = np.where(alpha == 255, frame, bg) |
|
|
|
|
|
process_time = (time.time() - process_start) * 1000 |
|
|
return output_frame, f"{process_time:.2f} ms" |
|
|
|
|
|
|
|
|
with gr.Blocks() as demo: |
|
|
gr.Markdown("# AI Background Remover") |
|
|
with gr.Row(): |
|
|
with gr.Column(): |
|
|
webcam = gr.Image(sources=["webcam"], streaming=True, label="Live Video") |
|
|
output_image = gr.Image(label="Processed Output") |
|
|
seg_enabled = gr.Checkbox(label="Enable Background Removal", value=True) |
|
|
blur_bg = gr.Checkbox(label="Blur Background") |
|
|
set_bg = gr.Checkbox(label="Custom Image Background") |
|
|
custom_image = gr.Image(label="Upload Custom Background", type="numpy") |
|
|
set_color = gr.Checkbox(label="Solid Color Background") |
|
|
bg_color = gr.Textbox(label="Background Color (R,G,B)", value="0,0,0") |
|
|
blur_intensity = gr.Slider(label="Blur Intensity", minimum=5, maximum=25, value=15, step=2) |
|
|
processing_time = gr.Textbox(label="Processing Time", value="0 ms") |
|
|
|
|
|
webcam.stream( |
|
|
fn=process_frame, |
|
|
inputs=[webcam, seg_enabled, blur_bg, set_bg, set_color, bg_color, blur_intensity, custom_image], |
|
|
outputs=[output_image, processing_time], |
|
|
_js="""async () => { |
|
|
// Ensure webcam permissions are requested |
|
|
await navigator.mediaDevices.getUserMedia({ video: true }); |
|
|
return true; |
|
|
}""" |
|
|
) |
|
|
|
|
|
if __name__ == "__main__": |
|
|
demo.launch(server_name="0.0.0.0", server_port=7860) |
