From 4619c122f18c67f1fc8cd5224fc98491fff05c7b Mon Sep 17 00:00:00 2001 From: ziesorx Date: Tue, 23 Sep 2025 20:32:08 +0700 Subject: [PATCH] feat: make tracking works --- core/communication/messages.py | 5 +-- core/communication/models.py | 1 - core/communication/websocket.py | 4 +- core/tracking/integration.py | 68 ++++++++++++++++++++++++++++++++- 4 files changed, 70 insertions(+), 8 deletions(-) diff --git a/core/communication/messages.py b/core/communication/messages.py index 7d3187d..5afde80 100644 --- a/core/communication/messages.py +++ b/core/communication/messages.py @@ -157,8 +157,7 @@ def create_state_report(cpu_usage: float, memory_usage: float, def create_image_detection(subscription_identifier: str, detection_data: Dict[str, Any], - model_id: int, model_name: str, - session_id: Optional[int] = None) -> ImageDetectionMessage: + model_id: int, model_name: str) -> ImageDetectionMessage: """ Create an image detection message. @@ -167,7 +166,6 @@ def create_image_detection(subscription_identifier: str, detection_data: Dict[st detection_data: Flat dictionary of detection results model_id: Model identifier model_name: Model name - session_id: Optional session ID Returns: ImageDetectionMessage object @@ -182,7 +180,6 @@ def create_image_detection(subscription_identifier: str, detection_data: Dict[st return ImageDetectionMessage( subscriptionIdentifier=subscription_identifier, - sessionId=session_id, data=data ) diff --git a/core/communication/models.py b/core/communication/models.py index eb7c39c..eb55cc6 100644 --- a/core/communication/models.py +++ b/core/communication/models.py @@ -108,7 +108,6 @@ class ImageDetectionMessage(BaseModel): type: Literal["imageDetection"] = "imageDetection" subscriptionIdentifier: str timestamp: str = Field(default_factory=lambda: datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%S.%fZ")) - sessionId: Optional[int] = None data: DetectionData diff --git a/core/communication/websocket.py b/core/communication/websocket.py index 71077f0..c40c912 100644 --- a/core/communication/websocket.py +++ b/core/communication/websocket.py @@ -304,9 +304,9 @@ class WebSocketHandler: # Get pipeline configuration for this model pipeline_parser = model_manager.get_pipeline_config(model_id) if pipeline_parser: - # Create tracking integration + # Create tracking integration with message sender tracking_integration = TrackingPipelineIntegration( - pipeline_parser, model_manager + pipeline_parser, model_manager, self._send_message ) # Initialize tracking model diff --git a/core/tracking/integration.py b/core/tracking/integration.py index 950a1dc..957e8a9 100644 --- a/core/tracking/integration.py +++ b/core/tracking/integration.py @@ -24,16 +24,18 @@ class TrackingPipelineIntegration: Manages tracking state transitions and pipeline execution triggers. """ - def __init__(self, pipeline_parser: PipelineParser, model_manager: Any): + def __init__(self, pipeline_parser: PipelineParser, model_manager: Any, message_sender=None): """ Initialize tracking-pipeline integration. Args: pipeline_parser: Pipeline parser with loaded configuration model_manager: Model manager for loading models + message_sender: Optional callback function for sending WebSocket messages """ self.pipeline_parser = pipeline_parser self.model_manager = model_manager + self.message_sender = message_sender # Initialize tracking components tracking_config = pipeline_parser.tracking_config.__dict__ if pipeline_parser.tracking_config else {} @@ -60,6 +62,11 @@ class TrackingPipelineIntegration: 'pipelines_executed': 0 } + # Test mode for mock detection + self.test_mode = True + self.test_detection_sent = False + self.start_time = time.time() + logger.info("TrackingPipelineIntegration initialized") async def initialize_tracking_model(self) -> bool: @@ -228,6 +235,9 @@ class TrackingPipelineIntegration: self.session_vehicles[session_id] = vehicle.track_id self.active_sessions[display_id] = session_id + # Send mock image detection message as per worker.md specification + await self._send_mock_detection(subscription_id, session_id) + # Execute detection pipeline (placeholder for Phase 5) pipeline_result = await self._execute_pipeline( frame, @@ -253,6 +263,13 @@ class TrackingPipelineIntegration: except Exception as e: logger.error(f"Error in tracking pipeline: {e}", exc_info=True) + # TEST MODE: Send mock detection after 10 seconds to test WebSocket communication + if self.test_mode and not self.test_detection_sent and (time.time() - self.start_time) > 10: + self.test_detection_sent = True + test_session_id = f"test-session-{int(time.time())}" + logger.info(f"[TEST MODE] Triggering mock detection with session {test_session_id}") + await self._send_mock_detection(subscription_id, test_session_id) + result['processing_time'] = time.time() - start_time return result @@ -295,6 +312,55 @@ class TrackingPipelineIntegration: return pipeline_result + async def _send_mock_detection(self, subscription_id: str, session_id: str): + """ + Send mock image detection message to backend following worker.md specification. + + Args: + subscription_id: Full subscription identifier (display-id;camera-id) + session_id: Session identifier for linking detection to user session + """ + try: + # Import here to avoid circular imports + from ..communication.messages import create_image_detection + + # Create flat detection data as required by the model + detection_data = { + "carModel": "Civic", + "carBrand": "Honda", + "carYear": 2023, + "bodyType": "Sedan", + "licensePlateText": "MOCK123", + "licensePlateConfidence": 0.95 + } + + # Get model info + model_id = 1 # Default model ID as integer + if self.tracking_model_id: + # Try to extract numeric ID from model string + try: + model_id = int(self.tracking_model_id.split('_')[-1].replace('v', '')) + except: + model_id = 1 + + # Create proper Pydantic message using the helper function + detection_message = create_image_detection( + subscription_identifier=subscription_id, + detection_data=detection_data, + model_id=model_id, + model_name="Vehicle Tracking Detection" + ) + + # Send to backend via WebSocket if sender is available + if self.message_sender: + await self.message_sender(detection_message) + logger.info(f"[MOCK DETECTION] Sent to backend: {detection_data}") + else: + logger.info(f"[MOCK DETECTION] No message sender available, would send: {detection_message}") + + except Exception as e: + logger.error(f"Error sending mock detection: {e}", exc_info=True) + def set_session_id(self, display_id: str, session_id: str): """ Set session ID for a display (from backend).