| import os |
| import cv2 |
| import torch |
| import gradio as gr |
| import numpy as np |
| from PIL import Image |
| from gfpgan import GFPGANer |
| from basicsr.archs.srvgg_arch import SRVGGNetCompact |
| from realesrgan.utils import RealESRGANer |
|
|
| |
| GFPGAN_MODEL = "GFPGANv1.4.pth" |
| GFPGAN_URL = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth" |
|
|
| if not os.path.exists(GFPGAN_MODEL): |
| import requests |
| r = requests.get(GFPGAN_URL, allow_redirects=True) |
| open(GFPGAN_MODEL, 'wb').write(r.content) |
|
|
| |
| ESRGAN_MODEL = "realesr-general-x4v3.pth" |
| ESRGAN_URL = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth" |
|
|
| if not os.path.exists(ESRGAN_MODEL): |
| import requests |
| r = requests.get(ESRGAN_URL, allow_redirects=True) |
| open(ESRGAN_MODEL, 'wb').write(r.content) |
|
|
| |
| esr_model = SRVGGNetCompact( |
| num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, |
| upscale=4, act_type='prelu' |
| ) |
|
|
| bg_upsampler = RealESRGANer( |
| scale=4, |
| model_path=ESRGAN_MODEL, |
| model=esr_model, |
| tile=0, |
| tile_pad=10, |
| pre_pad=0, |
| half=torch.cuda.is_available() |
| ) |
|
|
| |
| restorer = GFPGANer( |
| model_path=GFPGAN_MODEL, |
| upscale=2, |
| arch='clean', |
| channel_multiplier=2, |
| bg_upsampler=bg_upsampler |
| ) |
|
|
| |
| def enhance(image): |
| img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) |
|
|
| _, _, restored_img = restorer.enhance( |
| img, |
| has_aligned=False, |
| only_center_face=False, |
| paste_back=True |
| ) |
|
|
| restored_pil = Image.fromarray(cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB)) |
| return restored_pil |
|
|
| |
| iface = gr.Interface( |
| fn=enhance, |
| inputs=gr.Image(type="pil"), |
| outputs=gr.Image(type="pil"), |
| title="IMGEN - AI Photo Enhancer (Face + Outfit)", |
| description="Upload your photo (ID, CV, profile) and enhance both the face and outfit with AI using GFPGAN + RealESRGAN." |
| ) |
|
|
| if __name__ == "__main__": |
| iface.launch() |
|
|
