都有对应标注格式VOC (XML)选配Y0L0(TXT) .label| pic_ num| box_ num绑扎不规范: (3717 14510)并线线夹保护壳缺失: (3317 11285)耐张线夹保护壳缺失: (3748 16148)横杆腐蚀: (987 1556)塔头损坏: (972 1021)total: (9838 44520)电力电塔电线缺陷检测数据集统计表标签名称 图片数量 边框数量绑扎不规范 (绑扎不规范) 3717 14510并线线夹保护壳缺失 (并线线夹保护壳缺失) 3317 11285耐张线夹保护壳缺失 (耐张线夹保护壳缺失) 3748 16148横杆腐蚀 (横杆腐蚀) 987 1556塔头损坏 (塔头损坏) 972 1021总计 9838 44520项目结构深色版本power_line_defect_detection/├── dataset/│ ├── images/│ │ └── *.jpg│ ├── labels_voc/│ │ └── *.xml│ ├── labels_yolo/│ │ └── *.txt├── models/│ └── yolov8/├── src/│ ├── train.py│ ├── predict.py│ ├── utils.py│ ├── convert_voc_to_yolo.py├── weights/│ └── best_model.pt├── requirements.txt└── README.md安装依赖首先确保你已经安装了必要的库。创建一个requirements.txt文件内容如下深色版本torchtorchvisionnumpypandasmatplotlibtqdmpyyamlopencv-pythonultralyticslxml然后使用以下命令安装依赖bash深色版本pip install -r requirements.txt2. 数据集准备确保你的数据集已经按照以下结构组织深色版本dataset/├── images/│ └── *.jpg├── labels_voc/│ └── *.xml├── labels_yolo/│ └── *.txt每个文件夹中包含对应的图像文件和标签文件。确保所有图像文件都是.jpg格式标签文件是VOC格式的.xml文件和YOLO格式的.txt文件。数据集转换工具为了方便使用YOLOv8模型我们需要将VOC格式的标注文件转换为YOLO格式。可以使用以下脚本进行转换。3.1 src/convert_voc_to_yolo.pypython深色版本import osimport xml.etree.ElementTree as ETimport shutildef convert_voc_to_yolo(voc_dir, yolo_dir, class_names, image_size(640, 640)):os.makedirs(yolo_dir, exist_okTrue)for xml_file in os.listdir(voc_dir): if not xml_file.endswith(.xml): continue tree ET.parse(os.path.join(voc_dir, xml_file)) root tree.getroot() image_width int(root.find(size/width).text) image_height int(root.find(size/height).text) yolo_file os.path.join(yolo_dir, xml_file.replace(.xml, .txt)) with open(yolo_file, w) as f: for obj in root.findall(object): class_name obj.find(name).text if class_name not in class_names: continue class_id class_names.index(class_name) bndbox obj.find(bndbox) xmin int(bndbox.find(xmin).text) ymin int(bndbox.find(ymin).text) xmax int(bndbox.find(xmax).text) ymax int(bndbox.find(ymax).text) x_center (xmin xmax) / 2.0 / image_width y_center (ymin ymax) / 2.0 / image_height width (xmax - xmin) / image_width height (ymax - ymin) / image_height f.write(f{class_id} {x_center} {y_center} {width} {height}\n)ifname “main”:voc_dir “dataset/labels_voc”yolo_dir “dataset/labels_yolo”class_names [“绑扎不规范”, “并线线夹保护壳缺失”, “耐张线夹保护壳缺失”, “横杆腐蚀”, “塔头损坏”]convert_voc_to_yolo(voc_dir, yolo_dir, class_names)4. 数据集配置创建一个数据集类用于加载和预处理数据。4.1 src/utils.pypython深色版本import osimport torchfrom torch.utils.data import Dataset, DataLoaderfrom torchvision import transformsfrom PIL import Imageclass PowerLineDefectDetectionDataset(Dataset):definit(self, image_dir, label_dir, transformNone):self.image_dir image_dirself.label_dir label_dirself.transform transformself.image_files os.listdir(image_dir)self.label_files os.listdir(label_dir)def __len__(self): return len(self.image_files) def __getitem__(self, index): img_path os.path.join(self.image_dir, self.image_files[index]) label_path os.path.join(self.label_dir, self.image_files[index].replace(.jpg, .txt)) image Image.open(img_path).convert(RGB) if os.path.exists(label_path): with open(label_path, r) as f: labels f.readlines() labels [line.strip().split() for line in labels] labels [[int(label[0])] list(map(float, label[1:])) for label in labels] else: labels [] if self.transform: image self.transform(image) return image, labelsdef get_data_loaders(image_dir, label_dir, batch_size16, num_workers4):transform transforms.Compose([transforms.Resize((640, 640)),transforms.ToTensor(),transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]),])dataset PowerLineDefectDetectionDataset(image_dir, label_dir, transformtransform) train_size int(0.8 * len(dataset)) val_size len(dataset) - train_size train_dataset, val_dataset torch.utils.data.random_split(dataset, [train_size, val_size]) train_loader DataLoader(train_dataset, batch_sizebatch_size, shuffleTrue, num_workersnum_workers) val_loader DataLoader(val_dataset, batch_sizebatch_size, shuffleFalse, num_workersnum_workers) return train_loader, val_loader模型定义使用YOLOv8模型进行检测任务。这里我们使用预训练的YOLOv8模型并对其进行微调。5.1 src/train.pypython深色版本import torchimport torch.optim as optimfrom torch.utils.tensorboard import SummaryWriterfrom tqdm import tqdmfrom ultralytics import YOLOfrom src.utils import get_data_loadersimport torch.nn.functional as Fdef train_model(image_dir, label_dir, epochs100, batch_size16, learning_rate1e-4):device torch.device(“cuda” if torch.cuda.is_available() else “cpu”)model YOLO(yolov8n.pt) # Load a pretrained YOLOv8 model model model.to(device) train_loader, val_loader get_data_loaders(image_dir, label_dir, batch_sizebatch_size) optimizer optim.Adam(model.parameters(), lrlearning_rate) criterion torch.nn.MSELoss() writer SummaryWriter() for epoch in range(epochs): model.train() running_loss 0.0 for images, labels in tqdm(train_loader, descfEpoch {epoch 1}/{epochs}): images images.to(device) labels [torch.tensor(label, dtypetorch.float32).to(device) for label in labels] optimizer.zero_grad() outputs model(images) loss 0.0 for i in range(len(labels)): if len(labels[i]) 0: pred outputs[i, :, :, :5 * len(labels[i])] pred pred.permute(2, 0, 1).reshape(-1, 5) target labels[i][:, 1:] loss criterion(pred, target) loss.backward() optimizer.step() running_loss loss.item() train_loss running_loss / len(train_loader) writer.add_scalar(Training Loss, train_loss, epoch) model.eval() running_val_loss 0.0 with torch.no_grad(): for images, labels in val_loader: images images.to(device) labels [torch.tensor(label, dtypetorch.float32).to(device) for label in labels] outputs model(images) loss 0.0 for i in range(len(labels)): if len(labels[i]) 0: pred outputs[i, :, :, :5 * len(labels[i])] pred pred.permute(2, 0, 1).reshape(-1, 5) target labels[i][:, 1:] loss criterion(pred, target) running_val_loss loss.item() val_loss running_val_loss / len(val_loader) writer.add_scalar(Validation Loss, val_loss, epoch) print(fEpoch {epoch 1}/{epochs}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}) model.save(weights/best_model.pt) writer.close()ifname “main”:image_dir “dataset/images”label_dir “dataset/labels_yolo”train_model(image_dir, label_dir)6. 模型评估训练完成后可以通过验证集和测试集来评估模型的性能。示例如下6.1 src/predict.pypython深色版本import torchimport matplotlib.pyplot as pltfrom ultralytics import YOLOfrom src.utils import get_data_loadersimport numpy as npdef predict_and_plot(image_dir, label_dir, model_path, num_samples5):device torch.device(“cuda” if torch.cuda.is_available() else “cpu”)model YOLO(model_path) model model.to(device) model.eval() _, val_loader get_data_loaders(image_dir, label_dir) fig, axes plt.subplots(num_samples, 2, figsize(10, 5 * num_samples)) with torch.no_grad(): for i, (images, labels) in enumerate(val_loader): if i num_samples: break images images.to(device) labels [torch.tensor(label, dtypetorch.float32).to(device) for label in labels] outputs model(images) outputs outputs.cpu().numpy() images images.cpu().numpy().transpose((0, 2, 3, 1)) labels [label.cpu().numpy() for label in labels] for j in range(len(images)): ax axes[j] if num_samples 1 else axes ax[0].imshow(images[j]) ax[0].set_title(fTrue: {labels[j]}) ax[0].axis(off) ax[1].imshow(images[j]) ax[1].set_title(fPredicted: {outputs[j]}) ax[1].axis(off) plt.tight_layout() plt.show()ifname “main”:image_dir “dataset/images”label_dir “dataset/labels_yolo”model_path “weights/best_model.pt”predict_and_plot(image_dir, label_dir, model_path)7. 运行项目确保你的数据集已经放在相应的文件夹中。运行转换脚本将VOC格式的标注文件转换为YOLO格式bash深色版本python src/convert_voc_to_yolo.py在项目根目录下运行以下命令启动训练bash深色版本python src/train.py训练完成后运行以下命令进行评估和可视化bash深色版本python src/predict.py8. 功能说明数据集转换convert_voc_to_yolo.py脚本用于将VOC格式的标注文件转换为YOLO格式。训练模型train.py脚本用于训练YOLOv8模型使用均方误差损失函数和Adam优化器。评估模型predict.py脚本用于评估模型性能并可视化输入图像、真实标签和预测结果。9. 详细注释convert_voc_to_yolo.py数据集转换将VOC格式的标注文件转换为YOLO格式。utils.py数据集类定义了一个PowerLineDefectDetectionDataset类用于加载和预处理数据。数据加载器定义了一个get_data_loaders函数用于创建训练和验证数据加载器。train.py训练函数定义了一个train_model函数用于训练YOLOv8模型。训练过程在每个epoch中模型在训练集上进行前向传播和反向传播并在验证集上进行评估。predict.py预测和可视化定义了一个predict_and_plot函数用于在验证集上进行预测并可视化输入图像、真实标签和预测结果