yolo26训练自己的数据集
1.yolo26介绍
YOLO26 是一个统一的实时视觉模型,它引入了原生端到端推理、更轻量化的检测头、更新的训练方案，以及用于检测、分割、姿态估计、分类和有向检测的任务专用头。
在五个检测尺度上，YOLO26 在 COCO 上达到了 40.9-57.5 mAP，且 T4 TensorRT 延迟仅为 1.7-11.8 毫秒。该论文还报告称，与 YOLO11n 相比，YOLO26n 在 Intel Xeon CPU @ 2.00 GHz 上的 CPU ONNX 推理速度提升了高达 43%。具体参考YOLO 目标检测与分割 | Ultralytics。除了常见的检测，分割，分类，姿态，旋转框obb外，还支持了深度估计(通过模型计算每个像素到相机平面的垂直距离)，粗略的计算距离还可以。YOLO26 通过移除 DFL、实现无 NMS 端到端推理、引入 ProgLoss/STAL 训练策略等创新，在保持高精度的同时实现了极致的轻量化。凭借 2.4M 参数和 5.4B FLOPs 的优异效率，YOLO26 已成为工业检测、农业监测、智慧城市等边缘计算场景的首选基线模型之一。
2.标注数据与制作数据集
使用labelimg标注就行了，这里不罗嗦了，如果使用labelimg选择了yolo格式,直接整理弄成下面格式(待训练数据集)。
-
images/和labels/下的子目录名必须一一对应（train对train，val对val） -
每张图片（如
0001.jpg）在对应的labels/子目录中必须有一个同名的标注文件（0001.txt）
标注文件格式解释:
每个
.txt 文件包含该图片中所有目标的标注，每行一个目标，格式为：<class_id> <x_center> <y_center> <width> <height> 然后看上面待训练数据集结构图中，还少一个data.yaml。如下:
# 数据集根目录（可选，用于相对路径解析） path: /absolute/path/to/dataset # 训练/验证/测试图片路径（可以是目录，也可以是图片路径列表文件 .txt） train: images/train val: images/val test: images/test # 可选 # 类别数量 nc: 3 # 类别名称列表（索引 0 对应第一个名称） names: ['person', 'car', 'dog']
根据自己的实际去配置。names:也可以弄成字典形式，如下:
我这里使用lablelimg标注时选择了Pascal VOC 格式，得到了xml标注文件，需要转换为yolo格式，转换代码如下:
import os import xml.etree.ElementTree as ET from PIL import Image import shutil from tqdm import tqdm import random def convert_voc_to_yolo(xml_folder, image_folder, output_folder, class_list, train_ratio=0.8, val_ratio=0.1, test_ratio=0.1, random_seed=42): """ 将VOC格式的标注转换为YOLO格式（支持训练集、验证集、测试集划分） Args: xml_folder: XML文件所在的文件夹路径 image_folder: 图片文件所在的文件夹路径 output_folder: 输出文件夹路径 class_list: 类别名称列表，顺序与YOLO的class_id对应 train_ratio: 训练集比例 (默认0.8) val_ratio: 验证集比例 (默认0.1) test_ratio: 测试集比例 (默认0.1) random_seed: 随机种子，确保可重复性 """ # 验证比例和是否为1 total_ratio = train_ratio + val_ratio + test_ratio if abs(total_ratio - 1.0) > 0.001: raise ValueError(f"三个比例之和必须等于1，当前为: {total_ratio}") # 创建输出目录结构 images_train_dir = os.path.join(output_folder, 'images', 'train') images_val_dir = os.path.join(output_folder, 'images', 'val') images_test_dir = os.path.join(output_folder, 'images', 'test') labels_train_dir = os.path.join(output_folder, 'labels', 'train') labels_val_dir = os.path.join(output_folder, 'labels', 'val') labels_test_dir = os.path.join(output_folder, 'labels', 'test') for dir_path in [images_train_dir, images_val_dir, images_test_dir, labels_train_dir, labels_val_dir, labels_test_dir]: os.makedirs(dir_path, exist_ok=True) # 获取所有XML文件 xml_files = [f for f in os.listdir(xml_folder) if f.endswith('.xml')] if len(xml_files) == 0: print("错误: 未找到任何XML文件") return # 随机打乱文件列表 random.seed(random_seed) random.shuffle(xml_files) # 计算各集合的大小 total_count = len(xml_files) train_count = int(total_count * train_ratio) val_count = int(total_count * val_ratio) # 剩余的给测试集（处理舍入误差） test_count = total_count - train_count - val_count # 划分数据集 train_files = xml_files[:train_count] val_files = xml_files[train_count:train_count + val_count] test_files = xml_files[train_count + val_count:] # 创建类别到ID的映射 class_to_id = {name: idx for idx, name in enumerate(class_list)} print(f"数据集划分:") print(f" 总样本数: {total_count}") print(f" 训练集: {len(train_files)} ({len(train_files) / total_count * 100:.1f}%)") print(f" 验证集: {len(val_files)} ({len(val_files) / total_count * 100:.1f}%)") print(f" 测试集: {len(test_files)} ({len(test_files) / total_count * 100:.1f}%)") print(f" 类别列表: {class_list}") print("=" * 50) # 处理训练集 print("正在处理训练集...") for xml_file in tqdm(train_files): process_single_file( xml_file, xml_folder, image_folder, images_train_dir, labels_train_dir, class_to_id ) # 处理验证集 print("正在处理验证集...") for xml_file in tqdm(val_files): process_single_file( xml_file, xml_folder, image_folder, images_val_dir, labels_val_dir, class_to_id ) # 处理测试集 print("正在处理测试集...") for xml_file in tqdm(test_files): process_single_file( xml_file, xml_folder, image_folder, images_test_dir, labels_test_dir, class_to_id ) # 生成data.yaml配置文件 create_data_yaml(output_folder, class_list) # 生成数据集划分信息文件 create_split_info(output_folder, train_files, val_files, test_files) print("=" * 50) print("转换完成！") print(f"输出目录: {output_folder}") print(f"data.yaml配置文件已生成: {os.path.join(output_folder, 'data.yaml')}") print(f"数据集划分信息已保存: {os.path.join(output_folder, 'dataset_split_info.txt')}") def process_single_file(xml_file, xml_folder, image_folder, images_output_dir, labels_output_dir, class_to_id): """处理单个XML文件并生成对应的YOLO标注""" # 解析XML文件 xml_path = os.path.join(xml_folder, xml_file) try: tree = ET.parse(xml_path) root = tree.getroot() except Exception as e: print(f"警告: 解析XML文件失败 {xml_path}: {e}") return # 获取图片文件名（处理多种可能的情况） image_filename = None # 尝试从filename标签获取 filename_elem = root.find('filename') if filename_elem is not None: image_filename = filename_elem.text # 如果filename标签不存在，尝试从path标签获取 if not image_filename: path_elem = root.find('path') if path_elem is not None: image_filename = os.path.basename(path_elem.text) # 如果还是没有，从XML文件名推断 if not image_filename: # 尝试常见的图片扩展名 for ext in ['.jpg', '.jpeg', '.png', '.bmp', '.tiff']: potential_filename = xml_file.replace('.xml', ext) if os.path.exists(os.path.join(image_folder, potential_filename)): image_filename = potential_filename break if not image_filename: image_filename = xml_file.replace('.xml', '.jpg') # 默认使用jpg # 获取图片尺寸 size = root.find('size') if size is not None: try: img_width = int(size.find('width').text) img_height = int(size.find('height').text) except (ValueError, AttributeError): # 如果无法解析，从图片文件获取 img_width, img_height = None, None else: img_width, img_height = None, None # 如果XML中没有尺寸信息或解析失败，从图片文件中获取 if img_width is None or img_height is None: img_path = os.path.join(image_folder, image_filename) if os.path.exists(img_path): try: with Image.open(img_path) as img: img_width, img_height = img.size except Exception as e: print(f"警告: 无法读取图片 {img_path}: {e}") return else: print(f"警告: 找不到图片 {image_filename}，跳过 {xml_file}") return # 复制图片到输出目录 src_img_path = os.path.join(image_folder, image_filename) dst_img_path = os.path.join(images_output_dir, image_filename) if os.path.exists(src_img_path): try: shutil.copy2(src_img_path, dst_img_path) except Exception as e: print(f"警告: 复制图片失败 {src_img_path}: {e}") return else: print(f"警告: 图片文件不存在: {src_img_path}") return # 创建YOLO格式的标注文件 label_filename = os.path.splitext(image_filename)[0] + '.txt' label_path = os.path.join(labels_output_dir, label_filename) # 检查是否有目标需要标注 objects = root.findall('object') if len(objects) == 0: # 如果没有目标，创建一个空的标注文件 open(label_path, 'w').close() return with open(label_path, 'w') as f: # 遍历所有目标 for obj in objects: # 获取类别名称 class_name_elem = obj.find('name') if class_name_elem is None: print(f"警告: 在 {xml_file} 中未找到类别名称，跳过该目标") continue class_name = class_name_elem.text if class_name not in class_to_id: print(f"警告: 未找到类别 '{class_name}'，请检查CLASS_LIST，跳过该目标") continue class_id = class_to_id[class_name] # 获取边界框坐标 (VOC格式: xmin, ymin, xmax, ymax) bbox = obj.find('bndbox') if bbox is None: print(f"警告: 在 {xml_file} 中未找到边界框，跳过该目标") continue try: xmin = float(bbox.find('xmin').text) ymin = float(bbox.find('ymin').text) xmax = float(bbox.find('xmax').text) ymax = float(bbox.find('ymax').text) except (ValueError, AttributeError) as e: print(f"警告: 解析边界框坐标失败 {xml_file}: {e}") continue # 转换为YOLO格式 (中心点坐标和宽高，归一化) x_center = (xmin + xmax) / (2.0 * img_width) y_center = (ymin + ymax) / (2.0 * img_height) width = (xmax - xmin) / img_width height = (ymax - ymin) / img_height # 确保坐标在[0,1]范围内（处理标注误差） x_center = max(0, min(1, x_center)) y_center = max(0, min(1, y_center)) width = max(0, min(1, width)) height = max(0, min(1, height)) # 检查无效标注（宽度或高度为0） if width == 0 or height == 0: print(f"警告: 在 {xml_file} 中发现无效边界框，跳过") continue # 写入YOLO格式 f.write(f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}\n") def create_data_yaml(output_folder, class_names): """生成YOLO需要的data.yaml配置文件""" yaml_content = f"""# YOLO数据集配置文件 # 自动生成 - 划分比例为 8:1:1 path: {os.path.abspath(output_folder)} # 数据集根目录 train: images/train # 训练集图片路径 val: images/val # 验证集图片路径 test: images/test # 测试集图片路径 nc: {len(class_names)} # 类别数量 names: {class_names} # 类别名称列表 # 类别ID映射 (自动生成) # {dict(zip(range(len(class_names)), class_names))} """ yaml_path = os.path.join(output_folder, 'data.yaml') with open(yaml_path, 'w', encoding='utf-8') as f: f.write(yaml_content) def create_split_info(output_folder, train_files, val_files, test_files): """生成数据集划分信息文件""" info_path = os.path.join(output_folder, 'dataset_split_info.txt') with open(info_path, 'w', encoding='utf-8') as f: f.write("=" * 60 + "\n") f.write("数据集划分信息\n") f.write("=" * 60 + "\n\n") f.write(f"训练集文件列表 ({len(train_files)}个):\n") for file in train_files: f.write(f" - {file}\n") f.write(f"\n验证集文件列表 ({len(val_files)}个):\n") for file in val_files: f.write(f" - {file}\n") f.write(f"\n测试集文件列表 ({len(test_files)}个):\n") for file in test_files: f.write(f" - {file}\n") # 使用示例 if __name__ == "__main__": # 配置路径 XML_FOLDER = r"C:\data\xml" # 替换为你的XML文件夹路径 IMAGE_FOLDER = r"C:\data\images" # 替换为你的图片文件夹路径 OUTPUT_FOLDER = r"C:\data\data_set" # 替换为输出文件夹路径 # 定义类别列表（必须与标注时使用的类别名称完全一致） # 注意：列表顺序决定了YOLO的class_id，0对应第一个类别，1对应第二个，以此类推 CLASS_LIST = ["person"] # 根据你的实际类别修改 # 执行转换（使用8:1:1划分） convert_voc_to_yolo( xml_folder=XML_FOLDER, image_folder=IMAGE_FOLDER, output_folder=OUTPUT_FOLDER, class_list=CLASS_LIST, train_ratio=0.8, # 训练集80% val_ratio=0.1, # 验证集10% test_ratio=0.1, # 测试集10% random_seed=42 # 固定随机种子，确保结果可重复 )
这样将xml变为了txt文件，变为yolo格式了。
这里再提供一个通过xml得到类别的方法:
import xml.etree.ElementTree as ET from pathlib import Path def get_classes_from_xml(xml_dir): """ 从XML文件中提取所有类别名称 Args: xml_dir: XML文件目录路径 Returns: list: 排序后的类别名称列表 """ xml_dir = Path(xml_dir) classes = set() for xml_file in xml_dir.glob('*.xml'): try: tree = ET.parse(xml_file) root = tree.getroot() for obj in root.findall('object'): class_name = obj.find('name') if class_name is not None and class_name.text: classes.add(class_name.text) except Exception as e: print(f"警告: 处理 {xml_file.name} 时出错: {e}") return sorted(list(classes)) # 使用示例 if __name__ == "__main__": xml_directory = r"C\data\xml" # 修改为你的路径 class_list = get_classes_from_xml(xml_directory) print("所有类别:") print(class_list) print(f"\n类别数量: {len(class_list)}") # 打印每个类别及其索引（YOLO格式） print("\nYOLO类别映射:") for idx, class_name in enumerate(class_list): print(f" {idx}: {class_name}")
通过类别自动生成data.yaml
import yaml from pathlib import Path def generate_data_yaml(classes, output_path="data.yaml", dataset_path="./"): """ 生成YOLO训练所需的data.yaml配置文件 Args: classes: 类别名称列表，如 ['cat', 'dog', 'person'] output_path: 保存路径，默认 "data.yaml" dataset_path: 数据集根目录，默认 "./" Example: >>> generate_data_yaml(['cat', 'dog', 'person'], "data.yaml", "./dataset") """ # 构建配置 config = { 'path': dataset_path, 'train': 'train/images', 'val': 'val/images', 'test': 'test/images', 'nc': len(classes), 'names': classes } # 保存文件 with open(output_path, 'w', encoding='utf-8') as f: yaml.dump(config, f, default_flow_style=False, allow_unicode=True, sort_keys=False) print(f"✅ data.yaml 已生成: {output_path}") print(f" 类别数量: {len(classes)}") print(f" 类别列表: {classes}") # 使用示例 if __name__ == "__main__": # 方式1: 直接指定类别 # classes = ['bengbian', 'fanbai', 'qipao', 'quejiao', 'tuojiao', 'xiansun', 'yiwu', 'zangwu'] # 例如: ['cat', 'dog', 'person'] # generate_data_yaml(classes, "data.yaml", "./dataset") # 方式2: 从XML自动提取类别 import xml.etree.ElementTree as ET def get_classes_from_xml(xml_dir): classes = set() for xml_file in Path(xml_dir).glob('*.xml'): tree = ET.parse(xml_file) for obj in tree.findall('object'): name = obj.find('name') if name is not None and name.text: classes.add(name.text.strip()) return sorted(list(classes)) # 从XML提取并生成 classes = get_classes_from_xml(r"C:/data/xml") print(f"提取到类别: {classes}") generate_data_yaml(classes, "data.yaml", "./dataset")
3.预训练权重下载:
去下载预训练权重（不下载也没关系，训练时会自动检测，没有就自动下载）,如图，根据自己需要下载。
4.训练
最简单的训练脚本如下,根据需要可以在model.train配置其他参数，这里不赘述了。
from ultralytics import YOLO if __name__ == '__main__': # Load a model # model = YOLO("yolo26n.yaml") # build a new model from YAML model = YOLO("./pre_weights/yolo26n.pt") # load a pretrained model (recommended for training) # model = YOLO("yolo26n.yaml").load("yolo26n.pt") # build from YAML and transfer weights # Train the model results = model.train(data="data_set2/data.yaml", epochs=500, imgsz=640)
5.pt推理
训练完会得到.pt权重文件，一般使用last.pt或者best.pt。推理脚本如下:
import torch import cv2 import gc import time from ultralytics import YOLO def warmup_model(model, device, image_path=None, warmup_iterations=5): """ 预热模型 Args: model: YOLO模型 device: 设备 ('cuda' 或 'cpu') image_path: 图片路径，如果为None则创建虚拟图片 warmup_iterations: 预热次数 """ print("预热模型中...") # 准备预热图片 if image_path: warmup_image = cv2.imread(image_path) if warmup_image is None: print("⚠️ 无法读取图片，使用虚拟图片") warmup_image = torch.zeros((640, 640, 3), dtype=torch.uint8).numpy() else: # 创建虚拟图片 warmup_image = torch.zeros((640, 640, 3), dtype=torch.uint8).numpy() # 多次预热 for i in range(warmup_iterations): print(f" 预热 {i+1}/{warmup_iterations}...", end='') _ = model(warmup_image, verbose=False) if device == 'cuda': torch.cuda.synchronize() # 等待GPU完成 print(" ✅") # 清理显存缓存 if device == 'cuda': torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats() gc.collect() # 垃圾回收 print("✅ 预热完成") # 主程序 if __name__ == "__main__": # 1. 检查GPU if torch.cuda.is_available(): print(f"✅ 使用GPU: {torch.cuda.get_device_name(0)}") device = 'cuda' print(f" 显存总量: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.2f} GB") else: print("⚠️ GPU不可用，使用CPU") device = 'cpu' # 2. 加载模型 model = YOLO(r"C:\workspace\yolo_demo\runs\detect\train-5\weights\best.pt") model.to(device) # 3. 预热模型 image_path = "1.jpg" warmup_model(model, device, image_path, warmup_iterations=5) # 4. 正式推理 print("\n开始正式推理...") image = cv2.imread(image_path) t1 = time.time() results = model(image) if device == 'cuda': torch.cuda.synchronize() t2 = time.time() print(f"推理耗时: {(t2-t1)*1000:.2f}ms") # 5. 显示显存使用 if device == 'cuda': memory_allocated = torch.cuda.memory_allocated() / 1024**2 memory_cached = torch.cuda.memory_reserved() / 1024**2 print(f"显存使用: {memory_allocated:.2f} MB (分配) / {memory_cached:.2f} MB (缓存)") # 6. 打印检测结果 for r in results: boxes = r.boxes if boxes is not None: print(f"\n检测到 {len(boxes)} 个目标:") for i, box in enumerate(boxes, 1): cls_id = int(box.cls[0]) cls_name = model.names[cls_id] conf = float(box.conf[0]) xyxy = box.xyxy[0].cpu().numpy() print(f"目标 {i}: {cls_name}, 置信度: {conf:.3f}") print(f" 位置: ({xyxy[0]:.1f}, {xyxy[1]:.1f}) -> ({xyxy[2]:.1f}, {xyxy[3]:.1f})")
测试没得问题，这里就不贴图了。
5.onnx推理
5.1 pt格式转onnx格式
由于是pt格式，需要转换为.onnx格式。转换代码如下:
from ultralytics import YOLO # Load a model # model = YOLO("yolo26n.pt") # load an official model model = YOLO(r"C:\workspace\yolo_demo\runs\detect\train-5\weights\best.pt") # load a custom-trained model # Export the model model.export(format="onnx")
就得到了best.onnx
5.2 cpu推理
import time import cv2 from ultralytics import YOLO # 加载模型 model = YOLO(r"C:\workspace\yolo_demo\runs\detect\train-5\weights\best.onnx") # 读取并推理 image = cv2.imread("./data_set/yolo_dataset/test/images/21_1-rectangle.jpg") # 预热（忽略第一次加载） _ = model(image) t1=time.time() results = model(image, save=True, project="runs/detect", name="inference") t2=time.time() print("耗时:",t2-t1) #耗时: 0.045999765396118164 # 保存结果 cv2.imwrite("my_result.jpg", results[0].plot()) print("✅ 结果已保存到: my_result.jpg") # 打印检测框信息 boxes = results[0].boxes if boxes: print(f"\n检测到 {len(boxes)} 个目标:") print("-" * 60) for i, box in enumerate(boxes, 1): cls_id = int(box.cls[0]) conf = float(box.conf[0]) xyxy = box.xyxy[0].cpu().numpy() xywh = box.xywh[0].cpu().numpy() print(f"目标 {i}:") print(f" 类别: {model.names[cls_id]} (ID: {cls_id})") print(f" 置信度: {conf:.4f}") print(f" 边界框 (xyxy): x1={xyxy[0]:.1f}, y1={xyxy[1]:.1f}, x2={xyxy[2]:.1f}, y2={xyxy[3]:.1f}") print(f" 边界框 (xywh): cx={xywh[0]:.1f}, cy={xywh[1]:.1f}, w={xywh[2]:.1f}, h={xywh[3]:.1f}") print(f" 宽度: {xywh[2]:.1f}px, 高度: {xywh[3]:.1f}px") print("-" * 60) else: print("未检测到任何目标")
5.3 gpu推理
import time from typing import Optional import cv2 import numpy as np import onnxruntime as ort class YOLO26Detector: """YOLO26 ONNX GPU 推理类（端到端无NMS）""" def __init__( self, model_path: str, providers: Optional[list[str]] = None, # providers: list[str] | None = None, conf_threshold: float = 0.25, input_size: tuple[int, int] = (640, 640), ): """ Args: model_path: ONNX 模型文件路径 providers: ONNX Runtime 执行提供者列表，默认自动选择 GPU conf_threshold: 置信度阈值 input_size: 模型输入尺寸 (height, width) """ if providers is None: providers = [ "CUDAExecutionProvider", "CPUExecutionProvider", ] self.conf_threshold = conf_threshold self.input_size = input_size # (h, w) # 创建推理会话 sess_options = ort.SessionOptions() sess_options.graph_optimization_level = ( ort.GraphOptimizationLevel.ORT_ENABLE_ALL ) self.session = ort.InferenceSession( model_path, sess_options=sess_options, providers=providers ) # 获取实际使用的 provider self.provider = self.session.get_providers()[0] print(f"[YOLO26] 模型加载完成 | Provider: {self.provider}") # 获取输入信息 self.input_name = self.session.get_inputs()[0].name self.input_shape = self.session.get_inputs()[0].shape print(f"[YOLO26] 输入名称: {self.input_name}, 形状: {self.input_shape}") # 预热模型 self.warmup() def warmup(self, iterations: int = 10): """预热模型，确保 GPU 上下文初始化完成""" dummy_input = np.zeros( (1, 3, self.input_size[0], self.input_size[1]), dtype=np.float32 ) print(f"[YOLO26] 开始预热 ({iterations} 次)...") start = time.perf_counter() for _ in range(iterations): self.session.run(None, {self.input_name: dummy_input}) elapsed = (time.perf_counter() - start) * 1000 print(f"[YOLO26] 预热完成 | 平均耗时: {elapsed / iterations:.2f} ms/帧") def preprocess(self, image: np.ndarray) -> tuple[np.ndarray, float, int, int]: """ 图像预处理：Letterbox 缩放 + 归一化 + CHW 转换 Returns: input_tensor: 模型输入张量 (1, 3, H, W) scale: 缩放比例 pad_w: 水平方向 padding 像素 pad_h: 垂直方向 padding 像素 """ h, w = image.shape[:2] target_h, target_w = self.input_size # 计算缩放比例（保持宽高比） scale = min(target_h / h, target_w / w) new_h, new_w = int(h * scale), int(w * scale) # 缩放 resized = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_LINEAR) # Letterbox 填充（灰色 114） pad_h = target_h - new_h pad_w = target_w - new_w top, left = pad_h // 2, pad_w // 2 padded = np.full((target_h, target_w, 3), 114, dtype=np.uint8) padded[top : top + new_h, left : left + new_w] = resized # BGR -> RGB, HWC -> CHW, 归一化到 [0, 1] blob = padded[:, :, ::-1].transpose(2, 0, 1).astype(np.float32) / 255.0 input_tensor = np.expand_dims(blob, axis=0) return input_tensor, scale, pad_w, pad_h def postprocess( self, output: np.ndarray, scale: float, pad_w: int, pad_h: int ) -> list[dict]: """ 后处理：解析端到端输出 (N, 300, 6) 格式: [x1, y1, x2, y2, confidence, class_id] Returns: detections: 检测结果列表 """ detections = [] # output shape: (1, 300, 6) 或 (300, 6) if output.ndim == 3: output = output[0] for det in output: x1, y1, x2, y2, conf, cls_id = det # 置信度过滤 if conf < self.conf_threshold: continue # 将坐标映射回原图（去除 letterbox padding 和缩放） x1 = (x1 - pad_w / 2) / scale y1 = (y1 - pad_h / 2) / scale x2 = (x2 - pad_w / 2) / scale y2 = (y2 - pad_h / 2) / scale detections.append( { "bbox": [float(x1), float(y1), float(x2), float(y2)], "confidence": float(conf), "class_id": int(cls_id), } ) return detections def detect(self, image: np.ndarray) -> list[dict]: """ 对单张图像执行目标检测 Args: image: BGR 格式的 OpenCV 图像 (H, W, 3) Returns: detections: 检测结果列表，每项包含 bbox, confidence, class_id """ # 预处理 input_tensor, scale, pad_w, pad_h = self.preprocess(image) # 推理 outputs = self.session.run(None, {self.input_name: input_tensor}) output = outputs[0] # 后处理 detections = self.postprocess(output, scale, pad_w, pad_h) return detections def detect_batch(self, images: list[np.ndarray]) -> list[list[dict]]: """ 批量检测（要求所有图像尺寸一致） Args: images: BGR 图像列表 Returns: 每张图像的检测结果列表 """ batch_tensors = [] preprocess_info = [] for img in images: tensor, scale, pad_w, pad_h = self.preprocess(img) batch_tensors.append(tensor) preprocess_info.append((scale, pad_w, pad_h)) batch_input = np.concatenate(batch_tensors, axis=0) outputs = self.session.run(None, {self.input_name: batch_input}) all_detections = [] for i, (scale, pad_w, pad_h) in enumerate(preprocess_info): dets = self.postprocess(outputs[0][i : i + 1], scale, pad_w, pad_h) all_detections.append(dets) return all_detections # ===================== 使用示例 ===================== if __name__ == "__main__": # 初始化检测器（自动使用 CUDA GPU） detector = YOLO26Detector( model_path=r"C:\workspace\yolo_demo\runs\detect\train-5\weights\best.onnx", conf_threshold=0.25, input_size=(640, 640), ) # 读取图像 image = cv2.imread(r"C:\workspace\yolo_demo\data_set\yolo_dataset\test\images\21_1-rectangle.jpg") # 单张检测 detections = detector.detect(image) for det in detections: x1, y1, x2, y2 = [int(v) for v in det["bbox"]] conf = det["confidence"] cls_id = det["class_id"] print(f"类别: {cls_id}, 置信度: {conf:.4f}, 坐标: [{x1}, {y1}, {x2}, {y2}]") # 绘制检测框 cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2) cv2.putText( image, f"{cls_id}: {conf:.2f}", (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2, ) cv2.imwrite("result_26.jpg", image) print(f"共检测到 {len(detections)} 个目标")
贴个效果图吧，如下:
小结: 本文介绍了yolo26的从标注到训练到推理的全部流程，别给出实际脚本，关于NMS这一块yolo26是去除了的，如果自己的项目目标检测多个有重叠，且不需要，可以自己写一个NSM，这里不再详述。
本文若存在不做或错误之处，欢迎指出与评论。
浙公网安备 33010602011771号