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
from rich.panel import Panel
from rich.columns import Columns
from rich.console import Console

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 = "image restore"
description = " The model was trained to remove the moire patterns from images! Max input image resolution is \
up to 4K resolution\
<br /> \
It takes time to perform inference \
<br />"

article = "Reference:\n Towards Efficient and Scale-Robust Ultra-High-Definition Image Demoiréing. CVMI Lab, Nov. 21, 2022. Accessed: Nov. 21, 2022. [Online]. Available: https://github.com/CVMI-Lab/UHDM"


files = [
    '0001_01.jpg',
    '0002_01.jpg',
    '0003_01.jpg',
    '0004_01.jpg',
    '0005_01.jpg',
    '0006_01.jpg',
    '0007_01.jpg',
    '0008_01.jpg',
    '0009_01.jpg',
    '0010_01.jpg',
    '0011_01.jpg',
    '0012_01.jpg',
    '0013_04.jpg',
    '0014_01.jpg',
    '0015_01.jpg',
    '0016_01.jpg',
    '0017_01.jpg',
    '0018_01.jpg',
    '0019_01.jpg',
    '0020_01.jpg',
    '0021_01.jpg',
    '0022_01.jpg',
    '0023_01.jpg',
    '0024_01.jpg',
    '0025_01.jpg',
    '0026_01.jpg',
    '0027_01.jpg',
    '0028_01.jpg',
    '0029_01.jpg',
    '0030_01.jpg',
    '0031_01.jpg',
    '0032_01.jpg',
    '0033_01.jpg',
    '0034_01.jpg',
    '0035_01.jpg',
    '0036_01.jpg',
    '0037_07.jpg',
    '0038_01.jpg',
    '0039_01.jpg',
    '0040_01.jpg',
    '0041_01.jpg',
    '0042_01.jpg',
    '0043_01.jpg',
    '0044_01.jpg',
    '0045_01.jpg',
    '0046_01.jpg',
    '0047_01.jpg',
    '0048_01.jpg',
    '0049_01.jpg',
    '0050_01.jpg',
    '0051_01.jpg',
    '0052_01.jpg',
    '0053_01.jpg',
    '0054_01.jpg',
    '0055_01.jpg',
    '0056_01.jpg',
    '0057_01.jpg',
    '0058_01.jpg',
    '0059_01.jpg',
    '0060_01.jpg',
    '0061_01.jpg',
    '0062_01.jpg',
    '0063_02.jpg',
    '0064_01.jpg',
    '0065_01.jpg',
    '0066_01.jpg',
    '0067_01.jpg',
    '0068_01.jpg',
    '0069_01.jpg',
    '0070_01.jpg',
    '0071_01.jpg',
    '0072_01.jpg',
    '0073_01.jpg',
    '0074_01.jpg',
    '0075_01.jpg',
    '0076_01.jpg',
    '0077_01.jpg',
    '0078_01.jpg',
    '0079_01.jpg',
    '0080_01.jpg',
    '0081_01.jpg',
    '0082_01.jpg',
    '0083_01.jpg',
    '0084_01.jpg',
    '0085_02.jpg',
    '0086_01.jpg',
    '0087_02.jpg',
    '0088_01.jpg',
    '0089_01.jpg',
    '0090_02.jpg',
    '0091_01.jpg',
    '0092_01.jpg',
    '0093_01.jpg',
    '0094_01.jpg',
    '0095_01.jpg',
    '0096_01.jpg',
    '0097_01.jpg',
    '0098_01.jpg',
    '0099_01.jpg',
    '0100_01.jpg',
    '0101_01.jpg',
    '0102_01.jpg',
    '0103_01.jpg',
    '0104_01.jpg',
    '0105_01.jpg',
    '0106_01.jpg',
    '0107_01.jpg',
    '0108_01.jpg',
    '0109_01.jpg',
    '0110_01.jpg',
    '0111_01.jpg',
    '0112_01.jpg',
    '0113_01.jpg',
    '0114_01.jpg',
    '0115_01.jpg',
    '0116_01.jpg',
    '0117_01.jpg',
    '0118_01.jpg',
    '0119_01.jpg',
    '0120_01.jpg',
    '0121_01.jpg',
    '0122_01.jpg',
    '0123_01.jpg',
    '0124_01.jpg',
    '0125_01.jpg',
    '0126_01.jpg',
    '0127_01.jpg',
    '0128_01.jpg',
    '0129_01.jpg',
    '0130_01.jpg',
    '0131_01.jpg',
    '0132_02.jpg',
    '0133_01.jpg',
    '0134_01.jpg',
    '0135_01.jpg',
    '0136_01.jpg',
    '0137_01.jpg',
    '0138_02.jpg',
    '0139_01.jpg',
    '0140_01.jpg',
    '0141_01.jpg',
    '0142_01.jpg',
    '0143_01.jpg',
    '0144_01.jpg',
    '0145_01.jpg',
    '0146_01.jpg',
    '0147_01.jpg',
    '0148_01.jpg',
    '0149_01.jpg',
    '0150_02.jpg',
    '0151_01.jpg',
    '0152_01.jpg',
    '0153_01.jpg',
    '0154_01.jpg',
    '0155_01.jpg',
    '0156_01.jpg',
    '0157_01.jpg',
    '0158_01.jpg',
    '0159_01.jpg',
    '0160_01.jpg',
    '0161_01.jpg',
    '0162_01.jpg',
    '0163_01.jpg',
    '0164_01.jpg',
    '0165_01.jpg',
    '0166_01.jpg',
    '0167_01.jpg',
    '0168_01.jpg',
    '0169_01.jpg',
    '0170_01.jpg',
    '0171_01.jpg',
    '0172_01.jpg',
    '0173_01.jpg',
    '0174_01.jpg',
    '0175_01.jpg',
    '0176_01.jpg',
    '0177_03.jpg',
    '0178_01.jpg',
    '0179_01.jpg',
    '0180_01.jpg',
    '0181_01.jpg',
    '0182_01.jpg',
    '0183_01.jpg',
    '0184_01.jpg',
    '0185_01.jpg',
    '0186_01.jpg',
    '0187_01.jpg',
    '0188_01.jpg',
    '0189_01.jpg',
    '0190_01.jpg',
    '0191_01.jpg',
    '0192_01.jpg',
    '0193_01.jpg',
    '0194_02.jpg',
    '0195_01.jpg',
    '0196_01.jpg',
    '0197_02.jpg',
    '0198_01.jpg',
    '0199_02.jpg',
    '0200_01.jpg',
    '0201_01.jpg',
    '0202_02.jpg',
    '0203_01.jpg',
    '0204_01.jpg',
    '0205_01.jpg',
    '0206_01.jpg',
    '0207_01.jpg',
    '0208_01.jpg',
    '0209_01.jpg',
    '0210_01.jpg',
    '0211_01.jpg',
    '0212_01.jpg',
    '0213_01.jpg',
    '0214_01.jpg',
    '0215_01.jpg',
    '0216_01.jpg',
    '0217_02.jpg',
    '0218_01.jpg',
    '0219_01.jpg',
    '0220_01.jpg',
    '0221_01.jpg',
    '0222_01.jpg',
    '0223_01.jpg',
    '0224_01.jpg',
    '0225_01.jpg',
    '0226_02.jpg',
    '0227_01.jpg',
    '0228_02.jpg',
    '0229_01.jpg',
    '0230_01.jpg',
    '0231_01.jpg',
    '0232_01.jpg',
    '0233_02.jpg',
    '0234_01.jpg',
    '0235_01.jpg',
    '0236_01.jpg',
    '0237_01.jpg',
    '0238_01.jpg',
    '0239_01.jpg',
    '0240_02.jpg',
    '0241_01.jpg',
    '0242_01.jpg',
    '0243_01.jpg',
    '0244_01.jpg',
    '0245_01.jpg',
    '0246_01.jpg',
    '0247_01.jpg',
    '0248_01.jpg',
    '0249_01.jpg',
    '0250_01.jpg',
    '0251_01.jpg',
    '0252_01.jpg',
    '0253_01.jpg',
    '0254_01.jpg',
    '0255_01.jpg',
    '0256_01.jpg',
    '0257_01.jpg',
    '0258_02.jpg',
    '0259_01.jpg',
    '0260_02.jpg',
    '0261_01.jpg',
    '0262_01.jpg',
    '0263_01.jpg',
    '0264_01.jpg',
    '0265_01.jpg',
    '0266_01.jpg',
    '0267_01.jpg',
    '0268_01.jpg',
    '0269_01.jpg',
    '0270_01.jpg',
    '0271_01.jpg',
    '0272_01.jpg',
    '0273_01.jpg',
    '0274_01.jpg',
    '0275_01.jpg',
    '0276_01.jpg',
    '0277_02.jpg',
    '0278_01.jpg',
    '0279_01.jpg',
    '0280_01.jpg',
    '0281_01.jpg',
    '0282_01.jpg',
    '0283_01.jpg',
    '0284_01.jpg',
    '0285_01.jpg',
    '0286_01.jpg',
    '0287_01.jpg',
    '0288_01.jpg',
    '0289_02.jpg',
    '0290_01.jpg',
    '0291_01.jpg',
    '0292_01.jpg',
    '0293_01.jpg',
    '0294_01.jpg',
    '0295_01.jpg',
    '0296_01.jpg',
    '0297_01.jpg',
    '0298_02.jpg',
    '0299_02.jpg',
    '0300_01.jpg',
    '0301_02.jpg',
    '0302_01.jpg',
    '0303_01.jpg',
    '0304_01.jpg',
    '0305_02.jpg',
    '0306_03.jpg',
    '0307_01.jpg',
    '0308_01.jpg',
    '0309_01.jpg',
    '0310_01.jpg',
    '0311_01.jpg',
    '0312_02.jpg',
    '0313_02.jpg',
    '0314_01.jpg',
    '0315_01.jpg',
    '0316_01.jpg',
    '0317_03.jpg',
    '0318_01.jpg',
    '0319_01.jpg',
    '0320_01.jpg',
    '0321_01.jpg',
    '0322_01.jpg',
    '0323_01.jpg',
    '0324_01.jpg',
    '0325_02.jpg',
    '0326_01.jpg',
    '0327_01.jpg',
    '0328_01.jpg',
    '0329_01.jpg',
    '0330_01.jpg',
    '0331_01.jpg',
    '0332_03.jpg',
    '0333_01.jpg',
    '0334_01.jpg',
    '0335_02.jpg',
    '0336_01.jpg',
    '0337_01.jpg',
    '0338_01.jpg',
    '0339_01.jpg',
    '0340_01.jpg',
    '0341_03.jpg',
    '0342_01.jpg',
    '0343_02.jpg',
    '0344_02.jpg',
    '0345_01.jpg',
    '0346_01.jpg',
    '0347_01.jpg',
    '0348_02.jpg',
    '0349_01.jpg',
    '0350_01.jpg',
    '0351_01.jpg',
    '0352_01.jpg',
    '0353_01.jpg',
    '0354_01.jpg',
    '0355_01.jpg',
    '0356_01.jpg',
    '0357_01.jpg',
    '0358_01.jpg',
    '0359_01.jpg',
    '0360_01.jpg',
    '0361_01.jpg',
    '0362_01.jpg',
    '0363_01.jpg',
    '0364_01.jpg',
    '0365_01.jpg',
    '0366_01.jpg',
    '0367_01.jpg',
    '0368_01.jpg',
    '0369_01.jpg',
    '0370_01.jpg',
    '0371_01.jpg',
    '0372_03.jpg',
    '0373_01.jpg',
    '0374_02.jpg',
    '0375_01.jpg',
    '0376_01.jpg',
    '0377_01.jpg',
    '0378_01.jpg',
    '0379_01.jpg',
    '0380_01.jpg',
    '0381_01.jpg',
    '0382_01.jpg',
    '0383_01.jpg',
    '0384_01.jpg',
    '0385_01.jpg',
    '0386_01.jpg',
    '0387_01.jpg',
    '0388_01.jpg',
    '0389_01.jpg',
    '0390_01.jpg',
    '0391_02.jpg',
    '0392_01.jpg',
    '0393_01.jpg',
    '0394_01.jpg',
    '0395_01.jpg',
    '0396_02.jpg',
    '0397_01.jpg',
    '0398_01.jpg',
    '0399_01.jpg',
    '0400_01.jpg',
    '0401_02.jpg',
    '0402_02.jpg',
    '0403_01.jpg',
    '0404_01.jpg',
    '0405_01.jpg',
    '0406_01.jpg',
    '0407_01.jpg',
    '0408_01.jpg',
    '0409_01.jpg',
    '0410_10.jpg',
    '0411_01.jpg',
    '0412_01.jpg',
    '0413_02.jpg',
    '0414_01.jpg',
    '0415_01.jpg',
    '0416_01.jpg',
    '0417_01.jpg',
    '0418_01.jpg',
    '0419_01.jpg',
    '0421_01.jpg',
    '0422_01.jpg',
    '0423_01.jpg',
    '0424_01.jpg',
    '0425_01.jpg',
    '0426_01.jpg',
    '0427_01.jpg',
    '0428_01.jpg',
    '0429_03.jpg',
    '0430_02.jpg',
    '0431_01.jpg',
    '0432_02.jpg',
    '0433_01.jpg',
    '0434_01.jpg',
    '0435_01.jpg',
    '0436_02.jpg',
    '0437_02.jpg',
    '0439_01.jpg',
    '0440_01.jpg',
    '0441_01.jpg',
    '0442_01.jpg',
    '0443_01.jpg',
    '0444_01.jpg',
    '0445_01.jpg',
    '0446_01.jpg',
    '0448_02.jpg',
    '0449_02.jpg',
    '0450_01.jpg',
    '0451_01.jpg',
    '0452_01.jpg',
    '0453_01.jpg',
    '0454_03.jpg',
    '0455_02.jpg',
    '0456_01.jpg',
    '0457_01.jpg',
    '0458_03.jpg',
    '0459_01.jpg',
    '0460_01.jpg',
    '0461_01.jpg',
    '0464_02.jpg',
    '0465_03.jpg',
    '0466_01.jpg',
    '0467_01.jpg',
    '0468_01.jpg',
    '0469_02.jpg',
    '0470_01.jpg',
    '0471_01.jpg',
    '0472_01.jpg',
    '0473_01.jpg',
    '0474_01.jpg',
    '0475_03.jpg',
    '0476_01.jpg',
    '0477_01.jpg',
    '0478_02.jpg',
    '0479_01.jpg',
    '0480_01.jpg',
    '0481_01.jpg',
    '0482_09.jpg',
    '0483_01.jpg',
    '0484_01.jpg',
    '0485_03.jpg',
    '0486_01.jpg',
    '0487_01.jpg',
    '0488_01.jpg',
    '0489_01.jpg',
    '0490_01.jpg',
    '0491_02.jpg',
    '0492_01.jpg',
    '0493_01.jpg',
    '0494_04.jpg',
    '0495_01.jpg',
    '0496_01.jpg',
    '0497_02.jpg',
    '0498_02.jpg',
    '0499_01.jpg',
    '0500_02.jpg',
    '0501_02.jpg',
    '0502_02.jpg',
    '0503_01.jpg',
    '0504_01.jpg',
    '0505_01.jpg',
    '0506_03.jpg',
    '0507_01.jpg',
    '0508_01.jpg',
    '0509_01.jpg',
    '0510_02.jpg',
    '0511_01.jpg',
    '0513_02.jpg',
    '0514_01.jpg',
    '0517_01.jpg',
    '0519_01.jpg',
    '0520_01.jpg',
    '0522_02.jpg',
    '0526_01.jpg',
    '0528_01.jpg',
    '0530_02.jpg',
    '0533_03.jpg',
    '0536_02.jpg',
    '0558_01.jpg',
    '0559_01.jpg',
    '0560_01.jpg',
    '0561_01.jpg',
    '0562_01.jpg',
    '0563_01.jpg',
    '0564_01.jpg',
    '0565_01.jpg',
    '0566_01.jpg',
    '0567_01.jpg',
    '0568_01.jpg',
    '0569_01.jpg',
    '0571_01.jpg',
    '0572_01.jpg',
    '0573_01.jpg',
    '0574_03.jpg',
    '0575_02.jpg',
    '0576_01.jpg',
    '0578_01.jpg',
    '0579_02.jpg',
    '0580_01.jpg',
    '0581_01.jpg',
    '0583_01.jpg',
    '0584_01.jpg',
    '0585_01.jpg',
    '0586_02.jpg',
    '0587_01.jpg',
    '0588_01.jpg',
    '0589_02.jpg',
    '0591_01.jpg',
    '0593_01.jpg',
    '0594_01.jpg',
    '0595_01.jpg',
    '0597_02.jpg',
    '0599_01.jpg',
    '0600_01.jpg',
    '0601_01.jpg',
    '0602_01.jpg',
    '0603_01.jpg',
    '0604_02.jpg',
    '0605_01.jpg',
    '0606_01.jpg',
    '0607_01.jpg',
    '0608_01.jpg',
    '0609_02.jpg',
    '0610_04.jpg',
    '0611_01.jpg',
    '0612_01.jpg',
    '0613_01.jpg',
    '0614_05.jpg',
    '0615_01.jpg',
    '0616_01.jpg',
    '0617_03.jpg',
    '0618_01.jpg',
    '0619_01.jpg',
    '0620_01.jpg',
    '0621_02.jpg',
    '0622_02.jpg',
    '0623_01.jpg',
    '0624_02.jpg',
    '0627_05.jpg',
    '0628_03.jpg',
    '0629_01.jpg',
    '0630_01.jpg',
    '0633_01.jpg',
    '0634_01.jpg',
    '0671_10.jpg'
]

files = [os.path.join("n000129", k) for k in files]

Console().print(Columns ([Panel.fit(f"{k}", style="red on black") for k in files]))

iface1 = gr.Interface(fn=predict1,
                     inputs=gr.Image(type="pil"),
                     outputs=gr.Image(type="pil"),
                     examples=files,
                    title = title,
                    # description = description,
                    article = article
                    )

iface1.queue(concurrency_count=3)
iface1.launch()