YOLO模型在物体识别跟踪中的应用

YOLO技术在物体跟踪方面具有强大的功能本程序根据YOLO的流程和原理进行实战验证。以帮助小白快速理解和自我编程。YOLO模型训练流程1.准备数据集..../dataset/images/trainimages/val(验证labels(特征标签/trainlabels/val2.在指定文件加下创建****.yaml文件定义上述数据集的地址bottle_dataset.yamlpath: E:/Projects/Python/Python test/bottles/datasettrain: images/trainval: images/valnames:0: bottle1: nut2: wrong_parts3.程序如何实现参照本程序结构以下是运用YOLO技术实现物体识别跟踪的实战应用程序可以直接运行检验效果import os from ultralytics import YOLO import yaml from IPython.display import display,Image import cv2 import torch import multiprocessing from ultralytics.data.utils import check_cls_dataset # 更新后的函数名 # ---------------------- # 1. 数据集准备和配置 # ---------------------- # 创建数据集目录结构 (如果手动准备数据可以跳过这部分) def create_dataset_structure(dataset_path): os.makedirs(dataset_path,exist_okTrue) os.makedirs(os.path.join(dataset_path,images,train),exist_okTrue) os.makedirs(os.path.join(dataset_path,images,val),exist_okTrue) os.makedirs(os.path.join(dataset_path,labels,train),exist_okTrue) os.makedirs(os.path.join(dataset_path,labels,val),exist_okTrue) # 创建数据集配置文件 def create_data_yaml(dataset_path,class_names): data{path:dataset_path,train:images/train,val:images/val, names:{i:name for i,name in enumerate(class_names)}} with open(os.path.join(dataset_path,data.yaml),w) as f: yaml.dump(data,f) return os.path.join(dataset_path,data.yaml) # 示例使用 dataset_pathE:/Projects/Python/Python test/YOLOmodel/dataset class_names[surface_scratch,miss_part,wrong_parts] # 替换为你的类别 create_dataset_structure(dataset_path) data_yaml_pathcreate_data_yaml(dataset_path,class_names) # ---------------------- # 2. 模型训练 # ---------------------- def train_yolov8(data_yaml,model_sizen,epochs50,imgsz640,batch16,device0,project_namecustom_train): 训练YOLOv8模型 参数: data_yaml: 数据集配置文件路径 model_size: 模型大小 (n, s, m, l, x) epochs: 训练轮数 imgsz: 图像大小 batch: 批量大小 device: 使用的设备 (0 for GPU 0, cpu for CPU) project_name: 项目名称 # 加载预训练模型 modelYOLO(fyolov8{model_size}.pt) # 训练参数配置 train_args{data:data_yaml,epochs:epochs,imgsz:imgsz,batch:batch,device:device,project:project_name, name:fyolov8{model_size}_{epochs}epochs,save:True,save_period:10, # 每10个epoch保存一次检查点 optimizer:auto, # SGD, Adam, AdamW, NAdam, RAdam等 lr0:0.01, # 初始学习率 lrf:0.01, # 最终学习率 lr0 * lrf momentum:0.937,weight_decay:0.0005,warmup_epochs:3.0,box:7.5, # box损失权重 cls:0.5, # cls损失权重 dfl:1.5, # dfl损失权重 close_mosaic:10, # 最后10个epoch禁用mosaic增强 amp:True # 自动混合精度训练 } # 开始训练 resultsmodel.train(**train_args) return model,results # 训练模型 model,resultstrain_yolov8(data_yamldata_yaml_path,model_sizen, # 使用nano模型 (yolov8n.pt) epochs50,imgsz640,batch16,device0 # 使用GPU 0 ) # ---------------------- # 3. 模型验证 # ---------------------- def validate_model(model,data_yaml,batch16,imgsz640,device0): 验证模型性能 参数: model: 训练好的模型或模型路径 data_yaml: 数据集配置文件路径 batch: 验证批量大小 imgsz: 图像大小 device: 使用的设备 if isinstance(model,str): modelYOLO(model) metricsmodel.val(datadata_yaml,batchbatch,imgszimgsz,devicedevice,splitval, # 验证集 conf0.001, # 对象置信度阈值 iou0.6, # NMS IoU阈值 halfTrue, # 使用半精度FP16推理 plotsTrue # 生成评估图表 ) print(fmAP50-95: {metrics.box.map}) # mAP0.5:0.95 print(fmAP50: {metrics.box.map50}) # mAP0.5 print(fmAP75: {metrics.box.map75}) # mAP0.75 print(fPrecision: {metrics.box.p}) # 精确度 print(fRecall: {metrics.box.r}) # 召回率 return metrics # 验证模型 best_model_pathos.path.join(custom_train,fyolov8n_50epochs,weights,best.pt) metricsvalidate_model(best_model_path,data_yaml_path) # ---------------------- # 4. 模型预测 # ---------------------- def predict_image(model,image_path,output_dirpredictions,conf0.25,iou0.7,saveTrue,showTrue): 使用训练好的模型进行预测 参数: model: 训练好的模型或模型路径 image_path: 要预测的图像路径 output_dir: 输出目录 conf: 置信度阈值 iou: IoU阈值 save: 是否保存结果 show: 是否显示结果 if isinstance(model,str): modelYOLO(model) os.makedirs(output_dir,exist_okTrue) # 执行预测 resultsmodel.predict(sourceimage_path,confconf,iouiou,savesave,save_txtFalse,save_confTrue,save_cropFalse, show_labelsTrue,show_confTrue,show_boxesTrue,projectoutput_dir) # 显示结果 if show: for r in results: im_arrayr.plot() # 绘制预测结果的numpy数组 cv2.imshow(Prediction,im_array) cv2.waitKey(0) cv2.destroyAllWindows() return results # 示例预测 test_imageE:/Projects/Python/Python test/images/defect_19.jpg # 替换为你的测试图像 predictionspredict_image(best_model_path,test_image) # ---------------------- # 5. 模型导出 # ---------------------- def export_model(model,export_formatonnx,imgsz640,devicecpu,simplifyTrue): 导出模型为其他格式 参数: model: 训练好的模型或模型路径 export_format: 导出格式 (onnx, engine, tflite, tfjs, openvino, coreml) imgsz: 图像大小 device: 使用的设备 simplify: 是否简化ONNX模型 if isinstance(model,str): modelYOLO(model) export_args{format:export_format,imgsz:imgsz,device:device,simplify:simplify,opset:12 # ONNX opset版本 } # 导出模型 exported_model_pathmodel.export(**export_args) print(fModel exported to: {exported_model_path}) return exported_model_path # 导出为ONNX格式 onnx_model_pathexport_model(best_model_path,export_formatonnx) # ---------------------- # 6. 使用导出的模型进行推理 # ---------------------- def infer_with_exported_model(model_path,image_path,conf0.25): 使用导出的模型进行推理 参数: model_path: 导出的模型路径 image_path: 要推理的图像路径 conf: 置信度阈值 # 加载模型 modelYOLO(model_path) # 执行推理 resultsmodel(image_path,confconf) # 显示结果 for r in results: im_arrayr.plot() cv2.imshow(Exported Model Prediction,im_array) cv2.waitKey(0) cv2.destroyAllWindows() return results # 使用导出的ONNX模型进行推理 onnx_resultsinfer_with_exported_model(onnx_model_path,test_image)更多更详细实际应用及代码请关注公众号。