Spaces:
Runtime error
Runtime error
from chain_img_processor import ChainImgProcessor, ChainImgPlugin | |
import os | |
import gfpgan | |
import threading | |
from PIL import Image | |
from numpy import asarray | |
import cv2 | |
from roop.utilities import resolve_relative_path, conditional_download | |
modname = os.path.basename(__file__)[:-3] # calculating modname | |
model_gfpgan = None | |
THREAD_LOCK_GFPGAN = threading.Lock() | |
# start function | |
def start(core:ChainImgProcessor): | |
manifest = { # plugin settings | |
"name": "GFPGAN", # name | |
"version": "1.4", # version | |
"default_options": {}, | |
"img_processor": { | |
"gfpgan": GFPGAN | |
} | |
} | |
return manifest | |
def start_with_options(core:ChainImgProcessor, manifest:dict): | |
pass | |
class GFPGAN(ChainImgPlugin): | |
def init_plugin(self): | |
global model_gfpgan | |
if model_gfpgan is None: | |
model_path = resolve_relative_path('../models/GFPGANv1.4.pth') | |
model_gfpgan = gfpgan.GFPGANer(model_path=model_path, upscale=1, device=self.device) # type: ignore[attr-defined] | |
def process(self, frame, params:dict): | |
import copy | |
global model_gfpgan | |
if model_gfpgan is None: | |
return frame | |
if "face_detected" in params: | |
if not params["face_detected"]: | |
return frame | |
# don't touch original | |
temp_frame = copy.copy(frame) | |
if "processed_faces" in params: | |
for face in params["processed_faces"]: | |
start_x, start_y, end_x, end_y = map(int, face['bbox']) | |
padding_x = int((end_x - start_x) * 0.5) | |
padding_y = int((end_y - start_y) * 0.5) | |
start_x = max(0, start_x - padding_x) | |
start_y = max(0, start_y - padding_y) | |
end_x = max(0, end_x + padding_x) | |
end_y = max(0, end_y + padding_y) | |
temp_face = temp_frame[start_y:end_y, start_x:end_x] | |
if temp_face.size: | |
with THREAD_LOCK_GFPGAN: | |
_, _, temp_face = model_gfpgan.enhance( | |
temp_face, | |
paste_back=True | |
) | |
temp_frame[start_y:end_y, start_x:end_x] = temp_face | |
else: | |
with THREAD_LOCK_GFPGAN: | |
_, _, temp_frame = model_gfpgan.enhance( | |
temp_frame, | |
paste_back=True | |
) | |
if not "blend_ratio" in params: | |
return temp_frame | |
temp_frame = Image.blend(Image.fromarray(frame), Image.fromarray(temp_frame), params["blend_ratio"]) | |
return asarray(temp_frame) | |