### 实现YOLOv5中ONNX输出后处理 在YOLOv5中使用ONNX进行模型输出处理涉及几个重要步骤，包括环境配置、数据预处理以及最终的结果后处理。为了确保整个过程顺利，在WSL环境下完成这些工作是一个不错的选择[^2]。 #### 数据预处理 对于输入到ONNX模型的数据，通常需要做如下调整： - **图像格式转换**：这一步骤涉及到类型转换和归一化操作，目的是让图片适应网络的要求。例如，将BGR颜色空间转化为RGB，并缩放到0至1之间。 - **增加批次维度**：因为大多数深度学习框架期望接收四维张量作为输入（批量大小×通道数×高度×宽度），即使只有一张图片也需要构建这样的形状来满足API需求。 ```python import numpy as np from PIL import Image def preprocess(image_path, img_size=640): image = Image.open(image_path).convert('RGB') ratio = min(img_size / s for s in image.size) new_size = tuple(int(x * ratio) for x in image.size) resized_image = image.resize(new_size, Image.ANTIALIAS) padded_image = Image.new("RGB", (img_size, img_size), color=(114, 114, 114)) padded_image.paste(resized_image, ((img_size - new_size[0]) // 2, (img_size - new_size[1]) // 2)) input_data = np.array(padded_image, dtype=np.float32) / 255. input_data = input_data.transpose([2, 0, 1]) input_data = np.expand_dims(input_data, axis=0) return input_data ``` #### 执行模型推理 一旦完成了上述准备工作，则可以加载已保存好的`.onnx`文件来进行预测。需要注意的是，此时获得的输出可能不是标准形式，因此还需要进一步加工才能用于后续应用。 ```python import onnxruntime as ort session = ort.InferenceSession("model.onnx") input_name = session.get_inputs()[0].name output_names = [o.name for o in session.get_outputs()] results = session.run(output_names, {input_name: preprocessed_input}) ``` #### 结果后处理 最后也是最关键的一环就是对原始检测结果实施必要的变换，使之成为易于理解和使用的边界框列表。这部分主要包括坐标转换、非极大值抑制(NMS)，以及其他筛选机制。 - **坐标转换函数(xywh2xyxy)**：此方法负责将以中心点加宽高表示的目标位置转变为左上角右下角的形式，便于计算IoU和其他几何关系。 - **NMS算法**：通过设定阈值去除重叠过多的对象候选区域，从而提高识别精度。 - **过滤低置信度框(filter_boxes)**：移除那些得分低于指定水平的预测项，减少误报率。 ```python def xywh2xyxy(x): y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x) y[:, 0] = x[:, 0] - x[:, 2] / 2 # top left x y[:, 1] = x[:, 1] - x[:, 3] / 2 # top left y y[:, 2] = x[:, 0] + x[:, 2] / 2 # bottom right x y[:, 3] = x[:, 1] + x[:, 3] / 2 # bottom right y return y def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45): """Runs Non-Maximum Suppression (NMS)""" xc = prediction[..., 4] > conf_thres # candidates output = [] for xi, x in enumerate(prediction): # image index, image inference x = x[xc[xi]] # confidence filter boxes = xywh2xyxy(x[:, :4]) scores = x[:, 4] indices = cv2.dnn.NMSBoxes( bboxes=boxes.tolist(), scores=scores.tolist(), score_threshold=conf_thres, nms_threshold=iou_thres ) if len(indices) > 0: output.append(x[indices.flatten()]) return output ```