import os
from glob import glob
import cv2
import numpy as np
from PIL import Image
import torch
from torchvision import transforms
from transformers import AutoModelForImageSegmentation
import gradio as gr
import spaces
from gradio_imageslider import ImageSlider

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

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


def array_to_pil_image(image, size=(1024, 1024)):
    image = cv2.resize(image, size, interpolation=cv2.INTER_LINEAR)
    image = Image.fromarray(image).convert('RGB')
    return image


class ImagePreprocessor():
    def __init__(self, resolution=(1024, 1024)) -> None:
        self.transform_image = transforms.Compose([
            # transforms.Resize(resolution),    # 1. keep consistent with the cv2.resize used in training 2. redundant with that in path_to_image()
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ])

    def proc(self, image):
        image = self.transform_image(image)
        return image


from transformers import AutoModelForImageSegmentation
weights_path = 'BiRefNet'
birefnet = AutoModelForImageSegmentation.from_pretrained('/'.join(('zhengpeng7', weights_path)), trust_remote_code=True)
birefnet.to(device)
birefnet.eval()

usage_to_weights_file = {
    'General': 'BiRefNet',
    'Portrait': 'BiRefNet-portrait',
    'DIS': 'BiRefNet-DIS5K',
    'HRSOD': 'BiRefNet-HRSOD',
    'COD': 'BiRefNet-COD',
    'DIS-TR_TEs': 'BiRefNet-DIS5K-TR_TEs'
}

# def predict(image_1, image_2):
#     images = [image_1, image_2]
@spaces.GPU
def predict(image, resolution, weights_file):
    global weights_path
    global birefnet
    if weights_file != weights_path:
        # Load BiRefNet with chosen weights
        birefnet = AutoModelForImageSegmentation.from_pretrained('/'.join(('zhengpeng7', usage_to_weights_file[weights_file] if weights_file is not None else 'BiRefNet')), trust_remote_code=True)
        birefnet.to(device)
        birefnet.eval()
        weights_path = weights_file

    resolution = f"{image.shape[1]}x{image.shape[0]}" if resolution == '' else resolution
    # Image is a RGB numpy array.
    resolution = [int(int(reso)//32*32) for reso in resolution.strip().split('x')]
    images = [image]
    image_shapes = [image.shape[:2] for image in images]
    images = [array_to_pil_image(image, resolution) for image in images]

    image_preprocessor = ImagePreprocessor(resolution=resolution)
    images_proc = []
    for image in images:
        images_proc.append(image_preprocessor.proc(image))
    images_proc = torch.cat([image_proc.unsqueeze(0) for image_proc in images_proc])

    with torch.no_grad():
        scaled_preds_tensor = birefnet(images_proc.to(device))[-1].sigmoid()   # BiRefNet needs an sigmoid activation outside the forward.
    preds = []
    for image_shape, pred_tensor in zip(image_shapes, scaled_preds_tensor):
        if device == 'cuda':
            pred_tensor = pred_tensor.cpu()
        preds.append(torch.nn.functional.interpolate(pred_tensor.unsqueeze(0), size=image_shape, mode='bilinear', align_corners=True).squeeze().numpy())
    image_preds = []
    for image, pred in zip(images, preds):
        image = image.resize(pred.shape[::-1])
        pred = np.repeat(np.expand_dims(pred, axis=-1), 3, axis=-1)
        image_preds.append((pred * image).astype(np.uint8))

    return image, image_preds[0]


examples = [[_] for _ in glob('materials/examples/*')][:]

# Add the option of resolution in a text box.
for idx_example, example in enumerate(examples):
    examples[idx_example].append('1024x1024')
examples.append(examples[-1].copy())
examples[-1][1] = '512x512'

demo = gr.Interface(
    fn=predict,
    inputs=[
        'image',
        gr.Textbox(lines=1, placeholder="Type the resolution (`WxH`) you want, e.g., `512x512`. Higher resolutions can be much slower for inference.", label="Resolution"),
        gr.Radio(list(usage_to_weights_file.keys()), label="Weights", info="Choose the weights you want.")
    ],
    outputs=ImageSlider(),
    examples=examples,
    title='Online demo for `Bilateral Reference for High-Resolution Dichotomous Image Segmentation`',
    description=('Upload a picture, our model will extract a highly accurate segmentation of the subject in it. :)'
                 '\nThe resolution used in our training was `1024x1024`, which is thus the suggested resolution to obtain good results!\n Ours codes can be found at https://github.com/ZhengPeng7/BiRefNet.\n We also maintain the HF model of BiRefNet at https://huggingface.co/ZhengPeng7/birefnet for easier access.')
)
demo.launch(debug=True)