YOLOv8是YOLO(You Only Look Once)物体检测系统的最新版本(v8)。YOLO是一个实时的、一次性的物体检测系统,其目的是在网络的一次前向传递中进行物体检测,使其快速而高效。YOLOv8是以前的YOLO模型的改进版,具有更好的准确性和更快的推理速度。
ONNX(Open Neural Network Exchange)是一种代表深度学习模型的开放格式。要将YOLOv8模型转换为ONNX格式,你需要使用ONNX Runtime这样的工具,它提供了一个API来将不同框架的模型转换为ONNX格式。具体步骤取决于你用来开发和运行 YOLOv8 模型的编程框架和工具。
如何转换?
from ultralytics import YOLOmodel = YOLO('yolov8s-seg.pt')input_height,input_width = 640, 640model.export(format="onnx", imgsz=[input_height,input_width], opset=11)
如何使用ONNX Runtime加载推理?
onnxruntime 是一个开源的运行时引擎,用于执行以Open Neural Network Exchange(ONNX)格式表示的机器学习模型。ONNX是一个代表深度学习模型的开放标准,允许不同框架和工具之间的互操作性。ONNX模型可以从各种深度学习框架中导出,如PyTorch、TensorFlow和Keras,并且可以在不同的硬件平台上执行,如CPU、GPU和FPGA。
onnxruntime 统一的API为在不同的硬件平台和操作系统上执行ONNX模型提供了一个统一的API。它支持广泛的执行提供者,包括CPU、GPU和专门的加速器,如NVIDIA TensorRT和Intel OpenVINO。 还提供对模型优化和量化的支持,以提高模型性能,减少内存和存储需求。onnxruntime
onnxruntime 可用于各种应用,如计算机视觉、自然语言处理和语音识别。它提供了一个高性能和灵活的运行时引擎,可以在生产环境中高效执行深度学习模型。
安装onnx运行时
# For Nvidia GPU computers: pip install onnxruntime-gpu # Otherwise: pip install onnxruntime
import onnxruntimeopt_session = onnxruntime.SessionOptions()opt_session.enable_mem_pattern = Falseopt_session.enable_cpu_mem_arena = Falseopt_session.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL
你提供的代码设置了一个带有一些选项的onnxruntime.SessionOptions对象。让我们一个一个地去看它们:
- enable_mem_pattern:这个选项控制是否启用内存模式优化。当启用时,onnxruntime可以分析模型的内存使用模式,并更有效地分配内存,这可以减少内存使用并提高性能。将此选项设置为 "假",则禁用内存模式优化。
- enable_cpu_mem_arena:这个选项控制是否启用CPU内存竞技场。内存竞技场是onnxruntime用来更有效地管理内存的一种技术。当启用时,onnxruntime会分块分配内存并重用它们,这可以提高性能。将此选项设置为False,则禁用CPU内存arena。
- graph_optimization_level:该选项控制推理过程中对ONNX图进行的优化级别。onnxruntime提供了几种优化级别,从ORT_DISABLE_ALL(禁用所有优化)到ORT_ENABLE_EXTENDED(启用所有可用优化)。在你提供的代码中,使用了ORT_DISABLE_ALL,它禁用了所有优化。
model_path = 'yolov8s-seg.onnx'EP_list = ['CUDAExecutionProvider', 'CPUExecutionProvider']ort_session = onnxruntime.InferenceSession(model_path, providers=EP_list)
你提供的代码设置了一个onnxruntime.InferenceSession对象,用于加载ONNX模型并对其运行推理。让我们来看看所使用的参数:
- model_path:这个参数指定了你要加载的ONNX模型文件的路径。在你提供的例子中,路径被设置为 "model_name.onnx"。
- providers(提供者):该参数指定推理时应使用的执行提供者列表。执行提供者负责执行ONNX模型中的操作,onnxruntime提供了几个内置的执行提供者,包括CPUExecutionProvider和CUDAExecutionProvider(它们分别使用CPU和GPU进行推理)。在你提供的例子中,使用了两个执行提供者的列表:['CUDAExecutionProvider', 'CPUExecutionProvider']。这意味着onnxruntime将首先尝试使用CUDAExecutionProvider,如果它不可用(例如,因为没有GPU),它将退回到CPUExecutionProvider。
模型信息
model_inputs = ort_session.get_modelmeta()classes = eval(model_inputs.custom_metadata_map['names'])num_classes = len(classes)classes
InferenceSession对象的get_modelmeta()方法返回关于模型的输入和输出张量的元数据,以及导出模型时包含的任何自定义元数据。
model_inputs = ort_session.get_inputs()input_names = [model_inputs[i].name for i in range(len(model_inputs))]input_shape = model_inputs[0].shape
InferenceSession对象的get_inputs()方法返回模型的输入张量元数据的列表。在这段代码中,model_inputs通过调用olt_session.get_inputs()被分配了输入张量元数据的列表。
- 然后,input_names通过遍历model_inputs列表并获得每个输入张量的名称属性来分配一个输入张量名称的列表。
- input_shape通过获取model_inputs列表中第一个元素的shape属性来分配第一个输入张量的形状。这假设所有的输入张量都有相同的形状,这对深度学习模型来说通常是这样的。张量元数据对象的形状属性是一个元组,指定了张量的形状。例如,一个代表一批图像的4D张量,其尺寸为[batch_size, channels, height, width],其形状为(batch_size, channels, height, width)。
model_output = ort_session.get_outputs()output_names = [model_output[i].name for i in range(len(model_output))]output_names
InferenceSession对象的get_outputs()方法返回模型的输出张量元数据的列表。在这段代码中,通过调用olt_session.get_outputs(),model_output被分配了输出张量元数据的列表。
- 然后,output_names通过迭代model_output列表并获得每个输出张量的名称属性来分配一个输出张量名称的列表。
输入图像
import cv2from PIL import Image# Read Imageimage = cv2.imread('cat.jpg')# Image shapeimage_height, image_width = image.shape[:2]# Input shapeinput_height, input_width = input_shape[2:]# Convert image bgr to rgbimage_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)# Resize image to input shaperesized = cv2.resize(image_rgb, (input_width, input_height))# Scale input pixel value to 0 to 1input_image = resized / 255.0input_image = input_image.transpose(2,0,1)input_tensor = input_image[np.newaxis, :, :, :].astype(np.float32)input_tensor.shape
我们来看看你提供的代码:
- input_height和input_width:这些变量包含ONNX模型所需的输入张量的高度和宽度。高度和宽度是从input_shape变量中获得的,该变量是先前从ONNX模型中获取的。
- image_rgb:这个变量包含RGB颜色格式的输入图像数据。原始图像数据是用OpenCV的imread函数读取的,它默认以BGR颜色格式读取图像。cvtColor函数被用来将图像数据从BGR颜色格式转换为RGB颜色格式。
- 调整大小:该变量包含调整后的图像数据,其大小与ONNX模型所要求的输入张量的尺寸一致。resize函数用于调整图像数据的大小。
- input_image:该变量包含对图像数据进行归一化和重新排列后的输入张量数据。然后,图像数据的尺寸被重新排列,以符合 ONNX 模型所要求的输入张量格式,即(batch_size, channel, height, width)。
- input_tensor:这个变量包含最终的输入张量数据,它是一个NumPy数组,形状为(batch_size, channel, height, width)。批量大小被设置为1,因为只有一张图片被处理。数组的数据类型是float32。
一旦你准备好了输入的张量数据,你就可以把它传递给onnxruntime.InferenceSession对象的run方法来对模型运行推理。
outputs = ort_session.run(output_names, {input_names[0]: input_tensor})outputs
输出:该变量包含ONNX模型在推理过程中产生的输出数据。运行方法为模型返回一个输出的元组,每个输出以NumPy数组表示。在这种情况下,我们只对第一个输出感兴趣,这就是为什么我们使用[0]来访问它。
为了在模型上运行推理,我们将输入张量数据作为一个字典传递给run方法,其中键是模型的输入节点的名称(在本例中是input_names[0]),而值是输入张量数据。我们还提供了一个输出节点名称(output_names)的列表,以便运行方法知道要返回哪些输出。一旦运行方法被调用,它将对模型进行推理,并以NumPy数组的形式返回输出数据。
提取输出框
box_output = outputs[0]num_masks = 32predictions = np.squeeze(box_output).Tnum_classes = box_output.shape[1] - num_masks - 4
首先,box_output ,假设包含模型的输出,对应于输入图像中检测到的物体的边界盒。np.squeeze() 方法被用来从输出张量中去除任何尺寸为1的维度,有效地将其转换为一个二维numpy数组。然后用.T 方法对张量进行转置,使每一行对应于一个边界盒,每一列对应于边界盒的一个参数。num_masks 变量被假定为模型输出的分割掩码的数量。模型检测到的类的数量是通过从num_masks 和4减去box_output 中的列数来计算的。额外的4列对应于边界盒参数(x,y,w,h)。得到的predictions 数组是一个二维的numpy数组,其形状为(num_boxes, num_classes + 4) ,其中num_boxes 是在输入图像中检测到的边界盒的数量。predictions 的第一列num_classes ,对应于检测到的物体属于每个num_classes 类别的概率。最后四列对应的是边界盒参数(x, y, w, h)
# Filter out object confidence scores below thresholdconf_threshold = 0.25scores = np.max(predictions[:, 4:4+num_classes], axis=1)predictions = predictions[scores > conf_threshold, :]scores = scores[scores > conf_threshold]scores
每个检测到的物体的置信度分数是通过对该物体所有类别的最大概率来获得的。这是通过选择predictions 中对应于类概率的列(predictions[:, 4:4+num_classes] ),并使用np.max() 取每一行的最大值来实现的。然后,scores 变量被赋予这些最大的置信度分数。scores 接下来,通过只选择predictions 中对应的置信度分数高于阈值的行,过滤掉任何置信度分数低于阈值的边界框。最后,scores 变量包含一个一维的小数组,其中包含剩余边界框的置信度分数。
box_predictions = predictions[..., :num_classes+4]mask_predictions = predictions[..., num_classes+4:]box_predictions
box_predictions 变量被分配给predictions 的子集,该子集由第一列num_classes+4 组成,对应于边界盒参数(x, y, w, h) 和类别概率。
mask_predictions 变量被分配给predictions 的子集,该子集由其余列组成,对应于预测的分割掩码。box_predictions 和mask_predictions 都是二维numpy数组,形状分别为(num_boxes, num_classes+4) 和(num_boxes, num_masks) ,其中num_boxes 是应用置信度阈值得到的过滤边界盒的数量,num_classes 是模型检测到的物体类别的数量,num_masks 是模型输出的分割掩码的数量。
# Get the class with the highest confidenceclass_ids = np.argmax(box_predictions[:, 4:], axis=1)class_ids
class_ids 变量被分配为每个边界框的最高置信度分数的列的索引,该索引是通过对包含类别概率的box_predictions 子集(box_predictions[:, 4:] )使用np.argmax() 得到的。
得到的class_ids 数组是一个长度为num_boxes 的一维数字数组,包含每个边界框的预测类标签。请注意,类的ID是零索引的,所以它们的范围从0到1。num_classes - 1
def xywh2xyxy(x): # Convert bounding box (x, y, w, h) to bounding box (x1, y1, x2, y2) y = np.copy(x) y[..., 0] = x[..., 0] - x[..., 2] / 2 y[..., 1] = x[..., 1] - x[..., 3] / 2 y[..., 2] = x[..., 0] + x[..., 2] / 2 y[..., 3] = x[..., 1] + x[..., 3] / 2 return ydef extract_boxes(box_predictions): # Extract boxes from predictions boxes = box_predictions[:, :4] # Scale boxes to original image dimensions boxes = rescale_boxes(boxes, (input_height, input_width), (image_height, image_width)) # Convert boxes to xyxy format boxes = xywh2xyxy(boxes) # Check the boxes are within the image boxes[:, 0] = np.clip(boxes[:, 0], 0, image_width) boxes[:, 1] = np.clip(boxes[:, 1], 0, image_height) boxes[:, 2] = np.clip(boxes[:, 2], 0, image_width) boxes[:, 3] = np.clip(boxes[:, 3], 0, image_height) return boxesdef rescale_boxes(boxes, input_shape, image_shape): # Rescale boxes to original image dimensions input_shape = np.array([input_shape[1], input_shape[0], input_shape[1], input_shape[0]]) boxes = np.divide(boxes, input_shape, dtype=np.float32) boxes *= np.array([image_shape[1], image_shape[0], image_shape[1], image_shape[0]]) return boxes
# Get bounding boxes for each objectboxes = extract_boxes(box_predictions)boxes
extract_boxes 函数从模型中获取边界盒预测值,并以可用于可视化的格式提取盒子。该函数首先从预测中提取方框,然后将其缩放为原始图像尺寸。它将方框从(x, y, w, h)格式转换为(x1, y1, x2, y2)格式。它还检查方框是否在图像的边界内,如果有必要的话,将它们剪掉。该函数以(x1, y1, x2, y2)的格式返回方框。rescale_boxes 函数被用于extract_boxes ,以重新调整方框的大小。它接收(x, y, w, h)格式的方框、模型的输入形状(高度、宽度)和图像形状(高度、宽度),并返回重新缩放到原始图像尺寸的方框。
def nms(boxes, scores, iou_threshold): # Sort by score sorted_indices = np.argsort(scores)[::-1] keep_boxes = [] while sorted_indices.size > 0: # Pick the last box box_id = sorted_indices[0] keep_boxes.append(box_id) # Compute IoU of the picked box with the rest ious = compute_iou(boxes[box_id, :], boxes[sorted_indices[1:], :]) # Remove boxes with IoU over the threshold keep_indices = np.where(ious < iou_threshold)[0] # print(keep_indices.shape, sorted_indices.shape) sorted_indices = sorted_indices[keep_indices + 1] return keep_boxesdef compute_iou(box, boxes): # Compute xmin, ymin, xmax, ymax for both boxes xmin = np.maximum(box[0], boxes[:, 0]) ymin = np.maximum(box[1], boxes[:, 1]) xmax = np.minimum(box[2], boxes[:, 2]) ymax = np.minimum(box[3], boxes[:, 3]) # Compute intersection area intersection_area = np.maximum(0, xmax - xmin) * np.maximum(0, ymax - ymin) # Compute union area box_area = (box[2] - box[0]) * (box[3] - box[1]) boxes_area = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) union_area = box_area + boxes_area - intersection_area # Compute IoU iou = intersection_area / union_area return iou
# Apply non-maxima suppression to suppress weak, overlapping bounding boxesiou_threshold = 0.3indices = nms(boxes, scores, iou_threshold)indices
nms 函数执行非最大值抑制,以去除弱的和重叠的边界框。该函数需要三个参数:
boxes: 一个形状为(N, 4)的numpy数组,包含(x1, y1, x2, y2)格式的N界限框的坐标。scores: 一个形状为(N,)的numpy数组,包含每个N界定框的置信度分数。iou_threshold: 一个浮动值,指定用于重叠边界盒的IoU阈值。
该函数首先按照分数从高到低的顺序对边界盒进行排序。然后,它挑选出分数最高的盒子,并将其添加到要保留的盒子列表中。它计算这个盒子与所有其他盒子的物价指数,并删除任何物价指数大于阈值的盒子。该函数重复这一过程,直到没有箱子可以考虑。
该函数返回一个索引列表,该列表对应于在非最大值抑制后应保留的盒子。这些指数可以用来从原始数组中提取相应的盒子和分数。
在上面的代码中,nms 函数被应用于boxes 和scores 数组,IoU阈值为0.3,得到的索引被存储在indices 变量中。
boxes = boxes[indices]scores = scores[indices]class_ids = class_ids[indices]mask_predictions = mask_predictions[indices]
nms() 是一个函数,它接收 , ,和一个 作为输入。该函数首先根据分数的高低对盒子进行降序排序。然后选择得分最高的盒子作为起点,其索引被附加到 。然后,该函数计算所选盒子与 中所有其他盒子的IoU,并删除IoU超过阈值的盒子的索引。这个过程在剩下的盒子中得分最高的盒子中重复进行,直到没有剩下的盒子。boxes scores iou_threshold keep_boxes boxes
提取输出掩码
mask_output = np.squeeze(outputs[1])def sigmoid(x): return 1 / (1 + np.exp(-x))# Calculate the mask maps for each boxnum_mask, mask_height, mask_width = mask_output.shape # CHWmasks = sigmoid(mask_predictions @ mask_output.reshape((num_mask, -1)))masks = masks.reshape((-1, mask_height, mask_width))masks
首先,定义sigmoid函数,随后用于计算掩码图。来自YOLO模型的掩码输出被挤压以去除任何额外的尺寸。
接下来,提取面具输出的尺寸。mask_predictions ,再乘以压扁的mask_output ,得到每个面具的每个类别的分数。结果是一个二维数组,其形状为(num_mask, mask_height*mask_width) 。然后将其重塑为(num_mask, mask_height, mask_width) 。
最后,sigmoid函数被应用于这个数组,它为每个遮罩中的每个像素返回一个介于0和1之间的值。这代表该像素属于该物体的概率。得到的掩码被重塑为具有形状(-1, mask_height, mask_width) ,其中第一个维度对应于掩码的数量(即非最大抑制后检测到的物体的数量)。
# Downscale the boxes to match the mask sizescale_boxes = rescale_boxes(boxes, (image_height, image_width), (mask_height, mask_width))scale_boxes
# For every box/mask pair, get the mask mapmask_maps = np.zeros((len(scale_boxes), image_height, image_width))blur_size = (int(image_width / mask_width), int(image_height / mask_height))
import mathfor i in range(len(scale_boxes)): scale_x1 = int(math.floor(scale_boxes[i][0])) scale_y1 = int(math.floor(scale_boxes[i][1])) scale_x2 = int(math.ceil(scale_boxes[i][2])) scale_y2 = int(math.ceil(scale_boxes[i][3])) x1 = int(math.floor(boxes[i][0])) y1 = int(math.floor(boxes[i][1])) x2 = int(math.ceil(boxes[i][2])) y2 = int(math.ceil(boxes[i][3])) scale_crop_mask = masks[i][scale_y1:scale_y2, scale_x1:scale_x2] crop_mask = cv2.resize(scale_crop_mask, (x2 - x1, y2 - y1), interpolation=cv2.INTER_CUBIC) crop_mask = cv2.blur(crop_mask, blur_size) crop_mask = (crop_mask > 0.5).astype(np.uint8) mask_maps[i, y1:y2, x1:x2] = crop_mask
该代码块包含一个循环,在scale_boxes 中的每个边界盒上进行迭代。对于每个盒子,它计算出经过缩放和未经过缩放的盒子的整数坐标,scale_x1 、scale_y1 、scale_x2 、scale_y2 、x1 、y1 、x2 、y2 。然后,它使用经过缩放的坐标从masks ,提取与该盒子相对应的掩码,并使用双线性插值将其调整为与原始盒子的大小一致。使用一个正方形核对得到的遮罩进行模糊处理,其尺寸与原始图像中盒子的大小成正比。最后,通过0.5的阈值对遮罩进行二值化处理,并按照原始盒子的坐标存储在mask_maps 。得到的mask_maps 是一个三维数组,其中第一个维度对应于边界盒的索引,后两个维度对应于原始图像中的遮罩尺寸。
提取和获取点状多线条
def masks2segments(masks, strategy='largest'): """ It takes a list of masks(n,h,w) and returns a list of segments(n,xy) Args: masks (torch.Tensor): the output of the model, which is a tensor of shape (batch_size, 160, 160) strategy (str): 'concat' or 'largest'. Defaults to largest Returns: segments (List): list of segment masks """ segments = [] for x in masks.astype('uint8'): c = cv2.findContours(x, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0] if c: if strategy == 'concat': # concatenate all segments c = np.concatenate([x.reshape(-1, 2) for x in c]) elif strategy == 'largest': # select largest segment c = np.array(c[np.array([len(x) for x in c]).argmax()]).reshape(-1, 2) else: c = np.zeros((0, 2)) # no segments found segments.append(c.astype('float32')) return segmentsdef scale_coords(img1_shape, coords, img0_shape, ratio_pad=None, normalize=False): """ Rescale segment coordinates (xyxy) from img1_shape to img0_shape Args: img1_shape (tuple): The shape of the image that the coords are from. coords (torch.Tensor): the coords to be scaled img0_shape (tuple): the shape of the image that the segmentation is being applied to ratio_pad (tuple): the ratio of the image size to the padded image size. normalize (bool): If True, the coordinates will be normalized to the range [0, 1]. Defaults to False Returns: coords (torch.Tensor): the segmented image. """ if ratio_pad is None: # calculate from img0_shape gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]) # gain = old / new pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 # wh padding else: gain = ratio_pad[0][0] pad = ratio_pad[1] coords[..., 0] -= pad[0] # x padding coords[..., 1] -= pad[1] # y padding coords[..., 0] /= gain coords[..., 1] /= gain clip_coords(coords, img0_shape) if normalize: coords[..., 0] /= img0_shape[1] # width coords[..., 1] /= img0_shape[0] # height return coordsdef clip_coords(coords, shape): """ Clip line coordinates to the image boundaries. Args: coords (torch.Tensor) or (numpy.ndarray): A list of line coordinates. shape (tuple): A tuple of integers representing the size of the image in the format (height, width). Returns: (None): The function modifies the input `coordinates` in place, by clipping each coordinate to the image boundaries. """ # if isinstance(coords, torch.Tensor): # faster individually # coords[..., 0].clamp_(0, shape[1]) # x # coords[..., 1].clamp_(0, shape[0]) # y # else: # np.array (faster grouped) coords[..., 0] = coords[..., 0].clip(0, shape[1]) # x coords[..., 1] = coords[..., 1].clip(0, shape[0]) # y
这段代码定义了几个用于处理分割蒙版和坐标的实用函数:
masks2segments:接受一个遮罩列表(形状为(n, h, w))并返回一个分段遮罩列表。strategy参数可以设置为'concat'或'largest',以分别指定是否串联所有段或只选择最大的段。scale_coords:将段的坐标从一个图像形状重新调整为另一个。img1_shape参数是坐标所来自的原始图像的形状,coords是要缩放的坐标,img0_shape是被应用于分割的图像的形状。ratio_pad参数是可选的,是一个元组,包括图像大小与填充图像大小的比率以及在(width, height)方向的填充。normalize参数是可选的,如果设置为True,坐标将被规范化为[0, 1]。clip_coords:将线坐标夹在图像的边界上。这个函数接收一个线坐标列表和格式为(height, width)的图像形状,并通过将每个坐标剪切到图像边界来修改输入coordinates。
points = [scale_coords(mask_maps.shape[1:], x, (image_height, image_width), normalize=False) for x in masks2segments(mask_maps)]
这个代码段采用上一步生成的mask_maps ,并应用分割算法(masks2segments()),为每个遮罩获得一个segments 列表。每个段是一组对应于掩码边界的点。然后,应用scale_coords() 函数将这些段重新缩放为原始图像的大小。
由此产生的points 变量将包含每个遮罩的片段列表,其中每个片段是对应于遮罩边界的点(x, y)的列表,被缩放为原始图像的大小。
视觉化输出
image_draw = image.copy()for bbox, conf, label, seg in zip(boxes, scores, class_ids, points): # Draw segmentation pts = np.array(seg, np.int32) pts = pts.reshape((-1, 1, 2)) image_draw = cv2.polylines(image_draw, [pts],True, (86, 71, 255), 2) # Draw bounding box x_min, y_min, x_max, y_max = bbox.astype(np.int64) image_draw = cv2.rectangle(image_draw,(x_min, y_min), (x_max, y_max), (86, 255, 71), 2) # Draw label cls = classes[label] cv2.putText(image_draw, f'{cls}:{int(conf*100)}', (x_min, y_min - 2), cv2.FONT_HERSHEY_SIMPLEX, 0.60, [225, 255, 255],thickness=1)Image.fromarray(cv2.cvtColor(image_draw, cv2.COLOR_BGR2RGB))
输出
结论
onnxruntime 为在生产环境中执行深度学习模型提供了一个灵活和高性能的运行时引擎,并支持广泛的硬件平台和执行供应商。该代码还显示了如何处理模型的输出,以提取检测到的物体的边界框和类标签,这可用于各种下游应用,如跟踪、计数或识别。
参考资料: