钢铁焊缝缺陷数据集，支持目标检测与实例分割数据集

用于Yolo的实例分割与缺陷检测

。
数据集焊缝缺陷目标检测数据集6000张存在缺陷的图片。

4类缺陷类别：
1、Crack - 裂纹
2、Porosity - 孔洞
3、Spatters - 飞溅
4、Welding line - 焊接
 

1
 

如何使用YOLOv5进行目标检测和Mask R-CNN进行实例分割。YOLO系列模型主要针对目标检测任务，对于实例分割任务，Mask R-CNN是一个更为合适的选择。

焊缝缺陷检测（目标检测）

1. 环境准备

首先确保安装了必要的依赖项：

pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113
git clone https://github.com/ultralytics/yolov5.git
cd yolov5
pip install -r requirements.txt

2. 数据准备

创建一个data.yaml文件来定义数据集路径和类别信息：

train: ./path/to/train/images
val: ./path/to/val/images

nc: 4  # 类别数量
names: ['Crack', 'Porosity', 'Spatters', 'Welding line']

请根据实际情况调整路径。

3. 模型训练

使用YOLOv5进行训练：

python train.py --img 640 --batch 16 --epochs 100 --data path/to/data.yaml --weights yolov5s.pt

4. 推理及可视化

对新的图片进行推理并可视化结果：

python detect.py --weights runs/train/exp/weights/best.pt --source ./path/to/new/image.jpg --conf-thres 0.5 --save-txt --save-conf

焊缝缺陷分割（实例分割）

对于实例分割任务，我们将使用Detectron2框架中的Mask R-CNN模型。

1. 环境准备

安装Detectron2及相关依赖：

pip install torch torchvision
pip install 'git+https://github.com/facebookresearch/detectron2.git'

2. 数据准备

Detectron2需要COCO格式的数据集。假设您的数据已经转换为COCO格式（包括annotations.json和图像目录）。如果没有，请先转换数据格式。

3. 配置与训练

编写配置文件并开始训练：

from detectron2.config import get_cfg
from detectron2.engine import DefaultTrainer, default_argument_parser, default_setup, launch
from detectron2.data.datasets import register_coco_instances

# 注册自定义数据集
register_coco_instances("weld_defect_train", {}, "./path/to/train/annotations.json", "./path/to/train/images")
register_coco_instances("weld_defect_val", {}, "./path/to/val/annotations.json", "./path/to/val/images")

cfg = get_cfg()
cfg.merge_from_file("./detectron2/configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
cfg.DATASETS.TRAIN = ("weld_defect_train",)
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = "detectron2://COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x/137849600/model_final_f10217.pkl"  # 预训练权重
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025
cfg.SOLVER.MAX_ITER = 300    # 300次迭代。您可以根据需要调整
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128   # 较小的值更快，但可能不太准确。可以尝试不同的值
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 4  # 类别数量

os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg) 
trainer.resume_or_load(resume=False)
trainer.train()

4. 推理及可视化

加载训练好的模型并对新图片进行推理：

from detectron2.utils.visualizer import ColorMode, Visualizer
from detectron2.engine import DefaultPredictor
import cv2

cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")  # 加载最佳模型权重
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5  # 设置阈值
predictor = DefaultPredictor(cfg)

# 对新图片进行预测
im = cv2.imread("./path/to/new/image.jpg")
outputs = predictor(im)

# 可视化结果
v = Visualizer(im[:, :, ::-1], metadata=weld_defect_metadata, scale=0.8, instance_mode=ColorMode.IMAGE_BW)
v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
cv2.imshow('image', v.get_image()[:, :, ::-1])
cv2.waitKey(0)

总结

上述步骤分别为焊缝缺陷检测（使用YOLOv5）和焊缝缺陷分割（使用Mask R-CNN）完整的流程，包括环境设置、数据准备、模型训练、推理及结果可视化。可以做个参考，同学们。