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