from ultralytics import YOLO
import os
def main():
os.makedirs("runs/highway_damage", exist_ok=True)
model = YOLO('yolov8s.pt')
results = model.train(
data='highway_damage.yaml',
epochs=150,
imgsz=1024,
batch=8,
name='yolov8s_highway_1024',
project='runs/highway_damage',
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,
hsv_h=0.015,
hsv_s=0.7,
hsv_v=0.4,
degrees=10.0,
translate=0.1,
scale=0.5,
shear=2.0,
perspective=0.001,
flipud=0.0,
fliplr=0.5,
mosaic=1.0,
mixup=0.1
)
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 = ["Cracks", "Waterlogging", "Ravelling", "Muddy_road", "Road_side_garbage", "Potholes"]
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)
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
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(lines))
from ultralytics import YOLO
model = YOLO('runs/highway_damage/yolov8s_highway_1024/weights/best.pt')
results = model('test_image.jpg')
chinese_names = {
"Cracks": "裂缝",
"Waterlogging": "积水",
"Ravelling": "松散",
"Muddy_road": "泥泞道路",
"Road_side_garbage": "道路旁垃圾",
"Potholes": "坑洼"
}
for r in results:
boxes = r.boxes
for box in boxes:
cls_name = model.names[int(box.cls)]
print(f"检测到:{chinese_names.get(cls_name, cls_name)}")