|
from typing import Dict, List, Any |
|
from gfpgan import GFPGANer |
|
import cv2 |
|
from imageio import imread |
|
from basicsr.utils import imwrite |
|
import io |
|
import os |
|
import numpy as np |
|
import base64 |
|
|
|
|
|
class EndpointHandler(): |
|
def __init__(self, path=""): |
|
self.restorer = GFPGANer( |
|
model_path="./GFPGANv1.4.pth", |
|
upscale=2, |
|
arch="clean", |
|
channel_multiplier=2, |
|
bg_upsampler=None) |
|
|
|
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: |
|
""" |
|
data args: |
|
inputs (:obj: `str`) |
|
date (:obj: `str`) |
|
Return: |
|
A :obj:`list` | `dict`: will be serialized and returned |
|
""" |
|
|
|
inputs = data.pop("inputs",data) |
|
img = imread(io.BytesIO(base64.b64decode(inputs))) |
|
cropped_faces, restored_faces, restored_img = self.restorer.enhance( |
|
img, |
|
has_aligned=False, |
|
only_center_face=False, |
|
paste_back=True, |
|
weight=0.5) |
|
|
|
for idx, (cropped_face, restored_face) in enumerate(zip(cropped_faces, restored_faces)): |
|
retval, buffer = cv2.imencode('.png', restored_face) |
|
jpg_as_text = base64.b64encode(buffer) |
|
return jpg_as_text |
|
|