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
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:
parent
c94dfa10e7
commit
dc47eb8580
4 changed files with 34 additions and 49 deletions
|
@ -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:
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue