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

refactor: remove hardcoded modelid
All checks were successful
Build Worker Base and Application Images / check-base-changes (push) Successful in 7s
Details
Build Worker Base and Application Images / build-base (push) Has been skipped
Details
Build Worker Base and Application Images / build-docker (push) Successful in 2m46s
Details
Build Worker Base and Application Images / deploy-stack (push) Successful in 12s
Details

Browse source
This commit is contained in:
ziesorx 2025-09-25 00:18:02 +07:00
parent c94dfa10e7
commit dc47eb8580
4 changed files with 34 additions and 49 deletions
Download patch file Download diff file Expand all files Collapse all files

33
core/detection/pipeline.py
View file

@ -27,21 +27,23 @@ class DetectionPipeline:
Handles detection execution, branch coordination, and result aggregation.
"""
def __init__(self, pipeline_parser: PipelineParser, model_manager: Any, message_sender=None):
def __init__(self, pipeline_parser: PipelineParser, model_manager: Any, model_id: int, message_sender=None):
"""
Initialize detection pipeline.
Args:
pipeline_parser: Pipeline parser with loaded configuration
model_manager: Model manager for loading models
model_id: The model ID to use for loading models
message_sender: Optional callback function for sending WebSocket messages
"""
self.pipeline_parser = pipeline_parser
self.model_manager = model_manager
self.model_id = model_id
self.message_sender = message_sender
# Initialize components
self.branch_processor = BranchProcessor(model_manager)
self.branch_processor = BranchProcessor(model_manager, model_id)
self.redis_manager = None
self.db_manager = None
self.license_plate_manager = None
@ -150,23 +152,14 @@ class DetectionPipeline:
# Load detection model
logger.info(f"Loading detection model: {model_id} ({model_file})")
# Get the model ID from the ModelManager context
pipeline_models = list(self.model_manager.get_all_downloaded_models())
if pipeline_models:
actual_model_id = pipeline_models[0] # Use the first available model
self.detection_model = self.model_manager.get_yolo_model(actual_model_id, model_file)
else:
logger.error("No models available in ModelManager")
self.detection_model = self.model_manager.get_yolo_model(self.model_id, model_file)
if not self.detection_model:
logger.error(f"Failed to load detection model {model_file} from model {self.model_id}")
return False
self.detection_model_id = model_id
if self.detection_model:
logger.info(f"Detection model {model_id} loaded successfully")
return True
else:
logger.error(f"Failed to load detection model {model_id}")
return False
logger.info(f"Detection model {model_id} loaded successfully")
return True
except Exception as e:
logger.error(f"Error initializing detection model: {e}", exc_info=True)
@ -301,8 +294,8 @@ class DetectionPipeline:
"licensePlateText": license_text,
"licensePlateConfidence": confidence
},
modelId=52, # Default model ID
modelName="yolo11m" # Default model name
modelId=self.model_id,
modelName=self.pipeline_parser.pipeline_config.model_id if self.pipeline_parser.pipeline_config else "detection_model"
)
# Create imageDetection message
@ -342,8 +335,8 @@ class DetectionPipeline:
"licensePlateText": None,
"licensePlateConfidence": None
},
modelId=52, # Default model ID
modelName="yolo11m" # Default model name
modelId=self.model_id,
modelName=self.pipeline_parser.pipeline_config.model_id if self.pipeline_parser.pipeline_config else "detection_model"
)
# Create imageDetection message