from functools import lru_cache

from image_prediction.model_loader.database.connector import DatabaseConnector


class ModelLoader:
    def __init__(self, database_connector: DatabaseConnector):
        self.database_connector = database_connector

    @lru_cache(maxsize=None)
    def __get_object(self, identifier):
        return self.database_connector.get_object(identifier)

    def load_model(self, identifier):
        return self.__get_object(identifier)["model"]

    def load_classes(self, identifier):
        return self.__get_object(identifier)["classes"]