utils.py

# utils.py
# 工具模块：包含图像预处理、后处理和结果绘制函数

import cv2  # 图像处理
import numpy as np  # 数值计算
import random  # 随机数生成（用于颜色）
import torch  # PyTorch库（用于后处理）
import torchvision  # PyTorch视觉工具库（用于NMS）

# COCO数据集类别名称列表（80类）
coco_class = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
         'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
         'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
         'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
         'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
         'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
         'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
         'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
         'hair drier', 'toothbrush']

# 为每个类别随机生成一种颜色（用于绘制边界框）
colors = {name: [random.randint(0, 255) for _ in range(3)] for i, name in enumerate(coco_class)}


def img_process(img_path, size):
    """
    图像预处理：读取图像、等比缩放、归一化

    参数:
        img_path: 图像路径
        size: 模型输入尺寸（正方形）
    
    返回:
        frame: 原始图像（用于后续绘制）
        img_input: 处理后符合模型输入要求的图像
        ratio: 缩放比例（用于将预测框映射回原始图像）
    """
    # 读取原始图像（BGR格式）
    frame = cv2.imread(img_path)
    # 复制原始图像用于处理
    img_processed = np.copy(frame)
    # 获取图像高度和宽度
    [height, width, _] = img_processed.shape
    # 计算最长边（用于等比缩放）
    length = max((height, width))
    # 计算缩放比例（原始尺寸/目标尺寸）
    scale = length / size
    ratio = scale
    # 创建正方形画布（填充黑边保持比例）
    image = np.zeros((length, length, 3), np.uint8)
    # 将原始图像放置在画布左上角
    image[0:height, 0:width] = img_processed
    # 转换颜色空间（BGR->RGB，符合模型输入要求）
    img_input = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    # 缩放到模型输入尺寸
    img_input = cv2.resize(img_input, (size, size))
    # 保存缩放后的图像（调试用）
    cv2.imwrite("resize_football.jpg", img_input)
    # 归一化参数（均值和标准差）
    mean_data = [0, 0, 0]
    std_data = [255, 255, 255]
    # 执行归一化：(图像 - 均值) / 标准差
    img_input = (img_input - mean_data) / std_data  # 形状为HWC
    return frame, img_input, ratio


def out_process(data_local, data_conf, score_threshold, ratio):
    """
    模型输出后处理：过滤低置信度框、执行NMS、映射回原始图像尺寸

    参数:
        data_local: 边界框坐标数据 [1, 8400, 4]
        data_conf: 类别置信度数据 [1, 8400, 80]
        score_threshold: 置信度阈值
        ratio: 缩放比例（用于坐标映射）
    
    返回:
        detect: 处理后的检测结果（边界框、置信度、类别）
    """
    max_predictions = 300  # 最大保留检测框数量
    nms_result = []
    
    # 转换为PyTorch张量便于处理
    data_local = torch.tensor(data_local)
    data_conf = torch.tensor(data_conf)
    
    # 遍历每个样本（此处批次大小为1）
    for pred_bboxes, pred_scores in zip(data_local, data_conf):
        # 找到置信度大于阈值的索引
        i, j = (pred_scores > score_threshold).nonzero(as_tuple=False).T
        # 过滤出高置信度的边界框
        pred_bboxes = pred_bboxes[i]
        # 过滤出对应的置信度分数
        pred_cls_conf = pred_scores[i, j]
        # 过滤出对应的类别标签
        pred_cls_label = j[:]
    

    # 执行非极大值抑制（NMS）去除重叠框
    # 参数：边界框、分数、类别标签、IOU阈值（此处复用置信度阈值）
    idx_to_keep = torchvision.ops.boxes.batched_nms(
        boxes=pred_bboxes, 
        scores=pred_cls_conf, 
        idxs=pred_cls_label, 
        iou_threshold=score_threshold
    )

    # 根据NMS结果过滤
    pred_cls_conf = pred_cls_conf[idx_to_keep].unsqueeze(-1)  # 增加维度便于拼接
    pred_cls_label = pred_cls_label[idx_to_keep].unsqueeze(-1)
    pred_bboxes = pred_bboxes[idx_to_keep, :]

    # 拼接边界框、置信度和类别标签
    final_boxes = torch.cat([pred_bboxes, pred_cls_conf, pred_cls_label], dim=1)
    nms_result.append(final_boxes)
    
    # 限制最大检测框数量
    nms_result[:] = [
        im[:max_predictions] if (im is not None and im.shape[0] > max_predictions) else im 
        for im in nms_result
    ]
    
    # 转换为numpy数组并将边界框映射回原始图像尺寸
    detect = nms_result[0].numpy()
    detect[:, :4] = detect[:, :4] * ratio  # 前4列是坐标（x,y,x2,y2）

    return detect


def draw_res(img, boxes):
    """
    在图像上绘制检测结果

    参数:
        img: 原始图像
        boxes: 检测框列表，每个元素为[x, y, x2, y2, scores, class_ids]
    
    返回:
        img: 绘制后的图像
    """
    img = img.astype(np.uint8)  # 确保图像格式正确
    print(f"检测到 {len(boxes)} 个目标:")
    
    # 遍历每个检测框
    for i, [x, y, x2, y2, scores, class_ids] in enumerate(boxes):
        # 转换为整数坐标
        x = int(x)
        y = int(y)
        x2 = int(x2)
        y2 = int(y2)
        # 获取类别名称
        name = coco_class[int(class_ids)]
        # 打印检测信息（序号、坐标、置信度、类别）
        print(i + 1, [x, y, x2, y2], round(scores, 4), name)
        # 生成标签文本（类别 + 置信度）
        label = f'{name} ({scores:.2f})'
        # 计算文本尺寸（用于绘制文本背景框）
        W, H = cv2.getTextSize(label, 0, fontScale=1, thickness=2)[0]
        # 获取该类别的颜色
        color = colors[name]
        # 绘制边界框
        cv2.rectangle(img, (x, y), (int(x2), int(y2)), color, thickness=2)
        # 绘制文本背景框（半透明覆盖）
        cv2.rectangle(img, (x, int(y - H)), (int(x + W / 2), y), (0, 255,), -1, cv2.LINE_AA)
        # 绘制文本标签
        cv2.putText(img, label, (x, int(y) - 6), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1)
    
    return img

run-test.py

# run-test.py
# 导入必要的库
import sys
import time
import aidlite  # 用于模型推理的aidlite库
import cv2  # 用于图像处理
import numpy as np  # 用于数值计算
import argparse  # 用于解析命令行参数
import onnxruntime  # ONNX运行时（此处未直接使用，可能预留用于其他推理方式）

# 从工具模块导入后处理、绘图和图像预处理函数
from utils import out_process, draw_res, img_process


def run(input_shapes, output_shapes, is_quant=True):
    """
    运行模型推理的主函数

    参数:
        input_shapes: 模型输入形状列表
        output_shapes: 模型输出形状列表
        is_quant: 是否为量化模型，默认True
    """
    # 解析命令行参数
    args = parser_args()
    model_path = args.target_model  # 模型路径
    img_path = args.imgs  # 输入图像路径
    model_type = args.model_type  # 模型类型（QNN/SNPE等）
    invoke_nums = int(args.invoke_nums)  # 推理次数

    # 创建aidlite模型实例
    model = aidlite.Model.create_instance(model_path)
    # 设置模型输入输出属性：输入形状、数据类型（float32）、输出形状、数据类型（float32）
    model.set_model_properties(input_shapes, aidlite.DataType.TYPE_FLOAT32,
                               output_shapes, aidlite.DataType.TYPE_FLOAT32)
                               
    # 创建配置实例
    config = aidlite.Config.create_instance()
    # 设置实现类型为本地
    config.implement_type = aidlite.ImplementType.TYPE_LOCAL
    
    # 根据模型类型设置框架类型和加速类型
    if model_type.lower() == "qnn":
        config.framework_type = aidlite.FrameworkType.TYPE_QNN
        config.accelerate_type = aidlite.AccelerateType.TYPE_DSP  # 使用DSP加速
    elif model_type.lower() == "snpe2" or model_type.lower() == "snpe":
        config.framework_type = aidlite.FrameworkType.TYPE_SNPE2
        config.accelerate_type = aidlite.AccelerateType.TYPE_DSP  # 使用DSP加速
    
    # 设置线程数为4
    config.number_of_threads = 4
    # 如果是量化模型，设置量化模型标志
    if is_quant:
        config.is_quantify_model = 1
                               
    # 根据模型和配置构建解释器
    interpreter = aidlite.InterpreterBuilder.build_interpretper_from_model_and_config(model, config)
    if interpreter is None:
        print("构建解释器失败！")
        return False
        
    # 初始化解释器
    result = interpreter.init()
    if result != 0:
        print("解释器初始化失败！")
        return False
        
    # 加载模型
    result = interpreter.load_model()
    if result != 0:
        print("模型加载失败！")
        return False
    print("检测模型加载成功!")
        
    # 图像预处理：输入尺寸640x640
    size = 640
    # 处理图像：读取原图、生成模型输入图像、计算缩放比例
    frame, img_input, ratio = img_process(img_path, size)
    # 增加批次维度并转换为float32类型（模型输入要求）
    img_input = np.expand_dims(img_input, 0).astype(np.float32)
   

    # 记录每次推理时间
    invoke_time = []
    for i in range(invoke_nums):
        # 设置输入张量
        result = interpreter.set_input_tensor(0, img_input.data)
        if result != 0:
            print("设置输入张量失败")
        
        # 记录推理开始时间
        t1 = time.time()
        # 执行推理
        result = interpreter.invoke()
        # 计算推理耗时（毫秒）并记录
        cost_time = (time.time() - t1) * 1000
        invoke_time.append(cost_time)
    
    # 检查推理是否成功
    if result != 0:
        print("推理执行失败")
    
    # 获取输出张量并调整形状
    # 输出张量0: 置信度相关 [1,8400,80]
    qnn_conf = interpreter.get_output_tensor(0).reshape(*output_shapes[1])
    # 输出张量1: 边界框相关 [1,8400,4]
    qnn_local = interpreter.get_output_tensor(1).reshape(*output_shapes[0])
    
    # 检查输出是否有效
    if qnn_local is None:
        print("获取输出张量0失败!")
        
    # 销毁解释器释放资源
    result = interpreter.destory()

    # 计算推理时间统计值
    max_invoke_time = max(invoke_time)  # 最大耗时
    min_invoke_time = min(invoke_time)  # 最小耗时
    mean_invoke_time = sum(invoke_time) / invoke_nums  # 平均耗时
    var_invoketime = np.var(invoke_time)  # 方差（衡量耗时稳定性）
    
    # 打印性能统计结果
    print("====================================")
    print(f"QNN推理时间统计:\n --平均耗时: {mean_invoke_time} ms \n --最大耗时: {max_invoke_time} ms \n --最小耗时: {min_invoke_time} ms \n --耗时方差: {var_invoketime}")
    print("====================================")
    
    # 后处理：合并边界框和置信度数据
    results = np.concatenate((qnn_local, qnn_conf), axis=2)
    results = results[0]  # 去除批次维度

    # 应用置信度阈值过滤和NMS（非极大值抑制）
    score_threshold = 0.25  # 置信度阈值
    detect = out_process(qnn_local, qnn_conf, score_threshold, ratio)
    
    # 绘制检测结果
    res_img = draw_res(frame, list(detect))
    # 保存结果图像
    cv2.imwrite("python/results_img.jpg", res_img)   
    print("=======================================")


def parser_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description="运行模型基准测试")
    # 模型路径参数，默认使用yolo_nas_s量化模型
    parser.add_argument('--target_model', type=str, default='yolo_nas_s/yolo_nas_s_qcs8550_w8a8.qnn231.ctx.bin',
                        help="推理模型路径")
    # 输入图像路径，默认使用bus.jpg
    parser.add_argument('--imgs', type=str, default='bus.jpg', help="预测图像路径")
    # 推理次数，默认100次（用于统计性能）
    parser.add_argument('--invoke_nums', type=int, default=100, help="推理次数")
    # 模型类型，默认QNN（量化神经网络）
    parser.add_argument('--model_type', type=str, default='QNN', help="运行后端框架")
    args = parser.parse_args()
    return args


if __name__ == "__main__":
    # 定义模型输入输出形状
    input_shapes = [[1, 640, 640, 3]]  # [批次大小, 高度, 宽度, 通道数]
    output_shapes = [[1, 8400, 4], [1, 8400, 80]]  # 输出1:边界框坐标，输出2:类别置信度
    # 运行主函数
    run(input_shapes, output_shapes)