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\
\ 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.Image(type="pil"), outputs=gr.Image(type="pil"), examples=files, title = title, # description = description, article = article ) iface1.queue(concurrency_count=3) iface1.launch()