mirror of
https://github.com/deepfakes/faceswap.git
synced 2026-10-23 18:37:36 +08:00
torzdf
88352b0268
* requirements.txt: - Pin opencv to 4.1.1 (fixes cv2-dnn error) * lib.face_detect.DetectedFace: change LandmarksXY to landmarks_xy. Add left, right, top, bottom attributes * lib.model.session: Session manager for loading models into different graphs (for Nvidia + CPU) * plugins.extract._base: New parent class for all extract plugins * plugins.extract.pipeline. Remove MultiProcessing. Dynamically limit batchsize for Nvidia cards. Remove loglevel input * S3FD + FAN plugins. Standardise to Keras version for all backends * Standardize all extract plugins to new threaded codebase * Documentation. Start implementing Numpy style docstrings for Sphinx Documentation * Remove s3fd_amd. Change convert OTF to expect DetectedFace object * faces_detect - clean up and documentation * Remove PoolProcess * Migrate manual tool to new extract workflow * Remove AMD specific extractor code from cli and plugins * Sort tool to new extract workflow * Remove multiprocessing from project * Remove multiprocessing queues from QueueManager * Remove multiprocessing support from logger * Move face_filter to new extraction pipeline * Alignments landmarksXY > landmarks_xy and legacy handling * Intercept get_backend for sphinx doc build # Add Sphinx documentation
95 lines
3.9 KiB
Python
95 lines
3.9 KiB
Python
#!/usr/bin/env python3
|
|
""" Plugin loader for extract, training and model tasks """
|
|
|
|
import logging
|
|
import os
|
|
from importlib import import_module
|
|
|
|
logger = logging.getLogger(__name__) # pylint: disable=invalid-name
|
|
|
|
|
|
class PluginLoader():
|
|
""" Plugin loader for extract, training and model tasks """
|
|
@staticmethod
|
|
def get_detector(name, disable_logging=False):
|
|
""" Return requested detector plugin """
|
|
return PluginLoader._import("extract.detect", name, disable_logging)
|
|
|
|
@staticmethod
|
|
def get_aligner(name, disable_logging=False):
|
|
""" Return requested detector plugin """
|
|
return PluginLoader._import("extract.align", name, disable_logging)
|
|
|
|
@staticmethod
|
|
def get_model(name, disable_logging=False):
|
|
""" Return requested model plugin """
|
|
return PluginLoader._import("train.model", name, disable_logging)
|
|
|
|
@staticmethod
|
|
def get_trainer(name, disable_logging=False):
|
|
""" Return requested trainer plugin """
|
|
return PluginLoader._import("train.trainer", name, disable_logging)
|
|
|
|
@staticmethod
|
|
def get_converter(category, name, disable_logging=False):
|
|
""" Return the converter sub plugin """
|
|
return PluginLoader._import("convert.{}".format(category), name, disable_logging)
|
|
|
|
@staticmethod
|
|
def _import(attr, name, disable_logging):
|
|
""" Import the plugin's module """
|
|
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):
|
|
""" Return a list of available aligners/detectors """
|
|
extractpath = os.path.join(os.path.dirname(__file__),
|
|
"extract",
|
|
extractor_type)
|
|
extractors = sorted(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")
|
|
and item.name != "manual.py")
|
|
return extractors
|
|
|
|
@staticmethod
|
|
def get_available_models():
|
|
""" Return a list of available models """
|
|
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():
|
|
""" Return the default model """
|
|
models = PluginLoader.get_available_models()
|
|
return 'original' if 'original' in models else models[0]
|
|
|
|
@staticmethod
|
|
def get_available_convert_plugins(convert_category, add_none=True):
|
|
""" Return a list of available models """
|
|
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
|