Просмотр исходного кода

feat(main): add --mode polling with UltralyticsTracker and PollingTrackingCoordinator

wenhongquan 1 день назад
Родитель
Сommit
db7202975a
1 измененных файлов с 76 добавлено и 27 удалено
  1. 76 27
      dual_camera_system/main.py

+ 76 - 27
dual_camera_system/main.py

@@ -30,12 +30,15 @@ from config import (
     PANORAMA_CAMERA, PTZ_CAMERA, SDK_PATH,
     PANORAMA_CAMERA, PTZ_CAMERA, SDK_PATH,
     DETECTION_CONFIG, PTZ_CONFIG, COORDINATOR_CONFIG,
     DETECTION_CONFIG, PTZ_CONFIG, COORDINATOR_CONFIG,
     CALIBRATION_CONFIG, LOG_CONFIG, SYSTEM_CONFIG,
     CALIBRATION_CONFIG, LOG_CONFIG, SYSTEM_CONFIG,
-    CAMERA_GROUPS, get_enabled_groups
+    CAMERA_GROUPS, get_enabled_groups,
+    TRACKING_CONFIG,
 )
 )
 from dahua_sdk import DahuaSDK
 from dahua_sdk import DahuaSDK
 from panorama_camera import PanoramaCamera, ObjectDetector, DetectedObject
 from panorama_camera import PanoramaCamera, ObjectDetector, DetectedObject
 from ptz_camera import PTZCamera, PTZController
 from ptz_camera import PTZCamera, PTZController
 from coordinator import Coordinator, EventDrivenCoordinator, AsyncCoordinator, SequentialCoordinator
 from coordinator import Coordinator, EventDrivenCoordinator, AsyncCoordinator, SequentialCoordinator
+from tracker import UltralyticsTracker
+from polling_tracker import PollingTrackingCoordinator
 
 
 
 
 # 配置日志 - 使用LOG_CONFIG
 # 配置日志 - 使用LOG_CONFIG
@@ -235,34 +238,63 @@ class DualCameraSystem:
         self.ptz_camera = PTZCamera(self.sdk, ptz_config)
         self.ptz_camera = PTZCamera(self.sdk, ptz_config)
         
         
         # 初始化联动控制器
         # 初始化联动控制器
-        # 根据配置选择顺序模式或异步模式
-        sequential_mode = COORDINATOR_CONFIG.get('sequential_mode', {}).get('enabled', False)
-        
-        if sequential_mode:
-            logger.info("使用顺序联动控制器 (SequentialCoordinator)")
-            self.coordinator = SequentialCoordinator(
-                self.panorama_camera,
-                self.ptz_camera,
-                self.detector
-            )
-            # 应用顺序模式配置
-            seq_config = COORDINATOR_CONFIG.get('sequential_mode', {})
-            self.coordinator.set_capture_config(
-                ptz_stabilize_time=seq_config.get('ptz_stabilize_time', 1.0),
-                capture_wait_time=seq_config.get('capture_wait_time', 0.5),
-                return_to_panorama=seq_config.get('return_to_panorama', True),
-                default_pan=seq_config.get('default_pan', 0.0),
-                default_tilt=seq_config.get('default_tilt', 0.0),
-                default_zoom=seq_config.get('default_zoom', 1)
+        tracking_mode = self.config.get('tracking_mode', SYSTEM_CONFIG.get('tracking_mode', 'polling'))
+
+        if tracking_mode == 'polling':
+            logger.info("使用轮询跟踪抓拍协调器 (PollingTrackingCoordinator)")
+            tracking_cfg = TRACKING_CONFIG.copy()
+            tracking_cfg['tracker_type'] = self.config.get('tracker_type', tracking_cfg.get('tracker_type', 'bytetrack'))
+            tracking_cfg['max_tracking_targets'] = self.config.get('max_targets', tracking_cfg.get('max_tracking_targets', 4))
+
+            self.tracker = UltralyticsTracker(
+                model_path=tracking_cfg.get('model_path'),
+                model_type=tracking_cfg.get('model_type', 'auto'),
+                use_gpu=self.config.get('use_gpu', tracking_cfg.get('use_gpu', True)),
+                tracker_type=tracking_cfg['tracker_type'],
+                conf_threshold=tracking_cfg.get('conf_threshold', 0.5),
+                person_threshold=tracking_cfg.get('person_threshold', 0.5),
+                max_lost=tracking_cfg.get('max_lost', 30),
             )
             )
-        else:
-            logger.info("使用异步联动控制器 (AsyncCoordinator)")
-            self.coordinator = AsyncCoordinator(
+            self.coordinator = PollingTrackingCoordinator(
                 self.panorama_camera,
                 self.panorama_camera,
                 self.ptz_camera,
                 self.ptz_camera,
-                self.detector
+                self.tracker,
+                config=tracking_cfg,
+                calibrator=None,
             )
             )
-        
+        else:
+            # 根据配置选择顺序模式或异步模式
+            sequential_mode = COORDINATOR_CONFIG.get('sequential_mode', {}).get('enabled', False)
+
+            if sequential_mode:
+                logger.info("使用顺序联动控制器 (SequentialCoordinator)")
+                self.coordinator = SequentialCoordinator(
+                    self.panorama_camera,
+                    self.ptz_camera,
+                    self.detector
+                )
+                # 应用顺序模式配置
+                seq_config = COORDINATOR_CONFIG.get('sequential_mode', {})
+                self.coordinator.set_capture_config(
+                    ptz_stabilize_time=seq_config.get('ptz_stabilize_time', 1.0),
+                    capture_wait_time=seq_config.get('capture_wait_time', 0.5),
+                    return_to_panorama=seq_config.get('return_to_panorama', True),
+                    default_pan=seq_config.get('default_pan', 0.0),
+                    default_tilt=seq_config.get('default_tilt', 0.0),
+                    default_zoom=seq_config.get('default_zoom', 1)
+                )
+            else:
+                logger.info("使用异步联动控制器 (AsyncCoordinator)")
+                self.coordinator = AsyncCoordinator(
+                    self.panorama_camera,
+                    self.ptz_camera,
+                    self.detector
+                )
+
+        # 注入事件推送器
+        if getattr(self, 'event_pusher', None) is not None and hasattr(self.coordinator, 'set_event_pusher'):
+            self.coordinator.set_event_pusher(self.event_pusher)
+
         # 设置回调
         # 设置回调
         self._setup_callbacks()
         self._setup_callbacks()
         
         
@@ -297,7 +329,10 @@ class DualCameraSystem:
             if calibrator.load_calibration(calibration_file):
             if calibrator.load_calibration(calibration_file):
                 logger.info(f"已加载已有校准数据: {calibration_file}")
                 logger.info(f"已加载已有校准数据: {calibration_file}")
                 if self.coordinator and calibrator.is_calibrated():
                 if self.coordinator and calibrator.is_calibrated():
-                    self.coordinator.set_calibrator(calibrator)
+                    if hasattr(self.coordinator, 'set_calibrator'):
+                        self.coordinator.set_calibrator(calibrator)
+                    else:
+                        self.coordinator.calibrator = calibrator
                     logger.info("校准器已传递给联动控制器")
                     logger.info("校准器已传递给联动控制器")
                 self.calibrator = calibrator
                 self.calibrator = calibrator
             else:
             else:
@@ -483,7 +518,10 @@ class DualCameraSystem:
 
 
         # 将校准器传递给联动控制器,使其使用校准查找表而非线性映射
         # 将校准器传递给联动控制器,使其使用校准查找表而非线性映射
         if self.coordinator and self.calibrator and self.calibrator.is_calibrated():
         if self.coordinator and self.calibrator and self.calibrator.is_calibrated():
-            self.coordinator.set_calibrator(self.calibrator)
+            if hasattr(self.coordinator, 'set_calibrator'):
+                self.coordinator.set_calibrator(self.calibrator)
+            else:
+                self.coordinator.calibrator = self.calibrator
             logger.info("校准器已传递给联动控制器,跟踪将使用校准查找表")
             logger.info("校准器已传递给联动控制器,跟踪将使用校准查找表")
 
 
         return True
         return True
@@ -834,6 +872,14 @@ def main():
     parser.add_argument('--demo', action='store_true', help='演示模式(不连接实际摄像头)')
     parser.add_argument('--demo', action='store_true', help='演示模式(不连接实际摄像头)')
     parser.add_argument('--skip-calibration', action='store_true', help='跳过自动校准')
     parser.add_argument('--skip-calibration', action='store_true', help='跳过自动校准')
     parser.add_argument('--force-calibration', action='store_true', help='强制重新校准')
     parser.add_argument('--force-calibration', action='store_true', help='强制重新校准')
+    parser.add_argument('--mode', type=str, default='polling',
+                        choices=['polling', 'async', 'sequential'],
+                        help='跟踪联动模式 (polling/async/sequential)')
+    parser.add_argument('--tracker', type=str, default='bytetrack',
+                        choices=['bytetrack', 'botsort'],
+                        help='Ultralytics tracker 类型')
+    parser.add_argument('--max-targets', type=int, default=4,
+                        help='最大跟踪目标数')
     
     
     args = parser.parse_args()
     args = parser.parse_args()
     
     
@@ -948,6 +994,9 @@ def run_single_group_mode(args):
     
     
     config['model_size'] = args.model_size
     config['model_size'] = args.model_size
     config['use_gpu'] = not args.no_gpu
     config['use_gpu'] = not args.no_gpu
+    config['tracking_mode'] = args.mode
+    config['tracker_type'] = args.tracker
+    config['max_targets'] = args.max_targets
     
     
     # 创建系统实例
     # 创建系统实例
     system = DualCameraSystem(config)
     system = DualCameraSystem(config)