车辆检测

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


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

    model_xml = "/home/bhc/BHC/model/intel/vehicle-detection-adas-0002/FP16/vehicle-detection-adas-0002.xml"
    model_bin = "/home/bhc/BHC/model/intel/vehicle-detection-adas-0002/FP16/vehicle-detection-adas-0002.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("2.mp4")
    exec_net = ie.load_network(network=net, device_name="CPU")
    ret, frame = cap.read()
    cars = 0
    while True:
        ret, frame = cap.read()
        if ret is not True:
            break
        mask = np.zeros_like(frame)                                                 #mask：相同shape和type的array，全部为0值
        mh, mw, mc = mask.shape
        cv.line(mask, (0, mh//2), (mw, mh//2), (255, 255, 255), 3, 8, 0)            #mask：中间画线
        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]:                                                       #(1, 1, N, 7)
            if obj[2] > 0.5:                                                        #[image_id, label, conf, x_min, y_min, x_max, y_max]
                xmin = int(obj[3] * iw)
                ymin = int(obj[4] * ih)
                xmax = int(obj[5] * iw)
                ymax = int(obj[6] * ih)
                cv.rectangle(frame, (xmin, ymin), (xmax, ymax), (0, 255, 255), 2, 8)
                cx = xmin + (xmax - xmin) // 2
                cy = ymin + (ymax - ymin) // 2
                cv.circle(mask, (cx, cy), 3, (255, 255, 255), 3, 8, 0)              #mask：画圆圈
                cv.putText(frame, str(obj[2]), (xmin, ymin), cv.FONT_HERSHEY_PLAIN, 1.0, (0, 0, 255), 1)
        cv.putText(frame, "infer time(ms): %.3f, FPS: %.2f"%(inf_end*1000, 1/(inf_end+0.0001)), (10, 50),
                   cv.FONT_HERSHEY_SIMPLEX, 1.0, (255, 0, 255), 2, 8)
        contours, hiearchy = cv.findContours(mask[:, :, 0], cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)   #mask：寻找轮廓
        for cnt in range(len(contours)):
            bx, by, bw, bh = cv.boundingRect(contours[cnt])                                             #mask：寻找覆盖轮廓的矩形
            b_cx = bx + bw // 2
            b_cy = by + bh // 2
            dy = frame.shape[0] // 2 - b_cy                                                             #mask：寻找轮廓矩形（车）是否过了中间线
            if 0 < dy < 15:
                cars += 1
        cv.imshow("Pedestrian Detection", frame)
        cv.imshow("motion mask", mask)
        print("number of cars: ", cars)
        c = cv.waitKey(1)
        if c == 27:
            break
    cv.waitKey(0)
    cv.destroyAllWindows()


if __name__ == "__main__":
    ssd_video_demo()