|
|
import os |
|
|
import sys |
|
|
import types |
|
|
import cv2 |
|
|
import gradio as gr |
|
|
import numpy as np |
|
|
from gfpgan import GFPGANer |
|
|
from huggingface_hub import hf_hub_download |
|
|
|
|
|
|
|
|
try: |
|
|
import torchvision.transforms.functional_tensor |
|
|
except ModuleNotFoundError: |
|
|
|
|
|
sys.modules['torchvision.transforms.functional_tensor'] = types.ModuleType('functional_tensor') |
|
|
|
|
|
|
|
|
loaded_models = {} |
|
|
|
|
|
def load_model(version): |
|
|
if version not in loaded_models: |
|
|
if version == "v1.4 (original)": |
|
|
model_path = hf_hub_download("leonelhs/gfpgan", "GFPGANv1.4.pth") |
|
|
loaded_models[version] = GFPGANer( |
|
|
model_path=model_path, |
|
|
upscale=1, |
|
|
arch='original', |
|
|
channel_multiplier=2, |
|
|
bg_upsampler=None |
|
|
) |
|
|
elif version == "v1.3 (clean)": |
|
|
model_path = hf_hub_download("leonelhs/gfpgan", "GFPGANv1.3.pth") |
|
|
loaded_models[version] = GFPGANer( |
|
|
model_path=model_path, |
|
|
upscale=1, |
|
|
arch='clean', |
|
|
channel_multiplier=2, |
|
|
bg_upsampler=None |
|
|
) |
|
|
return loaded_models[version] |
|
|
|
|
|
|
|
|
def enhance_face(image, version): |
|
|
restorer = load_model(version) |
|
|
|
|
|
if isinstance(image, np.ndarray): |
|
|
img = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) |
|
|
else: |
|
|
raise ValueError("Invalid image format") |
|
|
|
|
|
_, _, restored_img = restorer.enhance( |
|
|
img, has_aligned=False, only_center_face=False, paste_back=True |
|
|
) |
|
|
|
|
|
restored_img = cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB) |
|
|
return restored_img |
|
|
|
|
|
|
|
|
iface = gr.Interface( |
|
|
fn=enhance_face, |
|
|
inputs=[ |
|
|
gr.Image(type="numpy", label="Upload your Anime/Cartoon image"), |
|
|
gr.Radio(["v1.3 (clean)", "v1.4 (original)"], value="v1.4 (original)", label="Choose GFPGAN Version") |
|
|
], |
|
|
outputs=gr.Image(type="numpy", label="Restored Cartoon Face"), |
|
|
title="Cartoon Face Restoration with GFPGAN", |
|
|
description="Upload a cartoonized image (e.g. from AnimeGAN) and restore it with GFPGAN." |
|
|
) |
|
|
|
|
|
if __name__ == "__main__": |
|
|
iface.launch() |
