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

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

View file

@ -21,14 +21,16 @@ class BranchProcessor:
Manages branch synchronization and result collection.
"""
def __init__(self, model_manager: Any):
def __init__(self, model_manager: Any, model_id: int):
"""
Initialize branch processor.
Args:
model_manager: Model manager for loading models
model_id: The model ID to use for loading models
"""
self.model_manager = model_manager
self.model_id = model_id
# Branch models cache
self.branch_models: Dict[str, YOLOWrapper] = {}
@ -123,22 +125,16 @@ class BranchProcessor:
# Load model
logger.info(f"Loading branch model: {model_id} ({model_file})")
# Get the first available model ID from ModelManager
pipeline_models = list(self.model_manager.get_all_downloaded_models())
if pipeline_models:
actual_model_id = pipeline_models[0] # Use the first available model
model = self.model_manager.get_yolo_model(actual_model_id, model_file)
# Load model using the proper model ID
model = self.model_manager.get_yolo_model(self.model_id, model_file)
if model:
self.branch_models[model_id] = model
self.stats['models_loaded'] += 1
logger.info(f"Branch model {model_id} loaded successfully")
return model
else:
logger.error(f"Failed to load branch model {model_id}")
return None
if model:
self.branch_models[model_id] = model
self.stats['models_loaded'] += 1
logger.info(f"Branch model {model_id} loaded successfully")
return model
else:
logger.error("No models available in ModelManager for branch loading")
logger.error(f"Failed to load branch model {model_id}")
return None
except Exception as e: