Spaces:
Runtime error
Runtime error
File size: 2,840 Bytes
f2dbf59 |
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 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 |
#!/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)
|