BTAD数据集实战指南：从数据预处理到模型训练的全流程解析

1次阅读

共计 2413 个字符，预计需要花费 7 分钟才能阅读完成。

BTAD（Bottle and Tablet Anomaly Detection）数据集是工业缺陷检测领域的重要基准数据集，包含瓶装和片剂两类产品的缺陷样本。其特点包括：

多模态数据：包含可见光、X 光等多种成像方式
精细标注：提供像素级缺陷掩膜和边界框（JSON/XML 格式）
真实场景：采集自实际生产线，包含光照变化、遮挡等挑战

该数据集在 2020 年 ICCV workshop 发布后，已成为评估工业缺陷检测算法的标准测试平台。

BTAD 的目录结构如下：

BTAD/
├── bottle/
│   ├── train/          # 训练集
│   │   ├── good/       # 正常样本
│   │   └── anomaly/    # 缺陷样本
│   ├── test/
│   └── annotations/    # JSON 标注
└── tablet/
    └── ...

加载图像和标注的示例代码（Python）：

import json
from PIL import Image

# 读取标注文件
def load_annotation(json_path):
    with open(json_path) as f:
        data = json.load(f)
    bboxes = []
    for obj in data['objects']:
        bbox = obj['bbox']  # [x_min, y_min, x_max, y_max]
        bboxes.append(bbox)
    return bboxes

# 示例使用
img = Image.open('BTAD/bottle/train/anomaly/001.png')
bboxes = load_annotation('BTAD/bottle/annotations/001.json')

使用 Matplotlib 展示缺陷样本：

import matplotlib.pyplot as plt
import matplotlib.patches as patches

def show_sample(img_path, annotation_path):
    img = Image.open(img_path)
    bboxes = load_annotation(annotation_path)

    fig, ax = plt.subplots(1)
    ax.imshow(img)

    for bbox in bboxes:
        rect = patches.Rectangle((bbox[0], bbox[1]),
            bbox[2]-bbox[0], bbox[3]-bbox[1],
            linewidth=2, edgecolor='r', facecolor='none')
        ax.add_patch(rect)
    plt.show()

工业图像的特殊处理方法：

光照归一化 ：

import cv2

def normalize_illumination(img):
    lab = cv2.cvtColor(img, cv2.COLOR_RGB2LAB)
    l, a, b = cv2.split(lab)
    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8,8))
    l_norm = clahe.apply(l)
    lab_norm = cv2.merge((l_norm, a, b))
    return cv2.cvtColor(lab_norm, cv2.COLOR_LAB2RGB)

ROI 提取 （基于先验知识定位产品区域）

PyTorch 训练示例（基于 ResNet18）：

import torch
import torchvision
from torch.utils.data import Dataset

class BTADDataset(Dataset):
    def __init__(self, img_dir, transform=None):
        self.img_paths = [...]  # 遍历目录获取图像路径
        self.transform = transform

    def __getitem__(self, idx):
        img = Image.open(self.img_paths[idx])
        label = 0 if 'good' in self.img_paths[idx] else 1

        if self.transform:
            img = self.transform(img)
        return img, label

# 数据增强
train_transform = torchvision.transforms.Compose([transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(brightness=0.2, contrast=0.2),
    transforms.ToTensor(),])

# 初始化模型
model = torchvision.models.resnet18(pretrained=True)
model.fc = torch.nn.Linear(512, 2)  # 二分类
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

常见问题及解决方案：