Hi there,
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, but it just shows darker image.
Here is my code, please let me know if you found anything wrong in it, otherwise, the model should be updated:
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
Thanks,
Winston
Link Copied
