adsist-cms/python-detector-worker
3
0
Fork 0
Code Issues Pull requests Projects Releases Packages 1 Wiki Activity Actions

feat: make tracking works

Browse source
This commit is contained in:
ziesorx 2025-09-23 20:32:08 +07:00
parent 255be78d43
commit 4619c122f1
4 changed files with 70 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

5
core/communication/messages.py
View file

@ -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
)

1
core/communication/models.py
View file

@ -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

4
core/communication/websocket.py
View file

@ -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

68
core/tracking/integration.py
View file

@ -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).