管道内裂缝检测数据集 2000张 管道裂缝 带标注voc yol

时间:2024-10-07 09:25:55

管道内裂缝检测数据集 (Pipeline Crack Detection Dataset)

数据集概述

该数据集是一个专门用于训练和评估管道内裂缝检测模型的数据集。数据集包含2000张图像,每张图像都带有标注信息,标注格式为VOC(Pascal VOC)和YOLO格式。这些图像涵盖了各种管道内的裂缝情况,适用于基于深度学习的目标检测任务。通过这个数据集,可以训练出能够准确检测和定位管道内部裂缝的模型,从而帮助进行管道维护和安全检查。

数据集特点
  • 高质量图像:数据集中的图像具有高分辨率,能够提供丰富的细节信息。
  • 带标注:每张图像都有详细的标注信息,包括裂缝的位置和大小。
  • 多格式标注:标注信息同时以VOC和YOLO格式提供,方便不同框架的使用。
  • 实际应用场景:适用于需要精确检测管道内部裂缝的场景,如城市供水管道、油气输送管道等。
数据集结构

pipeline_crack_dataset/
├── images/                            # 图像文件
│   ├── 00001.jpg                      # 示例图像
│   ├── 00002.jpg
│   └── ...
├── annotations/                       # 标注文件
│   ├── VOC/                           # Pascal VOC格式标注
│   │   ├── 00001.xml                  # 示例VOC标注文件
│   │   ├── 00002.xml
│   │   └── ...
│   ├── YOLO/                          # YOLO格式标注
│   │   ├── 00001.txt                  # 示例YOLO标注文件
│   │   ├── 00002.txt
│   │   └── ...
├── data.yaml                          # 类别描述文件
├── README.md                          # 数据集说明
数据集内容
  1. images/

    • 功能:存放图像文件。
    • 内容
      • 00001.jpg:示例图像。
      • 00002.jpg:另一张图像。
      • ...
  2. annotations/

    • 功能:存放标注文件。
    • 内容
      • VOC/:存放Pascal VOC格式的标注文件。
        • 00001.xml:示例VOC标注文件。
        • 00002.xml:另一张图像的VOC标注文件。
        • ...
      • YOLO/:存放YOLO格式的标注文件。
        • 00001.txt:示例YOLO标注文件。
        • 00002.txt:另一张图像的YOLO标注文件。
        • ...
  3. data.yaml

    • 功能:定义数据集的类别和其他相关信息。
    • 内容
      train: pipeline_crack_dataset/images
      val: pipeline_crack_dataset/images
      nc: 1
      names: ['Crack']
  4. README.md

    • 功能:数据集的详细说明文档。
    • 内容
      • 数据集的来源和用途。
      • 数据集的结构和内容。
      • 如何使用数据集进行模型训练和评估。
      • 其他注意事项和建议。
数据集统计
  • 总图像数量:2000张
  • 总标注框数量:2296个
  • 裂缝类别:1类(Cracks)
  • 平均每张图像的标注框数量:约1.15个
使用说明
  1. 环境准备

    • 安装依赖库:
      pip install -r requirements.txt
    • 确保安装了常用的深度学习库,例如torchtorchvisionnumpy等。
  2. 数据集路径设置

    • 将数据集解压到项目目录下,确保路径正确。
  3. 训练模型

     

    以下是一个使用PyTorch和torchvision库进行管道内裂缝检测的示例代码。我们将使用预训练的Faster R-CNN模型,并对其进行微调以适应我们的数据集。

    import torch
    import torchvision
    from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
    from torchvision.models.detection import FasterRCNN
    from torchvision.models.detection.rpn import AnchorGenerator
    from torchvision.transforms import functional as F
    from torch.utils.data import DataLoader, Dataset
    from PIL import Image
    import os
    import xml.etree.ElementTree as ET
    
    # 自定义数据集类
    class PipelineCrackDataset(Dataset):
        def __init__(self, root, transforms=None):
            self.root = root
            self.transforms = transforms
            self.imgs = list(sorted(os.listdir(os.path.join(root, "images"))))
            self.annotations = list(sorted(os.listdir(os.path.join(root, "annotations", "VOC"))))
    
        def __getitem__(self, idx):
            img_path = os.path.join(self.root, "images", self.imgs[idx])
            annotation_path = os.path.join(self.root, "annotations", "VOC", self.annotations[idx])
    
            img = Image.open(img_path).convert("RGB")
            annotation_root = ET.parse(annotation_path).getroot()
    
            boxes = []
            labels = []
            for obj in annotation_root.findall('object'):
                xmin, ymin, xmax, ymax = [int(obj.find('bndbox').find(tag).text) for tag in ('xmin', 'ymin', 'xmax', 'ymax')]
                boxes.append([xmin, ymin, xmax, ymax])
                labels.append(1)  # 裂缝类别
    
            boxes = torch.as_tensor(boxes, dtype=torch.float32)
            labels = torch.as_tensor(labels, dtype=torch.int64)
    
            target = {}
            target["boxes"] = boxes
            target["labels"] = labels
            target["image_id"] = torch.tensor([idx])
    
            if self.transforms is not None:
                img, target = self.transforms(img, target)
    
            return F.to_tensor(img), target
    
        def __len__(self):
            return len(self.imgs)
    
    # 数据预处理
    def get_transform(train):
        transforms = []
        # transforms.append(T.ToTensor())
        if train:
            transforms.append(torchvision.transforms.RandomHorizontalFlip(0.5))
        return torchvision.transforms.Compose(transforms)
    
    # 加载数据集
    dataset = PipelineCrackDataset(root='pipeline_crack_dataset', transforms=get_transform(train=True))
    dataset_test = PipelineCrackDataset(root='pipeline_crack_dataset', transforms=get_transform(train=False))
    
    indices = torch.randperm(len(dataset)).tolist()
    dataset = torch.utils.data.Subset(dataset, indices[:-50])
    dataset_test = torch.utils.data.Subset(dataset_test, indices[-50:])
    
    data_loader = DataLoader(dataset, batch_size=2, shuffle=True, num_workers=4, collate_fn=lambda x: tuple(zip(*x)))
    data_loader_test = DataLoader(dataset_test, batch_size=1, shuffle=False, num_workers=4, collate_fn=lambda x: tuple(zip(*x)))
    
    # 定义模型
    model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
    num_classes = 2  # 1类裂缝 + 背景
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
    
    # 设置设备
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    
    # 定义优化器
    params = [p for p in model.parameters() if p.requires_grad]
    optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005)
    
    # 训练模型
    num_epochs = 10
    for epoch in range(num_epochs):
        model.train()
        for images, targets in data_loader:
            images = list(image.to(device) for image in images)
            targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
    
            loss_dict = model(images, targets)
            losses = sum(loss for loss in loss_dict.values())
    
            optimizer.zero_grad()
            losses.backward()
            optimizer.step()
    
        print(f'Epoch {epoch+1}/{num_epochs}, Loss: {losses.item()}')
    
        # 验证模型
        model.eval()
        with torch.no_grad():
            for images, targets in data_loader_test:
                images = list(image.to(device) for image in images)
                targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
                outputs = model(images)
    
    # 保存模型
    torch.save(model.state_dict(), 'pipeline_crack_detection_model.pth')
注意事项
  • 数据格式:确保输入的数据格式正确,特别是图像文件和标注文件的格式。
  • 超参数调整:根据实际情况调整学习率、批大小等超参数,以获得最佳训练效果。
  • 硬件要求:建议使用GPU进行训练和推理,以加快处理速度。如果没有足够的计算资源,可以考虑使用云服务提供商的GPU实例。
  • 数据增强:可以通过数据增强技术(如随机翻转、旋转等)来增加模型的鲁棒性。
  • 模型选择:除了Faster R-CNN,还可以尝试其他目标检测模型,如YOLOv5、SSD等,以找到最适合当前任务的模型。
  • 类别平衡:虽然数据集中裂缝类别的样本数量相对均衡,但在实际应用中可能需要进一步检查并处理类别不平衡问题,例如通过过采样或欠采样方法。

通过上述步骤,你可以成功地使用这个高质量的管道内裂缝检测数据集进行模型训练和评估。该数据集不仅适用于学术研究,还可以应用于实际的管道维护和安全检查中,帮助提升裂缝检测的准确性和效率。希望这个数据集能帮助你更好地理解和应用最新的深度学习技术。