app.py

import os
import cv2
import numpy as np
import torch
import gradio as gr
import spaces
from gradio.themes.base import Base
from gradio.themes.utils import colors, fonts, sizes
from PIL import Image, ImageOps
from transformers import AutoModelForImageSegmentation
from torchvision import transforms

# Custom White Theme with Inter font
class WhiteTheme(Base):
    def __init__(
        self,
        *,
        primary_hue: colors.Color | str = colors.orange,
        font: fonts.Font | str | tuple[fonts.Font | str, ...] = (
            fonts.GoogleFont("Inter"),
            "ui-sans-serif",
            "system-ui",
            "sans-serif",
        ),
        font_mono: fonts.Font | str | tuple[fonts.Font | str, ...] = (
            fonts.GoogleFont("Inter"),
            "ui-monospace",
            "system-ui",
            "monospace",
        )
    ):
        super().__init__(
            primary_hue=primary_hue,
            font=font,
            font_mono=font_mono,
        )
        
        self.set(
            body_background_fill="white",
            block_background_fill="white",
            panel_background_fill="white",
            body_text_color="black",
            block_label_text_color="black",
            block_border_color="white",
            panel_border_color="white",
            input_border_color="lightgray",
            shadow_drop="none",
            button_primary_background_fill="*primary_500",
            button_primary_background_fill_hover="*primary_600",
            button_primary_text_color="white",
            button_secondary_background_fill="white",
            button_secondary_border_color="lightgray",
            slider_color="*primary_500",
            slider_track_color="lightgray"
        )

torch.set_float32_matmul_precision('high')
torch.jit.script = lambda f: f

device = "cuda" if torch.cuda.is_available() else "cpu"

def refine_foreground(image, mask, r=90):
    if mask.size != image.size:
        mask = mask.resize(image.size)
    image = np.array(image) / 255.0
    mask = np.array(mask) / 255.0
    estimated_foreground = FB_blur_fusion_foreground_estimator_2(image, mask, r=r)
    image_masked = Image.fromarray((estimated_foreground * 255.0).astype(np.uint8))
    return image_masked

def FB_blur_fusion_foreground_estimator_2(image, alpha, r=90):
    alpha = alpha[:, :, None]
    F, blur_B = FB_blur_fusion_foreground_estimator(
        image, image, image, alpha, r)
    return FB_blur_fusion_foreground_estimator(image, F, blur_B, alpha, r=6)[0]

def FB_blur_fusion_foreground_estimator(image, F, B, alpha, r=90):
    if isinstance(image, Image.Image):
        image = np.array(image) / 255.0
    blurred_alpha = cv2.blur(alpha, (r, r))[:, :, None]
    blurred_FA = cv2.blur(F * alpha, (r, r))
    blurred_F = blurred_FA / (blurred_alpha + 1e-5)
    blurred_B1A = cv2.blur(B * (1 - alpha), (r, r))
    blurred_B = blurred_B1A / ((1 - blurred_alpha) + 1e-5)
    F = blurred_F + alpha * (image - alpha * blurred_F - (1 - alpha) * blurred_B)
    F = np.clip(F, 0, 1)
    return F, blurred_B

class ImagePreprocessor():
    def __init__(self, resolution=(1024, 1024)) -> None:
        self.transform_image = transforms.Compose([
            transforms.Resize(resolution),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406],
                                 [0.229, 0.224, 0.225]),
        ])

    def proc(self, image: Image.Image) -> torch.Tensor:
        image = self.transform_image(image)
        return image

# Load the model
birefnet = AutoModelForImageSegmentation.from_pretrained(
    'zhengpeng7/BiRefNet-matting', trust_remote_code=True)
birefnet.to(device)
birefnet.eval()

def remove_background_wrapper(image):
    if image is None:
        raise gr.Error("Please upload an image.")
    image_ori = Image.fromarray(image).convert('RGB')
    foreground, background, pred_pil, reverse_mask = remove_background(image_ori)
    return foreground, background, pred_pil, reverse_mask

@spaces.GPU
def remove_background(image_ori):
    original_size = image_ori.size
    image_preprocessor = ImagePreprocessor(resolution=(1024, 1024))
    image_proc = image_preprocessor.proc(image_ori)
    image_proc = image_proc.unsqueeze(0)
    
    with torch.no_grad():
        preds = birefnet(image_proc.to(device))[-1].sigmoid().cpu()
    pred = preds[0].squeeze()
    
    pred_pil = transforms.ToPILImage()(pred)
    pred_pil = pred_pil.resize(original_size, Image.BICUBIC)
    
    reverse_mask = ImageOps.invert(pred_pil)
    
    foreground = image_ori.copy()
    foreground.putalpha(pred_pil)
    
    background = image_ori.copy()
    background.putalpha(reverse_mask)
    
    torch.cuda.empty_cache()
    
    return foreground, background, pred_pil, reverse_mask

# Custom CSS for button styling
custom_css = """
@keyframes gradient-animation {
    0% { background-position: 0% 50%; }
    50% { background-position: 100% 50%; }
    100% { background-position: 0% 50%; }
}

#submit-button {
    background: linear-gradient(
        135deg,
        #e0f7fa, #e8f5e9, #fff9c4, #ffebee,
        #f3e5f5, #e1f5fe, #fff3e0, #e8eaf6
    );
    background-size: 400% 400%;
    animation: gradient-animation 15s ease infinite;
    border-radius: 12px;
    color: black;
}
"""

with gr.Blocks(css=custom_css, theme=WhiteTheme()) as demo:
    # Interface setup with input and output
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="numpy", sources=['upload'], label="Upload Image")
            btn = gr.Button("Process Image", elem_id="submit-button")
        with gr.Column():
            output_foreground = gr.Image(type="pil", label="Foreground")
            output_background = gr.Image(type="pil", label="Background")
            output_foreground_mask = gr.Image(type="pil", label="Foreground Mask")
            output_background_mask = gr.Image(type="pil", label="Background Mask")

    # Link the button to the processing function
    btn.click(fn=remove_background_wrapper, inputs=image_input, outputs=[
        output_foreground, output_background, output_foreground_mask, output_background_mask])

    demo.launch(debug=True)