热轧钢带缺陷数据集，称为Xsteel表面缺陷数据集（X-SDD），其中包含七种典型的热轧带钢缺陷类型，共有1360个缺陷图像。与常用的NEU表面缺陷数据库（NEU-CLS）的六种缺陷类型相比，X-SD

热轧钢带缺陷数据集，称为Xsteel表面缺陷数据集（X-SDD），其中包含七种典型的热轧带钢缺陷类型，共有1360个缺陷图像。与常用的NEU表面缺陷数据库（NEU-CLS）的六种缺陷类型相比，X-SDD包含更多类型。

该数据集包含7种类型的1360张缺陷图像，包括238个渣夹杂物（简称“夹杂物”）、397个红色铁皮、122个铁皮灰、134个表面划痕（简称“划痕”）、63个板系氧化标度、203个整理辊印刷和203个温度系氧化标度。
Xsteel表面缺陷数据集（X-SDD），并探讨如何使用这个数据集进行深度学习模型的训练，特别是在目标检测任务中。

1. X-SDD 数据集介绍

数据集概述

• 数据集名称: Xsteel表面缺陷数据集（X-SDD）
• 数据集来源: 该数据集由工业界和学术界共同制作，旨在提供更全面的热轧钢带表面缺陷数据。
• 数据集内容: 包含1360张缺陷图像，涵盖7种典型的热轧带钢缺陷类型。
• 缺陷类别:
  • 渣夹杂物（夹杂物）
  • 红色铁皮
  • 铁皮灰
  • 表面划痕（划痕）
  • 板系氧化标度
  • 整理辊印刷
  • 温度系氧化标度

数据集目录结构

X-SDD/ ├── images/ │ ├── image1.jpg │ ├── image2.jpg │ └── ... └── labels/ ├── image1.txt ├── image2.txt └── ...

2. 数据集预处理

数据集划分

为了训练和验证模型，我们需要将数据集划分为训练集和验证集。通常情况下，可以将数据集按80%用于训练，20%用于验证。

importosimportrandomimportshutildefsplit_dataset(dataset_dir,output_dir,train_ratio=0.8):os.makedirs(os.path.join(output_dir,'images','train'),exist_ok=True)os.makedirs(os.path.join(output_dir,'images','val'),exist_ok=True)os.makedirs(os.path.join(output_dir,'labels','train'),exist_ok=True)os.makedirs(os.path.join(output_dir,'labels','val'),exist_ok=True)image_files=[fforfinos.listdir(os.path.join(dataset_dir,'images'))iff.endswith('.jpg')]random.shuffle(image_files)train_size=int(len(image_files)*train_ratio)train_images=image_files[:train_size]val_images=image_files[train_size:]forimage_fileintrain_images:shutil.copy(os.path.join(dataset_dir,'images',image_file),os.path.join(output_dir,'images','train',image_file))label_file=image_file.replace('.jpg','.txt')shutil.copy(os.path.join(dataset_dir,'labels',label_file),os.path.join(output_dir,'labels','train',label_file))forimage_fileinval_images:shutil.copy(os.path.join(dataset_dir,'images',image_file),os.path.join(output_dir,'images','val',image_file))label_file=image_file.replace('.jpg','.txt')shutil.copy(os.path.join(dataset_dir,'labels',label_file),os.path.join(output_dir,'labels','val',label_file))if__name__=="__main__":dataset_dir='X-SDD'output_dir='X-SDD_split'split_dataset(dataset_dir,output_dir)

3. 数据增强

数据增强可以帮助提高模型的泛化能力。我们将使用albumentations库来进行数据增强。

安装依赖

pipinstallalbumentations

数据增强代码

importalbumentationsasAfromalbumentations.pytorchimportToTensorV2importcv2importosimportshutildefaugment_image(image_path,label_path,output_dir,num_augmentations=10):image=cv2.imread(image_path)withopen(label_path,'r')asf:labels=[line.strip().split()forlineinf.readlines()]bboxes=[[float(x)forxinlabel[1:]]forlabelinlabels]class_labels=[int(label[0])forlabelinlabels]transform=A.Compose([A.HorizontalFlip(p=0.5),A.VerticalFlip(p=0.5),A.RandomRotate90(p=0.5),A.RandomBrightnessContrast(p=0.2),A.RandomGamma(p=0.2),A.Blur(blur_limit=3,p=0.2),A.CLAHE(p=0.2),A.HueSaturationValue(p=0.2),A.RandomResizedCrop(height=image.shape[0],width=image.shape[1],scale=(0.8,1.0),p=0.5),ToTensorV2()],bbox_params=A.BboxParams(format='yolo',label_fields=['class_labels']))foriinrange(num_augmentations):augmented=transform(image=image,bboxes=bboxes,class_labels=class_labels)augmented_image=augmented['image']augmented_bboxes=augmented['bboxes']augmented_class_labels=augmented['class_labels']output_image_path=os.path.join(output_dir,f'{os.path.basename(image_path)}_aug_{i}.jpg')cv2.imwrite(output_image_path,augmented_image.permute(1,2,0).numpy())output_label_path=os.path.join(output_dir,f'{os.path.basename(label_path)}_aug_{i}.txt')withopen(output_label_path,'w')asf:forjinrange(len(augmented_bboxes)):f.write(f"{augmented_class_labels[j]}{' '.join(map(str,augmented_bboxes[j]))}\n")defaugment_dataset(input_dir,output_dir,num_augmentations=10):ifos.path.exists(output_dir):shutil.rmtree(output_dir)os.makedirs(output_dir,exist_ok=True)forsplitin['train','val']:input_split_dir=os.path.join(input_dir,split)output_split_dir=os.path.join(output_dir,split)os.makedirs(output_split_dir,exist_ok=True)forimage_fileinos.listdir(os.path.join(input_split_dir,'images')):image_path=os.path.join(input_split_dir,'images',image_file)label_path=os.path.join(input_split_dir,'labels',image_file.replace('.jpg','.txt'))augment_image(image_path,label_path,output_split_dir,num_augmentations)if__name__=="__main__":input_dir='X-SDD_split'output_dir='X-SDD_augmented'augment_dataset(input_dir,output_dir,num_augmentations=10)

4. 训练YOLOv5模型

安装YOLOv5

gitclone https://github.com/ultralytics/yolov5cdyolov5 pipinstall-rrequirements.txt

数据集配置文件

创建一个data.yaml文件，配置数据集路径和类别信息。

# data.yamltrain:../X-SDD_augmented/images/trainval:../X-SDD_augmented/images/valnc:7names:['夹杂物','红色铁皮','铁皮灰','划痕','板系氧化标度','整理辊印刷','温度系氧化标度']

训练脚本

# src/train_yolov5.pyimporttorchfromyolov5.models.experimentalimportattempt_loadfromyolov5.utils.torch_utilsimportselect_devicefromyolov5.utils.generalimportcheck_img_sizefromyolov5.utils.datasetsimportcreate_dataloaderfromyolov5.utils.lossimportComputeLossfromyolov5.models.yoloimportModelimporttimeimportyamldeftrain_model(data_yaml_path,model_config,epochs,batch_size,img_size,device):withopen(data_yaml_path,'r')asf:data=yaml.safe_load(f)train_loader=create_dataloader(data['train'],img_size,batch_size,32)[0]val_loader=create_dataloader(data['val'],img_size,batch_size,32)[0]model=Model(model_config,ch=3,nc=data['nc']).to(device)model.train()optimizer=torch.optim.Adam(model.parameters(),lr=0.001)compute_loss=ComputeLoss(model)start_time=time.time()forepochinrange(epochs):epoch_start_time=time.time()fori,(imgs,targets,paths,_)inenumerate(train_loader):imgs=imgs.to(device)targets=targets.to(device)pred=model(imgs)loss,loss_items=compute_loss(pred,targets)optimizer.zero_grad()loss.backward()optimizer.step()ifi%10==0:print(f'Epoch [{epoch+1}/{epochs}], Step [{i+1}/{len(train_loader)}], Loss:{loss.item()}')epoch_end_time=time.time()epoch_duration=epoch_end_time-epoch_start_timeprint(f'Epoch [{epoch+1}/{epochs}] completed in{epoch_duration:.2f}seconds')torch.save(model.state_dict(),f'models/yolov5_custom_epoch_{epoch+1}.pth')end_time=time.time()total_duration=end_time-start_timeprint(f'Total training time:{total_duration:.2f}seconds')if__name__=="__main__":data_yaml_path='data.yaml'model_config='models/yolov5s.yaml'epochs=100batch_size=16img_size=640device=select_device('0')# 使用GPU，如果需要使用CPU，可以改为'cpu'train_model(data_yaml_path,model_config,epochs,batch_size,img_size,device)

5. 运行训练脚本

1. 数据增强：

   python src/utils/data_augmentation.py

2. 训练模型：

   python src/train_yolov5.py

6. 详细解释

数据增强

• augment_image: 对单张图像进行数据增强，并保存增强后的图像和标签。
• augment_dataset: 对整个数据集进行数据增强，生成新的数据集。

训练脚本

• 数据加载：使用create_dataloader创建训练和验证的数据加载器。
• 模型加载：加载YOLOv5模型，并设置为训练模式。
• 优化器和损失函数：定义Adam优化器和YOLOv5的损失函数。
• 训练循环：记录每个epoch的训练时间和总训练时间，进行前向传播、计算损失、反向传播和优化。
• 保存模型：每个epoch结束后保存模型的权重。

7. 注意事项

1. 数据集路径：确保数据集路径正确，特别是data.yaml文件中的路径。
2. 模型配置：确保模型配置文件路径正确。
3. 图像大小：img_size可以根据实际需求调整，通常使用640或1280。
4. 设备：确保设备（CPU或GPU）可用。

8. 总结

通过以上步骤，你可以构建一个完整的X-SDD热轧钢带表面缺陷检测数据集，并使用YOLOv5进行训练。