Mask R-CNN 工业缺陷检测模型

基于 Mask R-CNN 的工业缺陷检测和实例分割模型，支持像素级缺陷定位。

模型信息

模型权重

1. MVTec 训练模型

在 MVTec Anomaly Detection 数据集上训练，包含 15 个工业品类。

训练配置：

• Backbone: ResNet50-FPN (COCO 预训练)
• 学习率: 0.005 (Head) / 0.0005 (Backbone)
• 批次大小: 16
• 训练轮次: 23 (早停)
• 每轮时间: ~165 秒

2. VisA 微调模型

在 VisA (Visual Anomaly) 数据集上微调，包含 12 个工业品类。

训练配置：

• Backbone: ResNet50-FPN (MVTec 预训练)
• 学习率: 0.002 (Head) / 0.0002 (Backbone)
• 批次大小: 16
• 训练轮次: 60 (早停)
• 每轮时间: ~430 秒

使用方法

1. 安装依赖

pip install torch torchvision
pip install pycocotools
pip install opencv-python

2. 加载模型

import torch
from torchvision.models.detection import maskrcnn_resnet50_fpn
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor

def load_mask_rcnn(weights_path, num_classes=2):
    """加载 Mask R-CNN 模型"""
    # 创建模型
    model = maskrcnn_resnet50_fpn(pretrained=False)
    
    # 替换分类头
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
    
    # 替换 mask 头
    in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
    hidden_layer = 256
    model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask, hidden_layer, num_classes)
    
    # 加载权重
    checkpoint = torch.load(weights_path, map_location='cpu')
    model.load_state_dict(checkpoint['model_state_dict'])
    
    return model

# 加载 MVTec 模型
model = load_mask_rcnn('checkpoints/maskrcnn_mvtec/best.pth')

# 加载 VisA 模型
model = load_mask_rcnn('checkpoints/maskrcnn_visa/best.pth')

3. 推理

import cv2
import torch
import numpy as np
from torchvision import transforms

def predict(model, image_path, score_threshold=0.5):
    """推理单张图片"""
    model.eval()
    
    # 读取图片
    img = cv2.imread(image_path)
    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    
    # 转换为 tensor
    transform = transforms.Compose([
        transforms.ToTensor(),
    ])
    img_tensor = transform(img_rgb)
    
    # 推理
    with torch.no_grad():
        predictions = model([img_tensor])
    
    pred = predictions[0]
    
    # 过滤低分预测
    keep = pred['scores'] > score_threshold
    
    return {
        'boxes': pred['boxes'][keep].numpy(),
        'labels': pred['labels'][keep].numpy(),
        'scores': pred['scores'][keep].numpy(),
        'masks': pred['masks'][keep].numpy(),
    }

def visualize(image_path, predictions, output_path='result.jpg'):
    """可视化预测结果"""
    img = cv2.imread(image_path)
    
    for i in range(len(predictions['boxes'])):
        box = predictions['boxes'][i].astype(int)
        score = predictions['scores'][i]
        mask = predictions['masks'][i, 0] > 0.5
        
        # 绘制 bbox
        cv2.rectangle(img, (box[0], box[1]), (box[2], box[3]), (0, 0, 255), 2)
        cv2.putText(img, f'anomaly: {score:.2f}', (box[0], box[1] - 10),
                   cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
        
        # 绘制 mask
        colored_mask = np.zeros_like(img)
        colored_mask[mask] = [0, 0, 255]
        img = cv2.addWeighted(img, 1, colored_mask, 0.4, 0)
    
    cv2.imwrite(output_path, img)
    return img

# 使用示例
model = load_mask_rcnn('checkpoints/maskrcnn_mvtec/best.pth')
predictions = predict(model, 'test_image.jpg', score_threshold=0.3)
visualize('test_image.jpg', predictions, 'result.jpg')

评估结果

MVTec AD 数据集

VisA 数据集

模型架构

输入图像 (H, W, 3)
    │
    ▼
┌─────────────────┐
│  ResNet50-FPN    │  特征提取
│  (COCO 预训练)   │
└────────┬────────┘
         │
         ▼
┌─────────────────┐
│  Region Proposal │  候选区域生成
│  Network (RPN)   │
└────────┬────────┘
         │
         ▼
┌─────────────────┐
│  RoI Align       │  特征对齐
└────────┬────────┘
         │
    ┌────┴────┐
    │         │
    ▼         ▼
┌────────┐ ┌────────┐
│ Box    │ │ Mask   │  分类+回归 / 分割
│ Head   │ │ Head   │
└────────┘ └────────┘
    │         │
    ▼         ▼
  bbox +    mask
  class     (28×28)

训练细节

数据预处理

• 图片：原始尺寸，归一化到 [0, 1]
• 标注：COCO 格式 (bbox + segmentation polygon)
• 数据增强：RandomHorizontalFlip, ColorJitter, RandomErasing

训练策略

• 优化器：SGD (momentum=0.9, weight_decay=0.0005)
• 学习率：Cosine Annealing with Warmup
• 混合精度：AMP 启用
• 梯度裁剪：max_norm=10.0
• 早停：patience=15

硬件环境

• GPU: NVIDIA RTX 4090 (24GB)
• CPU: 16 vCPU Intel Xeon Platinum 8358P
• 训练时间：MVTec ~1小时，VisA ~3小时

参考

• Mask R-CNN Paper
• torchvision Mask R-CNN
• MVTec AD Dataset
• VisA Dataset

引用

@article{he2017mask,
  title={Mask R-CNN},
  author={He, Kaiming and Gkioxari, Georgia and Doll{\'a}r, Piotr and Girshick, Ross},
  journal={IEEE Transactions on Pattern Analysis and Machine Intelligence},
  year={2017}
}