Spaces:
Build error
Build error
| 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() |