TensorFlow.js中使用自定义YOLOv11n模型进行目标检测的预处理问题

在使用TensorFlow.js加载自定义训练的YOLOv11n模型进行目标检测时，开发者可能会遇到预测结果异常的问题。本文将深入分析这一问题的原因，并提供正确的预处理方法。

问题现象

当开发者尝试使用自己训练的YOLOv11n模型（而非标准的COCO-SSD模型）进行目标检测时，虽然程序能够正常运行，但预测结果明显错误。例如：

• 输入图像中包含公交车（类别5）和四个人（类别0）
• 模型却预测出大量手机（类别67）
• 即使有部分人物预测，其坐标位置也不正确

问题根源

经过分析，这类问题通常源于图像预处理步骤不正确。YOLO系列模型对输入图像的预处理有特定要求，包括：

1. 图像尺寸调整
2. 归一化处理
3. 填充(padding)策略
4. 维度扩展

正确的预处理方法

以下是TensorFlow.js中正确的预处理代码示例：

// 将图像转换为Tensor
const tensorBus = tf.browser.fromPixels(bus);

// 获取图像原始高度和宽度
const [h, w] = tensorBus.shape.slice(0, 2);

// 计算需要填充的最大尺寸
const maxSize = Math.max(w, h);

// 对图像进行填充，使其成为正方形
const imgPadded = tensorBus.pad([
  [0, maxSize - h],  // 高度方向填充
  [0, maxSize - w],  // 宽度方向填充
  [0, 0]             // 颜色通道不填充
]);

// 完整的预处理流程
const inputTensor = tf.image
    .resizeBilinear(imgPadded, [640, 640])  // 调整到模型期望的输入尺寸
    .div(255.0)                            // 归一化到0-1范围
    .expandDims(0);                        // 添加batch维度

关键预处理步骤解析

1. 图像填充：YOLO模型通常需要正方形输入，因此需要将矩形图像填充为正方形，保持原始图像内容不变形。

2. 尺寸调整：将填充后的图像调整到模型训练时使用的固定尺寸（如640x640）。

3. 归一化：将像素值从0-255范围归一化到0-1范围，这是大多数深度学习模型的输入要求。

4. 维度扩展：添加batch维度，因为模型期望的输入形状是[batch, height, width, channels]。

注意事项

1. 确保预处理步骤与模型训练时的预处理完全一致
2. 检查模型的输入形状要求
3. 验证类别列表(class list)是否正确
4. 在不同框架中测试模型预测结果，确保模型本身训练正确

通过正确的预处理方法，可以显著提高自定义YOLO模型在TensorFlow.js中的预测准确性。开发者应当仔细检查预处理流程，确保其与模型训练时的数据处理方式完全匹配。