MTCNN 人脸检测

demo.py

import cv2
from detection.mtcnn import MTCNN

# 检测图片中的人脸
def test_image(imgpath):
    mtcnn = MTCNN('./mtcnn.pb')
    img = cv2.imread(imgpath)

    bbox, landmarks, scores = mtcnn.detect_faces(img)

    print('total box:', len(bbox))
    for box, pts in zip(bbox, landmarks):
        box = box.astype('int32')
        img = cv2.rectangle(img, (box[1], box[0]), (box[3], box[2]), (255, 0, 0), 3)

        pts = pts.astype('int32')
        for i in range(5):
            img = cv2.circle(img, (pts[i + 5], pts[i]), 1, (0, 255, 0), 2)
    cv2.imshow('image', img)
    cv2.waitKey()

# 检测视频中的人脸
def test_camera():
    mtcnn = MTCNN('./mtcnn.pb')
    cap = cv2.VideoCapture('rtsp://admin:hik12345@192.168.3.160/Streaming/Channels/1')
    while True:
        ret, img = cap.read()
        if not ret:
            break
        bbox, landmarks, scores = mtcnn.detect_faces(img)
        print('total box:', len(bbox), scores)
        for box, pts in zip(bbox, landmarks):
            box = box.astype('int32')
            img = cv2.rectangle(img, (box[1], box[0]), (box[3], box[2]), (255, 0, 0), 3)
            pts = pts.astype('int32')
            for i in range(5):
                img = cv2.circle(img, (pts[i], pts[i + 5]), 1, (0, 255, 0), 2)
        cv2.imshow('img', img)
        cv2.waitKey(1)

if __name__ == '__main__':
    # test_image()
    test_camera()

mtcnn.py

import tensorflow as tf
from detection.align_trans import get_reference_facial_points, warp_and_crop_face
import numpy as np
import cv2
import detection.face_preprocess as face_preprocess

class MTCNN:

    def __init__(self, model_path, min_size=40, factor=0.709, thresholds=[0.7, 0.8, 0.8]):
        self.min_size = min_size
        self.factor = factor
        self.thresholds = thresholds

        graph = tf.Graph()
        with graph.as_default():
            with open(model_path, 'rb') as f:
                graph_def = tf.GraphDef.FromString(f.read())
                tf.import_graph_def(graph_def, name='')
        self.graph = graph
        config = tf.ConfigProto(
            allow_soft_placement=True,
            intra_op_parallelism_threads=4,
            inter_op_parallelism_threads=4)
        config.gpu_options.allow_growth = True
        self.sess = tf.Session(graph=graph, config=config)
        self.refrence = get_reference_facial_points(default_square=True)

    # 人脸检测
    def detect_faces(self, img):
        feeds = {
            self.graph.get_operation_by_name('input').outputs[0]: img,
            self.graph.get_operation_by_name('min_size').outputs[0]: self.min_size,
            self.graph.get_operation_by_name('thresholds').outputs[0]: self.thresholds,
            self.graph.get_operation_by_name('factor').outputs[0]: self.factor
        }
        fetches = [self.graph.get_operation_by_name('prob').outputs[0],
                  self.graph.get_operation_by_name('landmarks').outputs[0],
                  self.graph.get_operation_by_name('box').outputs[0]]
        prob, landmarks, box = self.sess.run(fetches, feeds)
        return box, landmarks, prob

    # 对齐获取单个人脸
    def align_face(self, img):
        ret = self.detect_faces(img)
        if ret is None:
            return None
        bbox, landmarks, prob = ret
        if bbox.shape[0] == 0:
            return None

        landmarks_copy = landmarks.copy()
        landmarks[:, 0:5] = landmarks_copy[:, 5:10]
        landmarks[:, 5:10] = landmarks_copy[:, 0:5]
        # print(landmarks[0, :])

        bbox = bbox[0, 0:4]
        bbox = bbox.astype(int)

        bbox = bbox[::-1]
        bbox_copy = bbox.copy()
        bbox[0:2] = bbox_copy[2:4]
        bbox[2:4] = bbox_copy[0:2]
        # print(bbox)

        points = landmarks[0, :].reshape((2, 5)).T
        # print(points) 

        '''
        face_img = cv2.rectangle(img, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 0, 255), 6)
        for i in range(5):
            pts = points[i, :]
            face_img = cv2.circle(face_img, (pts[0], pts[1]), 2, (0, 255, 0), 2)
        cv2.imshow('img', face_img)
        if cv2.waitKey(100000) & 0xFF == ord('q'):
            cv2.destroyAllWindows()
        '''

        warped_face = face_preprocess.preprocess(img, bbox, points, image_size='112,112')
        '''
        cv2.imshow('face', warped_face)

        if cv2.waitKey(100000) & 0xFF == ord('q'):
            cv2.destroyAllWindows()
        '''
        # warped_face = cv2.cvtColor(warped_face, cv2.COLOR_BGR2RGB)
        # aligned = np.transpose(warped_face, (2, 0, 1))
        # return aligned
        return warped_face

    # 对齐获取多个人脸
    def align_multi_faces(self, img, limit=None):
        boxes, landmarks, _ = self.detect_faces(img)
        if limit:
            boxes = boxes[:limit]
            landmarks = landmarks[:limit]

        landmarks_copy = landmarks.copy()
        landmarks[:, 0:5] = landmarks_copy[:, 5:10]
        landmarks[:, 5:10] = landmarks_copy[:, 0:5]

        # print('landmarks', landmark)
        faces = []
        for idx in range(len(landmarks)):
            '''
            landmark = landmarks[idx, :]
            facial5points = [[landmark[j], landmark[j + 5]] for j in range(5)]
            warped_face = warp_and_crop_face(np.array(img), facial5points, self.refrence, crop_size=(112, 112))
            faces.append(warped_face)
            '''
            bbox = boxes[idx, 0:4]
            bbox = bbox.astype(int)

            bbox = bbox[::-1]
            bbox_copy = bbox.copy()
            bbox[0:2] = bbox_copy[2:4]
            bbox[2:4] = bbox_copy[0:2]
            # print(bbox)

            points = landmarks[idx, :].reshape((2, 5)).T
            # print(points)
            warped_face = face_preprocess.preprocess(img, bbox, points, image_size='112,112')

            cv2.imshow('faces', warped_face)
            # warped_face = cv2.cvtColor(warped_face, cv2.COLOR_BGR2RGB)
            # aligned = np.transpose(warped_face, (2, 0, 1))
            faces.append(warped_face)

        # print('faces',faces)
        # print('boxes',boxes)
        return faces, boxes, landmarks