import time import cv2 import onnxruntime as ort from PIL import Image import numpy as np # 置信度 confidence_thres = 0.35 # iou阈值 iou_thres = 0.5 # 类别 classes = {0: 'herb'} # 随机颜色 color_palette = np.random.uniform(100, 255, size=(len(classes), 3)) # 判断是使用GPU或CPU providers = [ ('CUDAExecutionProvider', { 'device_id': 0, # 可以选择GPU设备ID，如果你有多个GPU }), 'CPUExecutionProvider', # 也可以设置CPU作为备选 ] def calculate_iou(box, other_boxes): """ 计算给定边界框与一组其他边界框之间的交并比（IoU）。 参数： - box: 单个边界框，格式为 [x1, y1, width, height]。 - other_boxes: 其他边界框的数组，每个边界框的格式也为 [x1, y1, width, height]。 返回值： - iou: 一个数组，包含给定边界框与每个其他边界框的IoU值。 """ # 计算交集的左上角坐标 x1 = np.maximum(box[0], np.array(other_boxes)[:, 0]) y1 = np.maximum(box[1], np.array(other_boxes)[:, 1]) # 计算交集的右下角坐标 x2 = np.minimum(box[0] + box[2], np.array(other_boxes)[:, 0] + np.array(other_boxes)[:, 2]) y2 = np.minimum(box[1] + box[3], np.array(other_boxes)[:, 1] + np.array(other_boxes)[:, 3]) # 计算交集区域的面积 intersection_area = np.maximum(0, x2 - x1) * np.maximum(0, y2 - y1) # 计算给定边界框的面积 box_area = box[2] * box[3] # 计算其他边界框的面积 other_boxes_area = np.array(other_boxes)[:, 2] * np.array(other_boxes)[:, 3] # 计算IoU值 iou = intersection_area / (box_area + other_boxes_area - intersection_area) return iou def custom_NMSBoxes(boxes, scores, confidence_threshold, iou_threshold): # 如果没有边界框，则直接返回空列表 if len(boxes) == 0: return [] # 将得分和边界框转换为NumPy数组 scores = np.array(scores) boxes = np.array(boxes) # 根据置信度阈值过滤边界框 mask = scores > confidence_threshold filtered_boxes = boxes[mask] filtered_scores = scores[mask] # 如果过滤后没有边界框，则返回空列表 if len(filtered_boxes) == 0: return [] # 根据置信度得分对边界框进行排序 sorted_indices = np.argsort(filtered_scores)[::-1] # 初始化一个空列表来存储选择的边界框索引 indices = [] # 当还有未处理的边界框时，循环继续 while len(sorted_indices) > 0: # 选择得分最高的边界框索引 current_index = sorted_indices[0] indices.append(current_index) # 如果只剩一个边界框，则结束循环 if len(sorted_indices) == 1: break # 获取当前边界框和其他边界框 current_box = filtered_boxes[current_index] other_boxes = filtered_boxes[sorted_indices[1:]] # 计算当前边界框与其他边界框的IoU iou = calculate_iou(current_box, other_boxes) # 找到IoU低于阈值的边界框，即与当前边界框不重叠的边界框 non_overlapping_indices = np.where(iou <= iou_threshold)[0] # 更新sorted_indices以仅包含不重叠的边界框 sorted_indices = sorted_indices[non_overlapping_indices + 1] # 返回选择的边界框索引 return indices def draw_detections(img, box, score, class_id): """ 在输入图像上绘制检测到的对象的边界框和标签。 参数: img: 要在其上绘制检测结果的输入图像。 box: 检测到的边界框。 score: 对应的检测得分。 class_id: 检测到的对象的类别ID。 返回: 无 """ # 提取边界框的坐标 x1, y1, w, h = box # 根据类别ID检索颜色 color = color_palette[class_id] # 在图像上绘制边界框 cv2.rectangle(img, (int(x1), int(y1)), (int(x1 + w), int(y1 + h)), color, 2) # 创建标签文本，包括类名和得分 label = f'{classes[class_id]}: {score:.2f}' # 计算标签文本的尺寸 (label_width, label_height), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1) # 计算标签文本的位置 label_x = x1 label_y = y1 - 10 if y1 - 10 > label_height else y1 + 10 # 绘制填充的矩形作为标签文本的背景 cv2.rectangle(img, (label_x, label_y - label_height), (label_x + label_width, label_y + label_height), color, cv2.FILLED) # 在图像上绘制标签文本 cv2.putText(img, label, (label_x, label_y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA) def preprocess(img, input_width, input_height): """ 在执行推理之前预处理输入图像。 返回: image_data: 为推理准备好的预处理后的图像数据。 """ # 获取输入图像的高度和宽度 img_height, img_width = img.shape[:2] # 将图像颜色空间从BGR转换为RGB img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # 将图像大小调整为匹配输入形状 img = cv2.resize(img, (input_width, input_height)) # 通过除以255.0来归一化图像数据 image_data = np.array(img) / 255.0 # 转置图像，使通道维度为第一维 image_data = np.transpose(image_data, (2, 0, 1)) # 通道首 # 扩展图像数据的维度以匹配预期的输入形状 image_data = np.expand_dims(image_data, axis=0).astype(np.float32) # 返回预处理后的图像数据 return image_data, img_height, img_width def postprocess(input_image, output, input_width, input_height, img_width, img_height): """ 对模型输出进行后处理，提取边界框、得分和类别ID。 参数: input_image (numpy.ndarray): 输入图像。 output (numpy.ndarray): 模型的输出。 input_width (int): 模型输入宽度。 input_height (int): 模型输入高度。 img_width (int): 原始图像宽度。 img_height (int): 原始图像高度。 返回: numpy.ndarray: 绘制了检测结果的输入图像。 """ # 转置和压缩输出以匹配预期的形状 outputs = np.transpose(np.squeeze(output[0])) # 获取输出数组的行数 rows = outputs.shape[0] # 用于存储检测的边界框、得分和类别ID的列表 boxes = [] scores = [] class_ids = [] # 计算边界框坐标的缩放因子 x_factor = img_width / input_width y_factor = img_height / input_height # 遍历输出数组的每一行 for i in range(rows): # 从当前行提取类别得分 classes_scores = outputs[i][4:] # 找到类别得分中的最大得分 max_score = np.amax(classes_scores) # 如果最大得分高于置信度阈值 if max_score >= confidence_thres: # 获取得分最高的类别ID class_id = np.argmax(classes_scores) # 从当前行提取边界框坐标 x, y, w, h = outputs[i][0], outputs[i][1], outputs[i][2], outputs[i][3] # 计算边界框的缩放坐标 left = int((x - w / 2) * x_factor) top = int((y - h / 2) * y_factor) width = int(w * x_factor) height = int(h * y_factor) # 将类别ID、得分和框坐标添加到各自的列表中 class_ids.append(class_id) scores.append(max_score) boxes.append([left, top, width, height]) # 应用非最大抑制过滤重叠的边界框 indices = custom_NMSBoxes(boxes, scores, confidence_thres, iou_thres) # 遍历非最大抑制后的选定索引 for i in indices: # 根据索引获取框、得分和类别ID box = boxes[i] score = scores[i] class_id = class_ids[i] # 在输入图像上绘制检测结果 draw_detections(input_image, box, score, class_id) # 返回修改后的输入图像 return input_image def init_detect_model(model_path): # 使用ONNX模型文件创建一个推理会话，并指定执行提供者 session = ort.InferenceSession(model_path, providers=providers) # 获取模型的输入信息 model_inputs = session.get_inputs() # 获取输入的形状，用于后续使用 input_shape = model_inputs[0].shape # 从输入形状中提取输入宽度 input_width = input_shape[2] # 从输入形状中提取输入高度 input_height = input_shape[3] # 返回会话、模型输入信息、输入宽度和输入高度 return session, model_inputs, input_width, input_height def detect_object(image, session, model_inputs, input_width, input_height): # 如果输入的图像是PIL图像对象，将其转换为NumPy数组 if isinstance(image, Image.Image): result_image = np.array(image) else: # 否则，直接使用输入的图像（假定已经是NumPy数组） result_image = image # 预处理图像数据，调整图像大小并可能进行归一化等操作 img_data, img_height, img_width = preprocess(result_image, input_width, input_height) # 使用预处理后的图像数据进行推理 outputs = session.run(None, {model_inputs[0].name: img_data}) # 对推理结果进行后处理，例如解码检测框，过滤低置信度的检测等 output_image = postprocess(result_image, outputs, input_width, input_height, img_width, img_height) # 返回处理后的图像 return output_image if __name__ == '__main__': # 模型文件的路径 model_path = 'model/detect/best.onnx' # 初始化检测模型，加载模型并获取模型输入节点信息和输入图像的宽度、高度 session, model_inputs, input_width, input_height = init_detect_model(model_path) # 三种模式 1为图片预测，并显示结果图片；2为摄像头检测，并实时显示FPS； 3为视频检测，并保存结果视频 mode = 2 if mode == 1: # 读取图像文件 image_data = cv2.imread("street.jpg") # 使用检测模型对读入的图像进行对象检测 result_image = detect_object(image_data, session, model_inputs, input_width, input_height) # 将检测后的图像保存到文件 cv2.imwrite("output_image.jpg", result_image) # 在窗口中显示检测后的图像 cv2.imshow('Output', result_image) # 等待用户按键，然后关闭显示窗口 cv2.waitKey(0) elif mode == 2: # 摄像头索引号，通常为0表示第一个摄像头 camera_index = 0 # 打开摄像头 cap = cv2.VideoCapture(camera_index, cv2.CAP_DSHOW) # 设置分辨率 cap.set(cv2.CAP_PROP_FRAME_WIDTH, 3840) # 宽度 cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 2160) # 高度 # 检查摄像头是否成功打开 if not cap.isOpened(): print("无法打开摄像头") exit() width = cap.get(cv2.CAP_PROP_FRAME_WIDTH) height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT) print("摄像头分辨率:", width, "x", height) # 初始化帧数计数器和起始时间 frame_count = 0 start_time = time.time() # 目标图像尺寸 target_width = 1024 target_height = 768 # 循环读取摄像头视频流 while True: # 读取一帧 ret, frame = cap.read() # 检查帧是否成功读取 if not ret: print("Error: Could not read frame.") break # 1920*1080的图像，中心裁剪640*480的区域 cropped_frame = frame[int(height / 2 - target_height / 2):int(height / 2 + target_height / 2), int(width / 2 - target_width / 2):int(width / 2 + target_width / 2)] # 调整图像尺寸 resized_frame = cv2.resize(cropped_frame, (target_width, target_height)) # 使用检测模型对读入的帧进行对象检测 output_image = detect_object(resized_frame, session, model_inputs, input_width, input_height) # 计算帧速率 frame_count += 1 end_time = time.time() elapsed_time = end_time - start_time fps = frame_count / elapsed_time print(f"FPS: {fps:.2f}") # 将FPS绘制在图像上 cv2.putText(output_image, f"FPS: {fps:.2f}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA) # 在窗口中显示当前帧 cv2.imshow("Video", output_image) # 按下 'q' 键退出循环 if cv2.waitKey(1) & 0xFF == ord('q'): break # 释放摄像头资源 cap.release() # 关闭窗口 cv2.destroyAllWindows() elif mode == 3: # 输入视频路径 input_video_path = 'kun.mp4' # 输出视频路径 output_video_path = 'kun_det.mp4' # 打开视频文件 cap = cv2.VideoCapture(input_video_path) # 检查视频是否成功打开 if not cap.isOpened(): print("Error: Could not open video.") exit() # 读取视频的基本信息 frame_width = int(cap.get(3)) frame_height = int(cap.get(4)) fps = cap.get(cv2.CAP_PROP_FPS) # 定义视频编码器和创建VideoWriter对象 fourcc = cv2.VideoWriter_fourcc(*'mp4v') # 根据文件名后缀使用合适的编码器 out = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height)) # 初始化帧数计数器和起始时间 frame_count = 0 start_time = time.time() while True: ret, frame = cap.read() if not ret: print("Info: End of video file.") break # 对读入的帧进行对象检测 output_image = detect_object(frame, session, model_inputs, input_width, input_height) # 计算并打印帧速率 frame_count += 1 end_time = time.time() elapsed_time = end_time - start_time if elapsed_time > 0: fps = frame_count / elapsed_time print(f"FPS: {fps:.2f}") # 将处理后的帧写入输出视频 out.write(output_image) #（可选）实时显示处理后的视频帧 cv2.imshow("Output Video", output_image) if cv2.waitKey(1) & 0xFF == ord('q'): break # 释放资源 cap.release() out.release() cv2.destroyAllWindows() else: print("输入错误，请检查mode的赋值")