73 lines
2.4 KiB
Python
73 lines
2.4 KiB
Python
import cv2
|
|
from loguru import logger
|
|
|
|
from lama_cleaner.helper import download_model
|
|
from lama_cleaner.plugins.base_plugin import BasePlugin
|
|
from lama_cleaner.schema import RunPluginRequest
|
|
|
|
|
|
class GFPGANPlugin(BasePlugin):
|
|
name = "GFPGAN"
|
|
|
|
def __init__(self, device, upscaler=None):
|
|
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}")
|
|
|
|
import facexlib
|
|
|
|
if hasattr(facexlib.detection.retinaface, "device"):
|
|
facexlib.detection.retinaface.device = device
|
|
|
|
# Use GFPGAN for face enhancement
|
|
self.face_enhancer = MyGFPGANer(
|
|
model_path=model_path,
|
|
upscale=1,
|
|
arch="clean",
|
|
channel_multiplier=2,
|
|
device=device,
|
|
bg_upsampler=upscaler.model if upscaler is not None else None,
|
|
)
|
|
self.face_enhancer.face_helper.face_det.mean_tensor.to(device)
|
|
self.face_enhancer.face_helper.face_det = (
|
|
self.face_enhancer.face_helper.face_det.to(device)
|
|
)
|
|
|
|
def __call__(self, rgb_np_img, req: RunPluginRequest):
|
|
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"
|
|
)
|