基于YOLOv8和Faster R-CNN的输电线路异物目标检测项目 检测 输电线异物数据集 输电线缺陷数据集 绝缘子 如何使用YOLOv8和Faster R-CNN训练输电线路异物目标检测数据集
电力篇-输电线路缺陷数据集 输电线路异物目标检测数据集
16000张5种检测目标：
‘burst’ - 爆裂
‘defect’ - 缺陷
‘foreign_obj’ - 异物
‘insulator’ - 绝缘体
‘nest’ - 窝（巢）
带标注 -YOLO格式 可直接用于YOLO系列目标检测算法模型训练
如何使用YOLOv8和Faster R-CNN训练输电线路异物目标检测数据集的详细步骤和代码。假设数据集包含16000张图片和5种检测目标：‘burst’（爆裂）、‘defect’（缺陷）、‘foreign_obj’（异物）、‘insulator’（绝缘体）、‘nest’（窝（巢））。数据集已经划分好，并且标签格式为txt和xml。
一、数据集介绍
数据集结构
总共有16000张图片。
标签类别：‘burst’（爆裂）、‘defect’（缺陷）、‘foreign_obj’（异物）、‘insulator’（绝缘体）、‘nest’（窝（巢））。
标签格式：txt格式和xml格式。
目录结构
深色版本
PowerLineDefectDetection/
├── data/
│ ├── images/
│ │ ├── train/
│ │ ├── val/
│ │ └── test/
│ ├── labels_txt/
│ │ ├── train/
│ │ ├── val/
│ │ └── test/
│ ├── labels_xml/
│ │ ├── train/
│ │ ├── val/
│ │ └── test/
├── models/
│ ├── yolov8/
│ └── faster_rcnn/
├── utils/
│ ├── data_loader.py
│ ├── metrics.py
│ └── plot.py
├── main.py
├── train_yolov8.py
├── train_faster_rcnn.py
├── infer_yolov8.py
├── infer_faster_rcnn.py
├── evaluate_yolov8.py
├── evaluate_faster_rcnn.py
└── README.md
二、数据集准备
-
数据集划分
假设数据集已经按照训练集、验证集和测试集划分好。 -
标签转换
如果标签是xml格式，需要将其转换为YOLO格式（txt）。可以使用以下脚本进行转换：
python
深色版本
import os
import xml.etree.ElementTree as ET
def convert_xml_to_yolo(xml_file, txt_file, class_names):
tree = ET.parse(xml_file)
root = tree.getroot()
width = int(root.find(‘size/width’).text)
height = int(root.find(‘size/height’).text)
with open(txt_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)
bbox = obj.find('bndbox')
x_min = float(bbox.find('xmin').text)
y_min = float(bbox.find('ymin').text)
x_max = float(bbox.find('xmax').text)
y_max = float(bbox.find('ymax').text)
x_center = (x_min + x_max) / 2.0 / width
y_center = (y_min + y_max) / 2.0 / height
w = (x_max - x_min) / width
h = (y_max - y_min) / height
f.write(f"{class_id} {x_center} {y_center} {w} {h}\n")
def convert_all_labels(xml_dir, txt_dir, class_names):
os.makedirs(txt_dir, exist_ok=True)
for xml_file in os.listdir(xml_dir):
if xml_file.endswith(‘.xml’):
txt_file = os.path.join(txt_dir, os.path.splitext(xml_file)[0] + ‘.txt’)
convert_xml_to_yolo(os.path.join(xml_dir, xml_file), txt_file, class_names)
if name == ‘main’:
xml_dir = ‘data/labels_xml/train/’
txt_dir = ‘data/labels_txt/train/’
class_names = [‘burst’, ‘defect’, ‘foreign_obj’, ‘insulator’, ‘nest’]
convert_all_labels(xml_dir, txt_dir, class_names)
三、安装依赖
- YOLOv8
bash
深色版本
cd models/yolov8
git clone https://github.com/ultralytics/ultralytics.git
cd ultralytics
pip install -r requirements.txt - Faster R-CNN
bash
深色版本
pip install torch torchvision
四、训练模型 - 使用YOLOv8
编写一个训练脚本train_yolov8.py：
python
深色版本
import os
from ultralytics import YOLO
def main():
# 设置训练参数
data_yaml = ‘…/data/power_line_defect.yaml’
model_yaml = ‘yolov8n.yaml’ # 选择合适的YOLOv8模型
initial_weights = ‘path/to/initial/weights.pt’ # 初始权重文件
epochs = 100
batch_size = 16
img_size = 640
device = ‘cuda’ if torch.cuda.is_available() else ‘cpu’
# 加载模型
model = YOLO(model_yaml)
model.load(initial_weights) # 加载初始权重
# 开始训练
results = model.train(
data=data_yaml,
epochs=epochs,
batch=batch_size,
imgsz=img_size,
device=device
)
if name == ‘main’:
main()
运行训练脚本：
bash
深色版本
python train_yolov8.py
2. 使用Faster R-CNN
编写一个训练脚本train_faster_rcnn.py：
import os
import torch
import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.transforms import functional as F
from torch.utils.data import DataLoader, Dataset
import xml.etree.ElementTree as ET
import cv2
import numpy as np
class PowerLineDefectDataset(Dataset):
def init(self, root, transforms=None):
self.root = root
self.transforms = transforms
self.imgs = sorted([f for f in os.listdir(os.path.join(root, ‘images’)) if f.endswith(‘.jpg’)])
self.annotations = sorted([f for f in os.listdir(os.path.join(root, ‘labels_xml’)) if f.endswith(‘.xml’)])
def __getitem__(self, idx):
img_path = os.path.join(self.root, 'images', self.imgs[idx])
annotation_path = os.path.join(self.root, 'labels_xml', self.annotations[idx])
img = cv2.imread(img_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
tree = ET.parse(annotation_path)
root = tree.getroot()
boxes = []
labels = []
for obj in root.findall('object'):
class_name = obj.find('name').text
bbox = obj.find('bndbox')
x_min = float(bbox.find('xmin').text)
y_min = float(bbox.find('ymin').text)
x_max = float(bbox.find('xmax').text)
y_max = float(bbox.find('ymax').text)
boxes.append([x_min, y_min, x_max, y_max])
labels.append(class_name)
boxes = torch.as_tensor(boxes, dtype=torch.float32)
labels = torch.as_tensor([class_names.index(label) for label in labels], dtype=torch.int64)
image_id = torch.tensor([idx])
area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
iscrowd = torch.zeros((len(boxes),), dtype=torch.int64)
target = {}
target["boxes"] = boxes
target["labels"] = labels
target["image_id"] = image_id
target["area"] = area
target["iscrowd"] = iscrowd
if self.transforms is not None:
img = self.transforms(img)
return img, target
def __len__(self):
return len(self.imgs)
def get_transform(train):
transforms = []
transforms.append(T.ToTensor())
if train:
transforms.append(T.RandomHorizontalFlip(0.5))
return T.Compose(transforms)
def collate_fn(batch):
return tuple(zip(*batch))
def main():
# 设置训练参数
root = ‘…/data/’
class_names = [‘burst’, ‘defect’, ‘foreign_obj’, ‘insulator’, ‘nest’]
num_classes = len(class_names) + 1 # +1 for background
batch_size = 4
num_epochs = 100
device = torch.device(‘cuda’) if torch.cuda.is_available() else torch.device(‘cpu’)
# 加载数据
dataset = PowerLineDefectDataset(root, get_transform(train=True))
data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=4, collate_fn=collate_fn)
# 加载模型
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
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)
# 开始训练
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()}')
torch.save(model.state_dict(), 'faster_rcnn_power_line_defect.pth')
if name == ‘main’:
main()
运行训练脚本：
bash
深色版本
python train_faster_rcnn.py
五、推理和可视化
- 使用YOLOv8
编写一个推理脚本infer_yolov8.py：
python
深色版本
import os
import torch
import cv2
import numpy as np
from ultralytics import YOLO
from ultralytics.yolo.utils.general import non_max_suppression, scale_coords
from ultralytics.yolo.utils.plots import Annotator
def load_model(weights_path, device):
model = YOLO(weights_path)
model.to(device)
return model
def infer_image(model, img_path, img_size, device):
img = cv2.imread(img_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = cv2.resize(img, (img_size, img_size))
img = img.transpose(2, 0, 1)
img = img.astype(np.float32) / 255.0
img = torch.from_numpy(img).unsqueeze(0).to(device)
with torch.no_grad():
pred = model(img)[0]
pred = non_max_suppression(pred, conf_thres=0.25, iou_thres=0.45)
return pred, img
def draw_boxes(img, pred, names):
annotator = Annotator(img, line_width=2, example=str(names))
for i, det in enumerate(pred):
if len(det):
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], img.shape[2:]).round()
for *xyxy, conf, cls in reversed(det):
c = int(cls)
label = f’{names[c]} {conf:.2f}’
annotator.box_label(xyxy, label, color=(0, 255, 0))
return annotator.result()
if name == ‘main’:
weights_path = ‘runs/train/exp/weights/best.pt’
img_path = ‘…/data/images/test/000001.jpg’
img_size = 640
device = ‘cuda’ if torch.cuda.is_available() else ‘cpu’
model = load_model(weights_path, device)
pred, img = infer_image(model, img_path, img_size, device)
# 读取原始图像
original_img = cv2.imread(img_path)
# 绘制预测框
annotated_img = draw_boxes(original_img, pred, names=['burst', 'defect', 'foreign_obj', 'insulator', 'nest'])
# 显示结果
cv2.imshow('Result', annotated_img)
cv2.waitKey(0)
cv2.destroyAllWindows()
- 使用Faster R-CNN
编写一个推理脚本infer_faster_rcnn.py：
python
深色版本
import os
import torch
import cv2
import numpy as np
from torchvision.models.detection import fasterrcnn_resnet50_fpn
from torchvision.transforms import functional as F
def load_model(weights_path, device):
model = fasterrcnn_resnet50_fpn(pretrained=False)
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, 6) # 6 classes including background
model.load_state_dict(torch.load(weights_path, map_location=device))
model.to(device)
model.eval()
return model
def infer_image(model, img_path, device):
img = cv2.imread(img_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img_tensor = F.to_tensor(img).unsqueeze(0).to(device)
with torch.no_grad():
pred = model(img_tensor)[0]
return pred, img
def draw_boxes(img, pred, names):
for box, label, score in zip(pred[‘boxes’], pred[‘labels’], pred[‘scores’]):
if score > 0.5:
x1, y1, x2, y2 = box.int().tolist()
cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
cv2.putText(img, f’{names[label]} {score:.2f}', (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
return img
if name == ‘main’:
weights_path = ‘faster_rcnn_power_line_defect.pth’
img_path = ‘…/data/images/test/000001.jpg’
device = ‘cuda’ if torch.cuda.is_available() else ‘cpu’
model = load_model(weights_path, device)
pred, img = infer_image(model, img_path, device)
# 读取原始图像
original_img = cv2.imread(img_path)
# 绘制预测框
annotated_img = draw_boxes(original_img, pred, names=['background', 'burst', 'defect', 'foreign_obj', 'insulator', 'nest'])
# 显示结果
cv2.imshow('Result', annotated_img)
cv2.waitKey(0)
cv2.destroyAllWindows()
六、评价指标
- 使用YOLOv8
编写一个评价脚本evaluate_yolov8.py：
python
深色版本
import os
import torch
from ultralytics import YOLO
from ultralytics.yolo.utils.general import check_img_size
from ultralytics.yolo.utils.datasets import create_dataloader
def main():
data_yaml = ‘…/data/power_line_defect.yaml’
weights_path = ‘runs/train/exp/weights/best.pt’
img_size = 640
batch_size = 16
device = ‘cuda’ if torch.cuda.is_available() else ‘cpu’
# 加载数据
dataloader = create_dataloader('../data/images/val/', img_size, batch_size, 32, single_cls=False, pad=0.5, rect=True)[0]
# 加载模型
model = YOLO(weights_path)
model.to(device)
# 验证模型
results = model.val(
data=data_yaml,
batch=batch_size,
imgsz=img_size,
device=device,
dataloader=dataloader
)
if name == ‘main’:
main()
运行评价脚本：
bash
深色版本
python evaluate_yolov8.py
2. 使用Faster R-CNN
编写一个评价脚本evaluate_faster_rcnn.py：
python
深色版本
import os
import torch
import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.transforms import functional as F
from torch.utils.data import DataLoader, Dataset
import xml.etree.ElementTree as ET
import cv2
import numpy as np
class PowerLineDefectDataset(Dataset):
def init(self, root, transforms=None):
self.root = root
self.transforms = transforms
self.imgs = sorted([f for f in os.listdir(os.path.join(root, ‘images’)) if f.endswith(‘.jpg’)])
self.annotations = sorted([f for f in os.listdir(os.path.join(root, ‘labels_xml’)) if f.endswith(‘.xml’)])
def __getitem__(self, idx):
img_path = os.path.join(self.root, 'images', self.imgs[idx])
annotation_path = os.path.join(self.root, 'labels_xml', self.annotations[idx])
img = cv2.imread(img_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
tree = ET.parse(annotation_path)
root = tree.getroot()
boxes = []
labels = []
for obj in root.findall('object'):
class_name = obj.find('name').text
bbox = obj.find('bndbox')
x_min = float(bbox.find('xmin').text)
y_min = float(bbox.find('ymin').text)
x_max = float(bbox.find('xmax').text)
y_max = float(bbox.find('ymax').text)
boxes.append([x_min, y_min, x_max, y_max])
labels.append(class_name)
boxes = torch.as_tensor(boxes, dtype=torch.float32)
labels = torch.as_tensor([class_names.index(label) for label in labels], dtype=torch.int64)
image_id = torch.tensor([idx])
area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
iscrowd = torch.zeros((len(boxes),), dtype=torch.int64)
target = {}
target["boxes"] = boxes
target["labels"] = labels
target["image_id"] = image_id
target["area"] = area
target["iscrowd"] = iscrowd
if self.transforms is not None:
img = self.transforms(img)
return img, target
def __len__(self):
return len(self.imgs)
def get_transform(train):
transforms = []
transforms.append(T.ToTensor())
if train:
transforms.append(T.RandomHorizontalFlip(0.5))
return T.Compose(transforms)
def collate_fn(batch):
return tuple(zip(*batch))
def evaluate_model(model, data_loader, device):
model.eval()
ious = []
with torch.no_grad():
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]
outputs = model(images)
for output, target in zip(outputs, targets):
pred_boxes = output['boxes']
true_boxes = target['boxes']
for pred_box, true_box in zip(pred_boxes, true_boxes):
iou = torchvision.ops.box_iou(pred_box.unsqueeze(0), true_box.unsqueeze(0)).item()
ious.append(iou)
avg_iou = sum(ious) / len(ious)
print(f'Average IoU: {avg_iou:.4f}')
def main():
root = ‘…/data/’
class_names = [‘burst’, ‘defect’, ‘foreign_obj’, ‘insulator’, ‘nest’]
num_classes = len(class_names) + 1 # +1 for background
batch_size = 4
device = torch.device(‘cuda’) if torch.cuda.is_available() else torch.device(‘cpu’)
# 加载数据
dataset = PowerLineDefectDataset(root, get_transform(train=False))
data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=4, collate_fn=collate_fn)
# 加载模型
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False)
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
model.load_state_dict(torch.load('faster_rcnn_power_line_defect.pth', map_location=device))
model.to(device)
# 评估模型
evaluate_model(model, data_loader, device)
if name == ‘main’:
main()
运行评价脚本：
bash
深色版本
python evaluate_faster_rcnn.py
七、运行项目
确保数据集和标签文件已经准备好，并放在相应的目录中。
运行YOLOv8训练脚本：
bash
深色版本
python train_yolov8.py
运行Faster R-CNN训练脚本：
bash
深色版本
python train_faster_rcnn.py
运行YOLOv8推理脚本：
bash
深色版本
python infer_yolov8.py
运行Faster R-CNN推理脚本：
bash
深色版本
python infer_faster_rcnn.py
运行YOLOv8评价脚本：
bash
深色版本
python evaluate_yolov8.py
运行Faster R-CNN评价脚本：
bash
深色版本
python evaluate_faster_rcnn.py
八、代码说明
数据配置文件：data/power_line_defect.yaml指定了数据集的路径和类别信息。
训练模型：train_yolov8.py和train_faster_rcnn.py分别负责加载数据、初始化模型、损失函数和优化器，并进行训练。
推理和可视化：infer_yolov8.py和infer_faster_rcnn.py用于加载训练好的模型并对新图像进行预测，显示预测结果。
评价指标：evaluate_yolov8.py和evaluate_faster_rcnn.py计算模型的评价指标，如IoU和mAP。
希望这些代码和说明能帮助你完成基于YOLOv8和Faster R-CNN的输电线路异物目标检测项目