Skip to content

Commit

Permalink
Adding plugins
Browse files Browse the repository at this point in the history
  • Loading branch information
@Clorr
Clorr committed Jan 3, 2018
1 parent 0dcf192 commit 3e2976a
Show file tree
Hide file tree
Showing 16 changed files with 327 additions and 425 deletions.
3 changes: 2 additions & 1 deletion .gitignore
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,4 +5,5 @@
!Dockerfile*
!requirements*
!lib/*.py
!scripts
!scripts
!plugins
2 changes: 1 addition & 1 deletion faceswap.py
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
#!/usr/bin/env python3
import argparse
from scripts.extract import ExtractTrainingData
from lib.cli import TrainingProcessor
from scripts.train import TrainingProcessor
from scripts.convert import ConvertImage

if __name__ == "__main__":
Expand Down
0 lib/pixel_shuffler.py → lib/PixelShuffler.py
View file
File renamed without changes.
84 changes: 0 additions & 84 deletions lib/aligner.py
View file

This file was deleted.

180 changes: 6 additions & 174 deletions lib/cli.py
View file
Original file line number Diff line number Diff line change
@@ -1,16 +1,8 @@
import argparse
import os
import cv2
import numpy
import time

from lib.utils import get_image_paths, get_folder, load_images, stack_images
from lib.faces_detect import crop_faces
from lib.training_data import get_training_data

from lib.model import autoencoder_A, autoencoder_B
from lib.model import encoder, decoder_A, decoder_B

from lib.utils import get_image_paths, get_folder, load_images

class FullPaths(argparse.Action):
"""Expand user- and relative-paths"""
Expand All @@ -19,154 +11,6 @@ def __call__(self, parser, namespace, values, option_string=None):
setattr(namespace, self.dest, os.path.abspath(
os.path.expanduser(values)))


class TrainingProcessor(object):
arguments = None

def __init__(self, subparser, command, description='default'):
self.parse_arguments(description, subparser, command)

def process_arguments(self, arguments):
self.arguments = arguments
print("Model A Directory: {}".format(self.arguments.input_A))
print("Model B Directory: {}".format(self.arguments.input_B))
print("Training data directory: {}".format(self.arguments.model_dir))
print('Starting, this may take a while...')

try:
encoder.load_weights(self.arguments.model_dir + '/encoder.h5')
decoder_A.load_weights(self.arguments.model_dir + '/decoder_A.h5')
decoder_B.load_weights(self.arguments.model_dir + '/decoder_B.h5')
except Exception as e:
print('Not loading existing training data.')
print(e)

self.process()

def parse_arguments(self, description, subparser, command):
parser = subparser.add_parser(
command,
help="This command trains the model for the two faces A and B.",
description=description,
epilog="Questions and feedback: \
https://github.com/deepfakes/faceswap-playground"
)

parser.add_argument('-A', '--input-A',
action=FullPaths,
dest="input_A",
default="input_A",
help="Input directory. A directory containing training images for face A.\
Defaults to 'input'")
parser.add_argument('-B', '--input-B',
action=FullPaths,
dest="input_B",
default="input_B",
help="Input directory. A directory containing training images for face B.\
Defaults to 'input'")
parser.add_argument('-m', '--model-dir',
action=FullPaths,
dest="model_dir",
default="models",
help="Model directory. This is where the training data will \
be stored. Defaults to 'model'")
parser.add_argument('-p', '--preview',
action="store_true",
dest="preview",
default=False,
help="Show preview output. If not specified, write progress \
to file.")
parser.add_argument('-v', '--verbose',
action="store_true",
dest="verbose",
default=False,
help="Show verbose output")
parser.add_argument('-s', '--save-interval',
type=int,
dest="save_interval",
default=100,
help="Sets the number of iterations before saving the model.")
parser.add_argument('-w', '--write-image',
action="store_true",
dest="write_image",
default=False,
help="Writes the training result to a file even on preview mode.")
parser = self.add_optional_arguments(parser)
parser.set_defaults(func=self.process_arguments)

def add_optional_arguments(self, parser):
# Override this for custom arguments
return parser

def save_model_weights(self):
encoder.save_weights(self.arguments.model_dir + '/encoder.h5')
decoder_A.save_weights(self.arguments.model_dir + '/decoder_A.h5')
decoder_B.save_weights(self.arguments.model_dir + '/decoder_B.h5')
print('save model weights')

def show_sample(self, test_A, test_B):
figure_A = numpy.stack([
test_A,
autoencoder_A.predict(test_A),
autoencoder_B.predict(test_A),
], axis=1)
figure_B = numpy.stack([
test_B,
autoencoder_B.predict(test_B),
autoencoder_A.predict(test_B),
], axis=1)

figure = numpy.concatenate([figure_A, figure_B], axis=0)
figure = figure.reshape((4, 7) + figure.shape[1:])
figure = stack_images(figure)

figure = numpy.clip(figure * 255, 0, 255).astype('uint8')

if self.arguments.preview is True:
cv2.imshow('', figure)
if not self.arguments.preview or self.arguments.write_image:
cv2.imwrite('_sample.jpg', figure)

def process(self):
images_A = get_image_paths(self.arguments.input_A)
images_B = get_image_paths(self.arguments.input_B)
images_A = load_images(images_A) / 255.0
images_B = load_images(images_B) / 255.0

images_A += images_B.mean(axis=(0, 1, 2)) - \
images_A.mean(axis=(0, 1, 2))

print('press "q" to stop training and save model')

BATCH_SIZE = 64

for epoch in range(1000000):
if self.arguments.verbose:
print("Iteration number {}".format(epoch + 1))
start_time = time.time()
warped_A, target_A = get_training_data(images_A, BATCH_SIZE)
warped_B, target_B = get_training_data(images_B, BATCH_SIZE)

loss_A = autoencoder_A.train_on_batch(warped_A, target_A)
loss_B = autoencoder_B.train_on_batch(warped_B, target_B)
print(loss_A, loss_B)

if epoch % self.arguments.save_interval == 0:
self.save_model_weights()
self.show_sample(target_A[0:14], target_B[0:14])

key = cv2.waitKey(1)
if key == ord('q'):
self.save_model_weights()
exit()
if self.arguments.verbose:
end_time = time.time()
time_elapsed = int(round((end_time - start_time)))
m, s = divmod(time_elapsed, 60)
h, m = divmod(m, 60)
print("Iteration done in {:02d}h{:02d}m{:02d}s".format(h, m, s))


class DirectoryProcessor(object):
'''
Abstract class that processes a directory of images
Expand Down Expand Up @@ -214,6 +58,11 @@ def process_directory(self):

self.finalize()

# for now, we limit this class responsability to the read of files. images and faces are processed outside this class
def process_image(self, filename):
# implement your image processing!
raise NotImplementedError()

def parse_arguments(self, description, subparser, command):
self.parser.add_argument('-i', '--input-dir',
action=FullPaths,
Expand Down Expand Up @@ -248,23 +97,6 @@ def add_optional_arguments(self, parser):
# Override this for custom arguments
return parser

def process_image(self, filename):
try:
image = cv2.imread(filename)
for (idx, face) in enumerate(crop_faces(image)):
if idx > 0 and self.arguments.verbose:
print('- Found more than one face!')
self.verify_output = True

self.process_face(face, idx, filename)
self.faces_detected = self.faces_detected + 1
except Exception as e:
print('Failed to extract from image: {}. Reason: {}'.format(filename, e))

def process_face(self, face, index, filename):
# implement your face processing!
raise NotImplementedError()

def finalize(self):
print('-------------------------')
print('Images found: {}'.format(self.images_found))
Expand Down
2 changes: 1 addition & 1 deletion lib/faces_detect.py
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import face_recognition
from .DetectedFace import DetectedFace

def crop_faces(frame):
def detect_faces(frame):
face_locations = face_recognition.face_locations(frame)
#face_encodings = face_recognition.face_encodings(frame, face_locations)

Expand Down
Loading

0 comments on commit 3e2976a

Please sign in to comment.