Model card

Swin-Defect-Detection

基于 Swin Transformer + YOLOv8 的工业缺陷检测模型。

模型简介

本项目结合 Swin Transformer 的强大特征提取能力与 YOLOv8 的高效检测头，实现工业产品表面缺陷的高精度检测。

模型列表

模型	说明	大小	适用场景
`checkpoints/dual_pretrained/best.pt`	双路 Swin 预训练模型	384 MB	高精度检测，需要较强特征提取能力
`checkpoints/yolov8_base/best.pt`	YOLOv8 基础模型	22 MB	轻量级部署，推理速度快
`checkpoints/yolov8_neck/best.pt`	YOLOv8 + FPN Neck 模型	314 MB	平衡精度与速度

环境要求

pip install torch torchvision
pip install ultralytics

使用方法

下载模型

from modelscope.hub.snapshot_download import snapshot_download

# 下载整个模型仓库
model_dir = snapshot_download('Brilliantccc/Swin-Defect-Detection')

# 或下载指定文件
model_dir = snapshot_download(
    'Brilliantccc/Swin-Defect-Detection',
    file_path='checkpoints/yolov8_base/best.pt'
)

推理示例

from ultralytics import YOLO

# 加载模型
model = YOLO('checkpoints/yolov8_base/best.pt')

# 推理
results = model('test_image.jpg')

# 显示结果
results[0].show()

训练

# YOLOv8 基础训练
python train_yolov8.py --data configs/mvtec.yaml --epochs 100

# YOLOv8 + Neck 训练
python train_yolov8_neck.py --data configs/mvtec.yaml --epochs 100

# 双路 Swin 预训练
python train_dual_pretrained.py --data configs/mvtec.yaml --epochs 100

数据集格式

支持 MVTec AD 和 VisA 数据集格式，使用 YOLO 格式的标注文件。

许可证

Apache 2.0