1. 人脸68个关键点推理演示代码

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

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import warnings

import cv2
import onnx
import torch
import numpy as np
import onnxruntime

warnings.filterwarnings("ignore")


class ONNXModel(object):
    def __init__(self, onnx_path):
        """
        :param onnx_path:
        """
        self.onnx_session = onnxruntime.InferenceSession(onnx_path)
        self.input_name = self.get_input_name(self.onnx_session)
        self.output_name = self.get_output_name(self.onnx_session)
        print("input_name:{}".format(self.input_name))
        print("output_name:{}".format(self.output_name))

    def get_output_name(self, onnx_session):
        """
        output_name = onnx_session.get_outputs()[0].name
        :param onnx_session:
        :return:
        """
        output_name = []
        for node in onnx_session.get_outputs():
            output_name.append(node.name)
        return output_name

    def get_input_name(self, onnx_session):
        """
        input_name = onnx_session.get_inputs()[0].name
        :param onnx_session:
        :return:
        """
        input_name = []
        for node in onnx_session.get_inputs():
            input_name.append(node.name)
        return input_name

    def get_input_feed(self, input_name, image_tensor):
        """
        input_feed={self.input_name: image_tensor}
        :param input_name:
        :param image_tensor:
        :return:
        """
        input_feed = {}
        for name in input_name:
            input_feed[name] = image_tensor
        return input_feed

    def forward(self, image_tensor):
        '''
        image_tensor = image.transpose(2, 0, 1)
        image_tensor = image_tensor[np.newaxis, :]
        onnx_session.run([output_name], {input_name: x})
        :param image_tensor:
        :return:
        '''
        # 输入数据的类型必须与模型一致,以下三种写法都是可以的
        # scores, boxes = self.onnx_session.run(None, {self.input_name: image_tensor})
        # scores, boxes = self.onnx_session.run(self.output_name, input_feed={self.input_name: image_tensor})
        input_feed = self.get_input_feed(self.input_name, image_tensor)
        output = self.onnx_session.run(self.output_name, input_feed=input_feed)
        output = np.array(output).reshape(-1, 2)
        return output

worker = ONNXModel("./pfld.onnx")

img = cv2.imread("./s_28.jpg")
# resized = cv2.resize(img, (112, 112), interpolation=cv2.INTER_LINEAR)
img_in = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
# cv2.imwrite("tmp.jpg", img_in)
img_in = np.transpose(img_in, (2, 0, 1)).astype(np.float32)
img_in = np.expand_dims(img_in, axis=0)
img_in /= 255.0

output = worker.forward(img_in)
# torch.Size([1, 59500, 4]) torch.Size([1, 59500, 2]) torch.Size([1, 59500, 10])
print(output)
for xy in output:
    x = xy[0] * 112
    y = xy[1] * 112
    cv2.circle(img, (int(x), int(y)), 2, (0,255,0), -1)

cv2.imwrite("out1.jpg", img)