IOPaint/lama_cleaner/plugins/gfpgan_plugin.py

72 lines
2.3 KiB
Python
Raw Normal View History

2023-03-26 07:39:09 +02:00
import cv2
from loguru import logger
from lama_cleaner.helper import download_model
from lama_cleaner.plugins.base_plugin import BasePlugin
class GFPGANPlugin(BasePlugin):
name = "GFPGAN"
2023-03-26 14:52:06 +02:00
def __init__(self, device, upscaler=None):
2023-03-26 07:39:09 +02:00
super().__init__()
from .gfpganer import MyGFPGANer
url = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth"
model_md5 = "94d735072630ab734561130a47bc44f8"
model_path = download_model(url, model_md5)
logger.info(f"GFPGAN model path: {model_path}")
2023-03-28 10:36:41 +02:00
import facexlib
2023-03-31 16:17:38 +02:00
2023-03-28 10:36:41 +02:00
if hasattr(facexlib.detection.retinaface, "device"):
2023-04-01 15:26:40 +02:00
facexlib.detection.retinaface.device = device
2023-03-28 10:36:41 +02:00
2023-03-26 07:39:09 +02:00
# Use GFPGAN for face enhancement
self.face_enhancer = MyGFPGANer(
model_path=model_path,
upscale=1,
2023-03-26 07:39:09 +02:00
arch="clean",
channel_multiplier=2,
device=device,
2023-03-26 14:52:06 +02:00
bg_upsampler=upscaler.model if upscaler is not None else None,
2023-03-26 07:39:09 +02:00
)
2023-04-01 15:26:40 +02:00
self.face_enhancer.face_helper.face_det.mean_tensor.to(device)
2023-03-31 16:17:38 +02:00
self.face_enhancer.face_helper.face_det = (
2023-04-01 15:26:40 +02:00
self.face_enhancer.face_helper.face_det.to(device)
2023-03-31 16:17:38 +02:00
)
2023-03-26 07:39:09 +02:00
def __call__(self, rgb_np_img, files, form):
weight = 0.5
bgr_np_img = cv2.cvtColor(rgb_np_img, cv2.COLOR_RGB2BGR)
logger.info(f"GFPGAN input shape: {bgr_np_img.shape}")
_, _, bgr_output = self.face_enhancer.enhance(
bgr_np_img,
has_aligned=False,
only_center_face=False,
paste_back=True,
weight=weight,
)
logger.info(f"GFPGAN output shape: {bgr_output.shape}")
# try:
# if scale != 2:
# interpolation = cv2.INTER_AREA if scale < 2 else cv2.INTER_LANCZOS4
# h, w = img.shape[0:2]
# output = cv2.resize(
# output,
# (int(w * scale / 2), int(h * scale / 2)),
# interpolation=interpolation,
# )
# except Exception as error:
# print("wrong scale input.", error)
return bgr_output
def check_dep(self):
try:
import gfpgan
except ImportError:
return (
"gfpgan is not installed, please install it first. pip install gfpgan"
)