import os
import cv2
import numpy as np
import torch
import gradio as gr
import spaces  # Required for @spaces.GPU

from PIL import Image, ImageOps
from transformers import AutoModelForImageSegmentation
from torchvision import transforms

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')
    # Call the processing function
    foreground, background, pred_pil, reverse_mask = remove_background(image_ori)
    return foreground, background, pred_pil, reverse_mask

@spaces.GPU  # Decorate the processing function
def remove_background(image_ori):
    original_size = image_ori.size

    # Preprocess the image
    image_preprocessor = ImagePreprocessor(resolution=(1024, 1024))
    image_proc = image_preprocessor.proc(image_ori)
    image_proc = image_proc.unsqueeze(0)

    # Prediction
    with torch.no_grad():
        preds = birefnet(image_proc.to(device))[-1].sigmoid().cpu()
    pred = preds[0].squeeze()

    # Process Results
    pred_pil = transforms.ToPILImage()(pred)
    pred_pil = pred_pil.resize(original_size, Image.BICUBIC)  # Resize mask to original size

    # Create reverse mask (background mask)
    reverse_mask = ImageOps.invert(pred_pil)

    # Create foreground image (object with transparent background)
    foreground = image_ori.copy()
    foreground.putalpha(pred_pil)

    # Create background image
    background = image_ori.copy()
    background.putalpha(reverse_mask)

    torch.cuda.empty_cache()

    # Return images in the specified order
    return foreground, background, pred_pil, reverse_mask

iface = gr.Interface(
    fn=remove_background_wrapper,
    inputs=gr.Image(type="numpy"),
    outputs=[
        gr.Image(type="pil", label="Foreground"),
        gr.Image(type="pil", label="Background"),
        gr.Image(type="pil", label="Foreground Mask"),
        gr.Image(type="pil", label="Background Mask")
    ],
    allow_flagging="never"
)

if __name__ == "__main__":
    iface.launch(debug=True)