def __init__(self, arguments):
logger.debug("Initializing %s: (args: %s", self.__class__.__name__, arguments)
self.args = arguments
self.output_dir = get_folder(self.args.output_dir)
logger.info("Output Directory: %s", self.args.output_dir)
self.images = Images(self.args)
self.alignments = Alignments(self.args, True, self.images.is_video)
self.plugins = Plugins(self.args)
self.post_process = PostProcess(arguments)
self.verify_output = False
self.save_interval = None
if hasattr(self.args, "save_interval"):
self.save_interval = self.args.save_interval
logger.debug("Initialized %s", self.__class__.__name__)
def __init__(self, arguments):
logger.debug("Initializing %s: (args: %s)", self.__class__.__name__,
arguments)
self.args = arguments
self.patch_threads = None
self.images = Images(self.args)
self.validate()
self.alignments = Alignments(self.args, False, self.images.is_video)
self.opts = OptionalActions(self.args, self.images.input_images,
self.alignments)
self.add_queues()
self.disk_io = DiskIO(self.alignments, self.images, arguments)
self.predictor = Predict(self.disk_io.load_queue, self.queue_size,
arguments)
configfile = self.args.configfile if hasattr(self.args,
"configfile") else None
self.converter = Converter(get_folder(self.args.output_dir),
self.predictor.output_size,
self.predictor.has_predicted_mask,
self.disk_io.draw_transparent,
self.disk_io.pre_encode,
arguments,
configfile=configfile)
logger.debug("Initialized %s", self.__class__.__name__)
def __init__(self, arguments):
self.args = arguments
self.output_dir = get_folder(self.args.output_dir)
self.extract_faces = False
self.faces_count = 0
self.images = Images(self.args)
self.alignments = Alignments(self.args, False)
# Update Legacy alignments
Legacy(self.alignments, self.args.verbose, self.images.input_images)
self.post_process = PostProcess(arguments)
self.verify_output = False
self.opts = OptionalActions(self.args, self.images.input_images)
def __init__(self, arguments):
logger.debug("Initializing %s: (args: %s)", self.__class__.__name__,
arguments)
self.args = arguments
Utils.set_verbosity(self.args.loglevel)
self.images = Images(self.args)
self.validate()
self.alignments = Alignments(self.args, False, self.images.is_video)
# Update Legacy alignments
Legacy(self.alignments, self.images.input_images,
arguments.input_aligned_dir)
self.opts = OptionalActions(self.args, self.images.input_images,
self.alignments)
self.add_queues()
self.disk_io = DiskIO(self.alignments, self.images, arguments)
self.predictor = Predict(self.disk_io.load_queue, self.queue_size,
arguments)
self.converter = Converter(get_folder(self.args.output_dir),
self.predictor.output_size,
self.predictor.has_predicted_mask,
self.disk_io.draw_transparent,
self.disk_io.pre_encode, arguments)
logger.debug("Initialized %s", self.__class__.__name__)
def __init__(self, arguments):
self.args = arguments
self.output_dir = get_folder(self.args.output_dir)
self.images = Images(self.args)
self.faces = Faces(self.args)
self.alignments = Alignments(self.args)
self.opts = OptionalActions(self.args, self.images.input_images)
def __init__(self, arguments):
self.args = arguments
self.images = Images(self.args)
self.faces = Faces(self.args)
self.alignments = Alignments(self.args)
self.output_dir = self.faces.output_dir
self.export_face = True
def __init__(self, arguments):
self.args = arguments
self.images = Images(self.args)
self.faces = Faces(self.args)
self.alignments = Alignments(self.args)
self.output_dir = self.faces.output_dir
self.export_face = True
self.save_interval = self.args.save_interval if hasattr(self.args, "save_interval") else None
def __init__(self, arguments):
logger.debug("Initializing %s: (args: %s", self.__class__.__name__,
arguments)
self.args = arguments
Utils.set_verbosity(self.args.loglevel)
self.output_dir = get_folder(self.args.output_dir)
logger.info("Output Directory: %s", self.args.output_dir)
self.images = Images(self.args)
self.alignments = Alignments(self.args, True, self.images.is_video)
self.post_process = PostProcess(arguments)
self.extractor = Extractor(self.args.detector, self.args.aligner,
self.args.loglevel, self.args.multiprocess,
self.args.rotate_images, self.args.min_size)
self.save_queue = queue_manager.get_queue("extract_save")
self.verify_output = False
self.save_interval = None
if hasattr(self.args, "save_interval"):
self.save_interval = self.args.save_interval
logger.debug("Initialized %s", self.__class__.__name__)
def __init__(self, arguments):
logger.debug("Initializing %s: (args: %s)", self.__class__.__name__,
arguments)
self.args = arguments
self.output_dir = get_folder(self.args.output_dir)
self.extract_faces = False
self.faces_count = 0
self.images = Images(self.args)
self.alignments = Alignments(self.args, False, self.images.is_video)
# Update Legacy alignments
Legacy(self.alignments, self.images.input_images,
arguments.input_aligned_dir)
self.post_process = PostProcess(arguments)
self.verify_output = False
self.opts = OptionalActions(self.args, self.images.input_images,
self.alignments)
logger.debug("Initialized %s", self.__class__.__name__)
def __init__(self, arguments):
logger.debug("Initializing %s: (args: %s)", self.__class__.__name__,
arguments)
self._args = arguments
self._patch_threads = None
self._images = ImagesLoader(self._args.input_dir, fast_count=True)
self._alignments = Alignments(self._args, False, self._images.is_video)
if self._alignments.version == 1.0:
logger.error(
"The alignments file format has been updated since the given alignments "
"file was generated. You need to update the file to proceed.")
logger.error(
"To do this run the 'Alignments Tool' > 'Extract' Job.")
sys.exit(1)
self._opts = OptionalActions(self._args, self._images.file_list,
self._alignments)
self._add_queues()
self._disk_io = DiskIO(self._alignments, self._images, arguments)
self._predictor = Predict(self._disk_io.load_queue, self._queue_size,
arguments)
self._validate()
get_folder(self._args.output_dir)
configfile = self._args.configfile if hasattr(self._args,
"configfile") else None
self._converter = Converter(self._predictor.output_size,
self._predictor.coverage_ratio,
self._predictor.centering,
self._disk_io.draw_transparent,
self._disk_io.pre_encode,
arguments,
configfile=configfile)
logger.debug("Initialized %s", self.__class__.__name__)
def __init__(self, arguments, sample_size, display, lock, trigger_patch):
logger.debug("Initializing %s: (arguments: '%s', sample_size: %s, display: %s, lock: %s, "
"trigger_patch: %s)", self.__class__.__name__, arguments, sample_size,
display, lock, trigger_patch)
self.sample_size = sample_size
self.display = display
self.lock = lock
self.trigger_patch = trigger_patch
self.input_images = list()
self.predicted_images = list()
self.images = Images(arguments)
self.alignments = Alignments(arguments,
is_extract=False,
input_is_video=self.images.is_video)
self.filelist = self.get_filelist()
self.indices = self.get_indices()
self.predictor = Predict(queue_manager.get_queue("preview_predict_in"),
sample_size,
arguments)
self.generate()
logger.debug("Initialized %s", self.__class__.__name__)
def __init__(self, arguments):
logger.debug("Initializing %s: (args: %s", self.__class__.__name__, arguments)
self.args = arguments
Utils.set_verbosity(self.args.loglevel)
self.output_dir = get_folder(self.args.output_dir)
logger.info("Output Directory: %s", self.args.output_dir)
self.images = Images(self.args)
self.alignments = Alignments(self.args, True, self.images.is_video)
self.post_process = PostProcess(arguments)
self.extractor = Extractor(self.args.detector,
self.args.aligner,
self.args.loglevel,
self.args.multiprocess,
self.args.rotate_images,
self.args.min_size)
self.save_queue = queue_manager.get_queue("extract_save")
self.verify_output = False
self.save_interval = None
if hasattr(self.args, "save_interval"):
self.save_interval = self.args.save_interval
logger.debug("Initialized %s", self.__class__.__name__)
class Samples():
""" Holds 5 random test faces """
def __init__(self, arguments, sample_size, display, lock, trigger_patch):
logger.debug("Initializing %s: (arguments: '%s', sample_size: %s, display: %s, lock: %s, "
"trigger_patch: %s)", self.__class__.__name__, arguments, sample_size,
display, lock, trigger_patch)
self.sample_size = sample_size
self.display = display
self.lock = lock
self.trigger_patch = trigger_patch
self.input_images = list()
self.predicted_images = list()
self.images = Images(arguments)
self.alignments = Alignments(arguments,
is_extract=False,
input_is_video=self.images.is_video)
self.filelist = self.get_filelist()
self.indices = self.get_indices()
self.predictor = Predict(queue_manager.get_queue("preview_predict_in"),
sample_size,
arguments)
self.generate()
logger.debug("Initialized %s", self.__class__.__name__)
@property
def random_choice(self):
""" Return for random indices from the indices group """
retval = [random.choice(indices) for indices in self.indices]
logger.debug(retval)
return retval
def get_filelist(self):
""" Return a list of files, filtering out those frames which do not contain faces """
logger.debug("Filtering file list to frames with faces")
if self.images.is_video:
filelist = ["{}_{:06d}.png".format(os.path.splitext(self.images.input_images)[0],
frame_no)
for frame_no in range(1, self.images.images_found + 1)]
else:
filelist = self.images.input_images
retval = [filename for filename in filelist
if self.alignments.frame_has_faces(os.path.basename(filename))]
logger.debug("Filtered out frames: %s", self.images.images_found - len(retval))
return retval
def get_indices(self):
""" Returns a list of 'self.sample_size' evenly sized partition indices
pertaining to the filtered file list """
# Remove start and end values to get a list divisible by self.sample_size
no_files = len(self.filelist)
crop = no_files % self.sample_size
top_tail = list(range(no_files))[
crop // 2:no_files - (crop - (crop // 2))]
# Partition the indices
size = len(top_tail)
retval = [top_tail[start:start + size // self.sample_size]
for start in range(0, size, size // self.sample_size)]
logger.debug("Indices pools: %s", ["{}: (start: {}, end: {}, size: {})".format(idx,
min(pool),
max(pool),
len(pool))
for idx, pool in enumerate(retval)])
return retval
def generate(self):
""" Generate a random test set """
self.load_frames()
self.predict()
self.trigger_patch.set()
def load_frames(self):
""" Load a sample of random frames """
self.input_images = list()
for selection in self.random_choice:
filename = os.path.basename(self.filelist[selection])
image = self.images.load_one_image(self.filelist[selection])
# Get first face only
face = self.alignments.get_faces_in_frame(filename)[0]
detected_face = DetectedFace()
detected_face.from_alignment(face, image=image)
self.input_images.append({"filename": filename,
"image": image,
"detected_faces": [detected_face]})
self.display.source = self.input_images
self.display.update_source = True
logger.debug("Selected frames: %s", [frame["filename"] for frame in self.input_images])
def predict(self):
""" Predict from the loaded frames """
with self.lock:
self.predicted_images = list()
for frame in self.input_images:
self.predictor.in_queue.put(frame)
idx = 0
while idx < self.sample_size:
logger.debug("Predicting face %s of %s", idx + 1, self.sample_size)
item = self.predictor.out_queue.get()
if item == "EOF":
logger.debug("Received EOF")
break
self.predicted_images.append(item)
logger.debug("Predicted face %s of %s", idx + 1, self.sample_size)
idx += 1
logger.debug("Predicted faces")
class Extract(object):
""" The extract process. """
def __init__(self, arguments):
self.args = arguments
self.images = Images(self.args)
self.faces = Faces(self.args)
self.alignments = Alignments(self.args)
self.output_dir = self.faces.output_dir
self.export_face = True
def process(self):
""" Perform the extraction process """
print('Starting, this may take a while...')
Utils.set_verbosity(self.args.verbose)
if hasattr(self.args, 'processes') and self.args.processes > 1:
self.extract_multi_process()
else:
self.extract_single_process()
self.alignments.write_alignments(self.faces.faces_detected)
images, faces = Utils.finalize(self.images.images_found,
self.faces.num_faces_detected,
self.faces.verify_output)
self.images.images_found = images
self.faces.num_faces_detected = faces
def extract_single_process(self):
""" Run extraction in a single process """
for filename in tqdm(self.images.input_images, file=sys.stdout):
filename, faces = self.process_single_image(filename)
self.faces.faces_detected[os.path.basename(filename)] = faces
def extract_multi_process(self):
""" Run the extraction on the correct number of processes """
for filename, faces in tqdm(pool_process(self.process_single_image,
self.images.input_images,
processes=self.args.processes),
total=self.images.images_found,
file=sys.stdout):
self.faces.num_faces_detected += 1
self.faces.faces_detected[os.path.basename(filename)] = faces
def process_single_image(self, filename):
""" Detect faces in an image. Rotate the image the specified amount
until at least one face is found, or until image rotations are
depleted.
Once at least one face has been detected, pass to
process_single_face to process the individual faces """
retval = filename, list()
try:
image = Utils.cv2_read_write('read', filename)
for angle in self.images.rotation_angles:
currentimage = Utils.rotate_image_by_angle(image, angle)
faces = self.faces.get_faces(currentimage, angle)
process_faces = [(idx, face) for idx, face in faces]
if process_faces and angle != 0 and self.args.verbose:
print("found face(s) by rotating image {} degrees".format(angle))
if process_faces:
break
final_faces = [self.process_single_face(idx, face, filename, currentimage)
for idx, face in process_faces]
retval = filename, final_faces
except Exception as err:
if self.args.verbose:
print("Failed to extract from image: {}. Reason: {}".format(filename, err))
return retval
def process_single_face(self, idx, face, filename, image):
""" Perform processing on found faces """
output_file = self.output_dir / Path(filename).stem if self.export_face else None
self.faces.draw_landmarks_on_face(face, image)
resized_face, t_mat = self.faces.extractor.extract(image,
face,
256,
self.faces.align_eyes)
blurry_file = self.faces.detect_blurry_faces(face, t_mat, resized_face, filename)
output_file = blurry_file if blurry_file else output_file
if self.export_face:
filename = "{}_{}{}".format(str(output_file), str(idx), Path(filename).suffix)
Utils.cv2_read_write('write', filename, resized_face)
return {"r": face.r,
"x": face.x,
"w": face.w,
"y": face.y,
"h": face.h,
"landmarksXY": face.landmarksAsXY()}
class Convert():
""" The convert process. """
def __init__(self, arguments):
logger.debug("Initializing %s: (args: %s)", self.__class__.__name__,
arguments)
self.args = arguments
self.output_dir = get_folder(self.args.output_dir)
self.extractor = None
self.faces_count = 0
self.images = Images(self.args)
self.alignments = Alignments(self.args, False, self.images.is_video)
# Update Legacy alignments
Legacy(self.alignments, self.images.input_images,
arguments.input_aligned_dir)
self.post_process = PostProcess(arguments)
self.verify_output = False
self.opts = OptionalActions(self.args, self.images.input_images,
self.alignments)
logger.debug("Initialized %s", self.__class__.__name__)
def process(self):
""" Original & LowMem models go with converter
Note: GAN prediction outputs a mask + an image, while other
predicts only an image. """
Utils.set_verbosity(self.args.loglevel)
if not self.alignments.have_alignments_file:
self.load_extractor()
model = self.load_model()
converter = self.load_converter(model)
batch = BackgroundGenerator(self.prepare_images(), 1)
for item in batch.iterator():
self.convert(converter, item)
if self.extractor:
queue_manager.terminate_queues()
Utils.finalize(self.images.images_found, self.faces_count,
self.verify_output)
def load_extractor(self):
""" Set on the fly extraction """
logger.warning("No Alignments file found. Extracting on the fly.")
logger.warning(
"NB: This will use the inferior dlib-hog for extraction "
"and dlib pose predictor for landmarks. It is recommended "
"to perfom Extract first for superior results")
extract_args = {
"detector": "dlib-hog",
"aligner": "dlib",
"loglevel": self.args.loglevel
}
self.extractor = Extractor(None, extract_args)
self.extractor.launch_detector()
self.extractor.launch_aligner()
def load_model(self):
""" Load the model requested for conversion """
logger.debug("Loading Model")
model_dir = get_folder(self.args.model_dir)
model = PluginLoader.get_model(self.args.trainer)(model_dir,
self.args.gpus,
predict=True)
logger.debug("Loaded Model")
return model
def load_converter(self, model):
""" Load the requested converter for conversion """
conv = self.args.converter
converter = PluginLoader.get_converter(conv)(model.converter(
self.args.swap_model),
model=model,
arguments=self.args)
return converter
def prepare_images(self):
""" Prepare the images for conversion """
filename = ""
if self.extractor:
load_queue = queue_manager.get_queue("load")
for filename, image in tqdm(self.images.load(),
total=self.images.images_found,
file=sys.stdout):
if (self.args.discard_frames
and self.opts.check_skipframe(filename) == "discard"):
continue
frame = os.path.basename(filename)
if self.extractor:
detected_faces = self.detect_faces(load_queue, filename, image)
else:
detected_faces = self.alignments_faces(frame, image)
faces_count = len(detected_faces)
if faces_count != 0:
# Post processing requires a dict with "detected_faces" key
self.post_process.do_actions(
{"detected_faces": detected_faces})
self.faces_count += faces_count
if faces_count > 1:
self.verify_output = True
logger.verbose("Found more than one face in "
"an image! '%s'", frame)
yield filename, image, detected_faces
def detect_faces(self, load_queue, filename, image):
""" Extract the face from a frame (If alignments file not found) """
inp = {"filename": filename, "image": image}
load_queue.put(inp)
faces = next(self.extractor.detect_faces())
landmarks = faces["landmarks"]
detected_faces = faces["detected_faces"]
final_faces = list()
for idx, face in enumerate(detected_faces):
detected_face = DetectedFace()
detected_face.from_dlib_rect(face)
detected_face.landmarksXY = landmarks[idx]
final_faces.append(detected_face)
return final_faces
def alignments_faces(self, frame, image):
""" Get the face from alignments file """
if not self.check_alignments(frame):
return list()
faces = self.alignments.get_faces_in_frame(frame)
detected_faces = list()
for rawface in faces:
face = DetectedFace()
face.from_alignment(rawface, image=image)
detected_faces.append(face)
return detected_faces
def check_alignments(self, frame):
""" If we have no alignments for this image, skip it """
have_alignments = self.alignments.frame_exists(frame)
if not have_alignments:
tqdm.write("No alignment found for {}, " "skipping".format(frame))
return have_alignments
def convert(self, converter, item):
""" Apply the conversion transferring faces onto frames """
try:
filename, image, faces = item
skip = self.opts.check_skipframe(filename)
if not skip:
for face in faces:
image = converter.patch_image(image, face)
filename = str(self.output_dir / Path(filename).name)
if self.args.draw_transparent:
filename = "{}.png".format(os.path.splitext(filename)[0])
logger.trace("Set extension to png: `%s`", filename)
cv2.imwrite(filename, image) # pylint: disable=no-member
except Exception as err:
logger.error("Failed to convert image: '%s'. Reason: %s", filename,
err)
raise
class Convert(object):
""" The convert process. """
def __init__(self, arguments):
self.args = arguments
self.output_dir = get_folder(self.args.output_dir)
self.images = Images(self.args)
self.faces = Faces(self.args)
self.alignments = Alignments(self.args)
self.opts = OptionalActions(self.args, self.images.input_images)
def process(self):
""" Original & LowMem models go with Adjust or Masked converter
Note: GAN prediction outputs a mask + an image, while other
predicts only an image. """
Utils.set_verbosity(self.args.verbose)
if not self.alignments.have_alignments_file:
self.generate_alignments()
self.faces.faces_detected = self.alignments.read_alignments()
model = self.load_model()
converter = self.load_converter(model)
batch = BackgroundGenerator(self.prepare_images(), 1)
for item in batch.iterator():
self.convert(converter, item)
Utils.finalize(self.images.images_found,
self.faces.num_faces_detected,
self.faces.verify_output)
def generate_alignments(self):
""" Generate an alignments file if one does not already
exist. Does not save extracted faces """
print('Alignments file not found. Generating at default values...')
extract = Extract(self.args)
extract.export_face = False
extract.process()
def load_model(self):
""" Load the model requested for conversion """
model_name = self.args.trainer
model_dir = get_folder(self.args.model_dir)
num_gpus = self.args.gpus
model = PluginLoader.get_model(model_name)(model_dir, num_gpus)
if not model.load(self.args.swap_model):
print("Model Not Found! A valid model "
"must be provided to continue!")
exit(1)
return model
def load_converter(self, model):
""" Load the requested converter for conversion """
args = self.args
conv = args.converter
converter = PluginLoader.get_converter(conv)(
model.converter(False),
trainer=args.trainer,
blur_size=args.blur_size,
seamless_clone=args.seamless_clone,
sharpen_image=args.sharpen_image,
mask_type=args.mask_type,
erosion_kernel_size=args.erosion_kernel_size,
match_histogram=args.match_histogram,
smooth_mask=args.smooth_mask,
avg_color_adjust=args.avg_color_adjust)
return converter
def prepare_images(self):
""" Prepare the images for conversion """
filename = ""
for filename in tqdm(self.images.input_images, file=sys.stdout):
if not self.check_alignments(filename):
continue
image = Utils.cv2_read_write('read', filename)
faces = self.faces.get_faces_alignments(filename, image)
if not faces:
continue
yield filename, image, faces
def check_alignments(self, filename):
""" If we have no alignments for this image, skip it """
have_alignments = self.faces.have_face(filename)
if not have_alignments:
tqdm.write("No alignment found for {}, "
"skipping".format(os.path.basename(filename)))
return have_alignments
def convert(self, converter, item):
""" Apply the conversion transferring faces onto frames """
try:
filename, image, faces = item
skip = self.opts.check_skipframe(filename)
if not skip:
for idx, face in faces:
image = self.convert_one_face(converter,
(filename, image, idx, face))
if skip != "discard":
filename = str(self.output_dir / Path(filename).name)
Utils.cv2_read_write('write', filename, image)
except Exception as err:
print("Failed to convert image: {}. "
"Reason: {}".format(filename, err))
def convert_one_face(self, converter, imagevars):
""" Perform the conversion on the given frame for a single face """
filename, image, idx, face = imagevars
if self.opts.check_skipface(filename, idx):
return image
image = self.images.rotate_image(image, face.r)
# TODO: This switch between 64 and 128 is a hack for now.
# We should have a separate cli option for size
size = 128 if (self.args.trainer.strip().lower()
in ('gan128', 'originalhighres')) else 64
image = converter.patch_image(image,
face,
size)
image = self.images.rotate_image(image, face.r, reverse=True)
return image
class Extract():
""" The extract process. """
def __init__(self, arguments):
self.args = arguments
self.output_dir = get_folder(self.args.output_dir)
print("Output Directory: {}".format(self.args.output_dir))
self.images = Images(self.args)
self.alignments = Alignments(self.args, True)
self.plugins = Plugins(self.args)
self.post_process = PostProcess(arguments)
self.export_face = True
self.verify_output = False
self.save_interval = None
if hasattr(self.args, "save_interval"):
self.save_interval = self.args.save_interval
def process(self):
""" Perform the extraction process """
print('Starting, this may take a while...')
Utils.set_verbosity(self.args.verbose)
# queue_manager.debug_monitor(1)
self.threaded_io("load")
save_thread = self.threaded_io("save")
self.run_extraction(save_thread)
self.alignments.save()
Utils.finalize(self.images.images_found,
self.alignments.faces_count,
self.verify_output)
self.plugins.process_detect.join()
self.plugins.process_align.join()
def threaded_io(self, task, io_args=None):
""" Load images in a background thread """
io_args = tuple() if io_args is None else (io_args, )
if task == "load":
func = self.load_images
elif task == "save":
func = self.save_faces
elif task == "reload":
func = self.reload_images
io_thread = MultiThread(thread_count=1)
io_thread.in_thread(func, *io_args)
return io_thread
def load_images(self):
""" Load the images """
load_queue = queue_manager.get_queue("load")
for filename, image in self.images.load():
imagename = os.path.basename(filename)
if imagename in self.alignments.data.keys():
continue
load_queue.put((filename, image))
load_queue.put("EOF")
def reload_images(self, detected_faces):
""" Reload the images and pair to detected face """
load_queue = queue_manager.get_queue("detect")
for filename, image in self.images.load():
detect_item = detected_faces.pop(filename, None)
if not detect_item:
continue
detect_item["image"] = image
load_queue.put(detect_item)
load_queue.put("EOF")
def save_faces(self):
""" Save the generated faces """
if not self.export_face:
return
save_queue = queue_manager.get_queue("save")
while True:
item = save_queue.get()
if item == "EOF":
break
filename, output_file, resized_face, idx = item
out_filename = "{}_{}{}".format(str(output_file),
str(idx),
Path(filename).suffix)
# pylint: disable=no-member
cv2.imwrite(out_filename, resized_face)
def run_extraction(self, save_thread):
""" Run Face Detection """
to_process = self.process_item_count()
frame_no = 0
size = self.args.size if hasattr(self.args, "size") else 256
align_eyes = self.args.align_eyes if hasattr(self.args, "align_eyes") else False
if self.plugins.is_parallel:
self.plugins.launch_aligner()
self.plugins.launch_detector()
if not self.plugins.is_parallel:
self.run_detection(to_process)
self.plugins.launch_aligner()
for faces in tqdm(self.plugins.detect_faces(extract_pass="******"),
total=to_process,
file=sys.stdout,
desc="Extracting faces"):
exception = faces.get("exception", False)
if exception:
exit(1)
filename = faces["filename"]
faces["output_file"] = self.output_dir / Path(filename).stem
self.align_face(faces, align_eyes, size)
self.post_process.do_actions(faces)
faces_count = len(faces["detected_faces"])
if self.args.verbose and faces_count == 0:
print("Warning: No faces were detected in image: "
"{}".format(os.path.basename(filename)))
if not self.verify_output and faces_count > 1:
self.verify_output = True
self.process_faces(filename, faces)
frame_no += 1
if frame_no == self.save_interval:
self.alignments.save()
frame_no = 0
if self.export_face:
queue_manager.get_queue("save").put("EOF")
save_thread.join_threads()
def process_item_count(self):
""" Return the number of items to be processedd """
processed = sum(os.path.basename(frame) in self.alignments.data.keys()
for frame in self.images.input_images)
if processed != 0 and self.args.skip_existing:
print("Skipping {} previously extracted frames".format(processed))
if processed != 0 and self.args.skip_faces:
print("Skipping {} frames with detected faces".format(processed))
to_process = self.images.images_found - processed
if to_process == 0:
print("No frames to process. Exiting")
queue_manager.terminate_queues()
exit(0)
return to_process
def run_detection(self, to_process):
""" Run detection only """
self.plugins.launch_detector()
detected_faces = dict()
for detected in tqdm(self.plugins.detect_faces(extract_pass="******"),
total=to_process,
file=sys.stdout,
desc="Detecting faces"):
exception = detected.get("exception", False)
if exception:
break
del detected["image"]
filename = detected["filename"]
detected_faces[filename] = detected
self.threaded_io("reload", detected_faces)
@staticmethod
def align_face(faces, align_eyes, size, padding=48):
""" Align the detected face """
final_faces = list()
image = faces["image"]
landmarks = faces["landmarks"]
detected_faces = faces["detected_faces"]
for idx, face in enumerate(detected_faces):
detected_face = DetectedFace()
detected_face.from_dlib_rect(face, image)
detected_face.landmarksXY = landmarks[idx]
detected_face.frame_dims = image.shape[:2]
detected_face.load_aligned(image,
size=size,
padding=padding,
align_eyes=align_eyes)
final_faces.append(detected_face)
faces["detected_faces"] = final_faces
def process_faces(self, filename, faces):
""" Perform processing on found faces """
final_faces = list()
save_queue = queue_manager.get_queue("save")
filename = faces["filename"]
output_file = faces["output_file"]
for idx, face in enumerate(faces["detected_faces"]):
if self.export_face:
save_queue.put((filename,
output_file,
face.aligned_face,
idx))
final_faces.append(face.to_alignment())
self.alignments.data[os.path.basename(filename)] = final_faces
class Convert():
""" The convert process. """
def __init__(self, arguments):
logger.debug("Initializing %s: (args: %s)", self.__class__.__name__,
arguments)
self.args = arguments
self.output_dir = get_folder(self.args.output_dir)
self.extract_faces = False
self.faces_count = 0
self.images = Images(self.args)
self.alignments = Alignments(self.args, False, self.images.is_video)
# Update Legacy alignments
Legacy(self.alignments, self.images.input_images,
arguments.input_aligned_dir)
self.post_process = PostProcess(arguments)
self.verify_output = False
self.opts = OptionalActions(self.args, self.images.input_images,
self.alignments)
logger.debug("Initialized %s", self.__class__.__name__)
def process(self):
""" Original & LowMem models go with Adjust or Masked converter
Note: GAN prediction outputs a mask + an image, while other
predicts only an image. """
Utils.set_verbosity(self.args.loglevel)
if not self.alignments.have_alignments_file:
self.load_extractor()
model = self.load_model()
converter = self.load_converter(model)
batch = BackgroundGenerator(self.prepare_images(), 1)
for item in batch.iterator():
self.convert(converter, item)
if self.extract_faces:
queue_manager.terminate_queues()
Utils.finalize(self.images.images_found, self.faces_count,
self.verify_output)
def load_extractor(self):
""" Set on the fly extraction """
logger.warning("No Alignments file found. Extracting on the fly.")
logger.warning(
"NB: This will use the inferior dlib-hog for extraction "
"and dlib pose predictor for landmarks. It is recommended "
"to perfom Extract first for superior results")
for task in ("load", "detect", "align"):
queue_manager.add_queue(task, maxsize=0)
detector = PluginLoader.get_detector("dlib_hog")(
loglevel=self.args.loglevel)
aligner = PluginLoader.get_aligner("dlib")(loglevel=self.args.loglevel)
d_kwargs = {
"in_queue": queue_manager.get_queue("load"),
"out_queue": queue_manager.get_queue("detect")
}
a_kwargs = {
"in_queue": queue_manager.get_queue("detect"),
"out_queue": queue_manager.get_queue("align")
}
d_process = SpawnProcess(detector.run, **d_kwargs)
d_event = d_process.event
d_process.start()
a_process = SpawnProcess(aligner.run, **a_kwargs)
a_event = a_process.event
a_process.start()
d_event.wait(10)
if not d_event.is_set():
raise ValueError("Error inititalizing Detector")
a_event.wait(10)
if not a_event.is_set():
raise ValueError("Error inititalizing Aligner")
self.extract_faces = True
def load_model(self):
""" Load the model requested for conversion """
model_name = self.args.trainer
model_dir = get_folder(self.args.model_dir)
num_gpus = self.args.gpus
model = PluginLoader.get_model(model_name)(model_dir, num_gpus)
if not model.load(self.args.swap_model):
logger.error("Model Not Found! A valid model "
"must be provided to continue!")
exit(1)
return model
def load_converter(self, model):
""" Load the requested converter for conversion """
args = self.args
conv = args.converter
converter = PluginLoader.get_converter(conv)(
model.converter(False),
trainer=args.trainer,
blur_size=args.blur_size,
seamless_clone=args.seamless_clone,
sharpen_image=args.sharpen_image,
mask_type=args.mask_type,
erosion_kernel_size=args.erosion_kernel_size,
match_histogram=args.match_histogram,
smooth_mask=args.smooth_mask,
avg_color_adjust=args.avg_color_adjust,
draw_transparent=args.draw_transparent)
return converter
def prepare_images(self):
""" Prepare the images for conversion """
filename = ""
for filename, image in tqdm(self.images.load(),
total=self.images.images_found,
file=sys.stdout):
if (self.args.discard_frames
and self.opts.check_skipframe(filename) == "discard"):
continue
frame = os.path.basename(filename)
if self.extract_faces:
detected_faces = self.detect_faces(filename, image)
else:
detected_faces = self.alignments_faces(frame, image)
faces_count = len(detected_faces)
if faces_count != 0:
# Post processing requires a dict with "detected_faces" key
self.post_process.do_actions(
{"detected_faces": detected_faces})
self.faces_count += faces_count
if faces_count > 1:
self.verify_output = True
logger.verbose("Found more than one face in "
"an image! '%s'", frame)
yield filename, image, detected_faces
@staticmethod
def detect_faces(filename, image):
""" Extract the face from a frame (If not alignments file found) """
queue_manager.get_queue("load").put((filename, image))
item = queue_manager.get_queue("align").get()
detected_faces = item["detected_faces"]
return detected_faces
def alignments_faces(self, frame, image):
""" Get the face from alignments file """
if not self.check_alignments(frame):
return None
faces = self.alignments.get_faces_in_frame(frame)
detected_faces = list()
for rawface in faces:
face = DetectedFace()
face.from_alignment(rawface, image=image)
detected_faces.append(face)
return detected_faces
def check_alignments(self, frame):
""" If we have no alignments for this image, skip it """
have_alignments = self.alignments.frame_exists(frame)
if not have_alignments:
tqdm.write("No alignment found for {}, " "skipping".format(frame))
return have_alignments
def convert(self, converter, item):
""" Apply the conversion transferring faces onto frames """
try:
filename, image, faces = item
skip = self.opts.check_skipframe(filename)
if not skip:
for face in faces:
image = self.convert_one_face(converter, image, face)
filename = str(self.output_dir / Path(filename).name)
cv2.imwrite(filename, image) # pylint: disable=no-member
except Exception as err:
logger.error("Failed to convert image: '%s'. Reason: %s", filename,
err)
raise
def convert_one_face(self, converter, image, face):
""" Perform the conversion on the given frame for a single face """
# TODO: This switch between 64 and 128 is a hack for now.
# We should have a separate cli option for size
size = 128 if (self.args.trainer.strip().lower()
in ('gan128', 'originalhighres')) else 64
image = converter.patch_image(image, face, size)
return image
class Extract():
""" The extract process. """
def __init__(self, arguments):
logger.debug("Initializing %s: (args: %s", self.__class__.__name__, arguments)
self.args = arguments
Utils.set_verbosity(self.args.loglevel)
self.output_dir = get_folder(self.args.output_dir)
logger.info("Output Directory: %s", self.args.output_dir)
self.images = Images(self.args)
self.alignments = Alignments(self.args, True, self.images.is_video)
self.post_process = PostProcess(arguments)
self.extractor = Extractor(self.args.detector,
self.args.aligner,
self.args.loglevel,
self.args.multiprocess,
self.args.rotate_images,
self.args.min_size)
self.save_queue = queue_manager.get_queue("extract_save")
self.verify_output = False
self.save_interval = None
if hasattr(self.args, "save_interval"):
self.save_interval = self.args.save_interval
logger.debug("Initialized %s", self.__class__.__name__)
@property
def skip_num(self):
""" Number of frames to skip if extract_every_n is passed """
return self.args.extract_every_n if hasattr(self.args, "extract_every_n") else 1
def process(self):
""" Perform the extraction process """
logger.info('Starting, this may take a while...')
# queue_manager.debug_monitor(3)
self.threaded_io("load")
save_thread = self.threaded_io("save")
self.run_extraction()
save_thread.join()
self.alignments.save()
Utils.finalize(self.images.images_found // self.skip_num,
self.alignments.faces_count,
self.verify_output)
def threaded_io(self, task, io_args=None):
""" Perform I/O task in a background thread """
logger.debug("Threading task: (Task: '%s')", task)
io_args = tuple() if io_args is None else (io_args, )
if task == "load":
func = self.load_images
elif task == "save":
func = self.save_faces
elif task == "reload":
func = self.reload_images
io_thread = MultiThread(func, *io_args, thread_count=1)
io_thread.start()
return io_thread
def load_images(self):
""" Load the images """
logger.debug("Load Images: Start")
load_queue = self.extractor.input_queue
idx = 0
for filename, image in self.images.load():
idx += 1
if load_queue.shutdown.is_set():
logger.debug("Load Queue: Stop signal received. Terminating")
break
if idx % self.skip_num != 0:
logger.trace("Skipping image '%s' due to extract_every_n = %s",
filename, self.skip_num)
continue
if image is None or not image.any():
logger.warning("Unable to open image. Skipping: '%s'", filename)
continue
imagename = os.path.basename(filename)
if imagename in self.alignments.data.keys():
logger.trace("Skipping image: '%s'", filename)
continue
item = {"filename": filename,
"image": image}
load_queue.put(item)
load_queue.put("EOF")
logger.debug("Load Images: Complete")
def reload_images(self, detected_faces):
""" Reload the images and pair to detected face """
logger.debug("Reload Images: Start. Detected Faces Count: %s", len(detected_faces))
load_queue = self.extractor.input_queue
for filename, image in self.images.load():
if load_queue.shutdown.is_set():
logger.debug("Reload Queue: Stop signal received. Terminating")
break
logger.trace("Reloading image: '%s'", filename)
detect_item = detected_faces.pop(filename, None)
if not detect_item:
logger.warning("Couldn't find faces for: %s", filename)
continue
detect_item["image"] = image
load_queue.put(detect_item)
load_queue.put("EOF")
logger.debug("Reload Images: Complete")
def save_faces(self):
""" Save the generated faces """
logger.debug("Save Faces: Start")
while True:
if self.save_queue.shutdown.is_set():
logger.debug("Save Queue: Stop signal received. Terminating")
break
item = self.save_queue.get()
logger.trace(item)
if item == "EOF":
break
filename, face = item
logger.trace("Saving face: '%s'", filename)
try:
with open(filename, "wb") as out_file:
out_file.write(face)
except Exception as err: # pylint: disable=broad-except
logger.error("Failed to save image '%s'. Original Error: %s", filename, err)
continue
logger.debug("Save Faces: Complete")
def process_item_count(self):
""" Return the number of items to be processedd """
processed = sum(os.path.basename(frame) in self.alignments.data.keys()
for frame in self.images.input_images)
logger.debug("Items already processed: %s", processed)
if processed != 0 and self.args.skip_existing:
logger.info("Skipping previously extracted frames: %s", processed)
if processed != 0 and self.args.skip_faces:
logger.info("Skipping frames with detected faces: %s", processed)
to_process = (self.images.images_found - processed) // self.skip_num
logger.debug("Items to be Processed: %s", to_process)
if to_process == 0:
logger.error("No frames to process. Exiting")
queue_manager.terminate_queues()
exit(0)
return to_process
def run_extraction(self):
""" Run Face Detection """
to_process = self.process_item_count()
size = self.args.size if hasattr(self.args, "size") else 256
align_eyes = self.args.align_eyes if hasattr(self.args, "align_eyes") else False
exception = False
for phase in range(self.extractor.passes):
if exception:
break
is_final = self.extractor.final_pass
detected_faces = dict()
self.extractor.launch()
for idx, faces in enumerate(tqdm(self.extractor.detected_faces(),
total=to_process,
file=sys.stdout,
desc="Running pass {} of {}: {}".format(
phase + 1,
self.extractor.passes,
self.extractor.phase.title()))):
exception = faces.get("exception", False)
if exception:
break
filename = faces["filename"]
if self.extractor.final_pass:
self.output_processing(faces, align_eyes, size, filename)
self.output_faces(filename, faces)
if self.save_interval and idx + 1 % self.save_interval == 0:
self.alignments.save()
else:
del faces["image"]
detected_faces[filename] = faces
if is_final:
logger.debug("Putting EOF to save")
self.save_queue.put("EOF")
else:
logger.debug("Reloading images")
self.threaded_io("reload", detected_faces)
def output_processing(self, faces, align_eyes, size, filename):
""" Prepare faces for output """
self.align_face(faces, align_eyes, size, filename)
self.post_process.do_actions(faces)
faces_count = len(faces["detected_faces"])
if faces_count == 0:
logger.verbose("No faces were detected in image: %s",
os.path.basename(filename))
if not self.verify_output and faces_count > 1:
self.verify_output = True
def align_face(self, faces, align_eyes, size, filename):
""" Align the detected face and add the destination file path """
final_faces = list()
image = faces["image"]
landmarks = faces["landmarks"]
detected_faces = faces["detected_faces"]
for idx, face in enumerate(detected_faces):
detected_face = DetectedFace()
detected_face.from_bounding_box(face, image)
detected_face.landmarksXY = landmarks[idx]
detected_face.load_aligned(image, size=size, align_eyes=align_eyes)
final_faces.append({"file_location": self.output_dir / Path(filename).stem,
"face": detected_face})
faces["detected_faces"] = final_faces
def output_faces(self, filename, faces):
""" Output faces to save thread """
final_faces = list()
for idx, detected_face in enumerate(faces["detected_faces"]):
output_file = detected_face["file_location"]
extension = Path(filename).suffix
out_filename = "{}_{}{}".format(str(output_file), str(idx), extension)
face = detected_face["face"]
resized_face = face.aligned_face
face.hash, img = hash_encode_image(resized_face, extension)
self.save_queue.put((out_filename, img))
final_faces.append(face.to_alignment())
self.alignments.data[os.path.basename(filename)] = final_faces
class Extract():
""" The extract process. """
def __init__(self, arguments):
logger.debug("Initializing %s: (args: %s", self.__class__.__name__,
arguments)
self.args = arguments
self.output_dir = get_folder(self.args.output_dir)
logger.info("Output Directory: %s", self.args.output_dir)
self.images = Images(self.args)
self.alignments = Alignments(self.args, True)
self.plugins = Plugins(self.args)
self.post_process = PostProcess(arguments)
self.export_face = True
self.verify_output = False
self.save_interval = None
if hasattr(self.args, "save_interval"):
self.save_interval = self.args.save_interval
logger.debug("Initialized %s", self.__class__.__name__)
def process(self):
""" Perform the extraction process """
logger.info('Starting, this may take a while...')
Utils.set_verbosity()
# queue_manager.debug_monitor(1)
self.threaded_io("load")
save_thread = self.threaded_io("save")
self.run_extraction(save_thread)
self.alignments.save()
Utils.finalize(self.images.images_found, self.alignments.faces_count,
self.verify_output)
def threaded_io(self, task, io_args=None):
""" Load images in a background thread """
logger.debug("Threading task: (Task: '%s')", task)
io_args = tuple() if io_args is None else (io_args, )
if task == "load":
func = self.load_images
elif task == "save":
func = self.save_faces
elif task == "reload":
func = self.reload_images
io_thread = MultiThread(func, *io_args, thread_count=1)
io_thread.start()
return io_thread
def load_images(self):
""" Load the images """
logger.debug("Load Images: Start")
load_queue = queue_manager.get_queue("load")
for filename, image in self.images.load():
if load_queue.shutdown.is_set():
logger.debug("Load Queue: Stop signal received. Terminating")
break
imagename = os.path.basename(filename)
if imagename in self.alignments.data.keys():
logger.trace("Skipping image: '%s'", filename)
continue
item = {"filename": filename, "image": image}
load_queue.put(item)
load_queue.put("EOF")
logger.debug("Load Images: Complete")
def reload_images(self, detected_faces):
""" Reload the images and pair to detected face """
logger.debug("Reload Images: Start. Detected Faces Count: %s",
len(detected_faces))
load_queue = queue_manager.get_queue("detect")
for filename, image in self.images.load():
if load_queue.shutdown.is_set():
logger.debug("Reload Queue: Stop signal received. Terminating")
break
logger.trace("Reloading image: '%s'", filename)
detect_item = detected_faces.pop(filename, None)
if not detect_item:
logger.warning("Couldn't find faces for: %s", filename)
continue
detect_item["image"] = image
load_queue.put(detect_item)
load_queue.put("EOF")
logger.debug("Reload Images: Complete")
def save_faces(self):
""" Save the generated faces """
logger.debug("Save Faces: Start")
if not self.export_face:
logger.debug("Not exporting faces")
logger.debug("Save Faces: Complete")
return
save_queue = queue_manager.get_queue("save")
while True:
if save_queue.shutdown.is_set():
logger.debug("Save Queue: Stop signal received. Terminating")
break
item = save_queue.get()
if item == "EOF":
break
filename, output_file, resized_face, idx = item
out_filename = "{}_{}{}".format(str(output_file), str(idx),
Path(filename).suffix)
logger.trace("Saving face: '%s'", out_filename)
try:
cv2.imwrite(out_filename, resized_face) # pylint: disable=no-member
except Exception as err: # pylint: disable=broad-except
logger.error("Failed to save image '%s'. Original Error: %s",
out_filename, err)
continue
logger.debug("Save Faces: Complete")
def run_extraction(self, save_thread):
""" Run Face Detection """
save_queue = queue_manager.get_queue("save")
to_process = self.process_item_count()
frame_no = 0
size = self.args.size if hasattr(self.args, "size") else 256
align_eyes = self.args.align_eyes if hasattr(self.args,
"align_eyes") else False
if self.plugins.is_parallel:
logger.debug("Using parallel processing")
self.plugins.launch_aligner()
self.plugins.launch_detector()
if not self.plugins.is_parallel:
logger.debug("Using serial processing")
self.run_detection(to_process)
self.plugins.launch_aligner()
for faces in tqdm(self.plugins.detect_faces(extract_pass="******"),
total=to_process,
file=sys.stdout,
desc="Extracting faces"):
filename = faces["filename"]
self.align_face(faces, align_eyes, size, filename)
self.post_process.do_actions(faces)
faces_count = len(faces["detected_faces"])
if faces_count == 0:
logger.verbose("No faces were detected in image: %s",
os.path.basename(filename))
if not self.verify_output and faces_count > 1:
self.verify_output = True
self.process_faces(filename, faces, save_queue)
frame_no += 1
if frame_no == self.save_interval:
self.alignments.save()
frame_no = 0
if self.export_face:
save_queue.put("EOF")
save_thread.join()
def process_item_count(self):
""" Return the number of items to be processedd """
processed = sum(
os.path.basename(frame) in self.alignments.data.keys()
for frame in self.images.input_images)
logger.debug("Items already processed: %s", processed)
if processed != 0 and self.args.skip_existing:
logger.info("Skipping previously extracted frames: %s", processed)
if processed != 0 and self.args.skip_faces:
logger.info("Skipping frames with detected faces: %s", processed)
to_process = self.images.images_found - processed
logger.debug("Items to be Processed: %s", to_process)
if to_process == 0:
logger.error("No frames to process. Exiting")
queue_manager.terminate_queues()
exit(0)
return to_process
def run_detection(self, to_process):
""" Run detection only """
self.plugins.launch_detector()
detected_faces = dict()
for detected in tqdm(self.plugins.detect_faces(extract_pass="******"),
total=to_process,
file=sys.stdout,
desc="Detecting faces"):
exception = detected.get("exception", False)
if exception:
break
del detected["image"]
filename = detected["filename"]
detected_faces[filename] = detected
self.threaded_io("reload", detected_faces)
def align_face(self, faces, align_eyes, size, filename, padding=48):
""" Align the detected face and add the destination file path """
final_faces = list()
image = faces["image"]
landmarks = faces["landmarks"]
detected_faces = faces["detected_faces"]
for idx, face in enumerate(detected_faces):
detected_face = DetectedFace()
detected_face.from_dlib_rect(face, image)
detected_face.landmarksXY = landmarks[idx]
detected_face.frame_dims = image.shape[:2]
detected_face.load_aligned(image,
size=size,
padding=padding,
align_eyes=align_eyes)
final_faces.append({
"file_location":
self.output_dir / Path(filename).stem,
"face":
detected_face
})
faces["detected_faces"] = final_faces
def process_faces(self, filename, faces, save_queue):
""" Perform processing on found faces """
final_faces = list()
filename = faces["filename"]
for idx, detected_face in enumerate(faces["detected_faces"]):
if self.export_face:
save_queue.put((filename, detected_face["file_location"],
detected_face["face"].aligned_face, idx))
final_faces.append(detected_face["face"].to_alignment())
self.alignments.data[os.path.basename(filename)] = final_faces
class Extract(object):
""" The extract process. """
def __init__(self, arguments):
self.args = arguments
self.images = Images(self.args)
self.faces = Faces(self.args)
self.alignments = Alignments(self.args)
self.output_dir = self.faces.output_dir
self.export_face = True
self.save_interval = self.args.save_interval if hasattr(
self.args, "save_interval") else None
def process(self):
""" Perform the extraction process """
print('Starting, this may take a while...')
Utils.set_verbosity(self.args.verbose)
if (hasattr(self.args, 'multiprocess') and self.args.multiprocess
and GPUStats().device_count == 0):
# TODO Checking that there is no available GPU is not
# necessarily an indicator of whether the user is actually
# using the CPU. Maybe look to implement further checks on
# dlib/tensorflow compilations
self.extract_multi_process()
else:
self.extract_single_process()
self.write_alignments()
images, faces = Utils.finalize(self.images.images_found,
self.faces.num_faces_detected,
self.faces.verify_output)
self.images.images_found = images
self.faces.num_faces_detected = faces
def write_alignments(self):
self.alignments.write_alignments(self.faces.faces_detected)
def extract_single_process(self):
""" Run extraction in a single process """
frame_no = 0
for filename in tqdm(self.images.input_images, file=sys.stdout):
filename, faces = self.process_single_image(filename)
self.faces.faces_detected[os.path.basename(filename)] = faces
frame_no += 1
if frame_no == self.save_interval:
self.write_alignments()
frame_no = 0
def extract_multi_process(self):
""" Run the extraction on the correct number of processes """
frame_no = 0
for filename, faces in tqdm(pool_process(self.process_single_image,
self.images.input_images),
total=self.images.images_found,
file=sys.stdout):
self.faces.num_faces_detected += 1
self.faces.faces_detected[os.path.basename(filename)] = faces
frame_no += 1
if frame_no == self.save_interval:
self.write_alignments()
frame_no = 0
def process_single_image(self, filename):
""" Detect faces in an image. Rotate the image the specified amount
until at least one face is found, or until image rotations are
depleted.
Once at least one face has been detected, pass to
process_single_face to process the individual faces """
retval = filename, list()
try:
image = Utils.cv2_read_write('read', filename)
for angle in self.images.rotation_angles:
currentimage = Utils.rotate_image_by_angle(image, angle)
faces = self.faces.get_faces(currentimage, angle)
process_faces = [(idx, face) for idx, face in faces]
if process_faces and angle != 0 and self.args.verbose:
print("found face(s) by rotating image "
"{} degrees".format(angle))
if process_faces:
break
final_faces = [
self.process_single_face(idx, face, filename, currentimage)
for idx, face in process_faces
]
retval = filename, final_faces
except Exception as err:
if self.args.verbose:
print("Failed to extract from image: "
"{}. Reason: {}".format(filename, err))
return retval
def process_single_face(self, idx, face, filename, image):
""" Perform processing on found faces """
output_file = self.output_dir / Path(
filename).stem if self.export_face else None
self.faces.draw_landmarks_on_face(face, image)
resized_face, t_mat = self.faces.extractor.extract(
image, face, 256, self.faces.align_eyes)
blurry_file = self.faces.detect_blurry_faces(face, t_mat, resized_face,
filename)
output_file = blurry_file if blurry_file else output_file
if self.export_face:
filename = "{}_{}{}".format(str(output_file), str(idx),
Path(filename).suffix)
Utils.cv2_read_write('write', filename, resized_face)
return {
"r": face.r,
"x": face.x,
"w": face.w,
"y": face.y,
"h": face.h,
"landmarksXY": face.landmarks_as_xy()
}
class Extract():
""" The extract process. """
def __init__(self, arguments):
logger.debug("Initializing %s: (args: %s", self.__class__.__name__,
arguments)
self.args = arguments
Utils.set_verbosity(self.args.loglevel)
self.output_dir = get_folder(self.args.output_dir)
logger.info("Output Directory: %s", self.args.output_dir)
self.images = Images(self.args)
self.alignments = Alignments(self.args, True, self.images.is_video)
self.post_process = PostProcess(arguments)
self.extractor = Extractor(self.args.detector, self.args.aligner,
self.args.loglevel, self.args.multiprocess,
self.args.rotate_images, self.args.min_size)
self.save_queue = queue_manager.get_queue("extract_save")
self.verify_output = False
self.save_interval = None
if hasattr(self.args, "save_interval"):
self.save_interval = self.args.save_interval
logger.debug("Initialized %s", self.__class__.__name__)
@property
def skip_num(self):
""" Number of frames to skip if extract_every_n is passed """
return self.args.extract_every_n if hasattr(self.args,
"extract_every_n") else 1
def process(self):
""" Perform the extraction process """
logger.info('Starting, this may take a while...')
# queue_manager.debug_monitor(3)
self.threaded_io("load")
save_thread = self.threaded_io("save")
self.run_extraction()
save_thread.join()
self.alignments.save()
Utils.finalize(self.images.images_found // self.skip_num,
self.alignments.faces_count, self.verify_output)
def threaded_io(self, task, io_args=None):
""" Perform I/O task in a background thread """
logger.debug("Threading task: (Task: '%s')", task)
io_args = tuple() if io_args is None else (io_args, )
if task == "load":
func = self.load_images
elif task == "save":
func = self.save_faces
elif task == "reload":
func = self.reload_images
io_thread = MultiThread(func, *io_args, thread_count=1)
io_thread.start()
return io_thread
def load_images(self):
""" Load the images """
logger.debug("Load Images: Start")
load_queue = self.extractor.input_queue
idx = 0
for filename, image in self.images.load():
idx += 1
if load_queue.shutdown.is_set():
logger.debug("Load Queue: Stop signal received. Terminating")
break
if idx % self.skip_num != 0:
logger.trace("Skipping image '%s' due to extract_every_n = %s",
filename, self.skip_num)
continue
if image is None or not image.any():
logger.warning("Unable to open image. Skipping: '%s'",
filename)
continue
imagename = os.path.basename(filename)
if imagename in self.alignments.data.keys():
logger.trace("Skipping image: '%s'", filename)
continue
item = {"filename": filename, "image": image}
load_queue.put(item)
load_queue.put("EOF")
logger.debug("Load Images: Complete")
def reload_images(self, detected_faces):
""" Reload the images and pair to detected face """
logger.debug("Reload Images: Start. Detected Faces Count: %s",
len(detected_faces))
load_queue = self.extractor.input_queue
for filename, image in self.images.load():
if load_queue.shutdown.is_set():
logger.debug("Reload Queue: Stop signal received. Terminating")
break
logger.trace("Reloading image: '%s'", filename)
detect_item = detected_faces.pop(filename, None)
if not detect_item:
logger.warning("Couldn't find faces for: %s", filename)
continue
detect_item["image"] = image
load_queue.put(detect_item)
load_queue.put("EOF")
logger.debug("Reload Images: Complete")
def save_faces(self):
""" Save the generated faces """
logger.debug("Save Faces: Start")
while True:
if self.save_queue.shutdown.is_set():
logger.debug("Save Queue: Stop signal received. Terminating")
break
item = self.save_queue.get()
logger.trace(item)
if item == "EOF":
break
filename, face = item
logger.trace("Saving face: '%s'", filename)
try:
with open(filename, "wb") as out_file:
out_file.write(face)
except Exception as err: # pylint: disable=broad-except
logger.error("Failed to save image '%s'. Original Error: %s",
filename, err)
continue
logger.debug("Save Faces: Complete")
def process_item_count(self):
""" Return the number of items to be processedd """
processed = sum(
os.path.basename(frame) in self.alignments.data.keys()
for frame in self.images.input_images)
logger.debug("Items already processed: %s", processed)
if processed != 0 and self.args.skip_existing:
logger.info("Skipping previously extracted frames: %s", processed)
if processed != 0 and self.args.skip_faces:
logger.info("Skipping frames with detected faces: %s", processed)
to_process = (self.images.images_found - processed) // self.skip_num
logger.debug("Items to be Processed: %s", to_process)
if to_process == 0:
logger.error("No frames to process. Exiting")
queue_manager.terminate_queues()
exit(0)
return to_process
def run_extraction(self):
""" Run Face Detection """
to_process = self.process_item_count()
size = self.args.size if hasattr(self.args, "size") else 256
align_eyes = self.args.align_eyes if hasattr(self.args,
"align_eyes") else False
exception = False
for phase in range(self.extractor.passes):
if exception:
break
is_final = self.extractor.final_pass
detected_faces = dict()
self.extractor.launch()
for idx, faces in enumerate(
tqdm(self.extractor.detected_faces(),
total=to_process,
file=sys.stdout,
desc="Running pass {} of {}: {}".format(
phase + 1, self.extractor.passes,
self.extractor.phase.title()))):
exception = faces.get("exception", False)
if exception:
break
filename = faces["filename"]
if self.extractor.final_pass:
self.output_processing(faces, align_eyes, size, filename)
self.output_faces(filename, faces)
if self.save_interval and (idx +
1) % self.save_interval == 0:
self.alignments.save()
else:
del faces["image"]
detected_faces[filename] = faces
if is_final:
logger.debug("Putting EOF to save")
self.save_queue.put("EOF")
else:
logger.debug("Reloading images")
self.threaded_io("reload", detected_faces)
def output_processing(self, faces, align_eyes, size, filename):
""" Prepare faces for output """
self.align_face(faces, align_eyes, size, filename)
self.post_process.do_actions(faces)
faces_count = len(faces["detected_faces"])
if faces_count == 0:
logger.verbose("No faces were detected in image: %s",
os.path.basename(filename))
if not self.verify_output and faces_count > 1:
self.verify_output = True
def align_face(self, faces, align_eyes, size, filename):
""" Align the detected face and add the destination file path """
final_faces = list()
image = faces["image"]
landmarks = faces["landmarks"]
detected_faces = faces["detected_faces"]
for idx, face in enumerate(detected_faces):
detected_face = DetectedFace()
detected_face.from_bounding_box(face, image)
detected_face.landmarksXY = landmarks[idx]
detected_face.load_aligned(image, size=size, align_eyes=align_eyes)
final_faces.append({
"file_location":
self.output_dir / Path(filename).stem,
"face":
detected_face
})
faces["detected_faces"] = final_faces
def output_faces(self, filename, faces):
""" Output faces to save thread """
final_faces = list()
for idx, detected_face in enumerate(faces["detected_faces"]):
output_file = detected_face["file_location"]
extension = Path(filename).suffix
out_filename = "{}_{}{}".format(str(output_file), str(idx),
extension)
face = detected_face["face"]
resized_face = face.aligned_face
face.hash, img = hash_encode_image(resized_face, extension)
self.save_queue.put((out_filename, img))
final_faces.append(face.to_alignment())
self.alignments.data[os.path.basename(filename)] = final_faces
class Convert(object):
""" The convert process. """
def __init__(self, arguments):
self.args = arguments
self.output_dir = get_folder(self.args.output_dir)
self.images = Images(self.args)
self.faces = Faces(self.args)
self.alignments = Alignments(self.args)
self.opts = OptionalActions(self.args, self.images.input_images)
def process(self):
""" Original & LowMem models go with Adjust or Masked converter
Note: GAN prediction outputs a mask + an image, while other
predicts only an image. """
Utils.set_verbosity(self.args.verbose)
if not self.alignments.have_alignments_file:
self.generate_alignments()
self.faces.faces_detected = self.alignments.read_alignments()
model = self.load_model()
converter = self.load_converter(model)
batch = BackgroundGenerator(self.prepare_images(), 1)
for item in batch.iterator():
self.convert(converter, item)
Utils.finalize(self.images.images_found, self.faces.num_faces_detected,
self.faces.verify_output)
def generate_alignments(self):
""" Generate an alignments file if one does not already
exist. Does not save extracted faces """
print('Alignments file not found. Generating at default values...')
extract = Extract(self.args)
extract.export_face = False
extract.process()
def load_model(self):
""" Load the model requested for conversion """
model_name = self.args.trainer
model_dir = get_folder(self.args.model_dir)
num_gpus = self.args.gpus
model = PluginLoader.get_model(model_name)(model_dir, num_gpus)
if not model.load(self.args.swap_model):
print(
"Model Not Found! A valid model must be provided to continue!")
exit(1)
return model
def load_converter(self, model):
""" Load the requested converter for conversion """
args = self.args
conv = args.converter
converter = PluginLoader.get_converter(conv)(
model.converter(False),
trainer=args.trainer,
blur_size=args.blur_size,
seamless_clone=args.seamless_clone,
sharpen_image=args.sharpen_image,
mask_type=args.mask_type,
erosion_kernel_size=args.erosion_kernel_size,
match_histogram=args.match_histogram,
smooth_mask=args.smooth_mask,
avg_color_adjust=args.avg_color_adjust)
return converter
def prepare_images(self):
""" Prepare the images for conversion """
filename = ""
for filename in tqdm(self.images.input_images, file=sys.stdout):
if not self.check_alignments(filename):
continue
image = Utils.cv2_read_write('read', filename)
faces = self.faces.get_faces_alignments(filename, image)
if not faces:
continue
yield filename, image, faces
def check_alignments(self, filename):
""" If we have no alignments for this image, skip it """
have_alignments = self.faces.have_face(filename)
if not have_alignments:
tqdm.write("No alignment found for {}, skipping".format(
os.path.basename(filename)))
return have_alignments
def convert(self, converter, item):
""" Apply the conversion transferring faces onto frames """
try:
filename, image, faces = item
skip = self.opts.check_skipframe(filename)
if not skip:
for idx, face in faces:
image = self.convert_one_face(converter,
(filename, image, idx, face))
if skip != "discard":
filename = str(self.output_dir / Path(filename).name)
Utils.cv2_read_write('write', filename, image)
except Exception as err:
print("Failed to convert image: {}. Reason: {}".format(
filename, err))
def convert_one_face(self, converter, imagevars):
""" Perform the conversion on the given frame for a single face """
filename, image, idx, face = imagevars
if self.opts.check_skipface(filename, idx):
return image
image = self.images.rotate_image(image, face.r)
# TODO: This switch between 64 and 128 is a hack for now.
# We should have a separate cli option for size
image = converter.patch_image(
image, face, 64 if "128" not in self.args.trainer else 128)
image = self.images.rotate_image(image, face.r, reverse=True)
return image
class Convert(): # pylint:disable=too-few-public-methods
""" The Faceswap Face Conversion Process.
The conversion process is responsible for swapping the faces on source frames with the output
from a trained model.
It leverages a series of user selected post-processing plugins, executed from
:class:`lib.convert.Converter`.
The convert process is self contained and should not be referenced by any other scripts, so it
contains no public properties.
Parameters
----------
arguments: :class:`argparse.Namespace`
The arguments to be passed to the convert process as generated from Faceswap's command
line arguments
"""
def __init__(self, arguments):
logger.debug("Initializing %s: (args: %s)", self.__class__.__name__,
arguments)
self._args = arguments
self._patch_threads = None
self._images = ImagesLoader(self._args.input_dir, fast_count=True)
self._alignments = Alignments(self._args, False, self._images.is_video)
if self._alignments.version == 1.0:
logger.error(
"The alignments file format has been updated since the given alignments "
"file was generated. You need to update the file to proceed.")
logger.error(
"To do this run the 'Alignments Tool' > 'Extract' Job.")
sys.exit(1)
self._opts = OptionalActions(self._args, self._images.file_list,
self._alignments)
self._add_queues()
self._disk_io = DiskIO(self._alignments, self._images, arguments)
self._predictor = Predict(self._disk_io.load_queue, self._queue_size,
arguments)
self._validate()
get_folder(self._args.output_dir)
configfile = self._args.configfile if hasattr(self._args,
"configfile") else None
self._converter = Converter(self._predictor.output_size,
self._predictor.coverage_ratio,
self._predictor.centering,
self._disk_io.draw_transparent,
self._disk_io.pre_encode,
arguments,
configfile=configfile)
logger.debug("Initialized %s", self.__class__.__name__)
@property
def _queue_size(self):
""" int: Size of the converter queues. 16 for single process otherwise 32 """
if self._args.singleprocess:
retval = 16
else:
retval = 32
logger.debug(retval)
return retval
@property
def _pool_processes(self):
""" int: The number of threads to run in parallel. Based on user options and number of
available processors. """
if self._args.singleprocess:
retval = 1
elif self._args.jobs > 0:
retval = min(self._args.jobs, total_cpus(), self._images.count)
else:
retval = min(total_cpus(), self._images.count)
retval = 1 if retval == 0 else retval
logger.debug(retval)
return retval
def _validate(self):
""" Validate the Command Line Options.
Ensure that certain cli selections are valid and won't result in an error. Checks:
* If frames have been passed in with video output, ensure user supplies reference
video.
* If "on-the-fly" and an NN mask is selected, output warning and switch to 'extended'
* If a mask-type is selected, ensure it exists in the alignments file.
* If a predicted mask-type is selected, ensure model has been trained with a mask
otherwise attempt to select first available masks, otherwise raise error.
Raises
------
FaceswapError
If an invalid selection has been found.
"""
if (self._args.writer == "ffmpeg" and not self._images.is_video
and self._args.reference_video is None):
raise FaceswapError(
"Output as video selected, but using frames as input. You must "
"provide a reference video ('-ref', '--reference-video').")
if (self._args.on_the_fly and self._args.mask_type
not in ("none", "extended", "components")):
logger.warning(
"You have selected an incompatible mask type ('%s') for On-The-Fly "
"conversion. Switching to 'extended'", self._args.mask_type)
self._args.mask_type = "extended"
if (not self._args.on_the_fly
and self._args.mask_type not in ("none", "predicted")
and not self._alignments.mask_is_valid(self._args.mask_type)):
msg = (
"You have selected the Mask Type `{}` but at least one face does not have this "
"mask stored in the Alignments File.\nYou should generate the required masks "
"with the Mask Tool or set the Mask Type option to an existing Mask Type.\nA "
"summary of existing masks is as follows:\nTotal faces: {}, Masks: "
"{}".format(self._args.mask_type, self._alignments.faces_count,
self._alignments.mask_summary))
raise FaceswapError(msg)
if self._args.mask_type == "predicted" and not self._predictor.has_predicted_mask:
available_masks = [
k for k, v in self._alignments.mask_summary.items()
if k != "none" and v == self._alignments.faces_count
]
if not available_masks:
msg = (
"Predicted Mask selected, but the model was not trained with a mask and no "
"masks are stored in the Alignments File.\nYou should generate the "
"required masks with the Mask Tool or set the Mask Type to `none`."
)
raise FaceswapError(msg)
mask_type = available_masks[0]
logger.warning(
"Predicted Mask selected, but the model was not trained with a "
"mask. Selecting first available mask: '%s'", mask_type)
self._args.mask_type = mask_type
def _add_queues(self):
""" Add the queues for in, patch and out. """
logger.debug("Adding queues. Queue size: %s", self._queue_size)
for qname in ("convert_in", "convert_out", "patch"):
queue_manager.add_queue(qname, self._queue_size)
def process(self):
""" The entry point for triggering the Conversion Process.
Should only be called from :class:`lib.cli.launcher.ScriptExecutor`
"""
logger.debug("Starting Conversion")
# queue_manager.debug_monitor(5)
try:
self._convert_images()
self._disk_io.save_thread.join()
queue_manager.terminate_queues()
finalize(self._images.count, self._predictor.faces_count,
self._predictor.verify_output)
logger.debug("Completed Conversion")
except MemoryError as err:
msg = (
"Faceswap ran out of RAM running convert. Conversion is very system RAM "
"heavy, so this can happen in certain circumstances when you have a lot of "
"cpus but not enough RAM to support them all."
"\nYou should lower the number of processes in use by either setting the "
"'singleprocess' flag (-sp) or lowering the number of parallel jobs (-j)."
)
raise FaceswapError(msg) from err
def _convert_images(self):
""" Start the multi-threaded patching process, monitor all threads for errors and join on
completion. """
logger.debug("Converting images")
save_queue = queue_manager.get_queue("convert_out")
patch_queue = queue_manager.get_queue("patch")
self._patch_threads = MultiThread(self._converter.process,
patch_queue,
save_queue,
thread_count=self._pool_processes,
name="patch")
self._patch_threads.start()
while True:
self._check_thread_error()
if self._disk_io.completion_event.is_set():
logger.debug("DiskIO completion event set. Joining Pool")
break
if self._patch_threads.completed():
logger.debug("All patch threads completed")
break
sleep(1)
self._patch_threads.join()
logger.debug("Putting EOF")
save_queue.put("EOF")
logger.debug("Converted images")
def _check_thread_error(self):
""" Monitor all running threads for errors, and raise accordingly. """
for thread in (self._predictor.thread, self._disk_io.load_thread,
self._disk_io.save_thread, self._patch_threads):
thread.check_and_raise_error()
class Extract(object):
""" The extract process. """
def __init__(self, arguments):
self.args = arguments
self.images = Images(self.args)
self.faces = Faces(self.args)
self.alignments = Alignments(self.args)
self.output_dir = self.faces.output_dir
self.export_face = True
def process(self):
""" Perform the extraction process """
print('Starting, this may take a while...')
Utils.set_verbosity(self.args.verbose)
if hasattr(self.args, 'processes') and self.args.processes > 1:
self.extract_multi_process()
else:
self.extract_single_process()
self.alignments.write_alignments(self.faces.faces_detected)
images, faces = Utils.finalize(self.images.images_found,
self.faces.num_faces_detected,
self.faces.verify_output)
self.images.images_found = images
self.faces.num_faces_detected = faces
def extract_single_process(self):
""" Run extraction in a single process """
for filename in tqdm(self.images.input_images):
filename, faces = self.process_single_image(filename)
self.faces.faces_detected[os.path.basename(filename)] = faces
def extract_multi_process(self):
""" Run the extraction on the correct number of processes """
for filename, faces in tqdm(pool_process(
self.process_single_image,
self.images.input_images,
processes=self.args.processes),
total=self.images.images_found):
self.faces.num_faces_detected += 1
self.faces.faces_detected[os.path.basename(filename)] = faces
def process_single_image(self, filename):
""" Detect faces in an image. Rotate the image the specified amount
until at least one face is found, or until image rotations are
depleted.
Once at least one face has been detected, pass to process_single_face
to process the individual faces """
retval = filename, list()
try:
image = Utils.cv2_read_write('read', filename)
for angle in self.images.rotation_angles:
image = Utils.rotate_image_by_angle(image, angle)
faces = self.faces.get_faces(image, angle)
process_faces = [(idx, face) for idx, face in faces]
if process_faces and angle != 0 and self.args.verbose:
print("found face(s) by rotating image {} degrees".format(
angle))
if process_faces:
break
final_faces = [
self.process_single_face(idx, face, filename, image)
for idx, face in process_faces
]
retval = filename, final_faces
except Exception as err:
if self.args.verbose:
print("Failed to extract from image: {}. Reason: {}".format(
filename, err))
return retval
def process_single_face(self, idx, face, filename, image):
""" Perform processing on found faces """
output_file = self.output_dir / Path(
filename).stem if self.export_face else None
self.faces.draw_landmarks_on_face(face, image)
resized_face, t_mat = self.faces.extractor.extract(
image, face, 256, self.faces.align_eyes)
blurry_file = self.faces.detect_blurry_faces(face, t_mat, resized_face,
filename)
output_file = blurry_file if blurry_file else output_file
if self.export_face:
filename = "{}_{}{}".format(str(output_file), str(idx),
Path(filename).suffix)
Utils.cv2_read_write('write', filename, resized_face)
return {
"r": face.r,
"x": face.x,
"w": face.w,
"y": face.y,
"h": face.h,
"landmarksXY": face.landmarksAsXY()
}
class Extract(): # pylint:disable=too-few-public-methods
""" The Faceswap Face Extraction Process.
The extraction process is responsible for detecting faces in a series of images/video, aligning
these faces and then generating a mask.
It leverages a series of user selected plugins, chained together using
:mod:`plugins.extract.pipeline`.
The extract process is self contained and should not be referenced by any other scripts, so it
contains no public properties.
Parameters
----------
arguments: argparse.Namespace
The arguments to be passed to the extraction process as generated from Faceswap's command
line arguments
"""
def __init__(self, arguments):
logger.debug("Initializing %s: (args: %s", self.__class__.__name__,
arguments)
self._args = arguments
self._output_dir = str(get_folder(self._args.output_dir))
logger.info("Output Directory: %s", self._args.output_dir)
self._images = ImagesLoader(self._args.input_dir,
load_with_hash=False,
fast_count=True)
self._alignments = Alignments(self._args, True, self._images.is_video)
self._existing_count = 0
self._set_skip_list()
self._post_process = PostProcess(arguments)
configfile = self._args.configfile if hasattr(self._args,
"configfile") else None
normalization = None if self._args.normalization == "none" else self._args.normalization
self._extractor = Extractor(self._args.detector,
self._args.aligner,
self._args.masker,
configfile=configfile,
multiprocess=not self._args.singleprocess,
rotate_images=self._args.rotate_images,
min_size=self._args.min_size,
normalize_method=normalization)
self._threads = list()
self._verify_output = False
logger.debug("Initialized %s", self.__class__.__name__)
@property
def _save_interval(self):
""" int: The number of frames to be processed between each saving of the alignments file if
it has been provided, otherwise ``None`` """
if hasattr(self._args, "save_interval"):
return self._args.save_interval
return None
@property
def _skip_num(self):
""" int: Number of frames to skip if extract_every_n has been provided """
return self._args.extract_every_n if hasattr(self._args,
"extract_every_n") else 1
def _set_skip_list(self):
""" Add the skip list to the image loader
Checks against `extract_every_n` and the existence of alignments data (can exist if
`skip_existing` or `skip_existing_faces` has been provided) and compiles a list of frame
indices that should not be processed, providing these to :class:`lib.image.ImagesLoader`.
"""
if self._skip_num == 1 and not self._alignments.data:
logger.debug("No frames to be skipped")
return
skip_list = []
for idx, filename in enumerate(self._images.file_list):
if idx % self._skip_num != 0:
logger.trace(
"Adding image '%s' to skip list due to extract_every_n = %s",
filename, self._skip_num)
skip_list.append(idx)
# Items may be in the alignments file if skip-existing[-faces] is selected
elif os.path.basename(filename) in self._alignments.data:
self._existing_count += 1
logger.trace("Removing image: '%s' due to previously existing",
filename)
skip_list.append(idx)
if self._existing_count != 0:
logger.info(
"Skipping %s frames due to skip_existing/skip_existing_faces.",
self._existing_count)
logger.debug("Adding skip list: %s", skip_list)
self._images.add_skip_list(skip_list)
def process(self):
""" The entry point for triggering the Extraction Process.
Should only be called from :class:`lib.cli.ScriptExecutor`
"""
logger.info('Starting, this may take a while...')
# from lib.queue_manager import queue_manager ; queue_manager.debug_monitor(3)
self._threaded_redirector("load")
self._run_extraction()
for thread in self._threads:
thread.join()
self._alignments.save()
finalize(self._images.process_count + self._existing_count,
self._alignments.faces_count, self._verify_output)
def _threaded_redirector(self, task, io_args=None):
""" Redirect image input/output tasks to relevant queues in background thread
Parameters
----------
task: str
The name of the task to be put into a background thread
io_args: tuple, optional
Any arguments that need to be provided to the background function
"""
logger.debug("Threading task: (Task: '%s')", task)
io_args = tuple() if io_args is None else (io_args, )
func = getattr(self, "_{}".format(task))
io_thread = MultiThread(func, *io_args, thread_count=1)
io_thread.start()
self._threads.append(io_thread)
def _load(self):
""" Load the images
Loads images from :class:`lib.image.ImagesLoader`, formats them into a dict compatible
with :class:`plugins.extract.Pipeline.Extractor` and passes them into the extraction queue.
"""
logger.debug("Load Images: Start")
load_queue = self._extractor.input_queue
for filename, image in self._images.load():
if load_queue.shutdown.is_set():
logger.debug("Load Queue: Stop signal received. Terminating")
break
item = ExtractMedia(filename, image[..., :3])
load_queue.put(item)
load_queue.put("EOF")
logger.debug("Load Images: Complete")
def _reload(self, detected_faces):
""" Reload the images and pair to detected face
When the extraction pipeline is running in serial mode, images are reloaded from disk,
paired with their extraction data and passed back into the extraction queue
Parameters
----------
detected_faces: dict
Dictionary of :class:`plugins.extract.pipeline.ExtractMedia` with the filename as the
key for repopulating the image attribute.
"""
logger.debug("Reload Images: Start. Detected Faces Count: %s",
len(detected_faces))
load_queue = self._extractor.input_queue
for filename, image in self._images.load():
if load_queue.shutdown.is_set():
logger.debug("Reload Queue: Stop signal received. Terminating")
break
logger.trace("Reloading image: '%s'", filename)
extract_media = detected_faces.pop(filename, None)
if not extract_media:
logger.warning("Couldn't find faces for: %s", filename)
continue
extract_media.set_image(image)
load_queue.put(extract_media)
load_queue.put("EOF")
logger.debug("Reload Images: Complete")
def _run_extraction(self):
""" The main Faceswap Extraction process
Receives items from :class:`plugins.extract.Pipeline.Extractor` and either saves out the
faces and data (if on the final pass) or reprocesses data through the pipeline for serial
processing.
"""
size = self._args.size if hasattr(self._args, "size") else 256
saver = ImagesSaver(self._output_dir, as_bytes=True)
exception = False
phase_desc = "Extraction"
for phase in range(self._extractor.passes):
if exception:
break
is_final = self._extractor.final_pass
detected_faces = dict()
self._extractor.launch()
self._check_thread_error()
if self._args.singleprocess:
phase_desc = self._extractor.phase.title()
desc = "Running pass {} of {}: {}".format(phase + 1,
self._extractor.passes,
phase_desc)
status_bar = tqdm(self._extractor.detected_faces(),
total=self._images.process_count,
file=sys.stdout,
desc=desc)
for idx, extract_media in enumerate(status_bar):
self._check_thread_error()
if is_final:
self._output_processing(extract_media, size)
self._output_faces(saver, extract_media)
if self._save_interval and (idx +
1) % self._save_interval == 0:
self._alignments.save()
else:
extract_media.remove_image()
# cache extract_media for next run
detected_faces[extract_media.filename] = extract_media
status_bar.update(1)
if not is_final:
logger.debug("Reloading images")
self._threaded_redirector("reload", detected_faces)
saver.close()
def _check_thread_error(self):
""" Check if any errors have occurred in the running threads and their errors """
for thread in self._threads:
thread.check_and_raise_error()
def _output_processing(self, extract_media, size):
""" Prepare faces for output
Loads the aligned face, perform any processing actions and verify the output.
Parameters
----------
extract_media: :class:`plugins.extract.pipeline.ExtractMedia`
Output from :class:`plugins.extract.pipeline.Extractor`
size: int
The size that the aligned face should be created at
"""
for face in extract_media.detected_faces:
face.load_aligned(extract_media.image, size=size)
self._post_process.do_actions(extract_media)
extract_media.remove_image()
faces_count = len(extract_media.detected_faces)
if faces_count == 0:
logger.verbose("No faces were detected in image: %s",
os.path.basename(extract_media.filename))
if not self._verify_output and faces_count > 1:
self._verify_output = True
def _output_faces(self, saver, extract_media):
""" Output faces to save thread
Set the face filename based on the frame name and put the face to the
:class:`~lib.image.ImagesSaver` save queue and add the face information to the alignments
data.
Parameters
----------
saver: lib.images.ImagesSaver
The background saver for saving the image
extract_media: :class:`~plugins.extract.pipeline.ExtractMedia`
The output from :class:`~plugins.extract.Pipeline.Extractor`
"""
logger.trace("Outputting faces for %s", extract_media.filename)
final_faces = list()
filename, extension = os.path.splitext(
os.path.basename(extract_media.filename))
for idx, face in enumerate(extract_media.detected_faces):
output_filename = "{}_{}{}".format(filename, str(idx), extension)
face.hash, image = encode_image_with_hash(face.aligned_face,
extension)
saver.save(output_filename, image)
final_faces.append(face.to_alignment())
self._alignments.data[os.path.basename(
extract_media.filename)] = dict(faces=final_faces)
del extract_media
def __init__(self, arguments):
logger.debug("Initializing %s: (args: %s)", self.__class__.__name__, arguments)
self._args = arguments
# load faces
faces_alignments = AlignmentsBase(self._args.faces_align_dir)
print()
print(f'Faces alignments: {len(faces_alignments._data.keys())}')
print(faces_alignments._data.keys())
self._faces = {}
faces_loader = ImagesLoader(self._args.faces_dir)
for filename, image in faces_loader.load():
face_name = os.path.basename(filename)
faces = faces_alignments.get_faces_in_frame(face_name)
detected_faces = list()
for rawface in faces:
face = DetectedFace()
face.from_alignment(rawface, image=image)
feed_face = AlignedFace(face.landmarks_xy,
image=image,
centering='face',
size=image.shape[0],
coverage_ratio=1.0,
dtype="float32")
detected_faces.append(feed_face)
self._faces[face_name] = (filename, image, detected_faces)
print('Faces:', len(self._faces))
print(self._faces.keys())
print()
self._patch_threads = None
self._images = ImagesLoader(self._args.input_dir, fast_count=True)
self._alignments = Alignments(self._args, False, self._images.is_video)
if self._alignments.version == 1.0:
logger.error("The alignments file format has been updated since the given alignments "
"file was generated. You need to update the file to proceed.")
logger.error("To do this run the 'Alignments Tool' > 'Extract' Job.")
sys.exit(1)
self._opts = OptionalActions(self._args, self._images.file_list, self._alignments)
self._add_queues()
self._disk_io = DiskIO(self._alignments, self._images, arguments)
self._predictor = Predict(self._disk_io.load_queue, self._queue_size, self._faces, arguments)
self._validate()
get_folder(self._args.output_dir)
configfile = self._args.configfile if hasattr(self._args, "configfile") else None
self._converter = Converter(self._predictor.output_size,
self._predictor.coverage_ratio,
self._predictor.centering,
self._disk_io.draw_transparent,
self._disk_io.pre_encode,
arguments,
configfile=configfile)
logger.debug("Initialized %s", self.__class__.__name__)
