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,这里不再详述。
本文若存在不做或错误之处,欢迎指出与评论。