import gradio as gr
from model.nets import my_model
import torch
import cv2
import torch.utils.data as data
import torchvision.transforms as transforms
import PIL
from PIL import Image
from PIL import ImageFile
import math
import os
import torch.nn.functional as F

os.environ["CUDA_VISIBLE_DEVICES"] = "1"
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model1 = my_model(en_feature_num=48,
                     en_inter_num=32,
                     de_feature_num=64,
                     de_inter_num=32,
                     sam_number=1,
                     ).to(device)

load_path1 = "./mix.pth"
model_state_dict1 = torch.load(load_path1, map_location=device)
model1.load_state_dict(model_state_dict1)


def default_toTensor(img):
    t_list = [transforms.ToTensor()]
    composed_transform = transforms.Compose(t_list)
    return composed_transform(img)

def predict1(img):
    in_img = transforms.ToTensor()(img).to(device).unsqueeze(0)
    b, c, h, w = in_img.size()
    # pad image such that the resolution is a multiple of 32
    w_pad = (math.ceil(w / 32) * 32 - w) // 2
    w_odd_pad = w_pad
    h_pad = (math.ceil(h / 32) * 32 - h) // 2
    h_odd_pad = h_pad

    if w % 2 == 1:
        w_odd_pad += 1
    if h % 2 == 1:
        h_odd_pad += 1

    in_img = img_pad(in_img, w_pad=w_pad, h_pad=h_pad, w_odd_pad=w_odd_pad, h_odd_pad=h_odd_pad)
    with torch.no_grad():
        out_1, out_2, out_3 = model1(in_img)
        if h_pad != 0:
            out_1 = out_1[:, :, h_pad:-h_odd_pad, :]
        if w_pad != 0:
            out_1 = out_1[:, :, :, w_pad:-w_odd_pad]
    out_1 = out_1.squeeze(0)
    out_1 = PIL.Image.fromarray(torch.clamp(out_1 * 255, min=0, max=255
    ).byte().permute(1, 2, 0).cpu().numpy())

    return out_1

def img_pad(x,  w_pad, h_pad, w_odd_pad, h_odd_pad):
    '''
    Here the padding values are determined by the average r,g,b values across the training set
    in FHDMi dataset. For the evaluation on the UHDM, you can also try the commented lines where
    the mean values are calculated from UHDM training set, yielding similar performance.
    '''
    x1 = F.pad(x[:, 0:1, ...], (w_pad, w_odd_pad, h_pad, h_odd_pad), value=0.3827)
    x2 = F.pad(x[:, 1:2, ...], (w_pad, w_odd_pad, h_pad, h_odd_pad), value=0.4141)
    x3 = F.pad(x[:, 2:3, ...], (w_pad, w_odd_pad, h_pad, h_odd_pad), value=0.3912)

    y = torch.cat([x1, x2, x3], dim=1)

    return y


title = "Clean Your Moire Images!"
description = " The model was trained to remove the moire patterns from your captured screen images! Specially, this model is capable of tackling \
images up to 4K resolution, which adapts to most of the modern mobile phones. \
<br /> \
(Note: It may cost 80s per 4K image (e.g., iPhone's resolution: 4032x3024) since this demo runs on the CPU. The model can run \
on a NVIDIA 3090 GPU 17ms per standard 4K image). \
<br /> \
The best way for a demo testing is using your mobile phone to capture a screen image, which may cause moire patterns. \
You can scan the [QR code](https://github.com/CVMI-Lab/UHDM/blob/main/figures/QR.jpg) to play on your mobile phone. "

article = "Check out the [ECCV 2022 paper](https://arxiv.org/abs/2207.09935) and the \
            [official training code](https://github.com/CVMI-Lab/UHDM) which the demo is based on.\
            <center><img src='https://visitor-badge.glitch.me/badge?page_id=Andyx_screen_image_demoire' alt='visitor badge'></center>"


iface1 = gr.Interface(fn=predict1,
                     inputs=gr.inputs.Image(type="pil"),
                     outputs=gr.inputs.Image(type="pil"),
                     examples=['001.jpg',
                            '002.jpg',
                            '005.jpg'],
                    title = title,
                    description = description,
                    article = article
                    )


iface1.launch()