Предварительно обученная модель Intel OpenVINO UNet-CamVid-Onnx-0001 не предсказывает правильно

Question

Я загрузил предварительно обученную модель Unet Intel из репозитория OpenVINO Model Zoo на Github без каких-либо изменений.

Но похоже, что это не работает, пожалуйста, взгляните на приведенный ниже прогноз справа. Я ожидаю увидеть правильную сегментацию с разными цветовыми отметками для дороги, неба, дерева и т. Д., Но это просто показывает более темное изображение.

Вот мой код, дайте мне знать, если вы обнаружили в нем что-то не так:

      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

Answer 1

Код в порядке. Результатом модели unet-camvid-onnx-0001 являются попиксельные вероятности каждого входного пикселя из 12 классов набора данных CamVid. Значение RGB 12 классов можно найти в этом каталоге:INSTALL_DIR\deployment_tools\open_model_zoo\ data \palettes

Вы можете ссылаться на классы для значения RGB в своем коде.