1
0
Fork 0
mirror of https://github.com/deepfakes/faceswap synced 2025-06-08 20:13:52 -04:00
faceswap/plugins/plugin_loader.py
torzdf 27a685383e
Convert refactor (#703)
* Separate predict and implement pool

* Add check and raise error to multithreading

Box functions to config. Add crop box option.

* All masks to mask module. Refactor convert masks

Predicted mask passed from model. Cli update

* Intesect box with mask and fixes

* Use raw NN output for convert

Use raw mask for face adjustments. Split adjustments to pre and post warp

* Separate out adjustments. Add unmask sharpen

* Set sensible defaults. Pre PR Testing

* Fix queue sizes. Move masked.py to lib

* Fix Skip Frames. Fix GUI Config popup

* Sensible queue limits. Add a less resource intensive single processing mode

* Fix predicted mask. Amend smooth box defaults

* Deterministic ordering for video output

* Video to Video convert implemented

* Fixups

- Remove defaults from folders across all stages
- Move match-hist and aca into color adjustments selectable
- Handle crashes properly for pooled processes
- Fix output directory does not exist error when creating a new output folder
- Force output to frames if input is not a video

* Add Color Transfer adjustment method

Wrap info text in GUI plugin configure popup

* Refactor image adjustments. Easier to create plugins

Start implementing config options for video encoding

* Add video encoding config options

Allow video encoding for frames input (must pass in a reference video)
Move video and image output writers to plugins

* Image writers config options

Move scaling to cli
Move draw_transparent to images config
Add config options for cv2 writer
Add Pillow image writer

* Add gif filetype to Pillow. Fix draw transparent for Pillow

* Add Animated GIF writer

standardize opencv/pillow defaults

* [speedup] Pre-encode supported writers in the convert pool (opencv, pillow)

Move scaling to convert pool
Remove dfaker mask

* Fix default writer

* Bugfixes

* Better custom argparse formatting
2019-04-21 19:19:06 +00:00

90 lines
3.4 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):
""" Return requested detector plugin """
return PluginLoader._import("extract.detect", name)
@staticmethod
def get_aligner(name):
""" Return requested detector plugin """
return PluginLoader._import("extract.align", name)
@staticmethod
def get_model(name):
""" Return requested model plugin """
return PluginLoader._import("train.model", name)
@staticmethod
def get_trainer(name):
""" Return requested trainer plugin """
return PluginLoader._import("train.trainer", name)
@staticmethod
def get_converter(category, name):
""" Return the converter sub plugin """
return PluginLoader._import("convert.{}".format(category), 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(("plugins", attr, name))
module = import_module(mod)
return getattr(module, ttl)
@staticmethod
def get_available_extractors(extractor_type):
""" Return a list of available models """
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 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 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 item.name.endswith(".py"))
if add_none:
converters.insert(0, "none")
return converters