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