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): 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.images.input_images) self.post_process = PostProcess(arguments) self.verify_output = False 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() 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 in tqdm(self.images.input_images, 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: convert_item = self.detect_faces(filename) else: convert_item = self.alignments_faces(filename, frame) if not convert_item: continue image, detected_faces = convert_item 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, filename): """ Extract the face from a frame (If not alignments file found) """ image = self.images.load_one_image(filename) queue_manager.get_queue("load").put((filename, image)) item = queue_manager.get_queue("align").get() detected_faces = item["detected_faces"] return image, detected_faces def alignments_faces(self, filename, frame): """ Get the face from alignments file """ if not self.check_alignments(frame): return None faces = self.alignments.get_faces_in_frame(frame) image = self.images.load_one_image(filename) detected_faces = list() for rawface in faces: face = DetectedFace() face.from_alignment(rawface, image=image) detected_faces.append(face) return image, 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 idx, face in enumerate(faces): image = self.convert_one_face(converter, (filename, image, idx, 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, 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 # 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