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 | |
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 | |
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 | |
if Path("input_image_gfpgan.jpg").exists(): | |
os.remove("input_image_gfpgan.jpg") | |
# save incoming PIL image to disk | |
img.save("input_image_gfpgan.jpg") | |
out = gfpgan_predict(img) | |
Console().print(out) | |
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\ | |
<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" | |
# ========================================================================== | |
# 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() | |