-2

I downloaded the Intel pre-trained Unet model from the OpenVINO Model Zoo Github repo without any modifications.

But it does not seem working, please have a look the below prediction on the right side. I am expecting to see correct segmentation with different color marks for road, sky, tree etc... but it just shows darker image. enter image description here

Here is my code, please let me know if you found anything wrong in it:

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
ArunJose
  • 1,999
  • 1
  • 10
  • 33
Franva
  • 6,565
  • 23
  • 79
  • 144

1 Answers1

0

The code seems fine. The model output of unet-camvid-onnx-0001 is the per-pixel probabilities of each input pixel of the 12 classes of the CamVid dataset. The RGB value of the 12 classes can be found in this directory: INSTALL_DIR\deployment_tools\open_model_zoo\data\palettes

You can refer to the classes for the RGB value in your code.

Rommel_Intel
  • 1,369
  • 1
  • 4
  • 8
  • thanks, I did saw the 12 classes. Appreciated if this post could be kept posted once the model has updated with correct performance. – Franva Aug 11 '21 at 10:55