头部姿态评估

from openvino.inference_engine import IECore
import numpy as np
import time
import cv2 as cv

# 人脸检测模型
model_xml = "/home/bhc/BHC/model/intel/face-detection-0200/FP16/face-detection-0200.xml"
model_bin = "/home/bhc/BHC/model/intel/face-detection-0200/FP16/face-detection-0200.bin"
# 头部姿态评估
head_xml = "/home/bhc/BHC/model/intel/head-pose-estimation-adas-0001/FP16/head-pose-estimation-adas-0001.xml"
head_bin = "/home/bhc/BHC/model/intel/head-pose-estimation-adas-0001/FP16/head-pose-estimation-adas-0001.bin"


def face_landmark_demo():
    ie = IECore()
    for device in ie.available_devices:
        print(device)

    net = ie.read_network(model=model_xml, weights=model_bin)
    input_blob = next(iter(net.input_info))
    out_blob = next(iter(net.outputs))

    n, c, h, w = net.input_info[input_blob].input_data.shape
    print(n, c, h, w)

    # cap = cv.VideoCapture(0)
    cap = cv.VideoCapture("1.mp4")
    exec_net = ie.load_network(network=net, device_name="CPU")

    head_net = ie.read_network(model=head_xml, weights=head_bin)
    em_input_blob = next(iter(head_net.input_info))
    head_it = iter(head_net.outputs)
    head_out_blob1 = next(head_it)  # angle_p_fc                                #多个输出 （1, 1）yaw角度
    head_out_blob2 = next(head_it)  # angle_r_fc                                #（1, 1）pitch角度
    head_out_blob3 = next(head_it)  # angle_y_fc                                #（1, 1）roll角度
    print(head_out_blob1, head_out_blob2, head_out_blob3)

    en, ec, eh, ew = head_net.input_info[em_input_blob].input_data.shape
    print(en, ec, eh, ew)

    em_exec_net = ie.load_network(network=head_net, device_name="CPU")

    while True:
        ret, frame = cap.read()
        if ret is not True:
            break
        image = cv.resize(frame, (w, h))
        image = image.transpose(2, 0, 1)
        inf_start = time.time()
        res = exec_net.infer(inputs={input_blob: [image]})
        inf_end = time.time() - inf_start
        # print("infer time(ms)：%.3f"%(inf_end*1000))
        ih, iw, ic = frame.shape
        res = res[out_blob]
        for obj in res[0][0]:
            if obj[2] > 0.75:
                xmin = int(obj[3] * iw)
                ymin = int(obj[4] * ih)
                xmax = int(obj[5] * iw)
                ymax = int(obj[6] * ih)
                if xmin < 0:
                    xmin = 0
                if ymin < 0:
                    ymin = 0
                if xmax >= iw:
                    xmax = iw - 1
                if ymax >= ih:
                    ymax = ih - 1
                roi = frame[ymin:ymax, xmin:xmax, :]
                roi_img = cv.resize(roi, (ew, eh))
                roi_img = roi_img.transpose(2, 0, 1)
                head_res = em_exec_net.infer(inputs={em_input_blob: [roi_img]})
                angle_p_fc = head_res[head_out_blob1][0][0]
                angle_r_fc = head_res[head_out_blob2][0][0]
                angle_y_fc = head_res[head_out_blob3][0][0]
                head_pose = ""
                if angle_p_fc > 20 or angle_p_fc < -20:
                    head_pose += "pitch, "
                if angle_r_fc > 20 or angle_r_fc < -20:
                    head_pose += "roll, "
                if angle_y_fc > 20 or angle_y_fc < -20:
                    head_pose += "yaw, "
                cv.rectangle(frame, (xmin, ymin), (xmax, ymax), (0, 255, 255), 2, 8)
                cv.putText(frame, head_pose, (xmin, ymin), cv.FONT_HERSHEY_SIMPLEX, 1.0, (255, 0, 255), 2, 8)

        cv.putText(frame, "infer time(ms): %.3f, FPS: %.2f" % (inf_end * 1000, 1/inf_end), (50, 50),
                   cv.FONT_HERSHEY_SIMPLEX, 1.0, (255, 0, 255), 2, 8)
        cv.imshow("Face+emotion Detection", frame)
        c = cv.waitKey(1)
        if c == 27:
            break
    cv.waitKey(0)
    cv.destroyAllWindows()


if __name__ == "__main__":
    face_landmark_demo()