72 lines
3.3 KiB
Python
72 lines
3.3 KiB
Python
# Copyright (C) 2020-2022 Intel Corporation
|
|
#
|
|
# SPDX-License-Identifier: MIT
|
|
|
|
import os
|
|
import numpy as np
|
|
from model_loader import ModelLoader
|
|
|
|
class FaceDetectorHandler:
|
|
def __init__(self):
|
|
base_dir = os.path.abspath(os.environ.get("DETECTOR_MODEL_PATH",
|
|
"/opt/nuclio/open_model_zoo/intel/face-detection-0205/FP32"))
|
|
model_xml = os.path.join(base_dir, "face-detection-0205.xml")
|
|
model_bin = os.path.join(base_dir, "face-detection-0205.bin")
|
|
self.model = ModelLoader(model_xml, model_bin)
|
|
|
|
def infer(self, image, threshold):
|
|
infer_res = self.model.infer(image)["boxes"]
|
|
infer_res = infer_res[infer_res[:,4] > threshold]
|
|
|
|
results = []
|
|
faces = []
|
|
h_scale = image.height / 416
|
|
w_scale = image.width / 416
|
|
for face in infer_res:
|
|
xmin = int(face[0] * w_scale)
|
|
ymin = int(face[1] * h_scale)
|
|
xmax = int(face[2] * w_scale)
|
|
ymax = int(face[3] * h_scale)
|
|
confidence = face[4]
|
|
|
|
faces.append(np.array(image)[ymin:ymax, xmin:xmax])
|
|
results.append({
|
|
"confidence": str(confidence),
|
|
"label": "face",
|
|
"points": [xmin, ymin, xmax, ymax],
|
|
"type": "rectangle",
|
|
"attributes": []
|
|
})
|
|
|
|
return results, faces
|
|
|
|
class AttributesExtractorHandler:
|
|
def __init__(self):
|
|
age_gender_base_dir = os.path.abspath(os.environ.get("AGE_GENDER_MODEL_PATH",
|
|
"/opt/nuclio/open_model_zoo/intel/age-gender-recognition-retail-0013/FP32"))
|
|
age_gender_model_xml = os.path.join(age_gender_base_dir, "age-gender-recognition-retail-0013.xml")
|
|
age_gender_model_bin = os.path.join(age_gender_base_dir, "age-gender-recognition-retail-0013.bin")
|
|
self.age_gender_model = ModelLoader(age_gender_model_xml, age_gender_model_bin)
|
|
emotions_base_dir = os.path.abspath(os.environ.get("EMOTIONS_MODEL_PATH",
|
|
"/opt/nuclio/open_model_zoo/intel/emotions-recognition-retail-0003/FP32"))
|
|
emotions_model_xml = os.path.join(emotions_base_dir, "emotions-recognition-retail-0003.xml")
|
|
emotions_model_bin = os.path.join(emotions_base_dir, "emotions-recognition-retail-0003.bin")
|
|
self.emotions_model = ModelLoader(emotions_model_xml, emotions_model_bin)
|
|
self.genders_map = ["female", "male"]
|
|
self.emotions_map = ["neutral", "happy", "sad", "surprise", "anger"]
|
|
|
|
def infer(self, image):
|
|
age_gender_request = self.age_gender_model.async_infer(image)
|
|
emotions_request = self.emotions_model.async_infer(image)
|
|
# Wait until both age_gender and emotion recognition async inferences finish
|
|
while not (age_gender_request.wait(0) == 0 and emotions_request.wait(0) == 0):
|
|
continue
|
|
age = int(np.squeeze(age_gender_request.output_blobs["age_conv3"].buffer) * 100)
|
|
gender = self.genders_map[np.argmax(np.squeeze(age_gender_request.output_blobs["prob"].buffer))]
|
|
emotion = self.emotions_map[np.argmax(np.squeeze(emotions_request.output_blobs['prob_emotion'].buffer))]
|
|
return {"attributes": [
|
|
{"name": "age", "value": str(age)},
|
|
{"name": "gender", "value": gender},
|
|
{"name": "emotion", "value": emotion}
|
|
]}
|