import multiprocessing import numpy as np import os from openvino.inference_engine import IECore from .. import _utils as utils UNIT_LOWER_LIMIT = 0 UNIT_UPPER_LIMIT = 1 FD_NAME = "detection" LM_NAME = "landmarks" QL_NAME = "quality" FR_NAME = "recognition" AT_NAME = "attributes" MD_NAME = "mask" BIN_EXT = ".bin" XML_EXT = ".xml" class Engine: def __init__(self, models_dir, settings): ie_core = IECore() num_threads = multiprocessing.cpu_count() try: num_threads = min( num_threads, max(int(os.getenv("PV_OPENVINO_THREADS_LIMIT")), 1) ) except (TypeError, ValueError): pass ie_core.set_config({"CPU_THREADS_NUM": str(num_threads)}, "CPU") ( fd_model_type, lm_model_type, ql_model_type, fr_model_type, at_model_type, md_model_type, ) = utils.get_model_types(settings) fd_net = ie_core.read_network( model=os.path.join(models_dir, FD_NAME, fd_model_type, FD_NAME + XML_EXT), weights=os.path.join(models_dir, FD_NAME, fd_model_type, FD_NAME + BIN_EXT), ) self.fd_input_name = next(iter(fd_net.input_info)) self.fd_input_shape = utils.read_fd_input_shape(models_dir, fd_model_type) self.fd_bboxes_name = "bboxes" self.fd_scores_name = "scores" self.fd_select_idxs_name = "selected_indices" self.fd_net = ie_core.load_network(network=fd_net, device_name="CPU") lm_net = ie_core.read_network( model=os.path.join(models_dir, LM_NAME, lm_model_type, LM_NAME + XML_EXT), weights=os.path.join(models_dir, LM_NAME, lm_model_type, LM_NAME + BIN_EXT), ) self.lm_input_name = next(iter(lm_net.input_info)) self.lm_input_shape = utils.read_lm_input_shape(models_dir) self.lm_landmarks_name = "landmarks" self.lm_net = ie_core.load_network(network=lm_net, device_name="CPU") ql_net = ie_core.read_network( model=os.path.join(models_dir, QL_NAME, ql_model_type, QL_NAME + XML_EXT), weights=os.path.join(models_dir, QL_NAME, ql_model_type, QL_NAME + BIN_EXT), ) self.ql_input_name = next(iter(ql_net.input_info)) self.ql_input_shape = utils.read_lm_input_shape(models_dir) self.ql_qualities_name = "qualities" self.ql_acceptabilities_name = "acceptabilities" self.ql_net = ie_core.load_network(network=ql_net, device_name="CPU") fr_net = ie_core.read_network( model=os.path.join(models_dir, FR_NAME, fr_model_type, FR_NAME + XML_EXT), weights=os.path.join(models_dir, FR_NAME, fr_model_type, FR_NAME + BIN_EXT), ) self.fr_input_name = next(iter(fr_net.input_info)) self.fr_input_shape = utils.read_fr_input_shape(models_dir) self.fr_output_name = next(iter(fr_net.outputs)) self.fr_output_shape = utils.read_fr_output_shape(models_dir) self.fr_net = ie_core.load_network(network=fr_net, device_name="CPU") at_net = ie_core.read_network( model=os.path.join(models_dir, AT_NAME, at_model_type, AT_NAME + XML_EXT), weights=os.path.join(models_dir, AT_NAME, at_model_type, AT_NAME + BIN_EXT), ) self.at_input_name = next(iter(at_net.input_info)) self.at_input_shape = utils.read_at_input_shape(models_dir) self.at_net = ie_core.load_network(network=at_net, device_name="CPU") if "mask" in settings: md_model_path = settings["mask"]["models_dir"] md_net = ie_core.read_network( model=os.path.join(md_model_path, md_model_type, MD_NAME + XML_EXT), weights=os.path.join(md_model_path, md_model_type, MD_NAME + BIN_EXT), ) self.md_input_name = next(iter(md_net.input_info)) self.md_input_shape = md_net.input_info[ self.md_input_name ].input_data.shape[2:] self.md_net = ie_core.load_network(network=md_net, device_name="CPU") self.mask_enabled = True else: self.mask_enabled = False def predict_bounding_boxes(self, np_imgs): """ Args: np_imgs: (list) list of images loaded in numpy, of format (1, H, W, C) Returns: bboxes: (list) list containing arrays of bboxes for each image in order [x1, y1, x2, y2], scaled between 0, 1 confs: (list) list containing arrays of confidences scores of the faces for each image """ all_bboxes, all_scores, all_face_counts = [], [], [] np_imgs = np.transpose(np_imgs, (0, 3, 1, 2)) for np_img in np_imgs: ie_out = self.fd_net.infer(inputs={self.fd_input_name: np_img}) bboxes = ie_out[self.fd_bboxes_name] scores = ie_out[self.fd_scores_name] select_idxs = ie_out[self.fd_select_idxs_name] # keep select_idxs until we see -1 i = 0 for idx in select_idxs[:, 0]: if idx == -1: break i += 1 select_idxs = select_idxs[:i] # filter bboxes and scores based on select_idxs for batch_idx, class_idx, idx in select_idxs: all_bboxes.append(bboxes[batch_idx][idx]) all_scores.append(scores[batch_idx][class_idx][idx].item()) all_face_counts.append(len(select_idxs)) img_idxs = [] for img, num in enumerate(all_face_counts): img_idxs += [img] * num return all_bboxes, all_scores, img_idxs def predict_landmarks(self, np_imgs): """ Args: np_imgs: (list) list of imgages loaded in numpy of format (1, C, H, W) Returns: qualities: (numpy array) qualities value between 0 and 1 lmks: (numpy array) landmarks in the shape of (N, 5, 2) """ np_imgs = np.transpose(np_imgs, (0, 3, 1, 2)) landmarks = [] for np_img in np_imgs: ie_out = self.lm_net.infer(inputs={self.lm_input_name: np_img}) lmks = np.squeeze(ie_out[self.lm_landmarks_name]) landmarks.append(lmks) return np.asarray(landmarks) def get_qualities(self, np_imgs): """ Args: np_imgs: (list) list of imgages loaded in numpy of format (1, C, H, W) Returns: qualities: (numpy array) qualities value between 0 and 1 """ np_imgs = np.transpose(np_imgs, (0, 3, 1, 2)) qualities, acceptabilities = [], [] for np_img in np_imgs: ie_out = self.ql_net.infer(inputs={self.ql_input_name: np_img}) quality = np.squeeze(ie_out[self.ql_qualities_name]) qualities.append(quality) acceptability = np.squeeze(ie_out[self.ql_acceptabilities_name]) acceptabilities.append(acceptability) return ( np.clip(qualities, UNIT_LOWER_LIMIT, UNIT_UPPER_LIMIT), np.clip(acceptabilities, UNIT_LOWER_LIMIT, UNIT_UPPER_LIMIT), ) def predict_embeddings(self, np_imgs): """ Args: np_imgs: (list) list of images loaded in numpy of format (1, C, H, W) Returns: embs: (numpy array) array of embedding arrays """ np_imgs = np.transpose(np_imgs, (0, 3, 1, 2)) embeddings = [] for np_img in np_imgs: ie_out = self.fr_net.infer(inputs={self.fr_input_name: np_img}) embeddings.append(np.squeeze(ie_out[self.fr_output_name])) return np.asarray(embeddings) def predict_attributes(self, np_imgs): """ Args: np_img: (numpy array) img loaded in numpy of format (1, C, H, W) Returns: ages: (numpy array) age probabilities in the shape of (N, 1, 7) genders: (numpy array) gender probabilities in the shape of (N, 1, 2) """ np_imgs = np.transpose(np_imgs, (0, 3, 1, 2)) ages, genders = [], [] for np_img in np_imgs: ie_out = self.at_net.infer(inputs={self.at_input_name: np_img}) ages.append(ie_out["age_probs"][0]) genders.append(ie_out["gender_probs"][0]) return ages, genders def check_for_masks(self, np_imgs): """ Args: np_img: (numpy array) img loaded in numpy of format (1, C, H, W) Returns: mask_probabilities: (numpy array) mask probabilities in the shape of (N, 1, 4) """ np_imgs = np.transpose(np_imgs, (0, 3, 1, 2)) mask_probabilities = [] for np_img in np_imgs: ie_out = self.md_net.infer(inputs={self.md_input_name: np_img}) mask_probabilities.append(list(ie_out.values())[0][0][0]) return mask_probabilities