Refactor: done phase 5
This commit is contained in:
parent
5176f99ba7
commit
476f19cabe
3 changed files with 740 additions and 109 deletions
|
@ -2,6 +2,7 @@
|
|||
Detection Pipeline Module.
|
||||
Main detection pipeline orchestration that coordinates detection flow and execution.
|
||||
"""
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
import uuid
|
||||
|
@ -15,6 +16,7 @@ from ..models.pipeline import PipelineParser
|
|||
from .branches import BranchProcessor
|
||||
from ..storage.redis import RedisManager
|
||||
from ..storage.database import DatabaseManager
|
||||
from ..storage.license_plate import LicensePlateManager
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
@ -42,6 +44,7 @@ class DetectionPipeline:
|
|||
self.branch_processor = BranchProcessor(model_manager)
|
||||
self.redis_manager = None
|
||||
self.db_manager = None
|
||||
self.license_plate_manager = None
|
||||
|
||||
# Main detection model
|
||||
self.detection_model: Optional[YOLOWrapper] = None
|
||||
|
@ -53,6 +56,12 @@ class DetectionPipeline:
|
|||
# Pipeline configuration
|
||||
self.pipeline_config = pipeline_parser.pipeline_config
|
||||
|
||||
# SessionId to subscriptionIdentifier mapping
|
||||
self.session_to_subscription = {}
|
||||
|
||||
# SessionId to processing results mapping (for combining with license plate results)
|
||||
self.session_processing_results = {}
|
||||
|
||||
# Statistics
|
||||
self.stats = {
|
||||
'detections_processed': 0,
|
||||
|
@ -90,6 +99,15 @@ class DetectionPipeline:
|
|||
logger.warning("Failed to create car_frontal_info table")
|
||||
logger.info("Database connection initialized")
|
||||
|
||||
# Initialize license plate manager (using same Redis config as main Redis manager)
|
||||
if self.pipeline_parser.redis_config:
|
||||
self.license_plate_manager = LicensePlateManager(self.pipeline_parser.redis_config.__dict__)
|
||||
if not await self.license_plate_manager.initialize(self._on_license_plate_result):
|
||||
logger.error("Failed to initialize license plate manager")
|
||||
return False
|
||||
logger.info("License plate manager initialized")
|
||||
|
||||
|
||||
# Initialize main detection model
|
||||
if not await self._initialize_detection_model():
|
||||
logger.error("Failed to initialize detection model")
|
||||
|
@ -154,6 +172,193 @@ class DetectionPipeline:
|
|||
logger.error(f"Error initializing detection model: {e}", exc_info=True)
|
||||
return False
|
||||
|
||||
async def _on_license_plate_result(self, session_id: str, license_data: Dict[str, Any]):
|
||||
"""
|
||||
Callback for handling license plate results from LPR service.
|
||||
|
||||
Args:
|
||||
session_id: Session identifier
|
||||
license_data: License plate data including text and confidence
|
||||
"""
|
||||
try:
|
||||
license_text = license_data.get('license_plate_text', '')
|
||||
confidence = license_data.get('confidence', 0.0)
|
||||
|
||||
logger.info(f"[LICENSE PLATE CALLBACK] Session {session_id}: "
|
||||
f"text='{license_text}', confidence={confidence:.3f}")
|
||||
|
||||
# Find matching subscriptionIdentifier for this sessionId
|
||||
subscription_id = self.session_to_subscription.get(session_id)
|
||||
|
||||
if not subscription_id:
|
||||
logger.warning(f"[LICENSE PLATE] No subscription found for sessionId '{session_id}' (type: {type(session_id)}), cannot send imageDetection")
|
||||
logger.warning(f"[LICENSE PLATE DEBUG] Current session mappings: {dict(self.session_to_subscription)}")
|
||||
|
||||
# Try to find by type conversion in case of type mismatch
|
||||
# Try as integer if session_id is string
|
||||
if isinstance(session_id, str) and session_id.isdigit():
|
||||
session_id_int = int(session_id)
|
||||
subscription_id = self.session_to_subscription.get(session_id_int)
|
||||
if subscription_id:
|
||||
logger.info(f"[LICENSE PLATE] Found subscription using int conversion: '{session_id}' -> {session_id_int} -> '{subscription_id}'")
|
||||
else:
|
||||
logger.error(f"[LICENSE PLATE] Failed to find subscription with int conversion")
|
||||
return
|
||||
# Try as string if session_id is integer
|
||||
elif isinstance(session_id, int):
|
||||
session_id_str = str(session_id)
|
||||
subscription_id = self.session_to_subscription.get(session_id_str)
|
||||
if subscription_id:
|
||||
logger.info(f"[LICENSE PLATE] Found subscription using string conversion: {session_id} -> '{session_id_str}' -> '{subscription_id}'")
|
||||
else:
|
||||
logger.error(f"[LICENSE PLATE] Failed to find subscription with string conversion")
|
||||
return
|
||||
else:
|
||||
logger.error(f"[LICENSE PLATE] Failed to find subscription with any type conversion")
|
||||
return
|
||||
|
||||
# Send imageDetection message with license plate data combined with processing results
|
||||
await self._send_license_plate_message(subscription_id, license_text, confidence, session_id)
|
||||
|
||||
# Update database with license plate information if database manager is available
|
||||
if self.db_manager and license_text:
|
||||
success = self.db_manager.execute_update(
|
||||
table='car_frontal_info',
|
||||
key_field='session_id',
|
||||
key_value=session_id,
|
||||
fields={
|
||||
'license_character': license_text,
|
||||
'license_type': 'LPR_detected' # Mark as detected by LPR service
|
||||
}
|
||||
)
|
||||
if success:
|
||||
logger.info(f"[LICENSE PLATE] Updated database for session {session_id}")
|
||||
else:
|
||||
logger.warning(f"[LICENSE PLATE] Failed to update database for session {session_id}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error in license plate result callback: {e}", exc_info=True)
|
||||
|
||||
|
||||
async def _send_license_plate_message(self, subscription_id: str, license_text: str, confidence: float, session_id: str = None):
|
||||
"""
|
||||
Send imageDetection message with license plate data plus any available processing results.
|
||||
|
||||
Args:
|
||||
subscription_id: Subscription identifier to send message to
|
||||
license_text: License plate text
|
||||
confidence: License plate confidence score
|
||||
session_id: Session identifier for looking up processing results
|
||||
"""
|
||||
try:
|
||||
if not self.message_sender:
|
||||
logger.warning("No message sender configured, cannot send imageDetection")
|
||||
return
|
||||
|
||||
# Import here to avoid circular imports
|
||||
from ..communication.models import ImageDetectionMessage, DetectionData
|
||||
|
||||
# Get processing results for this session from stored results
|
||||
car_brand = None
|
||||
body_type = None
|
||||
|
||||
# Find session_id from session mappings (we need session_id as key)
|
||||
session_id_for_lookup = None
|
||||
|
||||
# Try direct lookup first (if session_id is already the right type)
|
||||
if session_id in self.session_processing_results:
|
||||
session_id_for_lookup = session_id
|
||||
else:
|
||||
# Try to find by type conversion
|
||||
for stored_session_id in self.session_processing_results.keys():
|
||||
if str(stored_session_id) == str(session_id):
|
||||
session_id_for_lookup = stored_session_id
|
||||
break
|
||||
|
||||
if session_id_for_lookup and session_id_for_lookup in self.session_processing_results:
|
||||
branch_results = self.session_processing_results[session_id_for_lookup]
|
||||
logger.info(f"[LICENSE PLATE] Retrieved processing results for session {session_id_for_lookup}")
|
||||
|
||||
if 'car_brand_cls_v2' in branch_results:
|
||||
brand_result = branch_results['car_brand_cls_v2'].get('result', {})
|
||||
car_brand = brand_result.get('brand')
|
||||
if 'car_bodytype_cls_v1' in branch_results:
|
||||
bodytype_result = branch_results['car_bodytype_cls_v1'].get('result', {})
|
||||
body_type = bodytype_result.get('body_type')
|
||||
|
||||
# Clean up stored results after use
|
||||
del self.session_processing_results[session_id_for_lookup]
|
||||
logger.debug(f"[LICENSE PLATE] Cleaned up stored results for session {session_id_for_lookup}")
|
||||
else:
|
||||
logger.warning(f"[LICENSE PLATE] No processing results found for session {session_id}")
|
||||
|
||||
# Create detection data with combined information
|
||||
detection_data_obj = DetectionData(
|
||||
detection={
|
||||
"carBrand": car_brand,
|
||||
"carModel": None,
|
||||
"bodyType": body_type,
|
||||
"licensePlateText": license_text,
|
||||
"licensePlateConfidence": confidence
|
||||
},
|
||||
modelId=52, # Default model ID
|
||||
modelName="yolo11m" # Default model name
|
||||
)
|
||||
|
||||
# Create imageDetection message
|
||||
detection_message = ImageDetectionMessage(
|
||||
subscriptionIdentifier=subscription_id,
|
||||
data=detection_data_obj
|
||||
)
|
||||
|
||||
# Send message
|
||||
await self.message_sender(detection_message)
|
||||
logger.info(f"[COMBINED MESSAGE] Sent imageDetection with brand='{car_brand}', bodyType='{body_type}', license='{license_text}' to '{subscription_id}'")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error sending license plate imageDetection message: {e}", exc_info=True)
|
||||
|
||||
async def _send_initial_detection_message(self, subscription_id: str):
|
||||
"""
|
||||
Send initial imageDetection message when vehicle is first detected.
|
||||
|
||||
Args:
|
||||
subscription_id: Subscription identifier to send message to
|
||||
"""
|
||||
try:
|
||||
if not self.message_sender:
|
||||
logger.warning("No message sender configured, cannot send imageDetection")
|
||||
return
|
||||
|
||||
# Import here to avoid circular imports
|
||||
from ..communication.models import ImageDetectionMessage, DetectionData
|
||||
|
||||
# Create detection data with all fields as None (vehicle just detected, no classification yet)
|
||||
detection_data_obj = DetectionData(
|
||||
detection={
|
||||
"carBrand": None,
|
||||
"carModel": None,
|
||||
"bodyType": None,
|
||||
"licensePlateText": None,
|
||||
"licensePlateConfidence": None
|
||||
},
|
||||
modelId=52, # Default model ID
|
||||
modelName="yolo11m" # Default model name
|
||||
)
|
||||
|
||||
# Create imageDetection message
|
||||
detection_message = ImageDetectionMessage(
|
||||
subscriptionIdentifier=subscription_id,
|
||||
data=detection_data_obj
|
||||
)
|
||||
|
||||
# Send message
|
||||
await self.message_sender(detection_message)
|
||||
logger.info(f"[INITIAL DETECTION] Sent imageDetection for vehicle detection to '{subscription_id}'")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error sending initial detection imageDetection message: {e}", exc_info=True)
|
||||
|
||||
async def execute_detection_phase(self,
|
||||
frame: np.ndarray,
|
||||
display_id: str,
|
||||
|
@ -249,21 +454,20 @@ class DetectionPipeline:
|
|||
|
||||
result['detections'] = valid_detections
|
||||
|
||||
# If we have valid detections, send imageDetection message with empty detection
|
||||
# If we have valid detections, create session and send initial imageDetection
|
||||
if valid_detections:
|
||||
logger.info(f"Found {len(valid_detections)} valid detections, sending imageDetection message")
|
||||
logger.info(f"Found {len(valid_detections)} valid detections, storing session mapping")
|
||||
|
||||
# Send imageDetection with empty detection data
|
||||
message_sent = await self._send_image_detection_message(
|
||||
subscription_id=subscription_id,
|
||||
detection_context=detection_context
|
||||
)
|
||||
result['message_sent'] = message_sent
|
||||
# Store mapping from display_id to subscriptionIdentifier (for detection phase)
|
||||
# Note: We'll store session_id mapping later in processing phase
|
||||
self.session_to_subscription[display_id] = subscription_id
|
||||
logger.info(f"[SESSION MAPPING] Stored mapping: displayId '{display_id}' -> subscriptionIdentifier '{subscription_id}'")
|
||||
|
||||
if message_sent:
|
||||
logger.info(f"Detection phase completed - imageDetection message sent for {display_id}")
|
||||
else:
|
||||
logger.warning(f"Failed to send imageDetection message for {display_id}")
|
||||
# Send initial imageDetection message with empty detection data
|
||||
await self._send_initial_detection_message(subscription_id)
|
||||
|
||||
logger.info(f"Detection phase completed - {len(valid_detections)} detections found for {display_id}")
|
||||
result['message_sent'] = True
|
||||
else:
|
||||
logger.debug("No valid detections found in detection phase")
|
||||
|
||||
|
@ -341,6 +545,11 @@ class DetectionPipeline:
|
|||
'confidence': confidence
|
||||
}
|
||||
|
||||
# Store session mapping for license plate callback
|
||||
if session_id:
|
||||
self.session_to_subscription[session_id] = subscription_id
|
||||
logger.info(f"[SESSION MAPPING] Stored mapping: sessionId '{session_id}' -> subscriptionIdentifier '{subscription_id}'")
|
||||
|
||||
# Initialize database record with session_id
|
||||
if session_id and self.db_manager:
|
||||
success = self.db_manager.insert_initial_detection(
|
||||
|
@ -391,6 +600,11 @@ class DetectionPipeline:
|
|||
)
|
||||
result['actions_executed'].extend(executed_parallel_actions)
|
||||
|
||||
# Store processing results for later combination with license plate data
|
||||
if result['branch_results'] and session_id:
|
||||
self.session_processing_results[session_id] = result['branch_results']
|
||||
logger.info(f"[PROCESSING RESULTS] Stored results for session {session_id} for later combination")
|
||||
|
||||
logger.info(f"Processing phase completed for session {session_id}: "
|
||||
f"{len(result['branch_results'])} branches, {len(result['actions_executed'])} actions")
|
||||
|
||||
|
@ -402,57 +616,6 @@ class DetectionPipeline:
|
|||
result['processing_time'] = time.time() - start_time
|
||||
return result
|
||||
|
||||
async def _send_image_detection_message(self,
|
||||
subscription_id: str,
|
||||
detection_context: Dict[str, Any]) -> bool:
|
||||
"""
|
||||
Send imageDetection message with empty detection data to backend.
|
||||
|
||||
Args:
|
||||
subscription_id: Subscription identifier
|
||||
detection_context: Detection context data
|
||||
|
||||
Returns:
|
||||
True if message sent successfully, False otherwise
|
||||
"""
|
||||
try:
|
||||
if not self.message_sender:
|
||||
logger.warning("No message sender available for imageDetection")
|
||||
return False
|
||||
|
||||
# Import here to avoid circular imports
|
||||
from ..communication.messages import create_image_detection
|
||||
|
||||
# Create empty detection data as specified
|
||||
detection_data = {}
|
||||
|
||||
# Get model info from pipeline configuration
|
||||
model_id = 52 # Default model ID
|
||||
model_name = "yolo11m" # Default
|
||||
|
||||
if self.pipeline_config:
|
||||
model_name = getattr(self.pipeline_config, 'model_id', 'yolo11m')
|
||||
# Try to extract numeric model ID from pipeline context, fallback to default
|
||||
if hasattr(self.pipeline_config, 'model_id'):
|
||||
# For now, use default model ID since pipeline config stores string identifiers
|
||||
model_id = 52
|
||||
|
||||
# Create imageDetection message
|
||||
detection_message = create_image_detection(
|
||||
subscription_identifier=subscription_id,
|
||||
detection_data=detection_data,
|
||||
model_id=model_id,
|
||||
model_name=model_name
|
||||
)
|
||||
|
||||
# Send to backend via WebSocket
|
||||
await self.message_sender(detection_message)
|
||||
logger.info(f"[DETECTION PHASE] Sent imageDetection with empty detection: {detection_data}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error sending imageDetection message: {e}", exc_info=True)
|
||||
return False
|
||||
|
||||
async def execute_detection(self,
|
||||
frame: np.ndarray,
|
||||
|
@ -697,9 +860,9 @@ class DetectionPipeline:
|
|||
if action_type == 'postgresql_update_combined':
|
||||
result = await self._execute_postgresql_update_combined(action, context)
|
||||
|
||||
# Send imageDetection message with actual processing results after database update
|
||||
# Update session state with processing results after database update
|
||||
if result.get('status') == 'success':
|
||||
await self._send_processing_results_message(context)
|
||||
await self._update_session_with_processing_results(context)
|
||||
else:
|
||||
logger.warning(f"Unknown parallel action type: {action_type}")
|
||||
result = {'status': 'error', 'message': f'Unknown action type: {action_type}'}
|
||||
|
@ -889,76 +1052,49 @@ class DetectionPipeline:
|
|||
logger.error(f"Error resolving field template {template}: {e}")
|
||||
return None
|
||||
|
||||
async def _send_processing_results_message(self, context: Dict[str, Any]):
|
||||
async def _update_session_with_processing_results(self, context: Dict[str, Any]):
|
||||
"""
|
||||
Send imageDetection message with actual processing results after database update.
|
||||
Update session state with processing results from branch execution.
|
||||
|
||||
Args:
|
||||
context: Detection context containing branch results and subscription info
|
||||
context: Detection context containing branch results and session info
|
||||
"""
|
||||
try:
|
||||
branch_results = context.get('branch_results', {})
|
||||
session_id = context.get('session_id', '')
|
||||
subscription_id = context.get('subscription_id', '')
|
||||
|
||||
# Extract detection results from branch results
|
||||
detection_data = {
|
||||
"carBrand": None,
|
||||
"carModel": None,
|
||||
"bodyType": None,
|
||||
"licensePlateText": None,
|
||||
"licensePlateConfidence": None
|
||||
}
|
||||
if not session_id:
|
||||
logger.warning("No session_id in context for processing results")
|
||||
return
|
||||
|
||||
# Extract car brand from car_brand_cls_v2 results
|
||||
car_brand = None
|
||||
if 'car_brand_cls_v2' in branch_results:
|
||||
brand_result = branch_results['car_brand_cls_v2'].get('result', {})
|
||||
detection_data["carBrand"] = brand_result.get('brand')
|
||||
car_brand = brand_result.get('brand')
|
||||
|
||||
# Extract body type from car_bodytype_cls_v1 results
|
||||
body_type = None
|
||||
if 'car_bodytype_cls_v1' in branch_results:
|
||||
bodytype_result = branch_results['car_bodytype_cls_v1'].get('result', {})
|
||||
detection_data["bodyType"] = bodytype_result.get('body_type')
|
||||
body_type = bodytype_result.get('body_type')
|
||||
|
||||
# Create detection message
|
||||
subscription_id = context.get('subscription_id', '')
|
||||
# Get the actual numeric model ID from context
|
||||
model_id_value = context.get('model_id', 52)
|
||||
if isinstance(model_id_value, str):
|
||||
try:
|
||||
model_id_value = int(model_id_value)
|
||||
except (ValueError, TypeError):
|
||||
model_id_value = 52
|
||||
model_name = str(getattr(self.pipeline_config, 'model_id', 'unknown'))
|
||||
|
||||
logger.debug(f"Creating DetectionData with modelId={model_id_value}, modelName='{model_name}'")
|
||||
|
||||
from core.communication.models import ImageDetectionMessage, DetectionData
|
||||
detection_data_obj = DetectionData(
|
||||
detection=detection_data,
|
||||
modelId=model_id_value,
|
||||
modelName=model_name
|
||||
)
|
||||
detection_message = ImageDetectionMessage(
|
||||
subscriptionIdentifier=subscription_id,
|
||||
data=detection_data_obj
|
||||
)
|
||||
|
||||
# Send to backend via WebSocket
|
||||
if self.message_sender:
|
||||
await self.message_sender(detection_message)
|
||||
logger.info(f"[RESULTS] Sent imageDetection with processing results: {detection_data}")
|
||||
else:
|
||||
logger.warning("No message sender available for processing results")
|
||||
logger.info(f"[PROCESSING RESULTS] Completed for session {session_id}: "
|
||||
f"brand={car_brand}, bodyType={body_type}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error sending processing results message: {e}", exc_info=True)
|
||||
logger.error(f"Error updating session with processing results: {e}", exc_info=True)
|
||||
|
||||
def get_statistics(self) -> Dict[str, Any]:
|
||||
"""Get detection pipeline statistics."""
|
||||
branch_stats = self.branch_processor.get_statistics() if self.branch_processor else {}
|
||||
license_stats = self.license_plate_manager.get_statistics() if self.license_plate_manager else {}
|
||||
|
||||
return {
|
||||
'pipeline': self.stats,
|
||||
'branches': branch_stats,
|
||||
'license_plate': license_stats,
|
||||
'redis_available': self.redis_manager is not None,
|
||||
'database_available': self.db_manager is not None,
|
||||
'detection_model_loaded': self.detection_model is not None
|
||||
|
@ -978,4 +1114,7 @@ class DetectionPipeline:
|
|||
if self.branch_processor:
|
||||
self.branch_processor.cleanup()
|
||||
|
||||
if self.license_plate_manager:
|
||||
asyncio.create_task(self.license_plate_manager.close())
|
||||
|
||||
logger.info("Detection pipeline cleaned up")
|
Loading…
Add table
Add a link
Reference in a new issue