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