diff --git a/ai2/cv_video.py b/ai2/cv_video.py index e5f1b97..4dd6ccc 100644 --- a/ai2/cv_video.py +++ b/ai2/cv_video.py @@ -149,8 +149,6 @@ def process_frames(frame_queue, processed_frame_queue, ov_model, cls, confidence # 优化推理性能 try: - # 导入torch库 - import torch # 尝试启用ONNX Runtime加速 ov_model.to('cuda:0' if torch.cuda.is_available() else 'cpu') # 调整批处理大小为1,减少内存占用 @@ -215,14 +213,24 @@ def process_frames(frame_queue, processed_frame_queue, ov_model, cls, confidence results = ov_model(process_frame, classes=cls, show=False) last_results = results[0] # 保存检测结果用于后续帧 - # 如果尺寸调整过,需要将结果转换回原始尺寸 + # 统计数量 + num_objs = len(last_results.boxes) if hasattr(last_results, "boxes") else 0 + + # 统计各类别数量 + class_counts = {} + if hasattr(last_results, "boxes") and hasattr(last_results.boxes, "cls"): + for c in last_results.boxes.cls.cpu().numpy(): + class_counts[int(c)] = class_counts.get(int(c), 0) + 1 + + # # 绘制检测框 if resize_scale < 1.0: - # 绘制检测框 annotated_frame = cv2.resize(results[0].plot(conf=False, line_width=1, font_size=1.5), (frame.shape[1], frame.shape[0])) else: annotated_frame = results[0].plot(conf=False, line_width=1, font_size=1.5) - + # 在画面左上角写入统计信息 + text = f"Total: {num_objs} " + " ".join([f"cls{c}:{n}" for c, n in class_counts.items()]) + cv2.putText(annotated_frame, text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0,255,0), 2) # 在负载高时启用跳帧处理 if frame_queue.qsize() > frame_queue.maxsize * 0.5: skip_counter = max_skip