缺陷检测之Anomalib

缺陷检测的现状

工业缺陷数据有一个比较显著的特征：样本不平衡。绝大部分采集得到的工业数据都是没有缺陷的，这样一来，正样本的数据在模型训练中根本没有起到作用，负样本又太少，很难训练得到有效的模型。使用有监督学习的方法还有一个问题：负样本的出现是十分偶然的，可能在数据集中根本没有出现某一类型的负样本，如此训练得到的模型很有可能翻车，所以只能另寻他法。

Anomalib介绍

Anomalib属于无监督学习，无监督算法只使用正样本进行训练，网络经过大量的正样本学习，在遇到负样本时，就会知道负样本和正样本"长得不一样"，然后输出和原图尺寸相同的一张概率分布图，来表示某处是异常区域的概率大小。

附属项目连接

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.

模型训练方法

通过查看anomalib train -h

可以看到一些训练的指令，这里我常用的是:

python 复制代码

 anomalib train --config .\config.yaml

下面附上我的config.yaml(其中model下的./datasets/imagenette运行时会自动下载)

python 复制代码

# anomalib==1.1.0
seed_everything: true
trainer:
  accelerator: auto
  strategy: auto
  devices: 1
  num_nodes: 1
  precision: null
  logger: null
  callbacks: null
  fast_dev_run: false
  max_epochs: null  # 默认999
  min_epochs: null
  max_steps: -1
  min_steps: null
  max_time: null
  limit_train_batches: null
  limit_val_batches: null
  limit_test_batches: null
  limit_predict_batches: null
  overfit_batches: 0.0
  val_check_interval: null
  check_val_every_n_epoch: 1
  num_sanity_val_steps: null
  log_every_n_steps: null
  enable_checkpointing: null
  enable_progress_bar: null
  enable_model_summary: null
  accumulate_grad_batches: 1
  gradient_clip_val: null
  gradient_clip_algorithm: null
  deterministic: null
  benchmark: null
  inference_mode: true
  use_distributed_sampler: true
  profiler: null
  detect_anomaly: false
  barebones: false
  plugins: null
  sync_batchnorm: false
  reload_dataloaders_every_n_epochs: 0
normalization:
  normalization_method: MIN_MAX
task: SEGMENTATION
metrics:
  image:
  - F1Score
  - AUROC
  pixel: null
  threshold:
    class_path: anomalib.metrics.F1AdaptiveThreshold
    init_args:
      default_value: 0.5
      thresholds: null
      ignore_index: null
      validate_args: true
      compute_on_cpu: false
      dist_sync_on_step: false
      sync_on_compute: true
      compute_with_cache: true
logging:
  log_graph: false
default_root_dir: anomalib\train
ckpt_path: null
data:
  class_path: anomalib.data.Folder
  init_args:
    name: A-1/crops/hxjzq
    normal_dir: D:\python\anomalib-1.1.0\data\normal1
    root: crop_new/cxjzq
    abnormal_dir: null
    normal_test_dir: null
    mask_dir: null
    normal_split_ratio: 0.2
    extensions: null
    train_batch_size: 1
    eval_batch_size: 16
    num_workers: 4
    image_size: null
    transform: null
    train_transform: null
    eval_transform: null
    test_split_mode: synthetic
    test_split_ratio: 0.2
    val_split_mode: same_as_test
    val_split_ratio: 0.5
    seed: null

model:
  class_path: anomalib.models.EfficientAd
  init_args:
    imagenet_dir: './datasets/imagenette'
    teacher_out_channels: 384
    model_size: 'small'
    lr: 0.001
    weight_decay: 1e-05
    padding: False
    pad_maps: True



#model:
#  class_path: anomalib.models.Patchcore
#  init_args:
#    backbone: wide_resnet50_2
#    layers:
#    - layer2
#    - layer3
#    pre_trained: true
#    coreset_sampling_ratio: 0.1
#    num_neighbors: 9

训练完成后会得到一个模型

模型转为onnx方法

（官方给的是--export_model onn 这里我们改成 --export_type onnx）

python 复制代码

anomalib export --model EfficientAd --export_type onnx --ckpt_path D:\python\anomalib-1.1.0\anomalib\train\Patchcore\A-1\crops\hxjzq\v0\weights\lightning\model.ckpt

--model的出处：（根据自己的配置文件进行修改）

下面的都可供选择：（根据原版论文讲述的EfficientAd的效果最好）

运行：

调用onnx方法

python 复制代码

# infer onnx model
import onnxruntime as ort
import cv2
import numpy as np

def infer_onnx(onnx_path, img_path,img_name):
    ort_session = ort.InferenceSession(onnx_path)
    image = cv2.imread(img_path)
    image_blod = cv2.dnn.blobFromImage(image, scalefactor=1 / 255.0, size=(640,640), swapRB=True, )
    model_inputs = ort_session.get_inputs()
    output = ort_session.run(None, {model_inputs[0].name: image_blod})
    # print('output',output)
    # np.savetxt('res.csv', output[0].squeeze(), delimiter=',', fmt='%f')
    # print(output[0])
    
    # 计算放缩比例
    print('image_shape',image.shape)
    ratio_x = image.shape[1] / 224
    ratio_y = image.shape[0] / 224

    img_res = output[0].squeeze()
    img_res = img_res.astype(np.uint8)
    cv2.imwrite('img_res.jpg',img_res)
    
    # thresh 返回的二值化图像
    _, thresh = cv2.threshold(img_res, 45, 255, cv2.THRESH_BINARY) # 255白色
    cv2.imwrite('thresh.jpg',thresh)  # 异常图片显示（白色区域）

    # print('thresh',thresh) 
    # contours  轮廓列表 
    contours, _ = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)  # 只检测白色
    if output[1] > 50:
        for cnt in contours:
            x, y, w, h = cv2.boundingRect(cnt) # 获取矩形左上角坐标和宽高
            cv2.rectangle(image, (int(x*ratio_x), int(y*ratio_y)), (int((x + w)*ratio_x), int((y + h)*ratio_y)), (0, 0, 255), 2)

        cv2.imwrite(f'{img_name}', image)
    # print(f'{img_name}:', output[1])


if __name__ == '__main__':
    onnx_path = r"E:\0806\export\Patchcore\A-2\xjslh\weights\onnx\model.onnx"
    img_path = r"E:\0806\crop\A-2\crops\xjslh"
    import os
    from tqdm import tqdm
    img_list = os.listdir(img_path)
    for i in tqdm(img_list):
        img_path_ = os.path.join(img_path, i)
        infer_onnx(onnx_path, img_path_,i)