网站建设选择北京华网天下,网站开发框架 知乎,手机网站开发 html,wordpress付免签插件【YOLO-Ultralytics】 【核心引擎】【v8.3.235版本】 模型预测器基类代码predictor.py解析 文章目录【YOLO-Ultralytics】 【核心引擎】【v8.3.235版本】 模型预测器基类代码predictor.py解析前言YOLO 预测功能全解析输入源#xff08;source 参数#xff09;说明模型格式source 参数说明模型格式model 参数说明关键注意事项所需的库和模块BasePredictor 类整体概览1. BasePredictor 类属性说明表2. BasePredictor 类方法说明表初始化函数__init__图像预处理preprocess模型推理inference图像预变换pre_transform预测后处理基类占位postprocess预测执行入口\_\_call\_\_CLI预测执行predict_cli输入源设置setup_source流式推理核心stream_inference模型初始化setup_model结果写入write_results预测图像保存save_predicted_images结果可视化show回调执行run_callbacks回调注册add_callback完整代码通用预测框架的核心特性工程化核心优化扩展与易用性能力关键注意事项总结前言代码路径ultralytics\engine\predictor.pyBasePredictor 是 Ultralytics YOLO 框架中所有预测器的基类抽象出 YOLO 模型推理的通用核心流程为检测 / 分割 / 分类等任务的专用预测器如 DetectionPredictor提供标准化的底层能力。该类封装了「配置加载→模型初始化→数据源加载→预处理→推理→后处理→结果保存 / 可视化→资源释放」的全流程通用逻辑支持多输入源摄像头、视频、图片、流、张量等、多模型格式PyTorch/ONNX/TensorRT 等、流式推理避免内存溢出、线程安全推理、自定义回调等核心特性是 YOLO 推理体系的基础骨架。【YOLOv8-Ultralytics 系列文章目录】YOLO 预测功能全解析modepredict模式用于目标检测/分割/分类等预测任务核心参数为model模型文件和source输入源后续会详细说明。 Run prediction on images, videos, directories, globs, YouTube, webcam, streams, etc. Usage - sources: $ yolo modepredict modelyolo11n.pt source0 # webcam img.jpg # image vid.mp4 # video screen # screenshot path/ # directory list.txt # list of images list.streams # list of streams path/*.jpg # glob https://youtu.be/LNwODJXcvt4 # YouTube rtsp://example.com/media.mp4 # RTSP, RTMP, HTTP, TCP stream Usage - formats: $ yolo modepredict modelyolo11n.pt # PyTorch yolo11n.torchscript # TorchScript yolo11n.onnx # ONNX Runtime or OpenCV DNN with dnnTrue yolo11n_openvino_model # OpenVINO yolo11n.engine # TensorRT yolo11n.mlpackage # CoreML (macOS-only) yolo11n_saved_model # TensorFlow SavedModel yolo11n.pb # TensorFlow GraphDef yolo11n.tflite # TensorFlow Lite yolo11n_edgetpu.tflite # TensorFlow Edge TPU yolo11n_paddle_model # PaddlePaddle yolo11n.mnn # MNN yolo11n_ncnn_model # NCNN yolo11n_imx_model # Sony IMX yolo11n_rknn_model # Rockchip RKNN yolo11n.pte # PyTorch Executorch 输入源source 参数说明参数值类型详细说明使用示例0本地摄像头调用本地摄像头0为默认摄像头1/2为其他摄像头yolo modepredict modelyolo11n.pt source0img.jpg或png/bmp等单张图片本地图片文件路径支持常见图像格式yolo modepredict modelyolo11n.pt sourcetest_image.pngvid.mp4或avi/mov等本地视频本地视频文件路径支持常见视频格式yolo modepredict modelyolo11n.pt sourcedemo_video.aviscreen屏幕截图实时捕获电脑屏幕作为输入仅桌面端Windows/Linux/macOSyolo modepredict modelyolo11n.pt sourcescreenpath/目录批量预测指定目录下所有图片/视频文件递归遍历子目录yolo modepredict modelyolo11n.pt source./dataset/images/list.txt图片列表文件TXT文件每行写一个图片路径批量预测列表内所有图片yolo modepredict modelyolo11n.pt sourceimage_list.txtlist.streams流列表文件后缀为.streams的文本文件每行写一个网络流地址批量预测多个流yolo modepredict modelyolo11n.pt sourcestream_list.streamspath/*.jpgGlob通配符匹配指定路径下符合通配符规则的文件需加引号避免系统解析通配符yolo modepredict modelyolo11n.pt source./data/*.jpghttps://youtu.be/xxxYouTube视频直接通过YouTube链接预测视频内容需网络通畅支持公开视频yolo modepredict modelyolo11n.pt sourcehttps://youtu.be/LNwODJXcvt4rtsp://example.com/media.mp4网络流支持RTSP/RTMP/HTTP/TCP等协议的网络视频流如摄像头推流、直播流yolo modepredict modelyolo11n.pt sourcertsp://admin:123456192.168.1.100:554/stream1模型格式model 参数说明模型文件名称/后缀推理引擎/框架详细说明使用示例yolo11n.ptPyTorch原生PyTorch模型基础格式支持推理/再训练灵活性最高yolo modepredict modelyolo11n.pt source0yolo11n.torchscriptTorchScriptPyTorch导出的TorchScript格式适配生产环境跨平台兼容性更好yolo modepredict modelyolo11n.torchscript sourceimg.jpgyolo11n.onnxONNX Runtime/OpenCV DNN通用ONNX格式支持ONNX Runtime推理加dnnTrue可使用OpenCV DNN推理1. ONNX Runtimeyolo modepredict modelyolo11n.onnx sourcevid.mp42. OpenCV DNNyolo modepredict modelyolo11n.onnx sourcevid.mp4 dnnTrueyolo11n_openvino_modelOpenVINOIntel OpenVINO优化模型适配Intel CPU/GPU/VPU推理速度提升yolo modepredict modelyolo11n_openvino_model sourcepath/yolo11n.engineTensorRTNVIDIA TensorRT优化模型适配NVIDIA GPU极致推理性能需提前导出yolo modepredict modelyolo11n.engine source0yolo11n.mlpackageCoreMLApple CoreML格式仅支持macOS/iOS适配Apple SiliconM系列芯片yolo modepredict modelyolo11n.mlpackage sourceimg.jpgmacOS环境yolo11n_saved_modelTensorFlow SavedModelTensorFlow原生SavedModel格式适配TensorFlow生态yolo modepredict modelyolo11n_saved_model sourcelist.txtyolo11n.pbTensorFlow GraphDefTensorFlow传统GraphDef格式.pb文件适配老版本TensorFlow部署yolo modepredict modelyolo11n.pb sourcepath/*.jpgyolo11n.tfliteTensorFlow Lite轻量级TFLite格式适配移动端/嵌入式设备手机、单片机等yolo modepredict modelyolo11n.tflite source0嵌入式设备yolo11n_edgetpu.tfliteTensorFlow Edge TPU适配Google Coral Edge TPU加速棒的TFLite格式硬件加速推理yolo modepredict modelyolo11n_edgetpu.tflite sourcelist.streamsyolo11n_paddle_modelPaddlePaddle百度飞桨PaddlePaddle格式适配Paddle推理引擎yolo modepredict modelyolo11n_paddle_model sourcehttps://youtu.be/xxxyolo11n.mnnMNN阿里MNN格式轻量级跨平台推理适配移动端/嵌入式yolo modepredict modelyolo11n.mnn sourcertsp://example.com/media.mp4yolo11n_ncnn_modelNCNN腾讯NCNN格式专为移动端优化适配手机/嵌入式设备yolo modepredict modelyolo11n_ncnn_model sourcevid.mp4yolo11n_imx_modelSony IMX适配Sony IMX系列视觉芯片的模型面向索尼嵌入式硬件yolo modepredict modelyolo11n_imx_model sourcepath/yolo11n_rknn_modelRockchip RKNN瑞芯微RKNN格式适配RK3588等瑞芯微芯片嵌入式硬件优化yolo modepredict modelyolo11n_rknn_model source0yolo11n.ptePyTorch ExecutorchPyTorch轻量级部署格式适配移动端/边缘设备yolo modepredict modelyolo11n.pte sourceimg.jpg关键注意事项通配符使用sourcepath/*.jpg需加引号避免操作系统提前解析通配符导致参数错误环境依赖不同模型格式需安装对应依赖如TensorRT需装NVIDIA TensorRT库OpenVINO需装OpenVINO工具包硬件适配CoreML仅支持macOS/iOSTensorRT仅支持NVIDIA GPURKNN仅支持瑞芯微芯片网络流认证RTSP等流若需账号密码需在地址中拼接如rtsp://账号:密码IP:端口/流地址模型版本示例中yolo11n.pt的n代表nano轻量级还有s/m/l/x等版本用法完全一致。所需的库和模块# Ultralytics AGPL-3.0 License - https://ultralytics.com/license# 引入未来版本的类型注解支持提升代码静态类型检查和编辑器提示能力from__future__importannotations# 导入系统平台检测模块区分Windows/Linux/macOS适配视频编码格式、窗口显示逻辑importplatform# 导入正则表达式模块解析视频帧号、处理路径/日志字符串中的结构化信息importre# 导入线程管理模块提供线程锁保证多线程环境下推理的资源安全importthreading# 导入路径处理模块面向对象的路径操作统一管理推理结果保存路径、数据源路径frompathlibimportPath# 导入类型注解模块Any类型用于兼容推理流程中多类型的输入/输出参数fromtypingimportAny# 导入OpenCV核心模块图像/视频的读取、写入、可视化、格式转换、视频编码等importcv2# 导入数值计算模块图像数组处理、维度转换、归一化、数组拼接等核心数值操作importnumpyasnp# 导入PyTorch核心模块张量操作、设备管理、模型推理、数据类型转换等importtorch# 从ultralytics配置模块导入配置加载函数合并基础配置与自定义覆盖参数、保存目录生成函数自动创建runs/predict/exp*目录fromultralytics.cfgimportget_cfg,get_save_dir# 从ultralytics数据模块导入推理数据源加载函数统一处理摄像头/视频/图片/流/张量等多类型输入源fromultralytics.dataimportload_inference_source# 从ultralytics数据增强模块导入LetterBox类图像等比例缩放自动填充保证尺寸为模型步长倍数fromultralytics.data.augmentimportLetterBox# 从ultralytics神经网络后端模块导入AutoBackend类统一加载PyTorch/ONNX/TensorRT等多格式模型fromultralytics.nn.autobackendimportAutoBackend# 从ultralytics工具模块导入默认配置常量、日志器、系统标识macOS/WINDOWS、回调函数字典、彩色字符串函数、通用操作工具计时/缩放等fromultralytics.utilsimportDEFAULT_CFG,LOGGER,MACOS,WINDOWS,callbacks,colorstr,ops# 从ultralytics工具检查模块导入图像尺寸检查函数保证尺寸为模型步长倍数、可视化可用性检查函数兼容不同系统的窗口显示fromultralytics.utils.checksimportcheck_imgsz,check_imshow# 从ultralytics工具文件模块导入路径增量函数自动生成不重复的文件/目录名避免结果覆盖fromultralytics.utils.filesimportincrement_path# 从ultralytics工具PyTorch模块导入模型编译函数提升推理效率、设备选择函数自动选择CPU/GPU/TensorRT、推理模式装饰器禁用梯度计算fromultralytics.utils.torch_utilsimportattempt_compile,select_device,smart_inference_mode# 若在使用 Ultralytics 模型推理时未传入streamTrue参数推理结果会堆积在内存中# 可能导致处理大型数据源、长时间运行的流或视频时出现内存不足错误# 建议启用streamTrue生成 Results 对象生成器并遍历处理详见指定文档。STREAM_WARNING inference results will accumulate in RAM unless streamTrue is passed, causing potential out-of-memory errors for large sources or long-running streams and videos. See https://docs.ultralytics.com/modes/predict/ for help. Example: results model(source..., streamTrue) # generator of Results objects for r in results: boxes r.boxes # Boxes object for bbox outputs masks r.masks # Masks object for segment masks outputs probs r.probs # Class probabilities for classification outputs BasePredictor 类整体概览项目详情类名BasePredictor父类object核心定位所有YOLO预测器的基础类封装模型加载、输入源处理、图像预处理、推理执行、结果后处理、可视化/保存等通用逻辑支持多输入源图像/视频/流/摄像头等、多模型格式PyTorch/ONNX/TensorRT等核心依赖模块ultralytics.cfg配置解析、ultralytics.data输入源加载、ultralytics.nn.autobackend多模型后端适配、ultralytics.utils工具函数/性能分析/设备选择、cv2图像/视频处理、torch张量操作、threading线程安全关键特性1. 支持多输入源图像/视频/文件夹/摄像头/RTSP/YouTube/张量等2. 适配多模型格式PyTorch/ONNX/TensorRT等3. 流式推理避免内存溢出4. 线程安全的推理流程5. 完整的结果保存图像/视频/标签/TXT与可视化6. 性能分析预处理/推理/后处理耗时7. 回调机制支持自定义扩展8. 自动预热模型提升推理速度典型使用流程初始化预测器→配置参数→加载模型→设置输入源→预热模型→逐批次预处理图像→模型推理→后处理结果→可视化/保存结果→释放资源→输出性能统计1. BasePredictor 类属性说明表属性名类型说明argsSimpleNamespace推理配置参数conf/imgsz/save/show等save_dirPath结果保存目录done_warmupbool模型是否完成预热modeltorch.nn.ModuleAutoBackend封装的推理模型支持多格式datadict数据配置字典devicetorch.device推理设备CPU/GPU/TensorRT等datasetInferenceDataset加载的推理数据源vid_writerdict[str, cv2.VideoWriter]视频写入器字典{保存路径: 写入器}plotted_imgnp.ndarray最后一次可视化的图像带预测框/掩码source_typeSimpleNamespace输入源类型stream/video/image/tensor等seenint已处理的图像数量windowslist[str]可视化窗口名称列表batchtuple当前批次数据路径/原始图像/日志字符串resultslist[Any]当前批次的推理结果Results对象列表transformscallable分类任务专用图像变换callbacksdict[str, list[callable]]回调函数字典按事件分类txt_pathPath预测结果txt保存路径_lockthreading.Lock线程锁保证多线程推理安全2. BasePredictor 类方法说明表方法名功能说明preprocess图像预处理格式转换/归一化/设备迁移/数据类型转换pre_transform预处理前的LetterBox缩放保持比例自动填充inference模型推理支持增强/可视化/嵌入提取postprocess后处理基方法子类重写如检测的NMScall对外推理接口支持流式/非流式predict_cliCLI模式推理消费生成器避免内存累积setup_source加载推理数据源检查输入类型stream_inference核心流式推理循环全流程执行setup_model初始化模型AutoBackend加载多格式模型设备选择/编译write_results结果写入txt/图片/视频可视化save_predicted_images保存预测结果图片/视频/帧show可视化展示推理结果run_callbacks执行指定事件的所有回调函数add_callback注册自定义回调函数初始化函数initdef__init__(self,cfgDEFAULT_CFG,overrides:dict[str,Any]|NoneNone,_callbacks:dict[str,list[callable]]|NoneNone,): 初始化BasePredictor类 参数: cfg (str | dict): 基础配置文件路径/字典默认DEFAULT_CFG overrides (dict, optional): 配置覆盖参数如conf0.5, imgsz640 _callbacks (dict, optional): 自定义回调函数字典按事件分类 # 加载并合并配置基础配置 覆盖参数self.argsget_cfg(cfg,overrides)# 生成结果保存目录默认 runs/predict/exp*self.save_dirget_save_dir(self.args)# 设置默认置信度阈值未指定时为0.25ifself.args.confisNone:self.args.conf0.25# 标记模型未完成预热self.done_warmupFalse# 检查可视化可用性Windows/Linux/macOS兼容ifself.args.show:self.args.showcheck_imshow(warnTrue)# 后续初始化的属性setup_model/setup_source后赋值self.modelNone# 推理模型实例AutoBackend封装支持PyTorch/ONNX/TensorRT等格式setup_model后初始化self.dataself.args.data# 数据配置字典如数据集路径、类别名等来自推理配置参数self.imgszNone# 推理输入图像尺寸tuple如(640,640)setup_source中check_imgsz后确定self.deviceNone# 推理设备torch.device对象如cuda:0/cpusetup_model后从model.device同步self.datasetNone# 推理数据源实例InferenceDataset封装多类型输入源setup_source后初始化self.vid_writer{}# 视频写入器字典key保存路径valuecv2.VideoWriter用于多视频/流结果保存self.plotted_imgNone# 最后一次绘制的带预测结果的图像np.ndarray可视化/保存时使用self.source_typeNone# 输入源类型SimpleNamespace含stream/video/image/tensor等标识setup_source后赋值self.seen0# 已处理的图像/帧总数流式推理中逐批次累加用于统计和耗时计算self.windows[]# 可视化窗口名称列表Linux下创建可调整窗口时记录避免重复创建self.batchNone# 当前批次数据tuple含paths/im0s/s遍历dataset时赋值self.resultsNone# 当前批次后处理结果list[Results]postprocess后赋值self.transformsNone# 分类任务专用图像变换callable检测任务暂未使用self.callbacks_callbacksorcallbacks.get_default_callbacks()# 推理回调函数字典按事件分类如on_predict_startself.txt_pathNone# 单张图像预测结果txt保存路径Pathwrite_results时动态生成self._lockthreading.Lock()# 线程锁保证多线程环境下推理流程的资源安全如模型推理/视频写入callbacks.add_integration_callbacks(self)# 注册集成回调如WB/MLflow日志、ClearML监控等项目详情函数名__init__功能概述初始化预测器核心配置设置保存目录、默认参数、回调函数、线程锁等基础属性返回值无构造函数核心逻辑解析配置→设置保存目录→初始化默认置信度→检查可视化可用性→初始化核心属性→注册集成回调设计亮点1. 配置分层解析支持基础配置覆盖项灵活调整2. 自动适配系统可视化能力3. 线程锁保证多线程安全4. 默认回调机制便于扩展注意事项1. 需确保overrides中的参数与DEFAULT_CFG兼容否则可能解析失败2. 保存目录会自动增量命名exp1/exp2避免覆盖已有结果图像预处理preprocessdefpreprocess(self,im:torch.Tensor|list[np.ndarray])-torch.Tensor: 图像预处理将输入转换为模型可接受的张量格式BCHW/归一化/设备迁移/数据类型转换 参数: im (torch.Tensor | list[np.ndarray]): 输入图像 - Tensor: 形状(N, 3, H, W) - list[np.ndarray]: 形状[(H, W, 3) x N]N为批次大小 返回: (torch.Tensor): 预处理后的张量N, 3, H, W已迁移到指定设备归一化到0-1 # 判断输入是否为张量not_tensornotisinstance(im,torch.Tensor)ifnot_tensor:# 对列表中每张图像执行pre_transformLetterBox缩放imnp.stack(self.pre_transform(im))# BGR→RGBOpenCV加载的图像是BGR模型需要RGBifim.shape[-1]3:imim[...,::-1]# 维度转换BHWC批次高度宽度通道→ BCHW批次通道高度宽度imim.transpose((0,3,1,2))# 转换为连续内存数组提升张量操作效率imnp.ascontiguousarray(im)# NumPy→PyTorch张量imtorch.from_numpy(im)# 迁移到推理设备CPU/GPUimim.to(self.device)# 数据类型转换uint8→fp16/fp32模型若启用fp16则用half否则floatimim.half()ifself.model.fp16elseim.float()# 归一化0-255→0.0-1.0仅非张量输入需要张量输入已预处理ifnot_tensor:im/255returnim项目详情函数名preprocess功能概述将输入图像张量/列表格式转换为模型可接受的张量格式完成设备迁移、数据类型转换、归一化等操作返回值torch.Tensor预处理后的图像张量[N,3,H,W]float16/float32归一化到0-1核心逻辑格式转换列表→张量→通道转换BGR→RGB→维度调整BHWC→BCHW→设备迁移→数据类型转换→归一化设计亮点1. 兼容张量/列表两种输入格式2. 自动适配模型精度fp16/fp323. 连续内存优化ascontiguousarray提升推理速度注意事项1. 列表输入需保证每张图像为(H,W,3)格式否则np.stack会报错2. 张量输入假定已归一化无需重复除以255模型推理inferencedefinference(self,im:torch.Tensor,*args,**kwargs): 模型推理核心方法 参数: im (torch.Tensor): 预处理后的图像张量N, 3, H, W *args/ **kwargs: 额外推理参数 返回: (torch.Tensor | list): 模型原始预测结果检测为张量嵌入为列表 # 可视化特征图的保存路径仅非张量输入且启用visualize时visualize(increment_path(self.save_dir/Path(self.batch[0][0]).stem,mkdirTrue)ifself.args.visualizeand(notself.source_type.tensor)elseFalse)# 调用模型推理支持推理增强、特征图可视化、嵌入提取returnself.model(im,augmentself.args.augment,visualizevisualize,embedself.args.embed,*args,**kwargs)项目详情函数名inference功能概述执行模型前向推理可选启用数据增强、可视化特征图、提取嵌入特征返回值torch.Tensor / list[torch.Tensor]模型原始预测结果/嵌入特征核心逻辑配置可视化路径→调用模型前向传播→返回预测结果设计亮点1. 灵活支持推理增强/特征可视化/嵌入提取2. 可视化路径自动增量命名避免覆盖注意事项1. 推理增强会增加耗时但可提升小目标检测效果2. 特征可视化仅支持PyTorch模型其他格式ONNX/TensorRT不支持图像预变换pre_transformdefpre_transform(self,im:list[np.ndarray])-list[np.ndarray]: 预处理前的图像变换LetterBox缩放保持比例自动填充 参数: im (list[np.ndarray]): 原始图像列表[(H, W, 3) x N] 返回: (list[np.ndarray]): 缩放/填充后的图像列表 # 判断批次内所有图像是否形状相同若批次内所有图像形状完全相同x.shape 的值都一样集合中只会保留 1 个唯一元素same_shapeslen({x.shapeforxinim})1# 初始化LetterBox保持比例缩放自动填充letterboxLetterBox(self.imgsz,# 目标尺寸# 自动填充开关仅当图像形状相同启用rect模型为PyTorch或动态非IMX时开启autosame_shapesandself.args.rectand(self.model.ptor(getattr(self.model,dynamic,False)andnotself.model.imx)),strideself.model.stride,# 模型下采样步长保证尺寸是步长的倍数)# 对每张图像执行LetterBox变换return[letterbox(imagex)forxinim]项目详情函数名pre_transform功能概述对输入图像列表执行LetterBox缩放保证图像尺寸为模型步长的整数倍适配矩形推理返回值list[np.ndarray]缩放后的图像列表核心逻辑检查图像尺寸一致性→初始化LetterBox→逐张图像缩放设计亮点1. 自动适配矩形推理减少无效填充提升推理速度2. 兼容动态尺寸模型/IMX模型等特殊场景注意事项1. 矩形推理仅在所有图像尺寸相同时启用2. LetterBox缩放会保持长宽比不足部分填充灰色预测后处理基类占位postprocessdefpostprocess(self,preds,img,orig_imgs): 后处理基方法子类必须重写如检测的NMS、分割的掩码处理 参数: preds: 模型原始预测结果 img: 预处理后的图像张量 orig_imgs: 原始图像列表 返回: 后处理后的结果子类自定义 returnpreds项目详情函数名postprocess功能概述基类占位方法子类需重写以实现任务专属的预测结果后处理如检测任务的NMS返回值Any原始预测结果基类无处理核心逻辑直接返回原始预测结果无任何处理设计亮点为子类预留扩展接口保证框架的可扩展性注意事项子类必须重写该方法否则无法得到结构化的预测结果预测执行入口__call__def__call__(self,sourceNone,modelNone,stream:boolFalse,*args,**kwargs): 对外统一推理接口支持流式/非流式推理 参数: source: 推理数据源图片/视频/摄像头/流等 model: 推理模型路径/张量/模型实例 stream (bool): 是否流式推理True返回生成器False返回列表 *args/ **kwargs: 额外推理参数 返回: (list[Results] | generator): 推理结果列表/生成器 self.streamstreamifstream:# 流式推理返回生成器避免内存累积returnself.stream_inference(source,model,*args,**kwargs)else:# 非流式推理将生成器转为列表一次性返回所有结果returnlist(self.stream_inference(source,model,*args,**kwargs))项目详情函数名__call__功能概述预测器核心执行入口支持流式/非流式推理适配不同输入源和模型返回值list[Results] / generator非流式返回结果列表流式返回结果生成器核心逻辑设置流式标记→调用流式推理→非流式时转换为列表流式时返回生成器设计亮点1. 流式推理避免内存溢出适合视频/流等大输入源2. 统一入口兼容两种推理模式注意事项1. 非流式推理会将所有结果存入内存大输入源可能导致OOM2. 生成器需迭代消费否则无法释放资源CLI预测执行predict_clidefpredict_cli(self,sourceNone,modelNone): CLI模式专用推理方法消费生成器避免内存累积 参数: source: 推理数据源 model: 推理模型 注意: 必须消费生成器即使不使用结果否则会导致内存溢出 该方法专为CLI模式设计保证长视频/大流推理时的内存安全。 genself.stream_inference(source,model)# 消费生成器不存储结果for_ingen:pass项目详情函数名predict_cli功能概述适配命令行的预测执行方法消费流式推理生成器避免内存累积返回值无核心逻辑调用流式推理生成器→迭代消费生成器不存储结果避免内存溢出设计亮点专为CLI场景设计确保长时推理/大输入源不会内存溢出注意事项1. 不可移除生成器迭代逻辑否则会导致结果累积2. 该方法无返回值结果仅通过保存/可视化输出输入源设置setup_sourcedefsetup_source(self,source): 加载推理数据源检查输入类型初始化数据集 参数: source: 推理数据源图片/视频/摄像头/流等 # 检查并调整图像尺寸保证是模型步长的倍数self.imgszcheck_imgsz(self.args.imgsz,strideself.model.stride,min_dim2)# 加载推理数据源多源统一接口self.datasetload_inference_source(sourcesource,batchself.args.batch,# 批次大小vid_strideself.args.vid_stride,# 视频帧采样步长bufferself.args.stream_buffer,# 流缓冲开关channelsgetattr(self.model,ch,3),# 模型输入通道数默认3)# 获取输入源类型stream/video/image/tensor等self.source_typeself.dataset.source_type# 检测是否为大流/长视频/多图片输入警告流式推理必要性if(self.source_type.streamorself.source_type.screenshotorlen(self.dataset)1000# 大量图片orany(getattr(self.dataset,video_flag,[False]))# 长视频):# 对视频/流等大输入源警告非流式推理的OOM风险importtorchvisionifnotgetattr(self,stream,True):LOGGER.warning(STREAM_WARNING)# 重置视频写入器字典self.vid_writer{}项目详情函数名setup_source功能概述验证图像尺寸、加载推理输入源、检测输入源类型、初始化视频写入器返回值无核心逻辑验证图像尺寸→加载输入源→检测源类型→警告非流式推理风险→初始化视频写入器设计亮点1. 自动验证图像尺寸适配模型步长2. 统一加载所有输入源类型3. 主动警告潜在内存风险注意事项1. 需确保输入源路径/格式正确否则load_inference_source会报错2. 视频源需保证ffmpeg可用否则无法写入视频流式推理核心stream_inferencedefstream_inference(self,sourceNone,modelNone,*args,**kwargs): 核心流式推理循环执行「模型预热→批次遍历→预处理→推理→后处理→结果保存→资源释放」全流程 参数: source: 推理数据源 model: 推理模型 *args/ **kwargs: 额外推理参数 生成: (Results): 逐批次生成结构化推理结果 ifself.args.verbose:LOGGER.info()# 初始化模型未初始化时ifnotself.model:self.setup_model(model)# 线程锁保证多线程推理时的资源安全withself._lock:# 加载数据源每次predict调用时重新加载self.setup_source(sourceifsourceisnotNoneelseself.args.source)# 创建结果保存目录save/save_txt启用时ifself.args.saveorself.args.save_txt:(self.save_dir/labelsifself.args.save_txtelseself.save_dir).mkdir(parentsTrue,exist_okTrue)# 模型预热首次推理前避免首次推理耗时过长ifnotself.done_warmup:self.model.warmup(imgsz(1ifself.model.ptorself.model.tritonelseself.dataset.bs,self.model.ch,*self.imgsz))self.done_warmupTrue# 初始化批次计数、窗口列表、当前批次self.seen,self.windows,self.batch0,[],None# 初始化计时器预处理/推理/后处理profilers(ops.Profile(deviceself.device),ops.Profile(deviceself.device),ops.Profile(deviceself.device),)# 执行推理开始回调self.run_callbacks(on_predict_start)# 遍历数据集批次forself.batchinself.dataset:# 执行批次开始回调self.run_callbacks(on_predict_batch_start)# 解析当前批次数据路径列表、原始图像列表、日志字符串列表paths,im0s,sself.batch# 1. 预处理计时withprofilers[0]:imself.preprocess(im0s)# 2. 模型推理计时withprofilers[1]:predsself.inference(im,*args,**kwargs)# 若启用嵌入提取直接生成嵌入结果跳过后续处理ifself.args.embed:yieldfrom[preds]ifisinstance(preds,torch.Tensor)elsepredscontinue# 3. 后处理计时withprofilers[2]:self.resultsself.postprocess(preds,im,im0s)# 执行后处理结束回调self.run_callbacks(on_predict_postprocess_end)# 4. 结果可视化/保存/写入遍历批次内每张图像nlen(im0s)try:foriinrange(n):self.seen1# 累计处理图像数# 记录单张图像的耗时分摊批次耗时self.results[i].speed{preprocess:profilers[0].dt*1e3/n,inference:profilers[1].dt*1e3/n,postprocess:profilers[2].dt*1e3/n,}# 若启用verbose/save/save_txt/show写入结果并更新日志ifself.args.verboseorself.args.saveorself.args.save_txtorself.args.show:s[i]self.write_results(i,Path(paths[i]),im,s)exceptStopIteration:# 按下q键退出时终止循环break# 打印批次日志verbose模式ifself.args.verbose:LOGGER.info(\n.join(s))# 执行批次结束回调self.run_callbacks(on_predict_batch_end)# 生成当前批次的推理结果yieldfromself.results# 5. 资源释放# 释放视频写入器forvinself.vid_writer.values():ifisinstance(v,cv2.VideoWriter):v.release()# 销毁可视化窗口ifself.args.show:cv2.destroyAllWindows()# close any open windows# 打印最终耗时统计verbose模式ifself.args.verboseandself.seen:# 计算单张图像平均耗时ttuple(x.t/self.seen*1e3forxinprofilers)# speeds per imageLOGGER.info(fSpeed: %.1fms preprocess, %.1fms inference, %.1fms postprocess per image at shape f{(min(self.args.batch,self.seen),getattr(self.model,ch,3),*im.shape[2:])}%t)# 打印结果保存路径ifself.args.saveorself.args.save_txtorself.args.save_crop:nllen(list(self.save_dir.glob(labels/*.txt)))sf\n{nl}label{s*(nl1)}saved to{self.save_dir/labels}ifself.args.save_txtelseLOGGER.info(fResults saved to{colorstr(bold,self.save_dir)}{s})# 执行推理结束回调self.run_callbacks(on_predict_end)verbose 模式冗余 / 详细模式 是控制日志输出详细程度的核心配置用于平衡 “执行透明度” 与 “日志简洁性”。项目详情函数名stream_inference功能概述预测器核心执行逻辑完成模型预热、逐批次推理、结果处理、可视化/保存、性能统计等全流程返回值generator逐批次的Results对象生成器核心逻辑模型加载→输入源设置→模型预热→逐批次预处理→推理→后处理→结果保存/可视化→性能统计→释放资源设计亮点1. 线程安全的推理流程2. 逐批次处理避免内存累积3. 精细化性能统计4. 完整的资源释放机制5. 回调机制支持自定义扩展注意事项1. 模型预热仅执行一次提升后续推理速度2. 生成器需迭代消费否则会阻塞在批次循环中3. 视频写入器需手动释放否则会导致视频文件损坏模型初始化setup_modeldefsetup_model(self,model,verbose:boolTrue): 初始化推理模型支持多格式模型加载PyTorch/ONNX/TensorRT等 参数: model: 模型路径/实例默认使用self.args.model verbose (bool): 是否打印详细日志 # AutoBackend多模型格式统一加载后端self.modelAutoBackend(modelmodelorself.args.model,# 模型路径/实例deviceselect_device(self.args.device,verboseverbose),# 选择推理设备dnnself.args.dnn,# 是否使用OpenCV DNNONNX模型dataself.args.data,# 数据配置fp16self.args.half,# 是否启用fp16推理fuseTrue,# 是否融合ConvBN提升PyTorch模型效率verboseverbose,# 日志开关)# 更新设备/半精度配置同步模型实际状态self.deviceself.model.device self.args.halfself.model.fp16# 从导出元数据中复用图像尺寸非动态模型ifhasattr(self.model,imgsz)andnotgetattr(self.model,dynamic,False):self.args.imgszself.model.imgsz# 模型设为评估模式禁用Dropout/BatchNorm训练行为self.model.eval()# 尝试编译模型提升推理效率如TorchScriptself.modelattempt_compile(self.model,deviceself.device,modeself.args.compile)项目详情函数名setup_model功能概述加载多格式模型PyTorch/ONNX/TensorRT等配置设备、精度、编译等参数返回值无核心逻辑加载AutoBackend模型→更新设备/精度配置→设置模型为评估模式→可选编译模型设计亮点1. 一站式加载多格式模型无需手动适配2. 自动更新设备/精度配置3. 可选编译模型提升推理速度注意事项1. TensorRT/ONNX等模型需提前导出否则加载失败2. 模型编译仅支持PyTorch 2.0且仅在GPU上生效结果写入write_resultsdefwrite_results(self,i:int,p:Path,im:torch.Tensor,s:list[str])-str: 写入单张图像的推理结果保存txt/图片/视频可视化生成日志字符串 参数: i (int): 批次内图像索引 p (Path): 图像路径 im (torch.Tensor): 预处理后的图像张量 s (list[str]): 日志字符串列表 返回: (str): 更新后的日志字符串 string# 初始化日志字符串# 扩展批次维度单张图像时iflen(im.shape)3:imim[None]# 处理流/张量/从图像输入添加批次索引ifself.source_type.streamorself.source_type.from_imgorself.source_type.tensor:stringf{i}: frameself.dataset.count# 流/视频帧计数else:# 解析视频帧号从日志字符串中matchre.search(rframe (\d)/,s[i])frameint(match[1])ifmatchelseNone# 初始化txt保存路径图片文件名视频文件名_帧号self.txt_pathself.save_dir/labels/(p.stem(ifself.dataset.modeimageelsef_{frame}))# 添加图像尺寸到日志string{:g}x{:g} .format(*im.shape[2:])# 获取当前图像的推理结果resultself.results[i]result.save_dirself.save_dir.__str__()# 设置结果保存目录# 添加推理结果详情耗时到日志stringf{result.verbose()}{result.speed[inference]:.1f}ms# 可视化/保存结果启用save/show时ifself.args.saveorself.args.show:# 绘制预测结果框/掩码/标签/置信度self.plotted_imgresult.plot(line_widthself.args.line_width,# 框线宽度boxesself.args.show_boxes,# 是否显示框confself.args.show_conf,# 是否显示置信度labelsself.args.show_labels,# 是否显示标签im_gpuNoneifself.args.retina_maskselseim[i],# 掩码绘制优化)# 保存结果到txtsave_txt启用时ifself.args.save_txt:result.save_txt(f{self.txt_path}.txt,save_confself.args.save_conf)# 保存裁剪后的目标save_crop启用时ifself.args.save_crop:result.save_crop(save_dirself.save_dir/crops,file_nameself.txt_path.stem)# 可视化展示show启用时ifself.args.show:self.show(str(p))# 保存预测图像/视频save启用时ifself.args.save:self.save_predicted_images(self.save_dir/p.name,frame)returnstring项目详情函数名write_results功能概述生成结果字符串、保存标签TXT、保存裁剪图像、可视化结果、保存预测图像/视频返回值str当前图像的结果信息字符串核心逻辑生成结果字符串→绘制预测图像→保存TXT标签→保存裁剪→可视化→保存图像/视频设计亮点1. 统一管理所有结果输出方式2. 结果字符串包含关键信息尺寸/耗时/检测结果3. 适配图像/视频不同保存逻辑注意事项1. 保存裁剪需确保save_crop为True否则不会生成2. 可视化需系统支持否则会静默失败预测图像保存save_predicted_imagesdefsave_predicted_images(self,save_path:Path,frame:int0): 保存预测结果图片→JPG视频→MP4/AVI流→帧视频 参数: save_path (Path): 保存路径 frame (int): 视频帧号仅视频/流输入 # 带预测结果的图像imself.plotted_img# 处理视频/流输入ifself.dataset.modein{stream,video}:# 获取帧率视频用原帧率流默认30fpsself.dataset.fpsifself.dataset.modevideoelse30# 帧保存目录save_frames启用时frames_pathself.save_dir/f{save_path.stem}_frames# 初始化视频写入器新视频ifsave_pathnotinself.vid_writer:ifself.args.save_frames:Path(frames_path).mkdir(parentsTrue,exist_okTrue)# 适配不同系统的视频编码suffix,fourcc(.mp4,avc1)ifMACOSelse(.avi,WMV2)ifWINDOWSelse(.avi,MJPG)self.vid_writer[save_path]cv2.VideoWriter(filenamestr(Path(save_path).with_suffix(suffix)),fourcccv2.VideoWriter_fourcc(*fourcc),fpsfps,# 帧率必须整数frameSize(im.shape[1],im.shape[0]),# 视频尺寸宽度高度)# 写入视频帧self.vid_writer[save_path].write(im)# 保存单帧save_frames启用时ifself.args.save_frames:cv2.imwrite(f{frames_path}/{save_path.stem}_{frame}.jpg,im)# 处理图片输入保存为JPGselse:cv2.imwrite(str(save_path.with_suffix(.jpg)),im)项目详情函数名save_predicted_images功能概述按输入源类型图像/视频/流保存预测结果视频保存为MP4/AVI图像保存为JPG返回值无设计亮点1. 跨平台适配视频编码MAC/WIN/LINUX2. 支持视频帧单独保存3. 图像统一保存为JPG保证兼容性注意事项1. 视频编码需系统支持对应fourcc否则会导致视频无法播放2. 保存路径需有写入权限否则会报错结果可视化showdefshow(self,p:str): 可视化展示带预测结果的图像 参数: p (str): 窗口名称默认空 imself.plotted_img# Linux系统创建可调整大小的窗口ifplatform.system()Linuxandpnotinself.windows:self.windows.append(p)cv2.namedWindow(p,cv2.WINDOW_NORMAL|cv2.WINDOW_KEEPRATIO)cv2.resizeWindow(p,im.shape[1],im.shape[0])# 显示图像cv2.imshow(p,im)# 按键监听图片窗口停留300ms视频/流停留1ms按q键退出ifcv2.waitKey(300ifself.dataset.modeimageelse1)0xFFord(q):raiseStopIteration项目详情函数名show功能概述在窗口中显示预测结果图像支持图像/视频的实时可视化返回值无核心逻辑创建自适应窗口→显示图像→监听按键按q退出设计亮点1. Linux系统支持窗口自适应调整2. 图像/视频适配不同等待时间3. 按键退出机制友好注意事项1. 无显示器的环境如服务器禁用可视化否则会报错2. 视频可视化需快速迭代否则会卡顿回调执行run_callbacksdefrun_callbacks(self,event:str): 执行指定事件的所有回调函数 参数: event (str): 事件名称如on_predict_start/on_predict_batch_end forcallbackinself.callbacks.get(event,[]):callback(self)项目详情函数名run_callbacks功能概述执行指定事件的所有注册回调函数支持自定义扩展预测流程返回值无核心逻辑遍历指定事件的回调列表→逐个执行回调函数传入预测器自身设计亮点1. 松耦合的回调机制无需修改核心代码即可扩展功能2. 支持多回调函数注册注意事项1. 回调函数需接收预测器实例作为参数2. 回调函数异常会中断预测流程需做好异常处理回调注册add_callbackdefadd_callback(self,event:str,func:callable): 注册自定义回调函数到指定事件 参数: event (str): 事件名称 func (callable): 回调函数接收self作为参数 self.callbacks[event].append(func)项目详情函数名add_callback功能概述为指定事件注册新的回调函数扩展预测器功能返回值无核心逻辑将回调函数添加到指定事件的回调列表中设计亮点简单易用的回调注册接口支持动态扩展预测流程注意事项1. 需确保事件名正确如on_predict_start否则回调不会执行2. 避免注册耗时回调否则会降低推理速度完整代码# Ultralytics AGPL-3.0 License - https://ultralytics.com/license Run prediction on images, videos, directories, globs, YouTube, webcam, streams, etc. Usage - sources: $ yolo modepredict modelyolo11n.pt source0 # webcam img.jpg # image vid.mp4 # video screen # screenshot path/ # directory list.txt # list of images list.streams # list of streams path/*.jpg # glob https://youtu.be/LNwODJXcvt4 # YouTube rtsp://example.com/media.mp4 # RTSP, RTMP, HTTP, TCP stream Usage - formats: $ yolo modepredict modelyolo11n.pt # PyTorch yolo11n.torchscript # TorchScript yolo11n.onnx # ONNX Runtime or OpenCV DNN with dnnTrue yolo11n_openvino_model # OpenVINO yolo11n.engine # TensorRT yolo11n.mlpackage # CoreML (macOS-only) yolo11n_saved_model # TensorFlow SavedModel yolo11n.pb # TensorFlow GraphDef yolo11n.tflite # TensorFlow Lite yolo11n_edgetpu.tflite # TensorFlow Edge TPU yolo11n_paddle_model # PaddlePaddle yolo11n.mnn # MNN yolo11n_ncnn_model # NCNN yolo11n_imx_model # Sony IMX yolo11n_rknn_model # Rockchip RKNN yolo11n.pte # PyTorch Executorch # 引入未来版本的类型注解支持提升代码静态类型检查和编辑器提示能力from__future__importannotations# 导入系统平台检测模块区分Windows/Linux/macOS适配视频编码格式、窗口显示逻辑importplatform# 导入正则表达式模块解析视频帧号、处理路径/日志字符串中的结构化信息importre# 导入线程管理模块提供线程锁保证多线程环境下推理的资源安全importthreading# 导入路径处理模块面向对象的路径操作统一管理推理结果保存路径、数据源路径frompathlibimportPath# 导入类型注解模块Any类型用于兼容推理流程中多类型的输入/输出参数fromtypingimportAny# 导入OpenCV核心模块图像/视频的读取、写入、可视化、格式转换、视频编码等importcv2# 导入数值计算模块图像数组处理、维度转换、归一化、数组拼接等核心数值操作importnumpyasnp# 导入PyTorch核心模块张量操作、设备管理、模型推理、数据类型转换等importtorch# 从ultralytics配置模块导入配置加载函数合并基础配置与自定义覆盖参数、保存目录生成函数自动创建runs/predict/exp*目录fromultralytics.cfgimportget_cfg,get_save_dir# 从ultralytics数据模块导入推理数据源加载函数统一处理摄像头/视频/图片/流/张量等多类型输入源fromultralytics.dataimportload_inference_source# 从ultralytics数据增强模块导入LetterBox类图像等比例缩放自动填充保证尺寸为模型步长倍数fromultralytics.data.augmentimportLetterBox# 从ultralytics神经网络后端模块导入AutoBackend类统一加载PyTorch/ONNX/TensorRT等多格式模型fromultralytics.nn.autobackendimportAutoBackend# 从ultralytics工具模块导入默认配置常量、日志器、系统标识macOS/WINDOWS、回调函数字典、彩色字符串函数、通用操作工具计时/缩放等fromultralytics.utilsimportDEFAULT_CFG,LOGGER,MACOS,WINDOWS,callbacks,colorstr,ops# 从ultralytics工具检查模块导入图像尺寸检查函数保证尺寸为模型步长倍数、可视化可用性检查函数兼容不同系统的窗口显示fromultralytics.utils.checksimportcheck_imgsz,check_imshow# 从ultralytics工具文件模块导入路径增量函数自动生成不重复的文件/目录名避免结果覆盖fromultralytics.utils.filesimportincrement_path# 从ultralytics工具PyTorch模块导入模型编译函数提升推理效率、设备选择函数自动选择CPU/GPU/TensorRT、推理模式装饰器禁用梯度计算fromultralytics.utils.torch_utilsimportattempt_compile,select_device,smart_inference_mode STREAM_WARNING inference results will accumulate in RAM unless streamTrue is passed, causing potential out-of-memory errors for large sources or long-running streams and videos. See https://docs.ultralytics.com/modes/predict/ for help. Example: results model(source..., streamTrue) # generator of Results objects for r in results: boxes r.boxes # Boxes object for bbox outputs masks r.masks # Masks object for segment masks outputs probs r.probs # Class probabilities for classification outputs classBasePredictor: YOLO通用预测器基类抽象所有推理任务的公共逻辑 核心作用为检测/分割/分类等专用预测器提供标准化的推理框架包含 - 配置加载与初始化 - 多源数据加载摄像头/视频/图片/流/张量等 - 图像预处理缩放/归一化/通道转换/设备迁移 - 模型推理多格式模型兼容 - 结果后处理子类重写 - 结果保存/可视化/日志输出 - 回调机制推理各阶段自定义逻辑 - 流式推理避免内存溢出 - 线程安全推理 属性: args (SimpleNamespace): 推理配置参数conf/imgsz/save/show等 save_dir (Path): 结果保存目录 done_warmup (bool): 模型是否完成预热 model (torch.nn.Module): AutoBackend封装的推理模型支持多格式 data (dict): 数据配置字典 device (torch.device): 推理设备CPU/GPU/TensorRT等 dataset (InferenceDataset): 加载的推理数据源 vid_writer (dict[str, cv2.VideoWriter]): 视频写入器字典{保存路径: 写入器} plotted_img (np.ndarray): 最后一次可视化的图像带预测框/掩码 source_type (SimpleNamespace): 输入源类型stream/video/image/tensor等 seen (int): 已处理的图像数量 windows (list[str]): 可视化窗口名称列表 batch (tuple): 当前批次数据路径/原始图像/日志字符串 results (list[Any]): 当前批次的推理结果Results对象列表 transforms (callable): 分类任务专用图像变换 callbacks (dict[str, list[callable]]): 回调函数字典按事件分类 txt_path (Path): 预测结果txt保存路径 _lock (threading.Lock): 线程锁保证多线程推理安全 方法: preprocess: 图像预处理格式转换/归一化/设备迁移/数据类型转换 pre_transform: 预处理前的LetterBox缩放保持比例自动填充 inference: 模型推理支持增强/可视化/嵌入提取 postprocess: 后处理基方法子类重写如检测的NMS __call__: 对外推理接口支持流式/非流式 predict_cli: CLI模式推理消费生成器避免内存累积 setup_source: 加载推理数据源检查输入类型 stream_inference: 核心流式推理循环全流程执行 setup_model: 初始化模型AutoBackend加载多格式模型设备选择/编译 write_results: 结果写入txt/图片/视频可视化 save_predicted_images: 保存预测结果图片/视频/帧 show: 可视化展示推理结果 run_callbacks: 执行指定事件的所有回调函数 add_callback: 注册自定义回调函数 def__init__(self,cfgDEFAULT_CFG,overrides:dict[str,Any]|NoneNone,_callbacks:dict[str,list[callable]]|NoneNone,): 初始化BasePredictor类 参数: cfg (str | dict): 基础配置文件路径/字典默认DEFAULT_CFG overrides (dict, optional): 配置覆盖参数如conf0.5, imgsz640 _callbacks (dict, optional): 自定义回调函数字典按事件分类 # 加载并合并配置基础配置 覆盖参数self.argsget_cfg(cfg,overrides)# 生成结果保存目录默认 runs/predict/exp*self.save_dirget_save_dir(self.args)# 设置默认置信度阈值未指定时为0.25ifself.args.confisNone:self.args.conf0.25# 标记模型未完成预热self.done_warmupFalse# 检查可视化可用性Windows/Linux/macOS兼容ifself.args.show:self.args.showcheck_imshow(warnTrue)# 后续初始化的属性setup_model/setup_source后赋值self.modelNone# 推理模型实例AutoBackend封装支持PyTorch/ONNX/TensorRT等格式setup_model后初始化self.dataself.args.data# 数据配置字典如数据集路径、类别名等来自推理配置参数self.imgszNone# 推理输入图像尺寸tuple如(640,640)setup_source中check_imgsz后确定self.deviceNone# 推理设备torch.device对象如cuda:0/cpusetup_model后从model.device同步self.datasetNone# 推理数据源实例InferenceDataset封装多类型输入源setup_source后初始化self.vid_writer{}# 视频写入器字典key保存路径valuecv2.VideoWriter用于多视频/流结果保存self.plotted_imgNone# 最后一次绘制的带预测结果的图像np.ndarray可视化/保存时使用self.source_typeNone# 输入源类型SimpleNamespace含stream/video/image/tensor等标识setup_source后赋值self.seen0# 已处理的图像/帧总数流式推理中逐批次累加用于统计和耗时计算self.windows[]# 可视化窗口名称列表Linux下创建可调整窗口时记录避免重复创建self.batchNone# 当前批次数据tuple含paths/im0s/s遍历dataset时赋值self.resultsNone# 当前批次后处理结果list[Results]postprocess后赋值self.transformsNone# 分类任务专用图像变换callable检测任务暂未使用self.callbacks_callbacksorcallbacks.get_default_callbacks()# 推理回调函数字典按事件分类如on_predict_startself.txt_pathNone# 单张图像预测结果txt保存路径Pathwrite_results时动态生成self._lockthreading.Lock()# 线程锁保证多线程环境下推理流程的资源安全如模型推理/视频写入callbacks.add_integration_callbacks(self)# 注册集成回调如WB/MLflow日志、ClearML监控等defpreprocess(self,im:torch.Tensor|list[np.ndarray])-torch.Tensor: 图像预处理将输入转换为模型可接受的张量格式BCHW/归一化/设备迁移/数据类型转换 参数: im (torch.Tensor | list[np.ndarray]): 输入图像 - Tensor: 形状(N, 3, H, W) - list[np.ndarray]: 形状[(H, W, 3) x N]N为批次大小 返回: (torch.Tensor): 预处理后的张量N, 3, H, W已迁移到指定设备归一化到0-1 # 判断输入是否为张量not_tensornotisinstance(im,torch.Tensor)ifnot_tensor:# 对列表中每张图像执行pre_transformLetterBox缩放imnp.stack(self.pre_transform(im))# BGR→RGBOpenCV加载的图像是BGR模型需要RGBifim.shape[-1]3:imim[...,::-1]# 维度转换BHWC批次高度宽度通道→ BCHW批次通道高度宽度imim.transpose((0,3,1,2))# 转换为连续内存数组提升张量操作效率imnp.ascontiguousarray(im)# NumPy→PyTorch张量imtorch.from_numpy(im)# 迁移到推理设备CPU/GPUimim.to(self.device)# 数据类型转换uint8→fp16/fp32模型若启用fp16则用half否则floatimim.half()ifself.model.fp16elseim.float()# 归一化0-255→0.0-1.0仅非张量输入需要张量输入已预处理ifnot_tensor:im/255returnimdefinference(self,im:torch.Tensor,*args,**kwargs): 模型推理核心方法 参数: im (torch.Tensor): 预处理后的图像张量N, 3, H, W *args/ **kwargs: 额外推理参数 返回: (torch.Tensor | list): 模型原始预测结果检测为张量嵌入为列表 # 可视化特征图的保存路径仅非张量输入且启用visualize时visualize(increment_path(self.save_dir/Path(self.batch[0][0]).stem,mkdirTrue)ifself.args.visualizeand(notself.source_type.tensor)elseFalse)# 调用模型推理支持推理增强、特征图可视化、嵌入提取returnself.model(im,augmentself.args.augment,visualizevisualize,embedself.args.embed,*args,**kwargs)defpre_transform(self,im:list[np.ndarray])-list[np.ndarray]: 预处理前的图像变换LetterBox缩放保持比例自动填充 参数: im (list[np.ndarray]): 原始图像列表[(H, W, 3) x N] 返回: (list[np.ndarray]): 缩放/填充后的图像列表 # 判断批次内所有图像是否形状相同若批次内所有图像形状完全相同x.shape 的值都一样集合中只会保留 1 个唯一元素same_shapeslen({x.shapeforxinim})1# 初始化LetterBox保持比例缩放自动填充letterboxLetterBox(self.imgsz,# 目标尺寸# 自动填充开关仅当图像形状相同启用rect模型为PyTorch或动态非IMX时开启autosame_shapesandself.args.rectand(self.model.ptor(getattr(self.model,dynamic,False)andnotself.model.imx)),strideself.model.stride,# 模型下采样步长保证尺寸是步长的倍数)# 对每张图像执行LetterBox变换return[letterbox(imagex)forxinim]defpostprocess(self,preds,img,orig_imgs): 后处理基方法子类必须重写如检测的NMS、分割的掩码处理 参数: preds: 模型原始预测结果 img: 预处理后的图像张量 orig_imgs: 原始图像列表 返回: 后处理后的结果子类自定义 returnpredsdef__call__(self,sourceNone,modelNone,stream:boolFalse,*args,**kwargs): 对外统一推理接口支持流式/非流式推理 参数: source: 推理数据源图片/视频/摄像头/流等 model: 推理模型路径/张量/模型实例 stream (bool): 是否流式推理True返回生成器False返回列表 *args/ **kwargs: 额外推理参数 返回: (list[Results] | generator): 推理结果列表/生成器 self.streamstreamifstream:# 流式推理返回生成器避免内存累积returnself.stream_inference(source,model,*args,**kwargs)else:# 非流式推理将生成器转为列表一次性返回所有结果returnlist(self.stream_inference(source,model,*args,**kwargs))defpredict_cli(self,sourceNone,modelNone): CLI模式专用推理方法消费生成器避免内存累积 参数: source: 推理数据源 model: 推理模型 注意: 必须消费生成器即使不使用结果否则会导致内存溢出 该方法专为CLI模式设计保证长视频/大流推理时的内存安全。 genself.stream_inference(source,model)# 消费生成器不存储结果for_ingen:passdefsetup_source(self,source): 加载推理数据源检查输入类型初始化数据集 参数: source: 推理数据源图片/视频/摄像头/流等 # 检查并调整图像尺寸保证是模型步长的倍数self.imgszcheck_imgsz(self.args.imgsz,strideself.model.stride,min_dim2)# 加载推理数据源多源统一接口self.datasetload_inference_source(sourcesource,batchself.args.batch,# 批次大小vid_strideself.args.vid_stride,# 视频帧采样步长bufferself.args.stream_buffer,# 流缓冲开关channelsgetattr(self.model,ch,3),# 模型输入通道数默认3)# 获取输入源类型stream/video/image/tensor等self.source_typeself.dataset.source_type# 检测是否为大流/长视频/多图片输入警告流式推理必要性if(self.source_type.streamorself.source_type.screenshotorlen(self.dataset)1000# 大量图片orany(getattr(self.dataset,video_flag,[False]))# 长视频):# 对视频/流等大输入源警告非流式推理的OOM风险importtorchvisionifnotgetattr(self,stream,True):LOGGER.warning(STREAM_WARNING)# 重置视频写入器字典self.vid_writer{}smart_inference_mode()defstream_inference(self,sourceNone,modelNone,*args,**kwargs): 核心流式推理循环执行「模型预热→批次遍历→预处理→推理→后处理→结果保存→资源释放」全流程 参数: source: 推理数据源 model: 推理模型 *args/ **kwargs: 额外推理参数 生成: (Results): 逐批次生成结构化推理结果 ifself.args.verbose:LOGGER.info()# 初始化模型未初始化时ifnotself.model:self.setup_model(model)# 线程锁保证多线程推理时的资源安全withself._lock:# 加载数据源每次predict调用时重新加载self.setup_source(sourceifsourceisnotNoneelseself.args.source)# 创建结果保存目录save/save_txt启用时ifself.args.saveorself.args.save_txt:(self.save_dir/labelsifself.args.save_txtelseself.save_dir).mkdir(parentsTrue,exist_okTrue)# 模型预热首次推理前避免首次推理耗时过长ifnotself.done_warmup:self.model.warmup(imgsz(1ifself.model.ptorself.model.tritonelseself.dataset.bs,self.model.ch,*self.imgsz))self.done_warmupTrue# 初始化批次计数、窗口列表、当前批次self.seen,self.windows,self.batch0,[],None# 初始化计时器预处理/推理/后处理profilers(ops.Profile(deviceself.device),ops.Profile(deviceself.device),ops.Profile(deviceself.device),)# 执行推理开始回调self.run_callbacks(on_predict_start)# 遍历数据集批次forself.batchinself.dataset:# 执行批次开始回调self.run_callbacks(on_predict_batch_start)# 解析当前批次数据路径列表、原始图像列表、日志字符串列表paths,im0s,sself.batch# 1. 预处理计时withprofilers[0]:imself.preprocess(im0s)# 2. 模型推理计时withprofilers[1]:predsself.inference(im,*args,**kwargs)# 若启用嵌入提取直接生成嵌入结果跳过后续处理ifself.args.embed:yieldfrom[preds]ifisinstance(preds,torch.Tensor)elsepredscontinue# 3. 后处理计时withprofilers[2]:self.resultsself.postprocess(preds,im,im0s)# 执行后处理结束回调self.run_callbacks(on_predict_postprocess_end)# 4. 结果可视化/保存/写入遍历批次内每张图像nlen(im0s)try:foriinrange(n):self.seen1# 累计处理图像数# 记录单张图像的耗时分摊批次耗时self.results[i].speed{preprocess:profilers[0].dt*1e3/n,inference:profilers[1].dt*1e3/n,postprocess:profilers[2].dt*1e3/n,}# 若启用verbose/save/save_txt/show写入结果并更新日志ifself.args.verboseorself.args.saveorself.args.save_txtorself.args.show:s[i]self.write_results(i,Path(paths[i]),im,s)exceptStopIteration:# 按下q键退出时终止循环break# 打印批次日志verbose模式ifself.args.verbose:LOGGER.info(\n.join(s))# 执行批次结束回调self.run_callbacks(on_predict_batch_end)# 生成当前批次的推理结果yieldfromself.results# 5. 资源释放# 释放视频写入器forvinself.vid_writer.values():ifisinstance(v,cv2.VideoWriter):v.release()# 销毁可视化窗口ifself.args.show:cv2.destroyAllWindows()# close any open windows# 打印最终耗时统计verbose模式ifself.args.verboseandself.seen:# 计算单张图像平均耗时ttuple(x.t/self.seen*1e3forxinprofilers)# speeds per imageLOGGER.info(fSpeed: %.1fms preprocess, %.1fms inference, %.1fms postprocess per image at shape f{(min(self.args.batch,self.seen),getattr(self.model,ch,3),*im.shape[2:])}%t)# 打印结果保存路径ifself.args.saveorself.args.save_txtorself.args.save_crop:nllen(list(self.save_dir.glob(labels/*.txt)))sf\n{nl}label{s*(nl1)}saved to{self.save_dir/labels}ifself.args.save_txtelseLOGGER.info(fResults saved to{colorstr(bold,self.save_dir)}{s})# 执行推理结束回调self.run_callbacks(on_predict_end)defsetup_model(self,model,verbose:boolTrue): 初始化推理模型支持多格式模型加载PyTorch/ONNX/TensorRT等 参数: model: 模型路径/实例默认使用self.args.model verbose (bool): 是否打印详细日志 # AutoBackend多模型格式统一加载后端self.modelAutoBackend(modelmodelorself.args.model,# 模型路径/实例deviceselect_device(self.args.device,verboseverbose),# 选择推理设备dnnself.args.dnn,# 是否使用OpenCV DNNONNX模型dataself.args.data,# 数据配置fp16self.args.half,# 是否启用fp16推理fuseTrue,# 是否融合ConvBN提升PyTorch模型效率verboseverbose,# 日志开关)# 更新设备/半精度配置同步模型实际状态self.deviceself.model.device self.args.halfself.model.fp16# 从导出元数据中复用图像尺寸非动态模型ifhasattr(self.model,imgsz)andnotgetattr(self.model,dynamic,False):self.args.imgszself.model.imgsz# 模型设为评估模式禁用Dropout/BatchNorm训练行为self.model.eval()# 尝试编译模型提升推理效率如TorchScriptself.modelattempt_compile(self.model,deviceself.device,modeself.args.compile)defwrite_results(self,i:int,p:Path,im:torch.Tensor,s:list[str])-str: 写入单张图像的推理结果保存txt/图片/视频可视化生成日志字符串 参数: i (int): 批次内图像索引 p (Path): 图像路径 im (torch.Tensor): 预处理后的图像张量 s (list[str]): 日志字符串列表 返回: (str): 更新后的日志字符串 string# 初始化日志字符串# 扩展批次维度单张图像时iflen(im.shape)3:imim[None]# 处理流/张量/从图像输入添加批次索引ifself.source_type.streamorself.source_type.from_imgorself.source_type.tensor:stringf{i}: frameself.dataset.count# 流/视频帧计数else:# 解析视频帧号从日志字符串中matchre.search(rframe (\d)/,s[i])frameint(match[1])ifmatchelseNone# 初始化txt保存路径图片文件名视频文件名_帧号self.txt_pathself.save_dir/labels/(p.stem(ifself.dataset.modeimageelsef_{frame}))# 添加图像尺寸到日志string{:g}x{:g} .format(*im.shape[2:])# 获取当前图像的推理结果resultself.results[i]result.save_dirself.save_dir.__str__()# 设置结果保存目录# 添加推理结果详情耗时到日志stringf{result.verbose()}{result.speed[inference]:.1f}ms# 可视化/保存结果启用save/show时ifself.args.saveorself.args.show:# 绘制预测结果框/掩码/标签/置信度self.plotted_imgresult.plot(line_widthself.args.line_width,# 框线宽度boxesself.args.show_boxes,# 是否显示框confself.args.show_conf,# 是否显示置信度labelsself.args.show_labels,# 是否显示标签im_gpuNoneifself.args.retina_maskselseim[i],# 掩码绘制优化)# 保存结果到txtsave_txt启用时ifself.args.save_txt:result.save_txt(f{self.txt_path}.txt,save_confself.args.save_conf)# 保存裁剪后的目标save_crop启用时ifself.args.save_crop:result.save_crop(save_dirself.save_dir/crops,file_nameself.txt_path.stem)# 可视化展示show启用时ifself.args.show:self.show(str(p))# 保存预测图像/视频save启用时ifself.args.save:self.save_predicted_images(self.save_dir/p.name,frame)returnstringdefsave_predicted_images(self,save_path:Path,frame:int0): 保存预测结果图片→JPG视频→MP4/AVI流→帧视频 参数: save_path (Path): 保存路径 frame (int): 视频帧号仅视频/流输入 # 带预测结果的图像imself.plotted_img# 处理视频/流输入ifself.dataset.modein{stream,video}:# 获取帧率视频用原帧率流默认30fpsself.dataset.fpsifself.dataset.modevideoelse30# 帧保存目录save_frames启用时frames_pathself.save_dir/f{save_path.stem}_frames# 初始化视频写入器新视频ifsave_pathnotinself.vid_writer:ifself.args.save_frames:Path(frames_path).mkdir(parentsTrue,exist_okTrue)# 适配不同系统的视频编码suffix,fourcc(.mp4,avc1)ifMACOSelse(.avi,WMV2)ifWINDOWSelse(.avi,MJPG)self.vid_writer[save_path]cv2.VideoWriter(filenamestr(Path(save_path).with_suffix(suffix)),fourcccv2.VideoWriter_fourcc(*fourcc),fpsfps,# 帧率必须整数frameSize(im.shape[1],im.shape[0]),# 视频尺寸宽度高度)# 写入视频帧self.vid_writer[save_path].write(im)# 保存单帧save_frames启用时ifself.args.save_frames:cv2.imwrite(f{frames_path}/{save_path.stem}_{frame}.jpg,im)# 处理图片输入保存为JPGselse:cv2.imwrite(str(save_path.with_suffix(.jpg)),im)defshow(self,p:str): 可视化展示带预测结果的图像 参数: p (str): 窗口名称默认空 imself.plotted_img# Linux系统创建可调整大小的窗口ifplatform.system()Linuxandpnotinself.windows:self.windows.append(p)cv2.namedWindow(p,cv2.WINDOW_NORMAL|cv2.WINDOW_KEEPRATIO)cv2.resizeWindow(p,im.shape[1],im.shape[0])# 显示图像cv2.imshow(p,im)# 按键监听图片窗口停留300ms视频/流停留1ms按q键退出ifcv2.waitKey(300ifself.dataset.modeimageelse1)0xFFord(q):raiseStopIterationdefrun_callbacks(self,event:str): 执行指定事件的所有回调函数 参数: event (str): 事件名称如on_predict_start/on_predict_batch_end forcallbackinself.callbacks.get(event,[]):callback(self)defadd_callback(self,event:str,func:callable): 注册自定义回调函数到指定事件 参数: event (str): 事件名称 func (callable): 回调函数接收self作为参数 self.callbacks[event].append(func)通用预测框架的核心特性特性实现方式多输入源适配setup_source调用load_inference_source统一加载图像/视频/流/摄像头等多模型格式兼容setup_model通过AutoBackend加载PyTorch/ONNX/TensorRT等格式模型流式推理防OOMstream_inference生成器逐批次返回结果predict_cli消费生成器性能精细化统计基于ops.Profile统计预处理/推理/后处理耗时输出单图像平均耗时跨平台视频保存根据系统MAC/WIN/LINUX选择不同视频编码avc1/WMV2/MJPG线程安全推理stream_inference使用self._lock保证多线程下推理安全模型预热优化首次推理前执行model.warmup提升后续推理速度工程化核心优化优化点实现方式内存效率优化流式推理避免结果累积LetterBox矩形推理减少无效填充数据类型适配自动转换图像为模型所需的fp16/fp32非张量输入自动归一化资源自动释放推理结束后释放视频写入器、关闭可视化窗口避免资源泄漏配置灵活解析__init__通过get_cfg合并基础配置与覆盖项支持动态参数调整错误友好提示非流式推理大输入源时输出OOM警告可视化不可用时自动禁用代码可扩展性基类占位postprocess方法子类重写实现任务专属后处理扩展与易用性能力扩展项用途回调机制run_callbacks/add_callback支持自定义扩展预测流程如日志/监控多结果输出方式支持保存图像/视频/TXT标签/裁剪图像可视化窗口显示矩形推理优化pre_transform的LetterBox根据图像尺寸一致性启用auto模式动态尺寸验证check_imgsz保证图像尺寸为模型步长整数倍避免推理错误统一结果封装子类返回Results对象包含检测框/掩码/类别等所有信息关键注意事项流式推理使用视频/流/大量图像必须启用streamTrue否则会导致内存溢出生成器需迭代消费不可直接转为列表除非确认数据量小。模型格式适配ONNX/TensorRT等模型需提前导出AutoBackend加载时需保证依赖库如onnxruntime/tensorrt已安装。可视化限制服务器/无显示器环境需禁用showTrue否则会触发OpenCV错误Linux系统需安装X11才能显示窗口。视频保存依赖需保证系统安装ffmpeg否则视频写入器初始化失败可降级为保存单帧图像。回调函数开发自定义回调函数需接收预测器实例作为参数避免修改预测器核心属性做好异常捕获。多线程推理多线程调用预测器时需确保每个线程使用独立实例或依赖self._lock保证线程安全。图像输入格式列表输入需为RGB通道的(H,W,3)格式OpenCV读取的BGR图像会自动转换无需手动处理。总结详细介绍了 Ultralytics 框架中作为 YOLO 全任务推理通用基础的 BasePredictor 类。