yoooger

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