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
