face_landmark.py 2.43 KB
import os
import cv2
import MNN
import logging
import numpy as np

class Landmark_Detector(object):

    def __init__(self, model_path):
        logging.info('******** Start Init Landmark Detector ********')
        self.interpreter = MNN.Interpreter(model_path)
        self.session = self.interpreter.createSession()
        self.input_tensor = self.interpreter.getSessionInput(self.session)
        logging.info('******** Success Init Landmark Detector ********')
    
    def extent_bbox(self, bbox, image_h, image_w):
        x_min, y_min, x_max, y_max = bbox[0], bbox[1], bbox[2], bbox[3]
        face_w = x_max - x_min
        face_h = y_max - y_min
        x_min = int(x_min - face_w * 0.05)
        x_max = int(x_max + face_w * 0.05)
        y_min = int(y_min - face_h * 0.0)
        y_max = int(y_max + face_h * 0.1)

        x_min = 0 if x_min < 0 else x_min
        y_min = 0 if y_min < 0 else y_min
        x_max = image_w if x_max > image_w else x_max
        y_max = image_h if y_max > image_h else y_max

        return [x_min, y_min, x_max, y_max]

    def detect(self, image, bbox):
        logging.info('******** Start Landmark Detect ********')
        
        image_h, image_w, image_c = image.shape
        bbox = self.extent_bbox(bbox, image_h, image_w)
        face_area = image[bbox[1]:bbox[3], bbox[0]:bbox[2], ]
 
        face_h, face_w, _ = face_area.shape    
        image_resize = cv2.resize(face_area, (112, 112), interpolation=cv2.INTER_CUBIC)
        image_norm = image_resize / 256.
        image_norm = image_norm.astype(np.float32)
            
        image_transpose = image_norm.transpose((2, 0, 1))
        input_tensor = MNN.Tensor((1, 3, 112, 112), MNN.Halide_Type_Float, image_transpose, MNN.Tensor_DimensionType_Caffe)
        self.input_tensor.copyFrom(input_tensor)
        self.interpreter.runSession(self.session)
        output_tensor = self.interpreter.getSessionOutput(self.session)
        landmarks = output_tensor.getData()
        # landmarks = landmarks[0]

        new_landmark = []
        for p in range(106):
            # x_p = landmarks[2 * p] * w_r + bbox[0]
            # y_p = landmarks[2 * p + 1] * h_r + bbox[1]
            x_p = landmarks[2 * p]
            y_p = landmarks[2 * p + 1]
            new_landmark.append((x_p, y_p))
            # cv2.circle(image_resize, (int(x_p), int(y_p)), 1, (0, 0, 255), -1)

        logging.info('******** Success Face Detect ********')

        return new_landmark, image_resize