|
|
- #!/usr/bin/env python3
-
- import cv2
- import numpy as np
- from scipy.spatial.transform import Rotation
- import deep_models_shared_python3 as dm
-
- class HeadPoseEstimator:
- def __init__(self, models_directory, use_neural_compute_stick=False):
- # Load the models
- models_dir = models_directory
- print('Using the following directory to load object detector models:', models_dir)
-
- # file with network architecture and other information
- head_detection_model_prototxt_filename = models_dir + '/head_detection/deploy.prototxt'
- # file with network weights
- head_detection_model_caffemodel_filename = models_dir + '/head_detection/res10_300x300_ssd_iter_140000.caffemodel'
- self.face_confidence_threshold = 0.2
-
- print('attempting to load neural network from files')
- print('prototxt file =', head_detection_model_prototxt_filename)
- print('caffemodel file =', head_detection_model_caffemodel_filename)
- self.head_detection_model = cv2.dnn.readNetFromCaffe(head_detection_model_prototxt_filename, head_detection_model_caffemodel_filename)
- dm.print_model_info(self.head_detection_model, 'head_detection_model')
-
- # attempt to use Neural Compute Stick 2
- if use_neural_compute_stick:
- print('HeadPoseEstimator.__init__: Attempting to use an Intel Neural Compute Stick 2 using the following command: self.head_detection_model.setPreferableTarget(cv2.dnn.DNN_TARGET_MYRIAD)')
- self.head_detection_model.setPreferableTarget(cv2.dnn.DNN_TARGET_MYRIAD)
-
- head_pose_model_dir = models_dir + '/open_model_zoo/head-pose-estimation-adas-0001/FP32/'
- head_pose_weights_filename = head_pose_model_dir + 'head-pose-estimation-adas-0001.bin'
- head_pose_config_filename = head_pose_model_dir + 'head-pose-estimation-adas-0001.xml'
- self.head_pose_model = cv2.dnn.readNet(head_pose_weights_filename, head_pose_config_filename)
-
- if use_neural_compute_stick:
- print('Not attempting to use a Intel Neural Compute Stick 2 for head pose estimation due to potential errors.')
-
- dm.print_model_info(self.head_pose_model, 'head_pose_model')
-
- landmarks_model_dir = models_dir + '/open_model_zoo/facial-landmarks-35-adas-0002/FP32/'
- landmarks_weights_filename = landmarks_model_dir + 'facial-landmarks-35-adas-0002.bin'
- landmarks_config_filename = landmarks_model_dir + 'facial-landmarks-35-adas-0002.xml'
- self.landmarks_model = cv2.dnn.readNet(landmarks_weights_filename, landmarks_config_filename)
-
- if use_neural_compute_stick:
- print('Not attempting to use a Intel Neural Compute Stick 2 for facial landmarks due to potential errors.')
-
- dm.print_model_info(self.head_pose_model, 'head_pose_model')
-
-
- dm.print_model_info(self.landmarks_model, 'landmarks_model')
-
- self.landmark_names = ['right_eye_left', 'right_eye_right',
- 'left_eye_right', 'left_eye_left', 'nose_tip',
- 'nose_bottom', 'nose_right', 'nose_left', 'mouth_right',
- 'mouth_left', 'mouth_top', 'mouth_bottom',
- 'right_eyebrow_right', 'right_eyebrow_middle', 'right_eyebrow_left',
- 'left_eyebrow_right', 'left_eyebrow_middle', 'left_eyebrow_left',
- 'right_cheek_18', 'right_cheek_19', 'right_cheek_20', 'right_cheek_21',
- 'right_cheek_22', 'right_cheek_23', 'right_cheek_24',
- 'chin_right', 'chin_middle', 'chin_left',
- 'left_cheek_28', 'left_cheek_29', 'left_cheek_30', 'left_cheek_31',
- 'left_cheek_32', 'left_cheek_33', 'left_cheek_34']
-
-
- def get_landmark_names(self):
- return self.landmark_names
-
- def get_landmark_colors(self):
- return None
-
- def get_landmark_color_dict(self):
- return None
-
- def detect_faces(self, rgb_image):
- orig_h, orig_w, c = rgb_image.shape
- face_image = rgb_image
- rot_h, rot_w, c = face_image.shape
- # Assumes that the width is smaller than the height, and crop
- # a width x width square image from the top.
- square_face_image = face_image[:rot_w, :, :]
- sqr_h, sqr_w, c = square_face_image.shape
- network_image = cv2.resize(square_face_image, (300, 300))
- # Some magic numbers came from
- # https://www.pyimagesearch.com/2018/02/26/face-detection-with-opencv-and-deep-learning/
- face_image_blob = cv2.dnn.blobFromImage(network_image, 1.0, (300, 300), (104.0, 177.0, 123.0))
- self.head_detection_model.setInput(face_image_blob)
- face_detections = self.head_detection_model.forward()[0,0,:,:]
- confidence_mask = face_detections[:, 2] > self.face_confidence_threshold
- face_detections = face_detections[confidence_mask]
- coordinates = face_detections[:, 3:7]
- # Scale and rotate coordinates to the original image
- coordinates = coordinates * np.array([sqr_w, sqr_h, sqr_w, sqr_h])
-
- face_id = 0
- boxes = []
-
- for x0, y0, x1, y1 in coordinates:
- orig_y0 = y0
- orig_y1 = y1
- orig_x0 = x0
- orig_x1 = x1
- face_id += 1
- bounding_box = [orig_x0, orig_y0, orig_x1, orig_y1]
- boxes.append(bounding_box)
-
- return boxes
-
-
- def get_sub_image(self, rgb_image, bounding_box, enlarge_box=True, enlarge_scale=1.15):
- if enlarge_box:
- scale = enlarge_scale
- orig_h, orig_w, c = rgb_image.shape
-
- x0 = bounding_box[0]
- y0 = bounding_box[1]
- x1 = bounding_box[2]
- y1 = bounding_box[3]
-
- m_x = (x1 + x0) / 2.0
- m_y = (y1 + y0) / 2.0
-
- b_w = x1 - x0
- b_h = y1 - y0
-
- b_w = scale * b_w
- b_h = scale * b_h
-
- x0 = int(round(m_x - (b_w/2.0)))
- x1 = int(round(m_x + (b_w/2.0)))
- y0 = int(round(m_y - (b_h/2.0)))
- y1 = int(round(m_y + (b_h/2.0)))
-
- x0 = max(0, x0)
- x1 = min(orig_w, x1)
- y0 = max(0, y0)
- y1 = min(orig_h, y1)
- else:
- x0 = int(round(bounding_box[0]))
- y0 = int(round(bounding_box[1]))
- x1 = int(round(bounding_box[2]))
- y1 = int(round(bounding_box[3]))
-
- actual_bounding_box = [x0, y0, x1, y1]
- image_to_crop = rgb_image
- sub_image = image_to_crop[y0:y1, x0:x1, :]
- return sub_image, actual_bounding_box
-
-
- def estimate_head_pose(self, rgb_image, bounding_box, enlarge_box=True, enlarge_scale=1.15):
- face_crop_image, actual_bounding_box = self.get_sub_image(rgb_image, bounding_box, enlarge_box=enlarge_box, enlarge_scale=enlarge_scale)
- sqr_h, sqr_w, c = face_crop_image.shape
-
- if (sqr_h > 0) and (sqr_w > 0):
- head_pose_image_blob = cv2.dnn.blobFromImage(face_crop_image,
- size=(60, 60),
- swapRB=False,
- crop=False,
- ddepth=cv2.CV_32F)
- self.head_pose_model.setInput(head_pose_image_blob)
- head_pose_out = self.head_pose_model.forward(['angle_r_fc', 'angle_p_fc', 'angle_y_fc'])
- rpy = head_pose_out
- roll = rpy[0][0][0]
- pitch = rpy[1][0][0]
- yaw = rpy[2][0][0]
- pitch = pitch * np.pi/180.0
- roll = roll * np.pi/180.0
- yaw = yaw * np.pi/180.0
-
- return yaw, pitch, roll
-
- return None, None, None
-
-
- def detect_facial_landmarks(self, rgb_image, bounding_box, enlarge_box=True, enlarge_scale=1.15):
- face_crop_image, actual_bounding_box = self.get_sub_image(rgb_image, bounding_box, enlarge_box=enlarge_box, enlarge_scale=enlarge_scale)
- sqr_h, sqr_w, c = face_crop_image.shape
-
- if (sqr_h > 0) and (sqr_w > 0):
- landmarks_image_blob = cv2.dnn.blobFromImage(face_crop_image,
- size=(60, 60),
- swapRB=False,
- crop=False,
- ddepth=cv2.CV_32F)
- self.landmarks_model.setInput(landmarks_image_blob)
- landmarks_out = self.landmarks_model.forward()
-
- s = landmarks_out.shape
- out = np.reshape(landmarks_out[0], (s[1]//2, 2))
- x0, y0, x1, y1 = actual_bounding_box
-
- landmarks = {}
- for n, v in enumerate(out):
- x = int(round((v[0] * sqr_w) + x0))
- y = int(round((v[1] * sqr_h) + y0))
- name = self.landmark_names[n]
- landmarks[name] = (x,y)
-
- return landmarks, self.landmark_names.copy()
- return None, None
-
-
- def draw_bounding_box(self, image, bounding_box):
- x0 = int(round(bounding_box[0]))
- y0 = int(round(bounding_box[1]))
- x1 = int(round(bounding_box[2]))
- y1 = int(round(bounding_box[3]))
- color = (0, 0, 255)
- thickness = 2
- cv2.rectangle(image, (x0, y0), (x1, y1), color, thickness)
-
-
- def draw_head_pose(self, image, yaw, pitch, roll, bounding_box):
- x0, y0, x1, y1 = bounding_box
- face_x = (x1 + x0) / 2.0
- face_y = (y1 + y0) / 2.0
- #
- # opencv uses right-handed coordinate system
- # x points to the right of the image
- # y points to the bottom of the image
- # z points into the image
- #
-
- h, w, c = image.shape
- camera_center = (w/2.0, h/2.0)
- #For rendering with an unknown camera
- focal_length = 50.0
- camera_matrix = np.array([[focal_length, 0.0, camera_center[0]],
- [0.0, focal_length, camera_center[1]],
- [0.0, 0.0, 1.0]])
- face_translation = np.array([0.0, 0.0, 3000.0])
- distortion_coefficients = np.array([0.0, 0.0, 0.0, 0.0])
- # negate the directions of the y and z axes
- axes = np.array([[2000.0, 0.0, 0.0 ],
- [0.0, -2000.0, 0.0 ],
- [0.0, 0.0, -2000.0],
- [0.0, 0.0, 0.0 ]])
- head_ypr = np.array([-yaw, pitch, roll])
- rotation_mat = Rotation.from_euler('yxz', head_ypr).as_dcm()
- rotation_vec, jacobian = cv2.Rodrigues(rotation_mat)
- image_points, jacobian = cv2.projectPoints(axes, rotation_vec, face_translation, camera_matrix, distortion_coefficients)
- face_pix = np.array([face_x, face_y])
-
- origin = image_points[3].ravel()
- x_axis = (image_points[0].ravel() - origin) + face_pix
- y_axis = (image_points[1].ravel() - origin) + face_pix
- z_axis = (image_points[2].ravel() - origin) + face_pix
-
- p0 = tuple(np.int32(np.round(face_pix)))
- p1 = tuple(np.int32(np.round(x_axis)))
- cv2.line(image, p0, p1, (0, 0, 255), 2)
- p1 = tuple(np.int32(np.round(y_axis)))
- cv2.line(image, p0, p1, (0, 255, 0), 2)
- p1 = tuple(np.int32(np.round(z_axis)))
- cv2.line(image, p0, p1, (255, 0, 0), 2)
-
-
- def draw_landmarks(self, image, landmarks):
- for name, xy in landmarks.items():
- x = xy[0]
- y = xy[1]
- if 'mouth' in name:
- color = (255, 0, 0)
- elif 'nose' in name:
- color = (0, 255, 0)
- elif 'eyebrow' in name:
- color = (0, 0, 0)
- elif 'right_eye' in name:
- color = (255, 255, 0)
- elif 'left_eye' in name:
- color = (0, 255, 255)
- elif 'chin' in name:
- color = (255, 0, 255)
- else:
- color = (0, 0, 255)
- cv2.circle(image, (x,y), 2, color, 1)
- font_scale = 1.0
- line_color = [0, 0, 0]
- line_width = 1
- font = cv2.FONT_HERSHEY_PLAIN
-
-
- def apply_to_image(self, rgb_image, draw_output=False):
- if draw_output:
- output_image = rgb_image.copy()
- else:
- output_image = None
-
- heads = []
- boxes = self.detect_faces(rgb_image)
- facial_landmark_names = self.landmark_names.copy()
- for bounding_box in boxes:
- if draw_output:
- self.draw_bounding_box(output_image, bounding_box)
- yaw, pitch, roll = self.estimate_head_pose(rgb_image, bounding_box, enlarge_box=True, enlarge_scale=1.15)
- if yaw is not None:
- ypr = (yaw, pitch, roll)
- if draw_output:
- self.draw_head_pose(output_image, yaw, pitch, roll, bounding_box)
- else:
- ypr = None
- landmarks, landmark_names = self.detect_facial_landmarks(rgb_image, bounding_box, enlarge_box=True, enlarge_scale=1.15)
- if (landmarks is not None) and draw_output:
- self.draw_landmarks(output_image, landmarks)
- heads.append({'box':bounding_box, 'ypr':ypr, 'landmarks':landmarks})
-
- return heads, output_image
|
