File size: 1,265 Bytes
c4d374b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 |
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
"""
# get inputs
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
|