英特尔OpenVINO预训模型UNET-CamVid-Onnx-0001预测不正确

问题描述

我从OpenVINO Model Zoo Github Repo下载了英特尔pre-trained Unet model，未做任何修改。

但它似乎不起作用，请看看下面右侧的预测。 我期待看到正确的分割与不同的颜色标记的道路，天空，树木等…但它只是展示了更黑暗的形象。

这是我的代码，如果您发现其中有什么错误，请告诉我：

from logging import exception
import cv2
import numpy as np
from openvino.inference_engine import IECore

class ColorMap:        
    SKY=[28,51,71]
    BUILDING=[28,28,28]
    POLE=[60,60,60]
    ROAD=[50,25,50]
    PAVEMENT=[95,14,91]
    FENCE=[74,60,60]
    VEHICLE=[0,0,56]
    PEDESTRIAN=[86,8,23]
    BIKE=[47,5,13]
    UNLABELED = [17,18,21]
    TREE=[40,40,61]
    SIGNSYMBOL=[86,86,0]
    COLORS = []
    COLORS_BGR = []
    COLOR_MAP = {}

    # the sequence of colors in this arrar matters!!! as it maps to the prediction classes    
    COLORS.append(SKY)
    COLORS.append(BUILDING)
    COLORS.append(POLE)
    COLORS.append(ROAD)
    COLORS.append(PAVEMENT)
    COLORS.append(TREE)
    COLORS.append(SIGNSYMBOL)
    COLORS.append(FENCE)
    COLORS.append(VEHICLE)
    COLORS.append(PEDESTRIAN)
    COLORS.append(BIKE)
    COLORS.append(UNLABELED)

    for color in COLORS:
        np_color = np.array(color)
        COLORS_BGR.append(np_color[[2,1,0]])

    def crop_to_square(frame):
        height = frame.shape[0]
        width  = frame.shape[1]
        delta = int((width-height) / 2)
        return frame[0:height, delta:width-delta]


model_xml = 'unet-camvid-onnx-0001.xml'
model_bin = "unet-camvid-onnx-0001.bin"
shape = (480, 368)

ie = IECore()
print("Available devices:", ie.available_devices)
net = ie.read_network(model=model_xml, weights=model_bin)
input_blob = next(iter(net.input_info))
# You can select device_name="CPU" to run on CPU
# exec_net = ie.load_network(network=net, device_name='MYRIAD')
exec_net = ie.load_network(network=net, device_name='CPU')

# Get video from the computers webcam
cap = cv2.VideoCapture('/media/winstonfan/Workspace/Learning/Github/depthai/videos/CamVid.mp4')

while cap.isOpened():
    ret, raw_image = cap.read()
    if not ret:
        continue
    image = crop_to_square(raw_image)
    image = cv2.resize(image, shape)
    cv2.imshow('Video ', image)
    image = image.astype(np.float32)
    image = np.expand_dims(image, axis=0)
    image = image.transpose((0, 3, 1, 2))
    image = image / 127.5 - 1.0

    # Do the inference on the MYRIAD device
    output = exec_net.infer(inputs={input_blob: image})
    output = np.squeeze(output['206'])

    data = np.argmax(output, axis=0)
    if data.shape != (368, 480):
        raise exception('unexpected shape of data from decode() method in handler.py');

    class_colors = ColorMap.COLORS
    class_colors = np.asarray(class_colors, dtype=np.uint8)

    output_colors = np.take(class_colors, data, axis=0)
    max_value = output_colors.max()

    output_colors = (output_colors /(max_value/2.0) - 1.0).astype(np.float32)

    sqz_img = np.moveaxis(np.squeeze(image), 0, 2)

    overlayed = cv2.addWeighted(sqz_img, 1, output_colors, 0.2, 0)

    cv2.imshow('Output', overlayed)

    if cv2.waitKey(1) == ord('q'):
        break

---

解决方案

代码似乎没有问题。Unet-camvid-onnx-0001的模型输出是CamVid数据集的12类每个输入像素的每像素概率。这12个类的RGB值可以在以下目录中找到： INSTALL_DIRdeployment_toolsopen_model_zoodatapalettes

您可以在代码中引用RGB值的类。