From c6c58450bb8c871459e25b598f509bb0269228e5 Mon Sep 17 00:00:00 2001 From: Godfrey Tutu Date: Sat, 20 Jun 2020 23:54:27 +0100 Subject: [PATCH] feature: facial_landmarks - Add the code for locating facial landmarks - Ensure the left and right eyes are located and cropped correctly - Return the left and right images --- src/facial_landmarks_detection.py | 129 ++++++++++++++++++++++++++++++ 1 file changed, 129 insertions(+) create mode 100644 src/facial_landmarks_detection.py diff --git a/src/facial_landmarks_detection.py b/src/facial_landmarks_detection.py new file mode 100644 index 0000000..9c60c0a --- /dev/null +++ b/src/facial_landmarks_detection.py @@ -0,0 +1,129 @@ +''' +This is a sample class for a model. You may choose to use it as-is or make any changes to it. +This has been provided just to give you an idea of how to structure your model class. +''' +import os +import sys +import cv2 +from openvino.inference_engine import IENetwork, IECore + +class Model_X: + ''' + Class for the Face Detection Model. + ''' + def __init__(self, model_name, device='CPU', extensions=None): + ''' + TODO: Use this to set your instance variables. + ''' + self.model_name = model_name + self.device = device + self.extensions = extensions + + def load_model(self): + ''' + TODO: You will need to complete this method. + This method is for loading the model to the device specified by the user. + If your model requires any Plugins, this is where you can load them. + ''' + model_bin = self.model_name + '.bin' + model_xml = self.model_name + '.xml' + self.model = IENetwork(model_xml, model_bin) + self.plugin = IECore() + if self.extensions: + self.plugin.add_extension(self.extensions, self.device) + self.exec_net = self.plugin.load_network(self.model, device_name=self.device, num_requests=1) + self.check_model() + + def predict(self, image): + ''' + TODO: You will need to complete this method. + This method is meant for running predictions on the input image. + ''' + frame = self.preprocess_input(image) + self.exec_net.start_async(request_id=0, inputs={ self.input_name: frame }) + if self.exec_net.requests[0].wait(-1) == 0: + return self.exec_net.requests[0].outputs[self.output_name], frame + + def check_model(self): + self.input_name = next(iter(self.model.inputs)) + self.output_name = next(iter(self.model.outputs)) + self.input_shape = self.model.inputs[self.input_name].shape + self.output_shape = self.model.outputs[self.output_name].shape + + def preprocess_input(self, image): + ''' + Before feeding the data into the model for inference, + you might have to preprocess it. This function is where you can do that. + ''' + height, width = self.input_shape[2], self.input_shape[3] + frame = cv2.resize(image, (width, height)) + return frame.transpose((2, 0, 1)) + + + + def preprocess_output(self, outputs): + ''' + Before feeding the output of this model to the next model, + you might have to preprocess the output. This function is where you can do that. + ''' + # print(outputs.shape, outputs) + eyes_coords = { + 'left_eye': [outputs[0][0], outputs[0][1]], + 'right_eye':[outputs[0][2], outputs[0][3]] + } + return eyes_coords + + def get_eyes(self, coords, image): + height = image.shape[0] + width = image.shape[1] + # frame = cv2.resize(frame.transpose((1, 2, 0)), (width, height)) + eyes = {} + for key in coords.keys(): + x_coord = coords[key][0] + y_coord = coords[key][1] + x_min = int(x_coord * width - 30) + y_min = int(y_coord * height - 30) + x_max = int(x_coord * width + 30) + y_max = int(y_coord * height + 30) + croped_image = image[y_min:y_max, x_min:x_max] + eyes[key] = croped_image + return eyes + +def draw_landmarks(coords, frame, image): + height = image.shape[0] + width = image.shape[1] + frame = cv2.resize(frame.transpose((1, 2, 0)), (width, height)) + for key in coords.keys(): + x_coord = coords[key][0] + y_coord = coords[key][1] + frame = cv2.circle(frame, (x_coord*width, y_coord*height), radius=2, color=(0, 0, 225), thickness=3) + frame = cv2.rectangle(frame, + (x_coord*width - 30, y_coord*height - 30), + (x_coord*width + 30, y_coord*height + 30), + color=(0, 0, 225), thickness=3) + x_min = int(x_coord * width - 30) + y_min = int(y_coord * height - 30) + x_max = int(x_coord * width + 30) + y_max = int(y_coord * height + 30) + image = frame[y_min:y_max, x_min:x_max] + cv2.imwrite(key+'.jpg', image) + cv2.imwrite('result.jpg', frame) + +def main(): + CPU_EXTENSION_MAC = '/opt/intel/openvino_2019.3.376/deployment_tools/inference_engine/lib/intel64/libcpu_extension.dylib' + model_name = 'models/intel/landmarks-regression-retail-0009/FP16/landmarks-regression-retail-0009' + image = 'bin/test_image2.png' + model = Model_X(model_name=model_name, device='CPU', extensions=CPU_EXTENSION_MAC) + model.load_model() + image = cv2.imread(image) + pred, frame = model.predict(image) + pred = model.preprocess_output(pred) + eyes = model.get_eyes(pred, image) + for eye, cropped_image in eyes.items(): + cv2.imwrite(eye +'_image.jpg', cropped_image) + # draw_landmarks(pred, frame, image) + +if __name__ == '__main__': + main() + +# py src/facial_landmarks_detection.py