app.py

"""
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())