使用Anomalib项目的padim无监督算法 进行自制工业缺陷数据集的模型训练和ONNX部署(一)——模型训练篇

目录

前言 

一、无监督学习缺陷检测 Anomalib介绍

二、Anomalib代码结构

三、任务描述和模型训练推理

四、总结与展望


前言 

        本文专注于padim算法在自制数据集上的训练过程,博主水平有限,对神经网络模型秉持能用就行的态度,所以文中不涉及网络结构和论文细节的解读,想看这些的同学请另寻资料哈~

一、无监督学习缺陷检测 Anomalib介绍

        组里最近给的新任务,对金属材质表面的各种缺陷进行检测。之前使用的是有监督的yolov5网络,标数据集着实痛苦无比。而且工业缺陷数据有一个比较显著的特征:样本不平衡。绝大部分采集得到的工业数据都是没有缺陷的,这样一来,正样本的数据在模型训练中根本没有起到作用,负样本又太少,很难训练得到有效的模型。使用有监督学习的方法还有一个问题:负样本的出现是十分偶然的,可能在数据集中根本没有出现某一类型的负样本,如此训练得到的模型很有可能翻车,所以只能另寻他法。

        查阅资料,发现无监督的算法更适合工业缺陷检测的场景。无监督算法只使用正样本进行训练,网络经过大量的正样本学习,在遇到负样本时,就会知道负样本和正样本“长得不一样”,然后输出和原图尺寸相同的一张概率分布图,来表示某处是异常区域的概率大小。

        无监督学习方法虽好,但作为伸手党应该去哪里要代码呢?个人感觉最好的地方就是Anomalib,它使用Python实现,资料完备易懂,目前仍然在更新。链接如下:GitHub - openvinotoolkit/anomalib: An anomaly detection library comprising state-of-the-art algorithms and features such as experiment management, hyper-parameter optimization, and edge inference.https://github.com/openvinotoolkit/anomalib

使用Anomalib项目的padim无监督算法 进行自制工业缺陷数据集的模型训练和ONNX部署(一)——模型训练篇_第1张图片

        看到英文名字就可以知道:这是一个异常检测(Abnormal)的库(library),里面的内容的确十分丰富,集成了十余种近年来准确率较高的缺陷(异常)检测算法,基本都是无监督学习的方法,诸如padim算法、fastflow算法等。

使用Anomalib项目的padim无监督算法 进行自制工业缺陷数据集的模型训练和ONNX部署(一)——模型训练篇_第2张图片

        找到这个宝库,我们下一步需要探索它的代码结构和用法。

二、Anomalib代码结构

        我在本地使用Pycharm编程,新建项目后Anomalib的代码结构如下:

使用Anomalib项目的padim无监督算法 进行自制工业缺陷数据集的模型训练和ONNX部署(一)——模型训练篇_第3张图片

        若想理清楚该项目的代码结构,方便后面的训练和部署,需要一点点进行分析。其中,比较重要的几个部分是:

1. anomalib文件夹:事实上本文件夹就是该团队发布的库的源码,我们同样可以通过pip install的方法来进行安装,

        1.1 该文件夹下的models子文件夹包含了十余种缺陷检测的算法,可供读者任意调用;

        1.2 pre_processingpost_processing子文件夹分别是预处理和后处理功能,最终显示推理结果就使用了post_processing下的visualizer

        1.3 deploy文件夹中的inferencers则是各种推理器,想要使用pytorch的读者可以关注torch_inferencer,想要使用onnx推理的读者则应该使用openvino_inferencer

2. datasets文件夹:顾名思义,该文件夹中应该存放待训练的数据集,如工业数据集MVTec,或者我们自制的数据集。

3. results文件夹:该文件夹存放的是训练和推理的结果,只有在完成了训练或推理任务后才会出现。

4. tools文件夹:该文件夹中的inference子文件夹存放了一系列推理代码,它们分别调用了anomalib.deploy中的不同inferencer(推理器)。该文件夹下的train.py是训练模型的入口。

        总结:我们后面需要使用tools/train.py进行模型训练,使用tools/inference中的某个推理代码进行模型推理预测。有了以上预备知识,我们终于可以开始模型的训练了。

三、任务描述和模型训练推理

        若读者只想在pycharm中看一个算法的效果,那么按照官方的示例,使用train.py训练后再使用lightning_inference.py进行推理即可,在新出现的result文件夹中便可以看到推理结果,十分方便,此处不予赘述。

        但对于包括我在内,后续需要在其他平台部署的读者,我们需要训练自己的数据集并得到onnx模型,这样就需要对config.yaml文件进行修改,方法如下:

        首先依然按照官方的方案,看Readme中训练Custom Dataset的方法:

使用Anomalib项目的padim无监督算法 进行自制工业缺陷数据集的模型训练和ONNX部署(一)——模型训练篇_第4张图片        这里我们使用的是padim算法,于是我们将anomalib/models/padim/config.yaml的dataset部分按以上部分修改后在终端运行以下训练命令:

python tools/train.py --model padim --config anomalib/models/padim/config.yaml

        但会报错:找不到“normalization”等属性,此时不要着急,之前的config.yaml中实际上这些属性都是在的,这里官方竟然没有写进来,可能是疏忽了,我把自己的config.yaml的datasets部分放在这里,读者可以按照自己的自制数据集的路径进行改写,需要添加和注意的地方我都做了注释:

dataset:
  name: tube               # 数据集的名字,如MVTec等,这个不重要
  format: folder
  path: ./datasets/img_192 # 自制数据集路径
  normal_dir: normal       # 自制数据集正样本子文件夹
  abnormal_dir: abnormal   # 自制数据集负样本子文件夹
  mask_dir: null           # 二值掩膜路径,自制数据集一般没有,填null
  normal_test_dir: null # name of the folder containing normal test images.
  task: classification # classification or segmentation
  extensions: null
  normalization: imagenet  # 此处添加imagenet
  split_ratio: 0.2 # ratio of the normal images that will be used to create a test split
  image_size: 256
  train_batch_size: 32
  test_batch_size: 32
  num_workers: 8
  transform_config:
    train: null
    val: null
  test_split_mode: from_dir # 此处添加
  test_split_ratio: 0.2
  val_split_mode: same_as_test
  val_split_ratio: 0.5
  create_validation_set: true
  tiling:
    apply: false
    tile_size: null
    stride: null
    remove_border_count: 0
    use_random_tiling: False
    random_tile_count: 16

        需要特别注意的是:由于我们使用自制数据集,往往没有ground truth的二值化掩膜(mask),所以这里需要将task字段设置为classification而不是segmentation。接着将metrics部分的pixel部分删掉,否则会报错。原因是没有提供二值化掩膜,无法计算基于像素的测试准确度,只保留image部分。修改如下:

metrics:
  image:
    - F1Score
    - AUROC
#  pixel:
#    - F1Score
#    - AUROC
  threshold:
    method: adaptive #options: [adaptive, manual]
    manual_image: null
    manual_pixel: null

        我们还需要onnx模型,这里需要在config.yaml中的optimization部分加入onnx字段:

optimization:
  export_mode: onnx # options: torch, onnx, openvino

        至此,我们就可以再次运行训练命令,开始训练了:

python tools/train.py --model padim --config anomalib/models/padim/config.yaml

        padim算法的论文我并没有研读,但似乎其训练过程只需要一轮, 训练完成后,在results文件夹中可以看到如下结构:

使用Anomalib项目的padim无监督算法 进行自制工业缺陷数据集的模型训练和ONNX部署(一)——模型训练篇_第5张图片

         tube是博主本人数据集的名字,下面的images文件夹中有对正样本和负样本图片的测试结果,以概率热图的形式给出;weights文件夹中有pytorch-lightning的模型和我们所需的onnx模型,以及包含了训练数据集信息的metadata.json(对于部署十分重要,下一篇再说)。

        做完以上步骤,我们得到了onnx模型,但在部署之前,还是看一下转换后的模型有没有精度损失,在终端输入以下命令:

 python tools/inference/openvino_inference.py --weights results/padim/tube/run/weights/onnx/model.onnx --metadata results/padim/tube/run/weights/onnx/meta_data.json  --input datasets/img_192/abnormal/cam0_17_04_54.jpg  --output results/padim/tube/run/images --config src/anomalib/models/padim/config.yaml

        得到以下推理结果:

使用Anomalib项目的padim无监督算法 进行自制工业缺陷数据集的模型训练和ONNX部署(一)——模型训练篇_第6张图片

         效果不错,在工业缺陷检测这方面,终于可以和yolov5说拜拜了~

四、总结与展望

        我们使用padim算法在自制工业数据集上获得了良好的检测效果,通过训练得到了onnx模型;在这个过程中,最需要注意的就是配置文件config.yaml的修改,要按照官网教程、本博客和自己数据集的路径进行配置和修改,防止出错。

        在得到了onnx模型后,我们想要脱离开庞大繁杂的Anomalib项目和环境,完成自己的软件算法部署,但图像数据在送入onnx模型之前需要什么预处理步骤?模型输入输出分别是什么?输出的数据需要怎样的后处理?怎么画出Anomalib中十分简洁明了的热图呢?这些都是需要我们解决的问题。

        下一篇博客我们将继续从Anomalib的代码结构和流程进行分析,使用onnxruntime引擎完成onnx在C++中的部署,实现算法的落地。

你可能感兴趣的:(深度学习,python,视觉检测,算法)