自定义调用onnx模型结果不理想的原因

使用yolov5的cls模型进行图像分类，并导出onnx进行模型的部署，发现以下问题：

分类的onnx在官方的代码中效果很好，但换成自己的代码中结果就不同，检测结果不准。
检测的onnx效果就很好。

初始图像加载代码：

def preprocess_image(image_path, input_shape):

    # 读取图像
    image = cv2.imread(image_path)
    # 调整图像大小
    image = cv2.resize(image, (input_shape[2], input_shape[3]))
    # 转换图像格式为RGB
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    # 归一化
    image = image.astype(np.float32) / 255.0
    # 调整图像形状为 (batch_size, channels, height, width)
    image = np.transpose(image, (2, 0, 1))
    image = np.expand_dims(image, axis=0)
    return image

这个代码的图像处理结果在目标检测中检测效果挺好，在图像分类中效果挺差：

分析了官方源码，图像加载中存在部分差异

# 目标检测
dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)

# 图像分类
dataset = LoadImages(source, img_size=imgsz, transforms=classify_transforms(imgsz[0]), vid_stride=vid_stride)

分类算法中多了个transforms=classify_transforms(imgsz[0])，

def classify_transforms(size=224):
    # Transforms to apply if albumentations not installed
    assert isinstance(size, int), f'ERROR: classify_transforms size {size} must be integer, not (list, tuple)'
    # T.Compose([T.ToTensor(), T.Resize(size), T.CenterCrop(size), T.Normalize(IMAGENET_MEAN, IMAGENET_STD)])
    return T.Compose([CenterCrop(size), ToTensor(), T.Normalize(IMAGENET_MEAN, IMAGENET_STD)])

也就是说分类的图像需要经过标准化处理。

# 归一化
image = image.astype(np.float32) / 255.0
# 标准化
image = (image - np.array(IMAGENET_MEAN, dtype=np.float32)) / np.array(IMAGENET_STD, dtype=np.float32)

自定义调用onnx模型结果不理想的原因

猜你喜欢