Spaces:
Runtime error
Runtime error
| #!/usr/bin/env python | |
| # -*- coding: utf-8 -*- | |
| import os | |
| import sys | |
| import argparse | |
| import cv2 | |
| import numpy as np | |
| #import timeit | |
| import onnxruntime | |
| from facefusion.affine import create_box_mask, warp_face_by_landmark, paste_back | |
| class GFPGANOnnx: | |
| def __init__(self, model_path, providers): | |
| self.session = onnxruntime.InferenceSession(model_path, providers=providers) | |
| inputs = self.session.get_inputs() | |
| self.input_size = (inputs[0].shape[2], inputs[0].shape[3]) | |
| self.input_name = inputs[0].name | |
| self.affine = False | |
| def pre_process(self, image): | |
| img = cv2.resize(image, self.input_size) | |
| img = img/255.0 | |
| 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() | |
| output = cv2.resize(output, (width, height)) | |
| return output | |
| def run(self, image): | |
| height, width = image.shape[0], image.shape[1] | |
| img = self.pre_process(image) | |
| #t = timeit.default_timer() | |
| outputs = self.session.run(None, {self.input_name: img}) | |
| output = outputs[0][0] | |
| output = self.post_process(output, height, width) | |
| #print('infer time:',timeit.default_timer()-t) | |
| output = output.astype(np.uint8) | |
| return output | |
| if __name__ == "__main__": | |
| from yoloface_onnx import YoloFaceOnnx | |
| providers=['CPUExecutionProvider'] | |
| model_path = '/Users/wadahana/workspace/AI/sd/ComfyUI/models/facefusion/gfpgan_1.4.onnx' | |
| yolo_path = '/Users/wadahana/workspace/AI/sd/ComfyUI/models/facefusion/yoloface_8n.onnx' | |
| detector = YoloFaceOnnx(model_path=yolo_path, providers=providers) | |
| session = GFPGANOnnx(model_path=model_path, providers=providers) | |
| image = cv2.imread('/Users/wadahana/Desktop/anime-3.jpeg') | |
| face_list = detector.detect(image=image, conf=0.7) | |
| print(f'total of face: {len(face_list)}') | |
| output = image | |
| for index, face in enumerate(face_list): | |
| cropped, affine_matrix = warp_face_by_landmark(image, face.landmarks, session.input_size) | |
| box_mask = create_box_mask(session.input_size, 0.3, (0,0,0,0)) | |
| crop_mask = np.minimum.reduce([box_mask]).clip(0, 1) | |
| result = session.run(cropped) | |
| cv2.imwrite(f'/Users/wadahana/Desktop/output_{index}.jpg', result) | |
| output = paste_back(output, result, crop_mask, affine_matrix) | |
| cv2.imwrite(f'/Users/wadahana/Desktop/output.jpg', output) | |