OpenVINO + OpenCV实现车辆检测与道路分割

- road-segmentation-adas-0001
- vehicle-detection-adas-0002

BG, road, curb, mark， 输出格式[NCHW]=[1x4x512x896]

NCHW = [1x1xNx7],其中N表示检测到boxes数目

# 道路分割
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("D:/images/video/project_video.mp4")
exec_net = ie.load_network(network=net, device_name="CPU")

# 推理道路分割image = cv.resize(frame, (w, h))
image = image.transpose(2, 0, 1)
res = exec_net.infer(inputs={input_blob: [image]})

# 解析道路分割结果
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)

# 车辆检测
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")

# 车辆检测
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]})

# 解析车辆检测结果
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)