def launch_aligner(): """ Load the aligner plugin to retrieve landmarks """ out_queue = queue_manager.get_queue("out") kwargs = {"in_queue": queue_manager.get_queue("in"), "out_queue": out_queue} process = SpawnProcess() event = process.event for plugin in ("fan", "dlib"): aligner = PluginLoader.get_aligner(plugin)() process.in_process(aligner.align, **kwargs) # Wait for Aligner to take init # The first ever load of the model for FAN has reportedly taken # up to 3-4 minutes, hence high timeout. event.wait(300) if not event.is_set(): if plugin == "fan": process.join() print("Error initializing FAN. Trying Dlib") continue else: raise ValueError("Error inititalizing Aligner") if plugin == "dlib": return try: err = None err = out_queue.get(True, 1) except QueueEmpty: pass if not err: break process.join() print("Error initializing FAN. Trying Dlib")
def launch_aligner(self): """ Launch the face aligner """ out_queue = queue_manager.get_queue("align") kwargs = { "in_queue": queue_manager.get_queue("detect"), "out_queue": out_queue } align_process = SpawnProcess() event = align_process.event align_process.in_process(self.aligner.align, **kwargs) # Wait for Aligner to take it's VRAM # The first ever load of the model for FAN has reportedly taken # up to 3-4 minutes, hence high timeout. # TODO investigate why this is and fix if possible event.wait(300) if not event.is_set(): raise ValueError("Error inititalizing Aligner") try: err = None err = out_queue.get(True, 1) except QueueEmpty: pass if err: queue_manager.terminate_queues() print(err) exit(1)
def init_extractor(self, verbose): """ Initialize FAN """ out_queue = queue_manager.get_queue("out") d_kwargs = { "in_queue": queue_manager.get_queue("in"), "out_queue": queue_manager.get_queue("align") } a_kwargs = { "in_queue": queue_manager.get_queue("align"), "out_queue": out_queue } detect_process = SpawnProcess() align_process = SpawnProcess() d_event = detect_process.event a_event = align_process.event detector = PluginLoader.get_detector("manual")(verbose=verbose) detect_process.in_process(detector.detect_faces, **d_kwargs) for plugin in ("fan", "dlib"): aligner = PluginLoader.get_aligner(plugin)(verbose=verbose) align_process.in_process(aligner.align, **a_kwargs) # Wait for Aligner to take init # The first ever load of the model for FAN has reportedly taken # up to 3-4 minutes, hence high timeout. a_event.wait(300) if not a_event.is_set(): if plugin == "fan": align_process.join() print("Error initializing FAN. Trying Dlib") continue else: raise ValueError("Error inititalizing Aligner") if plugin == "dlib": break try: err = None err = out_queue.get(True, 1) except QueueEmpty: pass if not err: break align_process.join() print("Error initializing FAN. Trying Dlib") d_event.wait(10) if not d_event.is_set(): raise ValueError("Error inititalizing Detector") self.extractor["detect"] = detector self.extractor["align"] = aligner
def load_extractor(self): """ Set on the fly extraction """ print("\nNo Alignments file found. Extracting on the fly.\n" "NB: This will use the inferior dlib-hog for extraction " "and dlib pose predictor for landmarks.\nIt is recommended " "to perfom Extract first for superior results\n") for task in ("load", "detect", "align"): queue_manager.add_queue(task, maxsize=0) detector = PluginLoader.get_detector("dlib_hog")( verbose=self.args.verbose) aligner = PluginLoader.get_aligner("dlib")(verbose=self.args.verbose) 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() d_event = d_process.event a_process = SpawnProcess() a_event = a_process.event d_process.in_process(detector.detect_faces, **d_kwargs) a_process.in_process(aligner.align, **a_kwargs) 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
class Plugins(): """ Detector and Aligner Plugins and queues """ def __init__(self, arguments): self.args = arguments self.detector = self.load_detector() self.aligner = self.load_aligner() self.is_parallel = self.set_parallel_processing() self.process_detect = None self.process_align = None self.add_queues() def set_parallel_processing(self): """ Set whether to run detect and align together or seperately """ detector_vram = self.detector.vram aligner_vram = self.aligner.vram gpu_stats = GPUStats() if (detector_vram == 0 or aligner_vram == 0 or gpu_stats.device_count == 0): return True if hasattr(self.args, "multiprocess") and not self.args.multiprocess: print("\nNB: Parallel processing disabled.\nYou may get faster " "extraction speeds by enabling it with the -mp switch\n") return False required_vram = detector_vram + aligner_vram + 320 # 320MB buffer stats = gpu_stats.get_card_most_free() free_vram = int(stats["free"]) if self.args.verbose: print("{} - {}MB free of {}MB".format(stats["device"], free_vram, int(stats["total"]))) if free_vram <= required_vram: if self.args.verbose: print("Not enough free VRAM for parallel processing. " "Switching to serial") return False return True def add_queues(self): """ Add the required processing queues to Queue Manager """ for task in ("load", "detect", "align", "save"): size = 0 if task == "load" or (not self.is_parallel and task == "detect"): size = 100 queue_manager.add_queue(task, maxsize=size) def load_detector(self): """ Set global arguments and load detector plugin """ detector_name = self.args.detector.replace("-", "_").lower() # Rotation rotation = None if hasattr(self.args, "rotate_images"): rotation = self.args.rotate_images detector = PluginLoader.get_detector(detector_name)( verbose=self.args.verbose, rotation=rotation) return detector def load_aligner(self): """ Set global arguments and load aligner plugin """ aligner_name = self.args.aligner.replace("-", "_").lower() aligner = PluginLoader.get_aligner(aligner_name)( verbose=self.args.verbose) return aligner def launch_aligner(self): """ Launch the face aligner """ out_queue = queue_manager.get_queue("align") kwargs = {"in_queue": queue_manager.get_queue("detect"), "out_queue": out_queue} self.process_align = SpawnProcess() event = self.process_align.event self.process_align.in_process(self.aligner.align, **kwargs) # Wait for Aligner to take it's VRAM # The first ever load of the model for FAN has reportedly taken # up to 3-4 minutes, hence high timeout. # TODO investigate why this is and fix if possible event.wait(300) if not event.is_set(): raise ValueError("Error inititalizing Aligner") try: err = None err = out_queue.get(True, 1) except QueueEmpty: pass if err: if isinstance(err, str): queue_manager.terminate_queues() print(err) exit(1) else: queue_manager.get_queue("detect").put(err) def launch_detector(self): """ Launch the face detector """ out_queue = queue_manager.get_queue("detect") kwargs = {"in_queue": queue_manager.get_queue("load"), "out_queue": out_queue} if self.args.detector == "mtcnn": mtcnn_kwargs = self.detector.validate_kwargs( self.get_mtcnn_kwargs()) kwargs["mtcnn_kwargs"] = mtcnn_kwargs if self.detector.parent_is_pool: self.process_detect = PoolProcess(self.detector.detect_faces) else: self.process_detect = SpawnProcess() event = None if hasattr(self.process_detect, "event"): event = self.process_detect.event self.process_detect.in_process(self.detector.detect_faces, **kwargs) if not event: return event.wait(60) if not event.is_set(): raise ValueError("Error inititalizing Detector") def get_mtcnn_kwargs(self): """ Add the mtcnn arguments into a kwargs dictionary """ mtcnn_threshold = [float(thr.strip()) for thr in self.args.mtcnn_threshold] return {"minsize": self.args.mtcnn_minsize, "threshold": mtcnn_threshold, "factor": self.args.mtcnn_scalefactor} def detect_faces(self, extract_pass="******"): """ Detect faces from in an image """ if self.is_parallel or extract_pass == "align": out_queue = queue_manager.get_queue("align") if not self.is_parallel and extract_pass == "detect": out_queue = queue_manager.get_queue("detect") while True: try: faces = out_queue.get(True, 1) if faces == "EOF": break exception = faces.get("exception", None) if exception is not None: queue_manager.terminate_queues() yield faces break except QueueEmpty: continue yield faces