mirror of
https://github.com/deepfakes/faceswap
synced 2025-06-08 20:13:52 -04:00
34945cfcd7
* Making Models as plugins * Do not reload model on each image #39 + Adding FaceFilter #53 * Adding @lukaville PR for #43 and #44 (possibly) * Training done in a separate thread * Better log for plugin load * Adding a prefetch to train.py #49 (Note that we prefetch 2 batches of images, due to the queue behavior) + More compact logging with verbose info included * correction of DirectoryProcessor signature * adding missing import * Convert with parallel preprocessing of files * Added coverage var for trainer Added a var with comment. Feel free to add it as argument * corrections * Modifying preview and normalization of image + correction * Cleanup
31 lines
1.2 KiB
Python
31 lines
1.2 KiB
Python
import cv2
|
|
|
|
from pathlib import Path
|
|
from lib.cli import DirectoryProcessor
|
|
from plugins.PluginLoader import PluginLoader
|
|
|
|
class ExtractTrainingData(DirectoryProcessor):
|
|
def create_parser(self, subparser, command, description):
|
|
self.parser = subparser.add_parser(
|
|
command,
|
|
help="Extract the faces from a pictures.",
|
|
description=description,
|
|
epilog="Questions and feedback: \
|
|
https://github.com/deepfakes/faceswap-playground"
|
|
)
|
|
|
|
def process(self):
|
|
extractor_name = "Align" # TODO Pass as argument
|
|
extractor = PluginLoader.get_extractor(extractor_name)()
|
|
|
|
try:
|
|
for filename in self.read_directory():
|
|
print('Processing %s' % (filename))
|
|
image = cv2.imread(filename)
|
|
for idx, face in self.get_faces(image):
|
|
resized_image = extractor.extract(image, face, 256)
|
|
output_file = self.output_dir / Path(filename).stem
|
|
cv2.imwrite(str(output_file) + str(idx) + Path(filename).suffix, resized_image)
|
|
|
|
except Exception as e:
|
|
print('Failed to extract from image: {}. Reason: {}'.format(filename, e))
|