/
githubmirror
/
faceswap
Обзор
Документация
Войти
/
githubmirror
/
faceswap
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
v3.0.0
plugins/plugin_loader.py
303 строки
11 KB
torzdf
Faceswap 3 (#1516)
21 дек 2025, 05:45
Не верифицирован
21 дек 2025, 05:45
837bc2d
Код
Авторство
О чём код?
#!/usr/bin/env python3 """ Plugin loader for Faceswap extract, training and convert tasks """ from __future__ import annotations import logging import os import typing as T from importlib import import_module from lib.utils import get_module_objects if T.TYPE_CHECKING: from collections.abc import Callable from plugins.extract.detect._base import Detector from plugins.extract.align._base import Aligner from plugins.extract.mask._base import Masker from plugins.extract.recognition._base import Identity from plugins.train.model._base import ModelBase from plugins.train.trainer._base import TrainerBase logger = logging.getLogger(__name__) class PluginLoader(): """ Retrieve, or get information on, Faceswap plugins Return a specific plugin, list available plugins, or get the default plugin for a task. Example ------- >>> from plugins.plugin_loader import PluginLoader >>> align_plugins = PluginLoader.get_available_extractors('align') >>> aligner = PluginLoader.get_aligner('cv2-dnn') """ @staticmethod def get_detector(name: str, disable_logging: bool = False) -> type[Detector]: """ Return requested detector plugin Parameters ---------- name: str The name of the requested detector plugin disable_logging: bool, optional Whether to disable the INFO log message that the plugin is being imported. Default: `False` Returns ------- :class:`plugins.extract.detect` object: An extraction detector plugin """ return PluginLoader._import("extract.detect", name, disable_logging) @staticmethod def get_aligner(name: str, disable_logging: bool = False) -> type[Aligner]: """ Return requested aligner plugin Parameters ---------- name: str The name of the requested aligner plugin disable_logging: bool, optional Whether to disable the INFO log message that the plugin is being imported. Default: `False` Returns ------- :class:`plugins.extract.align` object: An extraction aligner plugin """ return PluginLoader._import("extract.align", name, disable_logging) @staticmethod def get_masker(name: str, disable_logging: bool = False) -> type[Masker]: """ Return requested masker plugin Parameters ---------- name: str The name of the requested masker plugin disable_logging: bool, optional Whether to disable the INFO log message that the plugin is being imported. Default: `False` Returns ------- :class:`plugins.extract.mask` object: An extraction masker plugin """ return PluginLoader._import("extract.mask", name, disable_logging) @staticmethod def get_recognition(name: str, disable_logging: bool = False) -> type[Identity]: """ Return requested recognition plugin Parameters ---------- name: str The name of the requested reccognition plugin disable_logging: bool, optional Whether to disable the INFO log message that the plugin is being imported. Default: `False` Returns ------- :class:`plugins.extract.recognition` object: An extraction recognition plugin """ return PluginLoader._import("extract.recognition", name, disable_logging) @staticmethod def get_model(name: str, disable_logging: bool = False) -> type[ModelBase]: """ Return requested training model plugin Parameters ---------- name: str The name of the requested training model plugin disable_logging: bool, optional Whether to disable the INFO log message that the plugin is being imported. Default: `False` Returns ------- :class:`plugins.train.model` object: A training model plugin """ return PluginLoader._import("train.model", name, disable_logging) @staticmethod def get_trainer(name: str, disable_logging: bool = False) -> type[TrainerBase]: """ Return requested training trainer plugin Parameters ---------- name: str The name of the requested training trainer plugin disable_logging: bool, optional Whether to disable the INFO log message that the plugin is being imported. Default: `False` Returns ------- :class:`plugins.train.trainer` object: A training trainer plugin """ return PluginLoader._import("train.trainer", name, disable_logging) @staticmethod def get_converter(category: str, name: str, disable_logging: bool = False) -> Callable: """ Return requested converter plugin Converters work slightly differently to other faceswap plugins. They are created to do a specific task (e.g. color adjustment, mask blending etc.), so multiple plugins will be loaded in the convert phase, rather than just one plugin for the other phases. Parameters ---------- name: str The name of the requested converter plugin disable_logging: bool, optional Whether to disable the INFO log message that the plugin is being imported. Default: `False` Returns ------- :class:`plugins.convert` object: A converter sub plugin """ return PluginLoader._import(f"convert.{category}", name, disable_logging) @staticmethod def _import(attr: str, name: str, disable_logging: bool): """ Import the plugin's module Parameters ---------- name: str The name of the requested converter plugin disable_logging: bool Whether to disable the INFO log message that the plugin is being imported. Returns ------- :class:`plugin` object: A plugin """ name = name.replace("-", "_") ttl = attr.split(".")[-1].title() if not disable_logging: logger.info("Loading %s from %s plugin...", ttl, name.title()) attr = "model" if attr == "Trainer" else attr.lower() mod = ".".join(("plugins", attr, name)) module = import_module(mod) return getattr(module, ttl) @staticmethod def get_available_extractors(extractor_type: T.Literal["align", "detect", "mask"], add_none: bool = False, extend_plugin: bool = False) -> list[str]: """ Return a list of available extractors of the given type Parameters ---------- extractor_type: {'align', 'detect', 'mask'} The type of extractor to return the plugins for add_none: bool, optional Append "none" to the list of returned plugins. Default: False extend_plugin: bool, optional Some plugins have configuration options that mean that multiple 'pseudo-plugins' can be generated based on their settings. An example of this is the bisenet-fp mask which, whilst selected as 'bisenet-fp' can be stored as 'bisenet-fp-face' and 'bisenet-fp-head' depending on whether hair has been included in the mask or not. ``True`` will generate each pseudo-plugin, ``False`` will generate the original plugin name. Default: ``False`` Returns ------- list: A list of the available extractor plugin names for the given type """ extractpath = os.path.join(os.path.dirname(__file__), "extract", extractor_type) extractors = [item.name.replace(".py", "").replace("_", "-") for item in os.scandir(extractpath) if not item.name.startswith("_") and not item.name.endswith("defaults.py") and item.name.endswith(".py")] extendable = ["bisenet-fp", "custom"] if extend_plugin and extractor_type == "mask" and any(ext in extendable for ext in extractors): for msk in extendable: extractors.remove(msk) extractors.extend([f"{msk}_face", f"{msk}_head"]) extractors = sorted(extractors) if add_none: extractors.insert(0, "none") return extractors @staticmethod def get_available_models() -> list[str]: """ Return a list of available training models Returns ------- list: A list of the available training model plugin names """ modelpath = os.path.join(os.path.dirname(__file__), "train", "model") models = sorted(item.name.replace(".py", "").replace("_", "-") for item in os.scandir(modelpath) if not item.name.startswith("_") and not item.name.endswith("defaults.py") and item.name.endswith(".py")) return models @staticmethod def get_default_model() -> str: """ Return the default training model plugin name Returns ------- str: The default faceswap training model """ models = PluginLoader.get_available_models() return 'original' if 'original' in models else models[0] @staticmethod def get_available_convert_plugins(convert_category: str, add_none: bool = True) -> list[str]: """ Return a list of available converter plugins in the given category Parameters ---------- convert_category: {'color', 'mask', 'scaling', 'writer'} The category of converter plugin to return the plugins for add_none: bool, optional Append "none" to the list of returned plugins. Default: True Returns ------- list A list of the available converter plugin names in the given category """ convertpath = os.path.join(os.path.dirname(__file__), "convert", convert_category) converters = sorted(item.name.replace(".py", "").replace("_", "-") for item in os.scandir(convertpath) if not item.name.startswith("_") and not item.name.endswith("defaults.py") and item.name.endswith(".py")) if add_none: converters.insert(0, "none") return converters __all__ = get_module_objects(__name__)