18 lines
605 B
Python
18 lines
605 B
Python
from collections import namedtuple
|
|
|
|
from image_prediction.locations import MLRUNS_DIR
|
|
from image_prediction.model_loader.loader import ModelLoader
|
|
from image_prediction.model_loader.loaders.mlflow import MlflowLoader
|
|
|
|
ModelClassesPair = namedtuple("ModelClassesPair", ["model", "classes"])
|
|
|
|
|
|
def load_model_and_classes(identifier, model_loader: ModelLoader = None) -> ModelClassesPair:
|
|
if not model_loader:
|
|
model_loader = MlflowLoader(MLRUNS_DIR)
|
|
|
|
model = model_loader.load_model(identifier)
|
|
classes = model_loader.load_classes(identifier)
|
|
|
|
return ModelClassesPair(model, classes)
|