from ultralytics import YOLO
import os
def main():
os.makedirs("runs/power_inspection", exist_ok=True)
model = YOLO('yolov8s.pt')
results = model.train(
data='power_data.yaml',
epochs=150,
imgsz=800,
batch=16,
name='yolov8s_power_800',
project='runs/power_inspection',
device=0,
workers=8,
cache=False,
optimizer='AdamW',
lr0=0.01,
lrf=0.01,
momentum=0.937,
weight_decay=0.0005,
warmup_epochs=3,
patience=30,
save=True,
save_period=10,
verbose=True,
plots=True
)
print(f"✅ 训练完成！最佳模型路径：{results.save_dir}/weights/best.pt")
if __name__ == '__main__':
main()
import os
import xml.etree.ElementTree as ET
from pathlib import Path
class_names = [
"铁塔顶部", "电缆头部异物", "爬电距离装置头部", "开关头部", "陶瓷绝缘子污秽",
"导线端头部", "地线头部", "电缆头部扭曲", "接地装置头部", "陶瓷绝缘子头部破损",
"电缆头部腐蚀", "负荷断路开关头部", "断路器头部", "复合绝缘子头部裂缝",
"闸刀开关头部", "电缆保护套头部", "电缆连接头部", "复合开关头部",
"聚合物绝缘子头部污秽", "聚合物绝缘子头部裂纹", "陶瓷绝缘子污秽",
"总承载体（瓷瓶）头部", "锚定装置头部", "复合绝缘子本体破损", "电缆护套头部",
"复合绝缘子头部污秽", "复合绝缘子头部破损", "复合绝缘子本体污秽",
"陶瓷绝缘子头部裂缝", "总承载体（瓷瓶）本体", "聚合物绝缘子本体污秽",
"聚合物绝缘子本体弯曲", "聚合物绝缘子本体裂纹", "聚合物绝缘子头部弯曲",
"铁塔本体", "输电线本体", "地线本体", "支撑件本体", "电缆本体",
"电缆本体扭曲", "电缆本体异物", "开关本体", "地线接地装置本体",
"交汇处本体", "锚定装置本体", "接地装置本体", "导线头部", "电缆头部腐蚀",
"熔断器头部"
]
def convert_voc_to_yolo(voc_dir, yolo_dir, image_dir):
os.makedirs(yolo_dir, exist_ok=True)
for xml_file in Path(voc_dir).glob("*.xml"):
tree = ET.parse(xml_file)
root = tree.getroot()
img_w = int(root.find('size/width').text)
img_h = int(root.find('size/height').text)
yolo_lines = []
for obj in root.findall('object'):
cls_name = obj.find('name').text
if cls_name not in class_names:
continue
cls_id = class_names.index(cls_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 / img_w
y_center = (ymin + ymax) / 2 / img_h
width = (xmax - xmin) / img_w
height = (ymax - ymin) / img_h
yolo_lines.append(f"{cls_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}")
with open(os.path.join(yolo_dir, xml_file.stem + '.txt'), 'w') as f:
f.write('\n'.join(yolo_lines))
