人工智能训练师（三级）实操考试 — 3.2.5 人脸AI智能检测系统交互流程设计

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 = [ for name in open('voc-model-labels.txt').readlines()] # 创建 ONNX Runtime 的推理会话，用于运行模型进行推理 2分 ort_session = ('version-RFB-320.onnx') # 获取模型输入的名称 2分 input_name = ()[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. # 获取指定目录下的所有文件和文件夹名称列表 listdir = os.listdir(path) # 遍历目录下的每个文件 for file_path in listdir: # 拼接图像文件的完整路径 img_path = os.path.join(path, file_path) # 使用 OpenCV 读取图像文件 2分 orig_image = # 将图像从 BGR 颜色空间转换为 RGB 颜色空间 image = cv2.cvtColor(orig_image, cv2.COLOR_BGR2RGB) # 将图像调整为 320x240 的尺寸（符合模型输入的尺寸要求） 2分 image = (, (320, 240)) # 定义图像归一化的均值数组 2分 image_mean = ([127, 127, 127]) # 对图像进行归一化处理，减去均值并除以 128 image = (image - image_mean) / 128 # 将图像的维度从 (高度, 宽度, 通道数) 转换为 (通道数, 高度, 宽度) image = np.transpose(image, [2, 0, 1]) # 在第一个维度上扩展一个维度 1分 image = (image, axis=0) # 将图像数据类型转换为 float32 类型 image = image.astype(np.float32) # 记录开始时间 time_time = time.time() # 使用 ONNX Runtime 运行模型，输入图像数据，得到输出 2分 confidences, boxes = (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}" # 在原始图像上绘制目标框 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))