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 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. \
\
It takes time to perform inference \
\
",
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.inputs.Image(type="pil"),
outputs=gr.inputs.Image(type="pil"),
examples=files,
title = title,
description = description,
article = article
)
iface1.queue(concurrency_count=3)
iface1.launch()