import gradio as gr
from models.demoire.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
from pathlib import Path
import shutil
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
# from models.gfpgan import gfpgan_predict
from models.llflow.inference import main
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 predict(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
def predict_gfpgan(img):
with Console().status("[red] using [green] GFP-GAN v1.4", spinner="aesthetic"):
# if image already exists with this name then delete it
return img
title = "image enhancement"
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"
# ==========================================================================
# example images
# ==========================================================================
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]))
# # --------------------------------------------------------------------------
# # making interfaces for models
# # --------------------------------------------------------------------------
iface1 = gr.Interface(fn=predict,
inputs=gr.Image(type="pil"),
outputs=gr.Image(type="pil"),
examples=files,
title = title,
# description = description,
article = article,
allow_flagging="auto"
)
gfpgan = gr.Interface(fn=predict_gfpgan,
inputs=gr.Image(type="pil"),
outputs=gr.Image(type="pil"),
examples=files,
title = "GFP-GAN v 1.4",
# description = description,
article = "Practical face restoration algorithm for old photos",
allow_flagging="auto"
)
# setting queues for all models
iface1.queue(concurrency_count=3)
gfpgan.queue(concurrency_count=3)
# launching the interface in parallel model
gr.Parallel(iface1, gfpgan).launch()