anomalib实战学习记录：安装及训练官方数据

liyun52

12644人浏览 · 2024-07-17 11:17:09

liyun52 · 2024-07-17 11:17:09 发布

源码下载地址：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-1.1.0

操作系统是：win11

代码的结构为：新版本的anomalib没有单独的train.py文件了，后面需要自己新建train.py文件，或者用命令行进行训练。

一、安装

1、需要通过conda新建一个python=3.10的虚拟环境：

conda create -n anomalib_env python=3.10

2、激活环境：

conda activate anomalib_env

3、通过下载源码进行安装：

git clone https://github.com/openvinotoolkit/anomalib.git
cd anomalib
pip install -e .

4、安装anomalib需要的相关库

anomalib install -h
anomalib install
anomalib install -v
anomalib install --option core
anomalib install --option openvino

后面三个的命令我没有截图，依次安装即可。

二、训练

我是使用pycharm进行训练的。使用pycharm打开源码，新建train.py文件。train.py内容如下（来自官网）：

# Import the required modules
from anomalib.data import MVTec
from anomalib.models import Patchcore
from anomalib.engine import Engine

# Initialize the datamodule, model and engine
datamodule = MVTec()
model = Patchcore()
engine = Engine()

# Train the model
engine.train(datamodule=datamodule, model=model)

开始训练即可。刚开始训练的时候，会下载MVTec数据集，可能会下载不下来，可以在报错信息里面复制下载链接，自行下载，下载好了之后，放在对应的文件夹中。

我在训练的过程中，报错如下：

参考：解决RuntimeError: An attempt has been made to start a new process before...办法 - 知乎 (zhihu.com)https://zhuanlan.zhihu.com/p/553407062

修改train.py：

datamodule = MVTec()
# 修改为：
datamodule = MVTec(num_workers=0)

然后就可以训练了。

补充一下，mvtec数据集训练时，更多的参数设置：

# Import the required modules
from anomalib.data import MVTec
from anomalib.models import Patchcore
from anomalib.engine import Engine
from anomalib import TaskType
from anomalib.utils.normalization import NormalizationMethod
from anomalib.data.utils import (
    TestSplitMode,
    ValSplitMode,
)

# Initialize the datamodule, model and engine
datamodule = MVTec(
    num_workers=0,
    category='hazelnut', # mvtec数据集里面的类别
    root = "./datasets/MVTec", # 数据集存放的位置
    train_batch_size = 16,
    eval_batch_size = 16,
    task = TaskType.CLASSIFICATION, # 任务的类型：CLASSIFICATION（分类），DETECTION（检测），SEGMENTATION（分割）
    image_size = None,
    transform = None,
    train_transform = None,
    eval_transform = None,
    test_split_mode = TestSplitMode.FROM_DIR,
    test_split_ratio = 0.2,
    val_split_mode = ValSplitMode.SAME_AS_TEST,
    val_split_ratio = 0.5,
    seed = None
)

model = Patchcore()
engine = Engine(
    default_root_dir = './result_mvtec2', # 结果保存的路径
    task = TaskType.CLASSIFICATION, # 任务的类型：CLASSIFICATION（分类），DETECTION（检测），SEGMENTATION（分割）
    callbacks = None,
    normalization = NormalizationMethod.MIN_MAX,
    threshold = "F1AdaptiveThreshold",
    image_metrics = None,
    pixel_metrics = None,
    logger = None,
)

# Train the model
engine.train(datamodule=datamodule, model=model)

三、推理

predictions = engine.predict(
    datamodule=datamodule,
    model=model,
    ckpt_path="model.ckpt", #修改成自己训练生成的model.ckpt文件路径 
)

训练的时候，会生成一个results文件夹，model.ckpt文件就在results文件夹里面。推理的代码可以直接放在train.py文件的最下面，即engine.train(datamodule=datamodule, model=model)后面。

四、数据集的划分

官方教程：

Custom Data — Anomalib 2022 documentationhttps://anomalib.readthedocs.io/en/latest/markdown/guides/how_to/data/custom_data.html通过官网给的链接，下载hazelnut_toy数据集放在datasets文件夹里面：

可以重新建一个train_toy.py

内容如下：

from anomalib.models import Patchcore
from anomalib.engine import Engine
from anomalib.data import Folder

# Create the datamodule
datamodule = Folder(
    name="hazelnut_toy",
    root="datasets/hazelnut_toy",
    normal_dir="good",
    abnormal_dir="colour",
    task="classification",
    num_workers=0,
    image_size=[256,256],
    val_split_ratio=0.5,
)

# Setup the datamodule
datamodule.setup()

i, train_data = next(enumerate(datamodule.train_dataloader()))
print(train_data.keys())

i, val_data = next(enumerate(datamodule.val_dataloader()))
print(val_data.keys())

i, test_data = next(enumerate(datamodule.test_dataloader()))
print(test_data.keys())

model = Patchcore()
engine = Engine(task="classification")
engine.train(datamodule=datamodule, model=model)

推理的代码为：

predictions = engine.predict(
    datamodule=datamodule,
    model=model,
    ckpt_path="model.ckpt",  # 存放model.ckpt文件的路径，需要对应修改
)

结果如下：

补充一下使用Folder构建数据集并训练时，更多的参数设置：

from anomalib.models import Patchcore
from anomalib.engine import Engine
from anomalib.data import Folder

from anomalib.data.utils import (
    TestSplitMode,
    ValSplitMode,
)

import torchvision.transforms.v2 as v2

data_transform = v2.Compose([
    v2.Resize((256,256)),
    v2.RandomHorizontalFlip(),
    v2.ColorJitter(brightness=(0.1,1)),
])

# task = "classification" or "detection" or "segmentation"
task = "classification"

# Create the datamodule
datamodule = Folder(
    name="mvtec",
    root="../datasets/bottle",
    normal_dir="./train/good",
    abnormal_dir=["./test/broken_large", "./test/broken_small", "./test/contamination"],
    task=task,
    train_batch_size=4,
    eval_batch_size=4,
    num_workers=0,
    image_size=[256,256],
    normal_test_dir='./test/good',
    mask_dir = None,
    normal_split_ratio = 0.2,
    extensions = None,
    transform = data_transform,
    train_transform = None,
    eval_transform = None,
    test_split_mode = TestSplitMode.FROM_DIR,
    test_split_ratio = 0.2,
    val_split_mode = ValSplitMode.FROM_TEST,
    val_split_ratio=0.5,
    seed = None,
)

# Setup the datamodule
datamodule.setup()

model = Patchcore()

engine = Engine(
    task=task,
    default_root_dir='./result_1114'
)
engine.train(datamodule=datamodule, model=model)

五、onnx模型导出

新建export.py文件，内容如下：

from anomalib.models import Patchcore
from anomalib.engine import Engine


model = Patchcore()
engine = Engine(task="classification")
onnx_model = engine.export(
    model= model,
    export_type='onnx',
    export_root=None,
    input_size=[256, 256],
    transform=None,
    compression_type=None,
    datamodule = None,
    metric = None,
    ov_args = None,
    ckpt_path='model.ckpt', # 存放model.ckpt文件的路径，需要对应修改
    )
print(onnx_model)

可以导出onnx模型，具体细节可以研究一下engine.py文件里面的export部分，有比较详细的说明。

六、onnx使用

在主文件夹里面，有一个tool的文件夹，里面有一个inference文件夹。

可以使用里面的openvino_inference.py进行onnx模型推理，有两种使用方式，一种是使用命令行进行推理，一种是直接右击运行py文件进行推理。

使用命令行进行推理：

python tools/inference/openvino_inference.py --weights ./results\weights\onnx\model.onnx --metadata ./results\weights\onnx\metadata.json  --input ./datasets\img_test  --output ./datasets\img_test_out --show

注意，weights参数是onnx模型的路径，metadata是和onnx在一个文件夹里面（export时自动生成的），input是要推理的图像文件夹或者一张图片，output是推理结果存放的位置，show参数是在推理完成之后，将推理的图像展示出来，也可以不添加这个参数。

直接右击.py的话，就是修改get_parser()里面参数，记得要把required参数都设置成False，不然会报下面的错误（这里是因为我把weights的required值设置为True）：

可以将参数按照下面进行修改，这里我只列出了几个必要的参数，其他的当然也可以修改：

    parser.add_argument("--weights", type=Path, default='./results\weights\onnx\model.onnx', required=False, help="Path to model weights")
    parser.add_argument("--metadata", type=Path, default='./results\weights\onnx\metadata.json', required=False, help="Path to a JSON file containing the metadata.")
    parser.add_argument("--input", type=Path, default=r'./datasets\img_test', required=False, help="Path to an image to infer.")
    parser.add_argument("--output", type=Path, default=r'./datasets\img_test_out', required=False, help="Path to save the output image.")

然后就可以直接右击运行py文件即可，推理的结果会存放到设置的output参数的文件夹中。