MTCNN python实现

#!/usr/bin/python3 # -*- coding: utf-8 -*- # MIT License # # Copyright (c) 2018 Iván de Paz Centeno # # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal # in the Software without restriction, including without limitation the rights # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell # copies of the Software, and to permit persons to whom the Software is # furnished to do so, subject to the following conditions: # # The above copyright notice and this permission notice shall be included in all # copies or substantial portions of the Software. # # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. # IMPORTANT: # # This code is derivated from the MTCNN implementation of David Sandberg for Facenet # (https://github.com/davidsandberg/facenet/) # It has been rebuilt from scratch, taking the David Sandberg's implementation as a reference. # The code improves the readibility, fixes several mistakes in the definition of the network (layer names) # and provides the keypoints of faces as outputs along with the bounding boxes. # import cv2 import numpy as np import pkg_resources import tensorflow as tf from mtcnn.layer_factory import LayerFactory from mtcnn.network import Network #from mtcnn.exceptions import InvalidImage __author__ = "Iván de Paz Centeno" class InvalidImage(Exception): pass class PNet(Network): """ Network to propose areas with faces. """ def _config(self): layer_factory = LayerFactory(self) layer_factory.new_feed(name='data', layer_shape=(None, None, None, 3)) layer_factory.new_conv(name='conv1', kernel_size=(3, 3), channels_output=10, stride_size=(1, 1), padding='VALID', relu=False) layer_factory.new_prelu(name='prelu1') layer_factory.new_max_pool(name='pool1', kernel_size=(2, 2), stride_size=(2, 2)) layer_factory.new_conv(name='conv2', kernel_size=(3, 3), channels_output=16, stride_size=(1, 1), padding='VALID', relu=False) layer_factory.new_prelu(name='prelu2') layer_factory.new_conv(name='conv3', kernel_size=(3, 3), channels_output=32, stride_size=(1, 1), padding='VALID', relu=False) layer_factory.new_prelu(name='prelu3') layer_factory.new_conv(name='conv4-1', kernel_size=(1, 1), channels_output=2, stride_size=(1, 1), relu=False) layer_factory.new_softmax(name='prob1', axis=3) layer_factory.new_conv(name='conv4-2', kernel_size=(1, 1), channels_output=4, stride_size=(1, 1), input_layer_name='prelu3', relu=False) def _feed(self, image): return self._session.run(['pnet/conv4-2/BiasAdd:0', 'pnet/prob1:0'], feed_dict={'pnet/input:0': image}) class RNet(Network): """ Network to refine the areas proposed by PNet """ def _config(self): layer_factory = LayerFactory(self) layer_factory.new_feed(name='data', layer_shape=(None, 24, 24, 3)) layer_factory.new_conv(name='conv1', kernel_size=(3, 3), channels_output=28, stride_size=(1, 1), padding='VALID', relu=False) layer_factory.new_prelu(name='prelu1') layer_factory.new_max_pool(name='pool1', kernel_size=(3, 3), stride_size=(2, 2)) layer_factory.new_conv(name='conv2', kernel_size=(3, 3), channels_output=48, stride_size=(1, 1), padding='VALID', relu=False) layer_factory.new_prelu(name='prelu2') layer_factory.new_max_pool(name='pool2', kernel_size=(3, 3), stride_size=(2, 2), padding='VALID') layer_factory.new_conv(name='conv3', kernel_size=(2, 2), channels_output=64, stride_size=(1, 1), padding='VALID', relu=False) layer_factory.new_prelu(name='prelu3') layer_factory.new_fully_connected(name='fc1', output_count=128, relu=False) # shouldn't the name be "fc1"? layer_factory.new_prelu(name='prelu4') layer_factory.new_fully_connected(name='fc2-1', output_count=2, relu=False) # shouldn't the name be "fc2-1"? layer_factory.new_softmax(name='prob1', axis=1) layer_factory.new_fully_connected(name='fc2-2', output_count=4, relu=False, input_layer_name='prelu4') def _feed(self, image): return self._session.run(['rnet/fc2-2/fc2-2:0', 'rnet/prob1:0'], feed_dict={'rnet/input:0': image}) class ONet(Network): """ Network to retrieve the keypoints """ def _config(self): layer_factory = LayerFactory(self) layer_factory.new_feed(name='data', layer_shape=(None, 48, 48, 3)) layer_factory.new_conv(name='conv1', kernel_size=(3, 3), channels_output=32, stride_size=(1, 1), padding='VALID', relu=False) layer_factory.new_prelu(name='prelu1') layer_factory.new_max_pool(name='pool1', kernel_size=(3, 3), stride_size=(2, 2)) layer_factory.new_conv(name='conv2', kernel_size=(3, 3), channels_output=64, stride_size=(1, 1), padding='VALID', relu=False) layer_factory.new_prelu(name='prelu2') layer_factory.new_max_pool(name='pool2', kernel_size=(3, 3), stride_size=(2, 2), padding='VALID') layer_factory.new_conv(name='conv3', kernel_size=(3, 3), channels_output=64, stride_size=(1, 1), padding='VALID', relu=False) layer_factory.new_prelu(name='prelu3') layer_factory.new_max_pool(name='pool3', kernel_size=(2, 2), stride_size=(2, 2)) layer_factory.new_conv(name='conv4', kernel_size=(2, 2), channels_output=128, stride_size=(1, 1), padding='VALID', relu=False) layer_factory.new_prelu(name='prelu4') layer_factory.new_fully_connected(name='fc1', output_count=256, relu=False) layer_factory.new_prelu(name='prelu5') layer_factory.new_fully_connected(name='fc2-1', output_count=2, relu=False) layer_factory.new_softmax(name='prob1', axis=1) layer_factory.new_fully_connected(name='fc2-2', output_count=4, relu=False, input_layer_name='prelu5') layer_factory.new_fully_connected(name='fc2-3', output_count=10, relu=False, input_layer_name='prelu5') def _feed(self, image): return self._session.run(['onet/fc2-2/fc2-2:0', 'onet/fc2-3/fc2-3:0', 'onet/prob1:0'], feed_dict={'onet/input:0': image}) class StageStatus(object): """ Keeps status between MTCNN stages """ def __init__(self, pad_result: tuple = None, width=0, height=0): self.width = width self.height = height self.dy = self.edy = self.dx = self.edx = self.y = self.ey = self.x = self.ex = self.tmpw = self.tmph = [] if pad_result is not None: self.update(pad_result) def update(self, pad_result: tuple): s = self s.dy, s.edy, s.dx, s.edx, s.y, s.ey, s.x, s.ex, s.tmpw, s.tmph = pad_result class MTCNN(object): """ Allows to perform MTCNN Detection -> a) Detection of faces (with the confidence probability) b) Detection of keypoints (left eye, right eye, nose, mouth_left, mouth_right) """ def __init__(self, weights_file: str = None, min_face_size