openvino系列 4. 图像分割基本案例
这里介绍一个基本的图像分割的例子,其实不管是图像分割还是物体检测,都基本遵循下面三个步骤:
- 首先,我们需要读取模型(
ie.read_model
)并且编译(ie.compile_model
); - 第二步,我们读取图片,并且reshape其大小以符合模型的输入;
- 第三部,模型推理(
compiled_model([input_image])[output_layer_ir]
)。得到的结果的尺寸和模型的输出尺寸相符。
环境描述:
- 本案例运行环境:Win10
- IDE:VSCode
- openvino版本:2022.1
- 代码链接,
2-basic-segmentation-detection-example
文章目录
1 图像分割
这里介绍如何使用OpenVINO进行图像分割。
我们使用来自 Open Model Zoo 的预训练模型 road-segmentation-adas-0001。该模型识别四类:背景、道路、路缘和标记。
- 首先,我们需要读取模型(
ie.read_model
)并且编译(ie.compile_model
); - 第二步,我们读取图片,并且reshape其大小以符合模型的输入;
- 第三部,模型推理(
compiled_model([input_image])[output_layer_ir]
)。得到的结果的尺寸和模型的输出尺寸相符。
代码如下:
import cv2
import matplotlib.pyplot as plt
import numpy as np
import sys
from openvino.runtime import Core
print("1 Load the model.")
ie = Core()
model = ie.read_model(model="model/road-segmentation-adas-0001.xml")
compiled_model = ie.compile_model(model=model, device_name="CPU")
input_layer_ir = compiled_model.input(0)
output_layer_ir = compiled_model.output(0)
print("- Input layer info: {}".format(input_layer_ir))
print("- Output layer info: {}".format(output_layer_ir))
print("2 Load the image, and reshape to the same size as model input.")
# The segmentation network expects images in BGR format
image = cv2.imread("data/empty_road_mapillary.jpg")
rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image_h, image_w, _ = image.shape
print("- Image original shape: {0}".format(image.shape))
# N,C,H,W = batch size, number of channels, height, width
N, C, H, W = input_layer_ir.shape
# OpenCV resize expects the destination size as (width, height)
resized_image = cv2.resize(image, (W, H))
# reshape to network input shape
input_image = np.expand_dims(
resized_image.transpose(2, 0, 1), 0
)
print("- Image size reshape into: {0}".format(input_image.shape))
print("3 Inference.")
# Run the inference
result = compiled_model([input_image])[output_layer_ir]
print("- Shape of inference result: {0}".format(result.shape))
# Prepare data for visualization
segmentation_mask = np.argmax(result, axis=1)
print("- Shape of segmentation mask: {0}".format(segmentation_mask.shape))
plt.imshow(segmentation_mask.transpose(1, 2, 0))
Terminal打印如下:
1 Load the model.
- Input layer info: <ConstOutput: names[data] shape{1,3,512,896} type: f32>
- Output layer info: <ConstOutput: names[L0317_ReWeight, L0317_ReWeight_SoftMax] shape{1,4,512,896} type: f32>
2 Load the image, and reshape to the same size as model input.
- Image original shape: (1080, 1920, 3)
- Image size reshape into: (1, 3, 512, 896)
3 Inference.
- Shape of inference result: (1, 4, 512, 896)
- Shape of segmentation mask: (1, 512, 896)
最后,我们可以可视化图像分割之后的效果。代码如下:
def segmentation_map_to_image(
result: np.ndarray, colormap: np.ndarray, remove_holes: bool = False
) -> np.ndarray:
"""
Convert network result of floating point numbers to an RGB image with
integer values from 0-255 by applying a colormap.
:param result: A single network result after converting to pixel values in H,W or 1,H,W shape.
:param colormap: A numpy array of shape (num_classes, 3) with an RGB value per class.
:param remove_holes: If True, remove holes in the segmentation result.
:return: An RGB image where each pixel is an int8 value according to colormap.
"""
if len(result.shape) != 2 and result.shape[0] != 1:
raise ValueError(
f"Expected result with shape (H,W) or (1,H,W), got result with shape {result.shape}"
)
if len(np.unique(result)) > colormap.shape[0]:
raise ValueError(
f"Expected max {colormap[0]} classes in result, got {len(np.unique(result))} "
"different output values. Please make sure to convert the network output to "
"pixel values before calling this function."
)
# reshape result from [1,H,W] to [H,W]
elif result.shape[0] == 1:
result = result.squeeze(0)
result = result.astype(np.uint8)
contour_mode = cv2.RETR_EXTERNAL if remove_holes else cv2.RETR_TREE
mask = np.zeros((result.shape[0], result.shape[1], 3), dtype=np.uint8)
for label_index, color in enumerate(colormap):
label_index_map = result == label_index
label_index_map = label_index_map.astype(np.uint8) * 255
contours, hierarchies = cv2.findContours(
label_index_map, contour_mode, cv2.CHAIN_APPROX_SIMPLE
)
cv2.drawContours(
mask,
contours,
contourIdx=-1,
color=color.tolist(),
thickness=cv2.FILLED,
)
return mask
# Define colormap, each color represents a class
colormap = np.array([[68, 1, 84], [48, 103, 141], [53, 183, 120], [199, 216, 52]])
# Define the transparency of the segmentation mask on the photo
alpha = 0.3
# Use function from notebook_utils.py to transform mask to an RGB image
mask = segmentation_map_to_image(segmentation_mask, colormap)
resized_mask = cv2.resize(mask, (image_w, image_h))
# Create image with mask put on
image_with_mask = cv2.addWeighted(resized_mask, alpha, rgb_image, 1 - alpha, 0)
# Define titles with images
data = {"Base Photo": rgb_image, "Segmentation": mask, "Masked Photo": image_with_mask}
# Create subplot to visualize images
fig, axs = plt.subplots(1, len(data.items()), figsize=(15, 10))
# Fill subplot
for ax, (name, image) in zip(axs, data.items()):
ax.axis('off')
ax.set_title(name)
ax.imshow(image)
# Display image
plt.show(fig)