gfpgan_foward.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import argparse
import cv2
import numpy as np
import time
from onnxruntime import InferenceSession
from util.model_utils import load_models
from util.image_utils import save_faces

class GFPGANFaceAugment:
    def __init__(self, modelSession):
        self.ort_session = modelSession
        self.face_size = 512
        self.upscale_factor = 2

    def pre_process(self, img):
        img = img / 255.0
        img = img.astype("float32")
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img[:, :, 0] = (img[:, :, 0] - 0.5) / 0.5
        img[:, :, 1] = (img[:, :, 1] - 0.5) / 0.5
        img[:, :, 2] = (img[:, :, 2] - 0.5) / 0.5
        img = np.float32(img[np.newaxis, :, :, :])
        img = img.transpose(0, 3, 1, 2)
        return img

    def post_process(self, output, height, width):
        output = output.clip(-1, 1)
        output = (output + 1) / 2
        output = output.transpose(1, 2, 0)
        output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
        output = (output * 255.0).round()
        inv_soft_mask = np.ones((height, width, 1), dtype=np.float32)
        output = cv2.resize(output, (width, height))
        
        return output, inv_soft_mask

    def forward(self, img):
        
        height, width = img.shape[0], img.shape[1]
        img = self.pre_process(img)
      
        start = int(round(time.time() * 1000))
        ort_inputs = {self.ort_session.get_inputs()[0].name: img}
        
        ort_outs = self.ort_session.run(None, ort_inputs)
        output = ort_outs[0][0]
       
        output, inv_soft_mask = self.post_process(output, height, width)
        end = int(round(time.time() * 1000))
        estimation_time = end - start
        logger.info(f"'GFPGAN infer time: {estimation_time} ms")
        output = output.astype(np.uint8)


        return output, inv_soft_mask