ssd对象检测

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

# 图像识别
def ssd_image_demo():
    ie = IECore()
    for device in ie.available_devices:
        print(device)

    with open('object_detection_classes_coco.txt') as f:
        labels = [line.strip() for line in f.readlines()]

    model_xml = "ssdv2_graph.xml"
    model_bin = "ssdv2_graph.bin"

    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            #输入BCHW
    print(n, c, h, w)

    src = cv.imread("001.jpg")
    image = cv.resize(src, (w, h))
    image = image.transpose(2, 0, 1)                                    #HWC->CHW

    exec_net = ie.load_network(network=net, device_name="CPU")
    res = exec_net.infer(inputs={input_blob:[image]})

    ih, iw, ic = src.shape
    res = res[out_blob]                                                 #输出(1, 1, N, 7),
    for obj in res[0][0]:
        if obj[2] > 0.25:                                               #[image_id, label, conf, x_min, y_min, x_max, y_max]
            index = int(obj[1])-1
            print(index)
            xmin = int(obj[3] * iw)
            ymin = int(obj[4] * ih)
            xmax = int(obj[5] * iw)
            ymax = int(obj[6] * ih)
            print(labels[index])
            cv.rectangle(src, (xmin, ymin), (xmax, ymax), (0, 255, 255), 2, 8)
            cv.putText(src, labels[index] + str(obj[2]), (xmin, ymin), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1, 8)
    cv.imshow("SSD Object Detection", src)
    cv.waitKey(0)

# 视频识别
def ssd_video_demo():
    ie = IECore()
    for device in ie.available_devices:
        print(device)

    with open('object_detection_classes_coco.txt') as f:
        labels = [line.strip() for line in f.readlines()]

    model_xml = "ssdv2_graph.xml"
    model_bin = "ssdv2_graph.bin"

    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("1.mp4")
    exec_net = ie.load_network(network=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.25:
                index = int(obj[1])-1
                xmin = int(obj[3] * iw)
                ymin = int(obj[4] * ih)
                xmax = int(obj[5] * iw)
                ymax = int(obj[6] * ih)
                print(labels[index])
                cv.rectangle(frame, (xmin, ymin), (xmax, ymax), (0, 255, 255), 2, 8)
                cv.putText(frame, "infer time(ms): %.3f"%(inf_end*1000), (50, 50), cv.FONT_HERSHEY_SIMPLEX, 1.0, (255, 0, 255),
                           2, 8)
                cv.putText(frame, labels[index] + str(obj[2]), (xmin, ymin), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1, 8)
        cv.imshow("SSD Object Detection", frame)
        c = cv.waitKey(1)
        if c == 27:
            break


if __name__ == "__main__":
    ssd_video_demo()