3.2.5

import os

import time

import cv2

import numpy as np

import vision.utils.box_utils_numpy as box_utils

import onnxruntime as ort

 

# 定义预测函数，对模型输出的边界框和置信度进行后处理

def predict(width, height, confidences, boxes, prob_threshold, iou_threshold=0.3, top_k=-1):

    boxes = boxes[0]

    confidences = confidences[0]

    picked_box_probs = []

    picked_labels = []

    for class_index in range(1, confidences.shape[1]):

        probs = confidences[:, class_index]

        mask = probs > prob_threshold

        probs = probs[mask]

        if probs.shape[0] == 0:

            continue

        subset_boxes = boxes[mask, :]

        box_probs = np.concatenate([subset_boxes, probs.reshape(-1, 1)], axis=1)

        box_probs = box_utils.hard_nms(box_probs,

                                       iou_threshold=iou_threshold,

                                       top_k=top_k,

                                       )

        picked_box_probs.append(box_probs)

        picked_labels.extend([class_index] * box_probs.shape[0])

    if not picked_box_probs:

        return np.array([]), np.array([]), np.array([])

    picked_box_probs = np.concatenate(picked_box_probs)

    picked_box_probs[:, 0] *= width

    picked_box_probs[:, 1] *= height

    picked_box_probs[:, 2] *= width

    picked_box_probs[:, 3] *= height

    return picked_box_probs[:, :4].astype(np.int32), np.array(picked_labels), picked_box_probs[:, 4]

 

# 从标签文件中读取每一行，并去除行首尾的空白字符，得到类别名称列表 2分

class_names = [name.strip() for name in open('voc-model-labels.txt').readlines()]

 

# 创建 ONNX Runtime 的推理会话，用于运行模型进行推理 2分

ort_session = ort.InferenceSession('version-RFB-320.onnx')

 

# 获取模型输入的名称 2分

input_name = ort_session.get_inputs()[0].name

 

# 定义保存检测结果图像的目录路径

result_path = "./detect_imgs_results_onnx"

 

# 定义置信度阈值，用于筛选出置信度较高的检测结果

threshold = 0.7

# 定义存储待检测图像的目录路径

path = "imgs"

# 用于统计所有图像中检测到的目标框总数，初始化为 0

sum = 0

 

# 如果保存结果的目录不存在，则创建该目录 2分

if not os.path.exists(result_path):

    os.mkdir(result_path)

   

# 获取指定目录下的所有文件和文件夹名称列表

listdir = os.listdir(path)

 

# 遍历目录下的每个文件

for file_path in listdir:

    # 拼接图像文件的完整路径

    img_path = os.path.join(path, file_path)

    # 使用 OpenCV 读取图像文件 2分

    orig_image = cv2.imread(img_path)

    # 将图像从 BGR 颜色空间转换为 RGB 颜色空间（许多模型要求输入为 RGB 格式）

    image = cv2.cvtColor(orig_image, cv2.COLOR_BGR2RGB)

    # 将图像调整为 320x240 的尺寸（符合模型输入的尺寸要求） 2分

    image = cv2.resize(image, (320, 240))

    # 定义图像归一化的均值数组 2分

    image_mean = np.array([127, 127, 127])

    # 对图像进行归一化处理，减去均值并除以 128

    image = (image - image_mean) / 128

    # 将图像的维度从 (高度, 宽度, 通道数) 转换为 (通道数, 高度, 宽度)

    image = np.transpose(image, [2, 0, 1])

    # 在第一个维度上扩展一个维度，将图像变为 (1, 通道数, 高度, 宽度)，以符合模型输入的维度要求  1分

    image = np.expand_dims(image, axis=0)

    # 将图像数据类型转换为 float32 类型

    image = image.astype(np.float32)

    # 记录开始时间，用于计算模型推理的耗时

    time_time = time.time()

    # 使用 ONNX Runtime 运行模型，输入图像数据，得到模型输出的置信度和边界框  2分

    confidences, boxes = ort_session.run(None, {input_name: image})

    # 计算并打印模型推理的耗时

    print("cost time:{}".format(time.time() - time_time))

    # 调用 predict 函数对模型输出的边界框和置信度进行后处理，得到最终的边界框、类别标签和置信度

    boxes, labels, probs = predict(orig_image.shape[1], orig_image.shape[0], confidences, boxes, threshold)

    # 遍历每个检测到的目标框

    for i in range(boxes.shape[0]):

        # 获取当前目标框的坐标

        box = boxes[i, :]

        # 生成当前目标框的标签字符串，包含类别名称和置信度

        label = f"{class_names[labels[i]]}: {probs[i]:.2f}"

 

        # 在原始图像上绘制目标框，颜色为 (255, 255, 0)，线条粗细为 4

        cv2.rectangle(orig_image, (box[0], box[1]), (box[2], box[3]), (255, 255, 0), 4)

        # 将绘制了目标框的图像保存到结果目录中

        cv2.imwrite(os.path.join(result_path, file_path), orig_image)

    # 累加当前图像中检测到的目标框数量到总数中

    sum += boxes.shape[0]

# 打印所有图像中检测到的目标框总数

print("sum:{}".format(sum))