YOLO钢轨缺口识别

YOLOv8在钢轨故障识别中的应用

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
from ultralytics import YOLO

if __name__ == "__main__":
    # 加载YOLOv8n模型
    model = YOLO("yolov8n.pt")
    
    # 开始训练
    model.train(
        data="/gap.yaml",       # 训练数据配置文件
        epochs=100,             # 训练轮次
        imgsz=640,              # 图片大小
        device="0"              # 设备可设置为GPU（0）或CPU（"cpu"）
    )

# 加载经过训练的模型进行预测
model = YOLO("/best.pt")  # 加载训练好的YOLOv8n模型

# 对图像进行批量推理
results = model(["/aug_prefix_0_1352.jpg"])  # 返回一系列结果对象

# 处理并显示结果
for result in results:
    boxes = result.boxes        # 获取边界框输出
    masks = result.masks        # 获取分割掩码输出
    keypoints = result.keypoints  # 获取关键点输出
    probs = result.probs        # 获取分类概率输出
    obb = result.obb            # 获取定向边界框（OBB）输出

    result.show()               # 显示结果
    result.save(filename="C:/Users/yusialone/Desktop/3d/its_rail/jieguo/result.jpg")  # 保存结果