mediapipe face_mesh测试

目录

安装

onnx测试

tensorflow预测tflite代码：

---

安装

pip install mediapipe --user

onnx测试

   img_path=r'D:\data\val_result\1212\test\1\2\02370_1.jpg'
    img_path=r'D:\data\face\1212\test\1\2\02674_1.jpg'
    img=cv2.imread(img_path)

    img_o = cv2.resize(img, (192, 192))
    # img=cv2.cvtColor(img_o,cv2.COLOR_BGR2RGB)
    img = img_o.astype(np.float32)
    # img=img-127.5
    img/=255.0

    # dummy_input = np.ones([1, 192, 192, 3], dtype=np.float32)
    dummy_input = np.array([img])
    model_path = r'E:\project\angle_net\xuanzhuan_detect\ncnn-20220420-windows-vs2017\x64\bin/model_float32.onnx'
    # model_path = r'C:\Users\Administrator\Pictures\mm\resources.tar\resources\face_landmark_with_attention_192x192/model_float32.onnx'
    model_file_name = model_path.split(".")[0]
    session_option = onnxruntime.SessionOptions()
    session_option.optimized_model_filepath = f"{model_file_name}_cudaopt.onnx"
    session_option.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_EXTENDED
    session = onnxruntime.InferenceSession(
        model_path,
        session_option,
        providers=['CUDAExecutionProvider']
    )

    input_name = session.get_inputs()[0].name
    output_names = [o.name for o in session.get_outputs()]
    input_shape = session.get_inputs()[0].shape
    # Warmup
    output = session.run(
        output_names,
        {input_name: dummy_input.transpose((0,3,1,2))}
    )
    # Inference

    for i in range(5):
        start = time.time()
        output = session.run(output_names,{input_name: dummy_input.transpose((0,3,1,2))})

        for i in [1,2,3,4,5,6]:
            points=output[i][0][0]

            if len(points[0])==1404:
                data_len = len(points[0]) // 3
                # print(data_len)
                for index in range(data_len):
                    cv2.circle(img_o, (int(points[0][index * 3]), int(points[0][index * 3 + 1])), 1, (0, 0, 213),
                               -1)  # x,y,r,color
            else:
                data_len = len(points[0]) // 2
                # print(data_len)
                for index in range(data_len):
                    cv2.circle(img_o, (int(points[0][index * 2]), int(points[0][index * 2 + 1])), 1, (0, 0, 213),
                               -1)  # x,y,r,color


        cv2.imshow("result",img_o)
        cv2.waitKey(0)
        print(f'elapsed_time : {(time.time()-start)*1000} ms')

tensorflow预测tflite代码：

import os

import cv2
import tensorflow as tf
import numpy as np
from tensorflow.lite.python.interpreter import Interpreter

class MeshDetector(object):

    def __init__(self, model_path):
        """Initialization"""
        self.target_image = None

        self.interpreter =Interpreter(model_path=model_path)

        # Set model input
        self.interpreter.allocate_tensors()
        self.input_details = self.interpreter.get_input_details()
        self.output_details = self.interpreter.get_output_details()

    def _preprocess(self, image_bgr):

        image_resized = cv2.resize(image_bgr, (192, 192))
        image_rgb = cv2.cvtColor(image_resized, cv2.COLOR_BGR2RGB)
        image_norm = (image_rgb-127.5)/127.5
        return image_norm

    def get_mesh(self, image):

        # Preprocess the image before sending to the network.
        image = self._preprocess(image)
        image = tf.image.convert_image_dtype(image, tf.float32)
        image = image[tf.newaxis, :]

        # The actual detection.
        self.interpreter.set_tensor(self.input_details[0]["index"], image)
        self.interpreter.invoke()

        # Save the results.
        mesh = self.interpreter.get_tensor(self.output_details[3]["index"])[0].reshape(468, 3) / 192
        score = self.interpreter.get_tensor(self.output_details[1]["index"])[0]

        return mesh, score

    def draw_mesh(self, image, mesh, mark_size=2, line_width=1):
        """Draw the mesh on an image"""
        # The mesh are normalized which means we need to convert it back to fit
        # the image size.
        image_size = image.shape[0]
        mesh = mesh * image_size
        for point in mesh:
            cv2.circle(image, (int(point[0]), int(point[1])),mark_size, (0, 255, 128), -1)

        # Draw the contours.
        # Eyes
        left_eye_contour = np.array([mesh[33][0:2],
                                     mesh[7][0:2],
                                     mesh[163][0:2],
                                     mesh[144][0:2],
                                     mesh[145][0:2],
                                     mesh[153][0:2],
                                     mesh[154][0:2],
                                     mesh[155][0:2],
                                     mesh[133][0:2],
                                     mesh[173][0:2],
                                     mesh[157][0:2],
                                     mesh[158][0:2],
                                     mesh[159][0:2],
                                     mesh[160][0:2],
                                     mesh[161][0:2],
                                     mesh[246][0:2], ]).astype(np.int32)
        right_eye_contour = np.array([mesh[263][0:2],
                                      mesh[249][0:2],
                                      mesh[390][0:2],
                                      mesh[373][0:2],
                                      mesh[374][0:2],
                                      mesh[380][0:2],
                                      mesh[381][0:2],
                                      mesh[382][0:2],
                                      mesh[362][0:2],
                                      mesh[398][0:2],
                                      mesh[384][0:2],
                                      mesh[385][0:2],
                                      mesh[386][0:2],
                                      mesh[387][0:2],
                                      mesh[388][0:2],
                                      mesh[466][0:2]]).astype(np.int32)
        # Lips
        cv2.polylines(image, [left_eye_contour, right_eye_contour], False,
                      (0, 0, 255), line_width, cv2.LINE_AA)


if __name__ == "__main__":
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    
    mesh_detector= MeshDetector(r'model_float32.tflite')
    img_path = r'D:\data\face\1212\test\1\2\02674_1.jpg'
    img = cv2.imread(img_path)
    img = cv2.resize(img, (192, 192))
    face_mesh, score= mesh_detector.get_mesh(img)
    mesh_detector.draw_mesh(img, face_mesh, mark_size=1)
    # cv2.imshow("Mesh", cv2.resize(img, (512, 512)))
    cv2.imshow("Mesh",img)
    cv2.waitKey()