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 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
