""" Background Removal - Image and Video - BRIA RMBG 1.4: Fast general-purpose background removal - BRIA RMBG 2.0: High-quality BiRefNet-based (slower, better quality) - IS-Net/DIS: Ultra-fine edge detection (hair, fur, transparency) - Robust Video Matting: Video background removal with ONNX Runtime (2x faster!) CPU-only HuggingFace Space """ import os import sys import tempfile import argparse from typing import Tuple, Optional # Set OMP threads to match cpu-basic (2 vCPU) os.environ["OMP_NUM_THREADS"] = "2" os.environ["OMP_WAIT_POLICY"] = "ACTIVE" import gradio as gr import numpy as np from PIL import Image import torch import torch.nn.functional as F from torchvision.transforms.functional import normalize from huggingface_hub import hf_hub_download import cv2 import onnxruntime as ort # Force CPU for PyTorch (image model) device = torch.device("cpu") # ============================================================================= # Image Background Removal (BRIA RMBG 1.4) # ============================================================================= class BriaRMBGModel: """Wrapper for BRIA-RMBG 1.4 model""" def __init__(self): self.model = None self.input_size = [1024, 1024] def load_model(self): """Load the BRIA-RMBG 1.4 model from HuggingFace""" if self.model is not None: return print("Loading BRIA-RMBG 1.4 model...") # Download all required files from BRIA repo repo_id = "briaai/RMBG-1.4" model_path = hf_hub_download(repo_id=repo_id, filename="model.pth") briarmbg_path = hf_hub_download(repo_id=repo_id, filename="briarmbg.py") myconfig_path = hf_hub_download(repo_id=repo_id, filename="MyConfig.py") # Setup package directory briarmbg_dir = os.path.dirname(briarmbg_path) # Create __init__.py to make it a proper package init_path = os.path.join(briarmbg_dir, "__init__.py") if not os.path.exists(init_path): with open(init_path, "w") as f: f.write("# Auto-generated for package import\n") # Add parent directory to path so we can import as package parent_dir = os.path.dirname(briarmbg_dir) package_name = os.path.basename(briarmbg_dir) if parent_dir not in sys.path: sys.path.insert(0, parent_dir) # Import using the package structure import importlib briarmbg_module = importlib.import_module(f"{package_name}.briarmbg") BriaRMBG = briarmbg_module.BriaRMBG self.model = BriaRMBG() self.model.load_state_dict(torch.load(model_path, map_location=device)) self.model.to(device) self.model.eval() print("BRIA-RMBG model loaded!") def preprocess(self, image: Image.Image): """Preprocess image for model input""" if image.mode != "RGB": image = image.convert("RGB") orig_size = image.size im = np.array(image) if len(im.shape) < 3: im = im[:, :, np.newaxis] if im.shape[2] == 1: im = np.repeat(im, 3, axis=2) im_tensor = torch.tensor(im, dtype=torch.float32).permute(2, 0, 1) im_tensor = F.interpolate( im_tensor.unsqueeze(0), self.input_size, mode="bilinear", align_corners=False ) image_tensor = im_tensor / 255.0 image_tensor = normalize(image_tensor, [0.5, 0.5, 0.5], [1.0, 1.0, 1.0]) return image_tensor, orig_size def postprocess(self, result, orig_size): """Postprocess model output to mask""" result = F.interpolate( result, size=(orig_size[1], orig_size[0]), mode="bilinear", align_corners=False ) ma = torch.max(result) mi = torch.min(result) result = (result - mi) / (ma - mi + 1e-8) mask = (result.squeeze().cpu().numpy() * 255).astype(np.uint8) return mask @torch.no_grad() def remove_background(self, image: Image.Image): """Remove background from image""" self.load_model() input_tensor, orig_size = self.preprocess(image) input_tensor = input_tensor.to(device) result = self.model(input_tensor) if isinstance(result, (list, tuple)): result = result[0] if isinstance(result, (list, tuple)): result = result[0] mask = self.postprocess(result, orig_size) pil_mask = Image.fromarray(mask).convert("L") if image.mode != "RGB": image = image.convert("RGB") output = image.copy() output.putalpha(pil_mask) return output, pil_mask # ============================================================================= # Image Background Removal (IS-Net/DIS - Ultra-fine edges) - ONNX Runtime # ============================================================================= class ISNetModel: """Wrapper for IS-Net/DIS ONNX model - better for hair, fur, transparency""" def __init__(self): self.session = None self.input_size = [1024, 1024] def load_model(self): """Load IS-Net ONNX model""" if self.session is not None: return print("Loading IS-Net/DIS ONNX model...") model_path = os.path.join(os.path.dirname(__file__), "isnet_dis_fp32.onnx") # Setup ONNX Runtime with optimizations sess_options = ort.SessionOptions() sess_options.intra_op_num_threads = 2 sess_options.inter_op_num_threads = 1 sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL self.session = ort.InferenceSession( model_path, sess_options, providers=['CPUExecutionProvider'] ) print("IS-Net/DIS ONNX model loaded!") def preprocess(self, image: Image.Image): """Preprocess image for model input""" if image.mode != "RGB": image = image.convert("RGB") orig_size = image.size # Resize to 1024x1024 im_resized = image.resize((1024, 1024), Image.BILINEAR) im = np.array(im_resized).astype(np.float32) # Normalize: (x/255 - 0.5) / 1.0 = x/255 - 0.5 im = im / 255.0 im = (im - 0.5) / 1.0 # HWC -> NCHW im = np.transpose(im, (2, 0, 1))[np.newaxis, ...] return im.astype(np.float32), orig_size def postprocess(self, result, orig_size): """Postprocess model output to mask""" # Result is already sigmoid (0-1), resize to original mask_1024 = result[0, 0] # 1024x1024 # Resize to original size mask_pil = Image.fromarray((mask_1024 * 255).astype(np.uint8)) mask_pil = mask_pil.resize(orig_size, Image.BILINEAR) return np.array(mask_pil) def remove_background(self, image: Image.Image): """Remove background from image using IS-Net ONNX""" self.load_model() input_tensor, orig_size = self.preprocess(image) # Run ONNX inference - d0 is the main output outputs = self.session.run(None, {'input': input_tensor}) result = outputs[0] # d0: main segmentation mask mask = self.postprocess(result, orig_size) pil_mask = Image.fromarray(mask).convert("L") if image.mode != "RGB": image = image.convert("RGB") output = image.copy() output.putalpha(pil_mask) return output, pil_mask # ============================================================================= # Image Background Removal (BRIA RMBG 2.0 - BiRefNet based, high quality) # ============================================================================= class RMBG2Model: """Wrapper for BRIA-RMBG 2.0 model (BiRefNet architecture) - higher quality, slower""" def __init__(self): self.model = None self.input_size = 1024 # Native resolution for best quality def load_model(self): """Load RMBG-2.0 model from HuggingFace""" if self.model is not None: return print("Loading BRIA-RMBG 2.0 model (BiRefNet)...") # Download model files from ungated repo repo_id = "1038lab/RMBG-2.0" model_path = hf_hub_download(repo_id=repo_id, filename="model.safetensors") birefnet_path = hf_hub_download(repo_id=repo_id, filename="birefnet.py") config_path = hf_hub_download(repo_id=repo_id, filename="BiRefNet_config.py") # Setup import path model_dir = os.path.dirname(birefnet_path) if model_dir not in sys.path: sys.path.insert(0, model_dir) # Import model classes from safetensors.torch import load_file import importlib.util # Load BiRefNet_config spec = importlib.util.spec_from_file_location("BiRefNet_config", config_path) config_module = importlib.util.module_from_spec(spec) sys.modules["BiRefNet_config"] = config_module spec.loader.exec_module(config_module) # Fix relative import in birefnet.py - read, fix, and exec with open(birefnet_path, 'r', encoding='utf-8') as f: birefnet_code = f.read() # Replace relative import with absolute birefnet_code = birefnet_code.replace( 'from .BiRefNet_config import BiRefNetConfig', 'from BiRefNet_config import BiRefNetConfig' ) # Execute the fixed code birefnet_module = type(sys)("birefnet") exec(compile(birefnet_code, birefnet_path, 'exec'), birefnet_module.__dict__) BiRefNet = birefnet_module.BiRefNet BiRefNetConfig = config_module.BiRefNetConfig # Create and load model config = BiRefNetConfig(bb_pretrained=False) self.model = BiRefNet(config) state_dict = load_file(model_path) self.model.load_state_dict(state_dict) self.model.to(device) self.model.eval() print(f"BRIA-RMBG 2.0 loaded! ({sum(p.numel() for p in self.model.parameters())/1e6:.0f}M params)") def preprocess(self, image: Image.Image): """Preprocess image for model input""" if image.mode != "RGB": image = image.convert("RGB") orig_size = image.size # Resize to model input size im_resized = image.resize((self.input_size, self.input_size), Image.BILINEAR) im = np.array(im_resized).astype(np.float32) # Normalize im = im / 255.0 im = (im - 0.5) / 1.0 # HWC -> NCHW im_tensor = torch.from_numpy(im).permute(2, 0, 1).unsqueeze(0) return im_tensor.float(), orig_size def postprocess(self, result, orig_size): """Postprocess model output to mask""" # Get final output (BiRefNet returns list for deep supervision) if isinstance(result, (list, tuple)): result = result[-1] # Last output is the refined mask # Sigmoid and resize mask = torch.sigmoid(result).squeeze().cpu().numpy() # Resize to original size mask_pil = Image.fromarray((mask * 255).astype(np.uint8)) mask_pil = mask_pil.resize(orig_size, Image.BILINEAR) return np.array(mask_pil) @torch.no_grad() def remove_background(self, image: Image.Image): """Remove background from image using RMBG-2.0""" self.load_model() input_tensor, orig_size = self.preprocess(image) input_tensor = input_tensor.to(device) result = self.model(input_tensor) mask = self.postprocess(result, orig_size) pil_mask = Image.fromarray(mask).convert("L") if image.mode != "RGB": image = image.convert("RGB") output = image.copy() output.putalpha(pil_mask) return output, pil_mask # ============================================================================= # Video Background Removal (Robust Video Matting - ONNX Runtime) # ============================================================================= rvm_session = None def load_rvm_onnx(): """Load RobustVideoMatting ONNX model for fast CPU inference.""" global rvm_session if rvm_session is not None: return rvm_session print("Loading Robust Video Matting ONNX model...") # Load ONNX model from local Space directory model_path = os.path.join(os.path.dirname(__file__), "rvm_mobilenetv3_fp32.onnx") # Setup ONNX Runtime with optimizations (2 vCPU on cpu-basic) sess_options = ort.SessionOptions() sess_options.intra_op_num_threads = 2 sess_options.inter_op_num_threads = 1 sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL rvm_session = ort.InferenceSession( model_path, sess_options, providers=['CPUExecutionProvider'] ) print("RVM ONNX model loaded! (2x faster than PyTorch)") return rvm_session # Output resolution options RESOLUTION_OPTIONS = { "Original": None, "1080p": 1080, "720p": 720, "480p": 480, } def process_video(input_video, output_type="green_screen", max_resolution="720p", downsample_ratio=0.25, progress=None): """ Process video to remove background using ONNX Runtime. Args: input_video: Path to input video output_type: "green_screen" or "foreground_only" max_resolution: Output video resolution (affects speed) downsample_ratio: Internal model quality (affects edge accuracy) progress: Gradio Progress tracker (optional, for web UI) Returns: Path to processed video """ if input_video is None: return None # Get max height from resolution option max_h = RESOLUTION_OPTIONS.get(max_resolution) session = load_rvm_onnx() output_names = [o.name for o in session.get_outputs()] cap = cv2.VideoCapture(input_video) fps = cap.get(cv2.CAP_PROP_FPS) orig_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) orig_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) # Resize based on quality preset if max_h and orig_height > max_h: scale = max_h / orig_height width = int(orig_width * scale) height = max_h print(f"Resizing: {orig_width}x{orig_height} -> {width}x{height}") else: width, height = orig_width, orig_height with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp: output_path = tmp.name fourcc = cv2.VideoWriter_fourcc(*"mp4v") out = cv2.VideoWriter(output_path, fourcc, fps, (width, height)) green_bg = np.zeros((height, width, 3), dtype=np.uint8) green_bg[:] = (0, 255, 0) # Initialize recurrent states (zeros for first frame) r1 = np.zeros((1, 1, 1, 1), dtype=np.float32) r2 = np.zeros((1, 1, 1, 1), dtype=np.float32) r3 = np.zeros((1, 1, 1, 1), dtype=np.float32) r4 = np.zeros((1, 1, 1, 1), dtype=np.float32) frame_idx = 0 while True: ret, frame = cap.read() if not ret: break # Resize if needed if frame.shape[0] != height or frame.shape[1] != width: frame = cv2.resize(frame, (width, height)) # Preprocess frame frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) src = frame_rgb.astype(np.float32) / 255.0 src = np.transpose(src, (2, 0, 1))[np.newaxis, ...] # NCHW # ONNX inference outputs = session.run(output_names, { 'src': src, 'r1i': r1, 'r2i': r2, 'r3i': r3, 'r4i': r4, 'downsample_ratio': np.array([downsample_ratio], dtype=np.float32) }) fgr, pha, r1, r2, r3, r4 = outputs # Postprocess fgr_np = np.transpose(fgr[0], (1, 2, 0)) # CHW -> HWC pha_np = pha[0, 0] if fgr_np.shape[:2] != (height, width): fgr_np = cv2.resize(fgr_np, (width, height)) pha_np = cv2.resize(pha_np, (width, height)) pha_3ch = np.stack([pha_np] * 3, axis=-1) if output_type == "green_screen": result = (fgr_np * pha_3ch + (green_bg / 255.0) * (1 - pha_3ch)) else: result = fgr_np * pha_3ch result_bgr = cv2.cvtColor((result * 255).astype(np.uint8), cv2.COLOR_RGB2BGR) out.write(result_bgr) frame_idx += 1 if frame_count > 0: if progress is not None: progress(frame_idx / frame_count, desc=f"Processing frame {frame_idx}/{frame_count}") elif frame_idx % 30 == 0: print(f"Processing frame {frame_idx}/{frame_count}") cap.release() out.release() return output_path # ============================================================================= # Gradio Interface # ============================================================================= # Initialize image models (lazy loading) bria_model = BriaRMBGModel() rmbg2_model = RMBG2Model() isnet_model = ISNetModel() IMAGE_MODELS = { "BRIA-RMBG 1.4 (fast, general)": bria_model, "BRIA-RMBG 2.0 (high quality, slow)": rmbg2_model, "IS-Net/DIS (fine edges, hair)": isnet_model, } IMAGE_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.webp', '.bmp', '.gif', '.tiff'} VIDEO_EXTENSIONS = {'.mp4', '.avi', '.mov', '.mkv', '.webm', '.flv', '.wmv'} def detect_file_type(file_path): """Detect if file is image or video based on extension.""" if file_path is None: return None ext = os.path.splitext(file_path)[1].lower() if ext in IMAGE_EXTENSIONS: return "image" elif ext in VIDEO_EXTENSIONS: return "video" return None def on_file_upload(file_path): """Handle file upload - detect type and update preview.""" if file_path is None: return gr.Image(visible=False), gr.Video(visible=False), None file_type = detect_file_type(file_path) if file_type == "image": return gr.Image(value=file_path, visible=True), gr.Video(visible=False), "image" elif file_type == "video": return gr.Image(visible=False), gr.Video(value=file_path, visible=True), "video" return gr.Image(visible=False), gr.Video(visible=False), None def process_media(file_path, file_type, image_model_choice, output_type, max_resolution, edge_quality, progress=gr.Progress()): """Process based on stored file type.""" if file_path is None: return None, None # Auto-detect file type if not provided (for API calls) if file_type is None: file_type = detect_file_type(file_path) if file_type is None: return None, None if file_type == "image": pil_image = Image.open(file_path) model = IMAGE_MODELS.get(image_model_choice, bria_model) output_img, mask = model.remove_background(pil_image) return output_img, None elif file_type == "video": output_path = process_video(file_path, output_type, max_resolution, edge_quality, progress) return None, output_path return None, None def build_ui(): """Build compact unified Gradio interface.""" with gr.Blocks(title="Background Removal") as demo: gr.Markdown("## Background Removal\nUpload image or video. Image models: RMBG 1.4 (fast) | RMBG 2.0 (quality) | IS-Net (fine edges) | Video: RVM") # State to track file type file_type_state = gr.State(value=None) with gr.Row(): with gr.Column(scale=1): input_file = gr.File( label="Upload Image or Video", file_types=["image", "video"], type="filepath" ) # Preview - only one visible at a time preview_image = gr.Image(label="Preview", interactive=False, visible=False) preview_video = gr.Video(label="Preview", interactive=False, visible=False) # Image model selector image_model_dropdown = gr.Dropdown( choices=list(IMAGE_MODELS.keys()), value="BRIA-RMBG 1.4 (fast, general)", label="Image Model", info="1.4: fast | 2.0: quality (slow) | IS-Net: fine edges" ) with gr.Accordion("Video Settings", open=False): output_type = gr.Radio( choices=["green_screen", "foreground_only"], value="green_screen", label="Output Type" ) max_res = gr.Dropdown( choices=list(RESOLUTION_OPTIONS.keys()), value="720p", label="Output Resolution" ) edge_quality = gr.Slider( minimum=0.1, maximum=1.0, value=0.25, step=0.05, label="Edge Quality" ) process_btn = gr.Button("Remove Background", variant="primary") with gr.Column(scale=1): # Output components output_image = gr.Image(label="Result (Image)", type="pil") output_video = gr.Video(label="Result (Video)") # Wrapper for examples (auto-detects file type) def process_example(file_path, image_model_choice, output_type, max_resolution, edge_quality): file_type = detect_file_type(file_path) return process_media(file_path, file_type, image_model_choice, output_type, max_resolution, edge_quality) # Example files with cached results gr.Examples( examples=[ ["examples/cat.jpg", "BRIA-RMBG 1.4 (fast, general)", "green_screen", "720p", 0.25], ], inputs=[input_file, image_model_dropdown, output_type, max_res, edge_quality], outputs=[output_image, output_video], fn=process_example, cache_examples=True, cache_mode="lazy", label="Examples" ) # On upload: detect type and show preview input_file.change( fn=on_file_upload, inputs=[input_file], outputs=[preview_image, preview_video, file_type_state] ) # On process: fill the already-visible output process_btn.click( fn=process_media, inputs=[input_file, file_type_state, image_model_dropdown, output_type, max_res, edge_quality], outputs=[output_image, output_video] ) return demo # ============================================================================= # CLI Interface # ============================================================================= def cli_main(): """CLI entry point.""" parser = argparse.ArgumentParser(description="Background Removal - Image and Video") subparsers = parser.add_subparsers(dest="command", help="Commands") # Image command img_parser = subparsers.add_parser("image", help="Remove background from image") img_parser.add_argument("input", type=str, help="Input image path") img_parser.add_argument("--output", "-o", type=str, default="output.png", help="Output path") img_parser.add_argument("--mask", "-m", type=str, help="Save mask to path") img_parser.add_argument("--model", type=str, default="BRIA-RMBG 1.4 (fast, general)", choices=list(IMAGE_MODELS.keys()), help="Image model: BRIA (fast) or IS-Net (fine edges)") # Video command vid_parser = subparsers.add_parser("video", help="Remove background from video") vid_parser.add_argument("input", type=str, help="Input video path") vid_parser.add_argument("--output", "-o", type=str, default="output.mp4", help="Output path") vid_parser.add_argument("--type", type=str, default="green_screen", choices=["green_screen", "foreground_only"]) vid_parser.add_argument("--downsample", type=float, default=0.25, help="Downsample ratio") args = parser.parse_args() if args.command == "image": print(f"Processing image: {args.input}") input_img = Image.open(args.input) # Use selected model (default: BRIA) model = IMAGE_MODELS.get(args.model, bria_model) output_img, mask = model.remove_background(input_img) output_img.save(args.output) print(f"Saved: {args.output}") if args.mask: mask.save(args.mask) print(f"Saved mask: {args.mask}") elif args.command == "video": print(f"Processing video: {args.input}") output_path = process_video(args.input, args.type, args.downsample) if output_path: import shutil shutil.move(output_path, args.output) print(f"Saved: {args.output}") else: parser.print_help() if __name__ == "__main__": if len(sys.argv) > 1 and sys.argv[1] in ["image", "video", "-h", "--help"]: cli_main() else: demo = build_ui() demo.launch(theme=gr.themes.Soft())