红外变电站分割数据集,标注为json格式,总共有5类
避雷器(289张),绝缘子(919张),电流互感器(413张),套管(161张),电压互感器(153张)
红外变电站分割数据集 (Infrared Substation Segmentation Dataset)
数据集概述
该数据集是一个专门用于训练和评估红外图像中变电站设备分割模型的数据集。数据集包含约1935张红外图像,每张图像都带有详细的标注信息,标注格式为JSON格式。这些图像涵盖了五种不同类型的变电站设备:避雷器、绝缘子、电流互感器、套管和电压互感器。适用于基于深度学习的语义分割任务,通过这个数据集,可以训练出能够在红外图像中准确分割和识别不同类型变电站设备的模型,从而帮助进行电力设施维护、故障检测等应用。
数据集特点
- 高质量红外图像:数据集中的图像均为红外成像,具有高分辨率,能够提供丰富的温度分布信息,特别适合变电站设备的特征分析。
- 带标注:每张图像都有详细的标注信息,包括设备的位置和边界。
- JSON格式标注:标注信息以JSON格式提供,方便直接使用于支持JSON格式的语义分割框架。
- 实际应用场景:适用于需要精确分割和识别红外图像中变电站设备的场景,如电力设施维护、故障检测系统等。
数据集结构
infrared_substation_segmentation_dataset/
├── images/ # 图像文件
│ ├── 00001.jpg # 示例图像
│ ├── 00002.jpg
│ └── ...
├── annotations/ # JSON格式标注文件
│ ├── 00001.json # 示例JSON标注文件
│ ├── 00002.json
│ └── ...
├── data.yaml # 类别描述文件
├── README.md # 数据集说明
└── model/ # 预训练模型文件夹(可选)└── infrared_substation_segmentation_model.pth # 预训练模型(如果有的话)
数据集内容
-
images/
- 功能:存放图像文件。
- 内容:
00001.jpg
:示例图像。00002.jpg
:另一张图像。- ...
-
annotations/
- 功能:存放JSON格式的标注文件。
- 内容:
00001.json
:示例JSON标注文件。00002.json
:另一张图像的JSON标注文件。- ...
-
data.yaml
- 功能:定义数据集的类别和其他相关信息。
- 内容: yaml
深色版本
train: infrared_substation_segmentation_dataset/images val: infrared_substation_segmentation_dataset/images nc: 5 names: ['arrester', 'insulator', 'current transformer', 'bushing', 'voltage transformer']
-
README.md
- 功能:数据集的详细说明文档。
- 内容:
- 数据集的来源和用途。
- 数据集的结构和内容。
- 如何使用数据集进行模型训练和评估。
- 其他注意事项和建议。
-
model/
(可选)- 功能:存放预训练模型文件。
- 内容:
infrared_substation_segmentation_model.pth
:预训练的模型文件(如果有的话)。
数据集统计
- 总图像数量:约1935张
- 类别:5类
- 类别列表及图像数量:
arrester
(避雷器):289张insulator
(绝缘子):919张current transformer
(电流互感器):413张bushing
(套管):161张voltage transformer
(电压互感器):153张
使用说明
- 环境准备:确保安装了常用的深度学习库,例如
torch
,torchvision
,numpy
等。 - 数据集路径设置:将数据集解压到项目目录下,并确保路径正确。
- 加载预训练模型:如果有预训练模型,可以直接加载并对其进行微调或直接使用。
- 数据增强:可以通过随机翻转、旋转等方法增加数据多样性,提高模型鲁棒性。
- 超参数调整:根据实际情况调整学习率、批大小等超参数,以获得最佳训练效果。
- 硬件要求:建议使用GPU进行训练和推理,以加快处理速度。如果没有足够的计算资源,可以考虑使用云服务提供商的GPU实例。
- 类别平衡:虽然数据集中各类别的样本数量相对均衡,但在实际应用中可能需要进一步检查并处理类别不平衡问题,例如通过过采样或欠采样方法。
关键代码示例
以下是一个使用PyTorch和torchvision
库进行红外变电站设备分割的示例代码。我们将使用预训练的DeepLab v3+模型,并对其进行微调以适应我们的数据集。
import torch
import torchvision
from torchvision.models.segmentation import deeplabv3_resnet101
from torchvision.transforms import functional as F
from torch.utils.data import DataLoader, Dataset
from PIL import Image
import os
import json
import numpy as np# 自定义数据集类
class InfraredSubstationSegmentationDataset(Dataset):def __init__(self, root, transforms=None):self.root = rootself.transforms = transformsself.imgs = list(sorted(os.listdir(os.path.join(root, "images"))))self.annotations = list(sorted(os.listdir(os.path.join(root, "annotations"))))def __getitem__(self, idx):img_path = os.path.join(self.root, "images", self.imgs[idx])annotation_path = os.path.join(self.root, "annotations", self.annotations[idx])img = Image.open(img_path).convert("RGB")with open(annotation_path, 'r') as f:annotation = json.load(f)mask = np.zeros((img.height, img.width), dtype=np.uint8)for obj in annotation['objects']:class_id = {'arrester': 1,'insulator': 2,'current transformer': 3,'bushing': 4,'voltage transformer': 5}[obj['class']]polygon = np.array(obj['polygon'])cv2.fillPoly(mask, [polygon], class_id)mask = Image.fromarray(mask)if self.transforms is not None:img, mask = self.transforms(img, mask)return F.to_tensor(img), F.to_tensor(mask)def __len__(self):return len(self.imgs)# 数据预处理
def get_transform(train):transforms = []if train:transforms.append(torchvision.transforms.RandomHorizontalFlip(0.5))return torchvision.transforms.Compose(transforms)# 加载数据集
dataset = InfraredSubstationSegmentationDataset(root='infrared_substation_segmentation_dataset', transforms=get_transform(train=True))
dataset_test = InfraredSubstationSegmentationDataset(root='infrared_substation_segmentation_dataset', transforms=get_transform(train=False))indices = torch.randperm(len(dataset)).tolist()
dataset = torch.utils.data.Subset(dataset, indices[:-193])
dataset_test = torch.utils.data.Subset(dataset_test, indices[-193:])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 = deeplabv3_resnet101(pretrained=True)
model.classifier[-1] = torch.nn.Conv2d(256, 6, kernel_size=(1, 1), stride=(1, 1)) # 5类 + 背景
model.aux_classifier[-1] = torch.nn.Conv2d(256, 6, kernel_size=(1, 1), stride=(1, 1)) # 5类 + 背景# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)# 定义优化器
params_to_optimize = [{"params": [p for p in model.backbone.parameters() if p.requires_grad]},{"params": [p for p in model.classifier.parameters() if p.requires_grad]},
]
optimizer = torch.optim.SGD(params_to_optimize, 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 = list(target.to(device) for target in targets)optimizer.zero_grad()outputs = model(images)loss = torch.nn.functional.cross_entropy(outputs['out'], targets)if 'aux' in outputs:loss += 0.5 * torch.nn.functional.cross_entropy(outputs['aux'], targets)loss.backward()optimizer.step()print(f'Epoch {epoch+1}/{num_epochs}, Loss: {loss.item()}')# 验证模型model.eval()with torch.no_grad():for images, targets in data_loader_test:images = list(image.to(device) for image in images)targets = list(target.to(device) for target in targets)outputs = model(images)# 保存模型
torch.save(model.state_dict(), 'infrared_substation_segmentation_model.pth')
注意事项
- 数据格式:确保输入的数据格式正确,特别是图像文件和标注文件的格式。
- 超参数调整:根据实际情况调整学习率、批大小等超参数,以获得最佳训练效果。
- 硬件要求:建议使用GPU进行训练和推理,以加快处理速度。如果没有足够的计算资源,可以考虑使用云服务提供商的GPU实例。
- 数据增强:可以通过数据增强技术(如随机翻转、旋转等)来增加模型的鲁棒性。
- 模型选择:除了DeepLab v3+,还可以尝试其他语义分割模型,如U-Net、FCN等,以找到最适合当前任务的模型。
- 类别平衡:虽然数据集中各类别的样本数量相对均衡,但在实际应用中可能需要进一步检查并处理类别不平衡问题,例如通过过采样或欠采样方法。
通过上述步骤,你可以成功地使用这个高质量的红外变电站分割数据集进行模型训练和评估。该数据集不仅适用于学术研究,还可以应用于实际的电力设施维护、故障检测系统等领域,帮助提升对红外图像中变电站设备的分割准确性和效率。希望这个数据集能帮助你更好地理解和应用最新的深度学习技术。