道路分割+车辆识别

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

model_xml = "/home/bhc/BHC/model/intel/road-segmentation-adas-0001/FP16/road-segmentation-adas-0001.xml"
model_bin = "/home/bhc/BHC/model/intel/road-segmentation-adas-0001/FP16/road-segmentation-adas-0001.bin"

vehicel_xml = "/home/bhc/BHC/model/intel/vehicle-detection-adas-0002/FP16/vehicle-detection-adas-0002.xml"
vehicel_bin = "/home/bhc/BHC/model/intel//vehicle-detection-adas-0002/FP16/vehicle-detection-adas-0002.bin"


def ssd_video_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("2.mp4")
    exec_net = ie.load_network(network=net, device_name="CPU")

    # 车辆检测
    vnet = ie.read_network(model=vehicel_xml, weights=vehicel_bin)
    vehicle_input_blob = next(iter(vnet.input_info))
    vehicle_out_blob = next(iter(vnet.outputs))

    vn, vc, vh, vw = vnet.input_info[vehicle_input_blob].input_data.shape
    print(n, c, h, w)
    vehicle_exec_net = ie.load_network(network=vnet, device_name="CPU")

    while True:
        ret, frame = cap.read()
        if ret is not True:
            break
        # 车辆检测
        inf_start = time.time()
        image = cv.resize(frame, (vw, vh))
        image = image.transpose(2, 0, 1)
        vec_res = vehicle_exec_net.infer(inputs={vehicle_input_blob:[image]})
        # 推理道路分割
        image = cv.resize(frame, (w, h))
        image = image.transpose(2, 0, 1)
        res = exec_net.infer(inputs={input_blob: [image]})

        # 解析车辆检测结果
        ih, iw, ic = frame.shape
        vec_res = vec_res[vehicle_out_blob]
        for obj in vec_res[0][0]:
            if obj[2] > 0.5:
                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, 0, 255), 2, 8)
                cv.putText(frame, str(obj[2]), (xmin, ymin), cv.FONT_HERSHEY_PLAIN, 1.0, (0, 0, 255), 1)

        # 解析道路分割结果
        res = res[out_blob]
        res = np.squeeze(res, 0)
        res = res.transpose(1, 2, 0)
        res = np.argmax(res, 2)
        print(res.shape)
        hh, ww = res.shape
        mask = np.zeros((hh, ww, 3), dtype=np.uint8)
        mask[np.where(res > 0)] = (0, 255, 255)
        mask[np.where(res > 1)] = (255, 0, 255)
        mask = cv.resize(mask, (frame.shape[1], frame.shape[0]))
        result = cv.addWeighted(frame, 0.5, mask, 0.5, 0)
        inf_end = time.time() - inf_start
        cv.putText(result, "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)
        # out.write(result)
        cv.imshow("Pedestrian Detection", result)
        c = cv.waitKey(1)
        if c == 27:
            break
    cv.waitKey(0)
    cv.destroyAllWindows()


if __name__ == "__main__":
    ssd_video_demo()