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

1次阅读

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

BTAD（Braille Textile Anomaly Detection）数据集是工业缺陷检测领域的重要基准数据集之一，主要用于纺织品表面的缺陷检测。该数据集包含三种不同类型的纺织品缺陷，每种缺陷都有数百张高分辨率图像，并提供了精确的边界框标注。BTAD 数据集在学术界被广泛用于评估缺陷检测算法的性能，尤其是在小目标检测和复杂背景下的缺陷识别方面具有挑战性。

BTAD 数据集的目录结构通常如下：

BTAD/
├── train/
│   ├── images/
│   ├── annotations_json/
├── test/
│   ├── images/
│   ├── annotations_json/

以下是使用 Python 加载 BTAD 数据集的代码示例：

import json
import os
from PIL import Image
import numpy as np

def load_btad_dataset(base_path):
    """
    加载 BTAD 数据集
    :param base_path: 数据集根目录
    :return: 图像和标注列表
    """
    images = []
    annotations = []

    # 遍历 train 和 test 目录
    for split in ['train', 'test']:
        image_dir = os.path.join(base_path, split, 'images')
        anno_dir = os.path.join(base_path, split, 'annotations_json')

        for img_file in os.listdir(image_dir):
            if not img_file.endswith('.png'):
                continue

            # 加载图像
            img_path = os.path.join(image_dir, img_file)
            image = Image.open(img_path)
            images.append(image)

            # 加载对应的标注文件
            anno_path = os.path.join(anno_dir, img_file.replace('.png', '.json'))
            with open(anno_path, 'r') as f:
                anno = json.load(f)
            annotations.append(anno)

    return images, annotations

为了更好地理解数据集，我们可以使用 Matplotlib 可视化一些样本：

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

def visualize_sample(image, annotation):
    """
    可视化样本图像和标注
    :param image: PIL 图像
    :param annotation: 标注字典
    """
    fig, ax = plt.subplots(1)
    ax.imshow(image)

    # 绘制边界框
    for obj in annotation['objects']:
        bbox = obj['bbox']
        rect = patches.Rectangle((bbox['xmin'], bbox['ymin']),
            bbox['xmax'] - bbox['xmin'],
            bbox['ymax'] - bbox['ymin'],
            linewidth=1,
            edgecolor='r',
            facecolor='none'
        )
        ax.add_patch(rect)
        ax.text(bbox['xmin'], bbox['ymin'],
            obj['category'],
            color='white',
            bbox=dict(facecolor='red', alpha=0.5)
        )

    plt.axis('off')
    plt.show()

# 示例用法
images, annotations = load_btad_dataset('path/to/BTAD')
visualize_sample(images[0], annotations[0])

工业图像通常需要特殊的预处理方法：

光照归一化：工业环境下光照条件可能不稳定，可以使用直方图均衡化或 CLAHE 方法

import cv2

def normalize_illumination(image):
    """
    光照归一化
    :param image: 输入图像
    :return: 归一化后的图像
    """
    # 转换为灰度图
    gray = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)

    # CLAHE 自适应直方图均衡化
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
    normalized = clahe.apply(gray)

    return Image.fromarray(normalized)

ROI 提取：关注可能包含缺陷的区域，减少计算量

def extract_roi(image, annotation, margin=20):
    """
    提取感兴趣区域 (ROI)
    :param image: 输入图像
    :param annotation: 标注信息
    :param margin: 边界扩展像素
    :return: ROI 图像
    """
    img_array = np.array(image)

    # 合并所有缺陷框
    boxes = [obj['bbox'] for obj in annotation['objects']]
    xmin = min(b['xmin'] for b in boxes) - margin
    ymin = min(b['ymin'] for b in boxes) - margin
    xmax = max(b['xmax'] for b in boxes) + margin
    ymax = max(b['ymax'] for b in boxes) + margin

    # 确保不超出图像边界
    xmin = max(0, xmin)
    ymin = max(0, ymin)
    xmax = min(img_array.shape[1], xmax)
    ymax = min(img_array.shape[0], ymax)

    roi = img_array[ymin:ymax, xmin:xmax]
    return Image.fromarray(roi)

以下是基于 PyTorch 的 ResNet 训练示例：

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import models, transforms
from torch.utils.data import Dataset, DataLoader

class BTADDataset(Dataset):
    """自定义数据集类"""
    def __init__(self, images, annotations, transform=None):
        self.images = images
        self.annotations = annotations
        self.transform = transform

    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        image = self.images[idx]
        annotation = self.annotations[idx]

        # 转换为张量
        if self.transform:
            image = self.transform(image)

        # 创建目标张量
        # 这里简化为二分类问题（有缺陷 / 无缺陷）target = 1 if len(annotation['objects']) > 0 else 0
        target = torch.tensor(target, dtype=torch.long)

        return image, target

# 数据增强
train_transform = transforms.Compose([transforms.Resize((256, 256)),
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(10),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 加载数据集
images, annotations = load_btad_dataset('path/to/BTAD')
train_dataset = BTADDataset(images[:800], annotations[:800], train_transform)
test_dataset = BTADDataset(images[800:], annotations[800:], train_transform)

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

# 初始化模型
model = models.resnet18(pretrained=True)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, 2)  # 二分类

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# 训练循环
for epoch in range(10):
    model.train()
    running_loss = 0.0

    for inputs, labels in train_loader:
        optimizer.zero_grad()

        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    print(f'Epoch {epoch+1}, Loss: {running_loss/len(train_loader)}')