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

torch.jit.script = lambda f: f

from models.baseline import BiRefNet
from config import Config


config = Config()
device = config.device



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


model = BiRefNet(bb_pretrained=False)
state_dict = './BiRefNet_ep580.pth'
if os.path.exists(state_dict):
    birefnet_dict = torch.load(state_dict, map_location="cpu")
    unwanted_prefix = '_orig_mod.'
    for k, v in list(birefnet_dict.items()):
        if k.startswith(unwanted_prefix):
            birefnet_dict[k[len(unwanted_prefix):]] = birefnet_dict.pop(k)
    model.load_state_dict(birefnet_dict)
model = model.to(device)
model.eval()


# def predict(image_1, image_2):
#     images = [image_1, image_2]
@spaces.GPU
def predict(image, resolution):
    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 = model(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")],
    outputs=ImageSlider(),
    examples=examples,
    title='Online demo for `Bilateral Reference for High-Resolution Dichotomous Image Segmentation`',
    description=('Upload a picture, our model will give you the binary maps of the highly accurate segmentation of the salient objects in it. :)'
                 '\nThe resolution used in our training was `1024x1024`, which is too much burden for the huggingface free spaces like this one (cost nearly 40s). Please set resolution as more than `768x768` for images with many texture details to obtain good results!\n Ours codes can be found at https://github.com/ZhengPeng7/BiRefNet.')
)
demo.launch(debug=True)