| |
| |
|
|
|
|
|
|
| """ Plugin loader for extract, training and model tasks """ |
|
|
| from utils import logger |
| import os |
| from importlib import import_module |
| from typing import Type |
| from trainer._base import TrainerBase |
| from torch.utils.data import Dataset |
| from model._base import ModelBase |
|
|
| class PluginLoader(): |
| """ |
| Plugin loader for extract, training and model tasks |
| function: get_{model_type} |
| args: {model_name} |
| will return the a class named model_type under model_type.model_name.py |
| |
| as it return a class you should also annotate the returning classtype to make |
| code linting avaliable in some IDE |
| """ |
| @staticmethod |
| def get_classifier(name) -> Type[ModelBase]: |
| """ Return requested attribute encoder plugin """ |
| return PluginLoader._import("model.classifier", name) |
|
|
| @staticmethod |
| def get_trainer(name) -> Type[TrainerBase]: |
| """ Return requested trainer plugin """ |
| return PluginLoader._import("trainer", name) |
|
|
| @staticmethod |
| def get_dataset(name) -> Type[Dataset]: |
| """ Return requested trainer plugin """ |
| return PluginLoader._import("dataset", name) |
|
|
| @staticmethod |
| def _import(attr, name): |
| """ Import the plugin's module """ |
| name = name.replace("-", "_") |
| ttl = attr.split(".")[-1].title() |
| logger.info("Loading %s from %s plugin...", ttl, name.title()) |
| attr = "model" if attr == "Trainer" else attr.lower() |
| mod = ".".join((attr, name)) |
| module = import_module(mod) |
| logger.info(str(module) + str(ttl)) |
| return getattr(module, ttl) |
|
|
| @staticmethod |
| def get_available_trainer(): |
| """ Return a list of available models """ |
| modelpath = os.path.join(os.path.dirname(__file__), "trainer") |
| models = sorted(item.name.replace(".py", "").replace("_", "-") |
| for item in os.scandir(modelpath) |
| if not item.name.startswith("_") |
| and item.name.endswith(".py")) |
| return models |
|
|
| @staticmethod |
| def get_default_model(): |
| """ Return the default model """ |
| models = PluginLoader.get_available_models() |
| return 'original' if 'original' in models else models[0] |
|
|
|
|
|
|
