def process_image(self, filename):
# TODO move the model load and the converter creation in a method called on init, but after the arg parsing
(face_A,
face_B) = ('/decoder_A.h5',
'/decoder_B.h5') if not self.arguments.swap_model else (
'/decoder_B.h5', '/decoder_A.h5')
model_dir = self.arguments.model_dir
encoder.load_weights(model_dir + "/encoder.h5")
decoder_A.load_weights(model_dir + face_A)
decoder_B.load_weights(model_dir + face_B)
converter = PluginLoader.get_converter("Masked")(autoencoder_B)
try:
image = cv2.imread(filename)
for (idx, face) in enumerate(detect_faces(image)):
if idx > 0 and self.arguments.verbose:
print('- Found more than one face!')
self.verify_output = True
image = converter.patch_image(image, face)
self.faces_detected = self.faces_detected + 1
output_file = self.output_dir / Path(filename).name
cv2.imwrite(str(output_file), image)
except Exception as e:
print('Failed to convert image: {}. Reason: {}'.format(
filename, e))
def process_arguments(self, arguments):
if not arguments.swap_model:
self.face_A, self.face_B = ('/decoder_A.h5', '/decoder_B.h5')
else:
self.face_A, self.face_B = ('/decoder_B.h5', '/decoder_A.h5')
model_dir = arguments.model_dir
encoder.load_weights(model_dir + "/encoder.h5")
decoder_A.load_weights(model_dir + self.face_A)
decoder_B.load_weights(model_dir + self.face_B)
self.converter = PluginLoader.get_converter("Masked")(autoencoder_B)
super().process_arguments(arguments)
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))
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 __init__(self, description='default'):
print('Initializing')
self.parse_arguments(description)
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()
import cv2
import numpy
from lib.utils import get_image_paths, load_images, stack_images
from lib.training_data import get_training_data
from lib.model import autoencoder_A
from lib.model import autoencoder_B
from lib.model import encoder, decoder_A, decoder_B
try:
encoder.load_weights('models/encoder.h5')
decoder_A.load_weights('models/decoder_A.h5')
decoder_B.load_weights('models/decoder_B.h5')
except:
pass
def save_model_weights():
encoder.save_weights('models/encoder.h5')
decoder_A.save_weights('models/decoder_A.h5')
decoder_B.save_weights('models/decoder_B.h5')
print('save model weights')
def show_sample(test_A, test_B):
figure_A = numpy.stack([
test_A,
autoencoder_A.predict(test_A),
autoencoder_B.predict(test_A),
],