diff --git a/lib/cli.py b/lib/cli.py index a925af29b2..cc68498797 100644 --- a/lib/cli.py +++ b/lib/cli.py @@ -438,29 +438,32 @@ def get_argument_list(): "rounding": 2, "type": float, "dest": "ref_threshold", - "default": 0.6, + "default": 0.4, "help": "Threshold for positive face recognition. For use with " - "nfilter or filter. Lower values are stricter."}) + "nfilter or filter. Lower values are stricter. NB: Using " + "face filter will significantly decrease extraction speed."}) argument_list.append({"opts": ("-n", "--nfilter"), "action": FilesFullPaths, "filetypes": "image", "dest": "nfilter", "nargs": "+", "default": None, - "help": "Reference image for the persons you do " - "not want to process. Should be a front " - "portrait. Multiple images can be added " - "space separated"}) + "help": "Reference image for the persons you do not want to " + "process. Should be a front portrait with a single person " + "in the image. Multiple images can be added space " + "separated. NB: Using face filter will significantly " + "decrease extraction speed."}) argument_list.append({"opts": ("-f", "--filter"), "action": FilesFullPaths, "filetypes": "image", "dest": "filter", "nargs": "+", "default": None, - "help": "Reference images for the person you " - "want to process. Should be a front " - "portrait. Multiple images can be added " - "space separated"}) + "help": "Reference images for the person you want to process. " + "Should be a front portrait with a single person in the " + "image. Multiple images can be added space separated. NB: " + "Using face filter will significantly decrease extraction " + "speed."}) return argument_list diff --git a/lib/face_filter.py b/lib/face_filter.py index a373fcb458..4dd92e3617 100644 --- a/lib/face_filter.py +++ b/lib/face_filter.py @@ -2,9 +2,16 @@ """ Face Filterer for extraction in faceswap.py """ import logging +import os +import sys -import face_recognition +import cv2 +import numpy as np +from lib.faces_detect import DetectedFace +from lib.logger import get_loglevel +from lib.utils import GetModel +from plugins.extract.pipeline import Extractor logger = logging.getLogger(__name__) # pylint: disable=invalid-name @@ -15,77 +22,206 @@ def avg(arr): class FaceFilter(): - """ Face filter for extraction """ - def __init__(self, reference_file_paths, nreference_file_paths, threshold=0.6): + """ Face filter for extraction + NB: we take only first face, so the reference file should only contain one face. """ + + def __init__(self, reference_file_paths, nreference_file_paths, detector, aligner, loglevel, + multiprocess=False, threshold=0.4): logger.debug("Initializing %s: (reference_file_paths: %s, nreference_file_paths: %s, " - "threshold: %s)", self.__class__.__name__, reference_file_paths, - nreference_file_paths, threshold) - images = list(map(face_recognition.load_image_file, reference_file_paths)) - nimages = list(map(face_recognition.load_image_file, nreference_file_paths)) - # Note: we take only first face, so the reference file should only contain one face. - self.encodings = list(map(lambda im: face_recognition.face_encodings(im)[0], images)) - self.nencodings = list(map(lambda im: face_recognition.face_encodings(im)[0], nimages)) + "detector: %s, aligner: %s. loglevel: %s, multiprocess: %s, threshold: %s)", + self.__class__.__name__, reference_file_paths, nreference_file_paths, + detector, aligner, loglevel, multiprocess, threshold) + git_model_id = 7 + model_filename = ["vgg_face_v1.caffemodel", "vgg_face_v1.prototxt"] + + self.input_size = 224 + self.numeric_loglevel = get_loglevel(loglevel) + self.vgg_face = self.get_model(git_model_id, model_filename) + self.filters = self.load_images(reference_file_paths, nreference_file_paths) + self.align_faces(detector, aligner, loglevel, multiprocess) + self.get_filter_encodings() self.threshold = threshold - logger.trace("encodings: %s", self.encodings) - logger.trace("nencodings: %s", self.nencodings) logger.debug("Initialized %s", self.__class__.__name__) + # <<< GET MODEL >>> # + @staticmethod + def get_model(git_model_id, model_filename): + """ Check if model is available, if not, download and unzip it """ + root_path = os.path.abspath(os.path.dirname(sys.argv[0])) + cache_path = os.path.join(root_path, "plugins", "extract", ".cache") + model = GetModel(model_filename, cache_path, git_model_id).model_path + vgg_face = cv2.dnn.readNetFromCaffe(model[1], model[0]) # pylint: disable=no-member + vgg_face.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU) # pylint: disable=no-member + return vgg_face + + @staticmethod + def load_images(reference_file_paths, nreference_file_paths): + """ Load the images """ + retval = dict() + for fpath in reference_file_paths: + retval[fpath] = {"image": cv2.imread(fpath), # pylint: disable=no-member + "type": "filter"} + for fpath in nreference_file_paths: + retval[fpath] = {"image": cv2.imread(fpath), # pylint: disable=no-member + "type": "nfilter"} + logger.debug("Loaded filter images: %s", {k: v["type"] for k, v in retval.items()}) + return retval + + # Extraction pipeline + def align_faces(self, detector_name, aligner_name, loglevel, multiprocess): + """ Use the requested detectors to retrieve landmarks for filter images """ + extractor = Extractor(detector_name, aligner_name, loglevel, multiprocess=multiprocess) + self.run_extractor(extractor) + del extractor + self.load_aligned_face() + + def run_extractor(self, extractor): + """ Run extractor to get faces """ + exception = False + for _ in range(extractor.passes): + self.queue_images(extractor) + if exception: + break + extractor.launch() + for faces in extractor.detected_faces(): + exception = faces.get("exception", False) + if exception: + break + filename = faces["filename"] + detected_faces = faces["detected_faces"] + + if len(detected_faces) > 1: + logger.warning("Multiple faces found in %s file: '%s'. Using first detected " + "face.", self.filters[filename]["type"], filename) + detected_faces = [detected_faces[0]] + self.filters[filename]["detected_faces"] = detected_faces + + # Aligner output + if extractor.final_pass: + landmarks = faces["landmarks"] + self.filters[filename]["landmarks"] = landmarks + + def queue_images(self, extractor): + """ queue images for detection and alignment """ + in_queue = extractor.input_queue + for fname, img in self.filters.items(): + logger.debug("Adding to filter queue: '%s' (%s)", fname, img["type"]) + feed_dict = dict(filename=fname, image=img["image"]) + if img.get("detected_faces", None): + feed_dict["detected_faces"] = img["detected_faces"] + logger.debug("Queueing filename: '%s' items: %s", + fname, list(feed_dict.keys())) + in_queue.put(feed_dict) + logger.debug("Sending EOF to filter queue") + in_queue.put("EOF") + + def load_aligned_face(self): + """ Align the faces for vgg_face input """ + for filename, face in self.filters.items(): + logger.debug("Loading aligned face: '%s'", filename) + bounding_box = face["detected_faces"][0] + image = face["image"] + landmarks = face["landmarks"][0] + + detected_face = DetectedFace() + detected_face.from_bounding_box(bounding_box, image) + detected_face.landmarksXY = landmarks + detected_face.load_aligned(image, size=224) + face["face"] = detected_face.aligned_face + del face["image"] + logger.debug("Loaded aligned face: ('%s', shape: %s)", + filename, face["face"].shape) + + def get_filter_encodings(self): + """ Return filter face encodings from Keras VGG Face """ + for filename, face in self.filters.items(): + logger.debug("Getting encodings for: '%s'", filename) + encodings = self.predict(face["face"]) + logger.debug("Filter Filename: %s, encoding shape: %s", filename, encodings.shape) + face["encoding"] = encodings + del face["face"] + + def predict(self, face): + """ Return encodings for given image from vgg_face """ + if face.shape[0] != self.input_size: + face = self.resize_face(face) + blob = cv2.dnn.blobFromImage(face, # pylint: disable=no-member + 1.0, + (self.input_size, self.input_size), + [104, 117, 123], + False, + False) + self.vgg_face.setInput(blob) + preds = self.vgg_face.forward()[0, :] + return preds + + def resize_face(self, face): + """ Resize incoming face to model_input_size """ + if face.shape[0] < self.input_size: + interpolation = cv2.INTER_CUBIC # pylint:disable=no-member + else: + interpolation = cv2.INTER_AREA # pylint:disable=no-member + + face = cv2.resize(face, # pylint:disable=no-member + dsize=(self.input_size, self.input_size), + interpolation=interpolation) + return face + def check(self, detected_face): - """ Check Face - we could use detected landmarks, but I did not manage to do so. - TODO The copy/paste below should help """ + """ Check the extracted Face """ logger.trace("Checking face with FaceFilter") - encodings = face_recognition.face_encodings(detected_face.image) - if not encodings: - logger.verbose("No face encodings found") - return False - - if self.encodings: - distances = list(face_recognition.face_distance(self.encodings, encodings[0])) - logger.trace("Distances: %s", distances) - distance = avg(distances) - logger.trace("Average Distance: %s", distance) - mindistance = min(distances) - logger.trace("Minimum Distance: %s", mindistance) - if distance > self.threshold: - logger.verbose("Distance above threshold: %f < %f", distance, self.threshold) - return False - if self.nencodings: - ndistances = list(face_recognition.face_distance(self.nencodings, encodings[0])) - logger.trace("nDistances: %s", ndistances) - ndistance = avg(ndistances) - logger.trace("Average nDistance: %s", ndistance) - nmindistance = min(ndistances) - logger.trace("Minimum nDistance: %s", nmindistance) - if not self.encodings and ndistance < self.threshold: - logger.verbose("nDistance below threshold: %f < %f", ndistance, self.threshold) - return False - if self.encodings: - if mindistance > nmindistance: - logger.verbose("Distance to negative sample is smaller") - return False - if distance > ndistance: - logger.verbose("Average distance to negative sample is smaller") - return False - # k-nn classifier - var_k = min(5, min(len(distances), len(ndistances)) + 1) - var_n = sum(list(map(lambda x: x[0], - list(sorted([(1, d) for d in distances] + - [(0, d) for d in ndistances], - key=lambda x: x[1]))[:var_k]))) - ratio = var_n/var_k - if ratio < 0.5: - logger.verbose("K-nn is %.2f", ratio) - return False - return True - - -# # Copy/Paste (mostly) from private method in face_recognition -# face_recognition_model = face_recognition_models.face_recognition_model_location() -# face_encoder = dlib.face_recognition_model_v1(face_recognition_model) - -# def convert(detected_face): -# return np.array(face_encoder.compute_face_descriptor(detected_face.image, -# detected_face.landmarks, -# 1)) -# # end of Copy/Paste + distances = {"filter": list(), "nfilter": list()} + encodings = self.predict(detected_face.aligned_face) + for filt in self.filters.values(): + similarity = self.find_cosine_similiarity(filt["encoding"], encodings) + distances[filt["type"]].append(similarity) + + avgs = {key: avg(val) if val else None for key, val in distances.items()} + mins = {key: min(val) if val else None for key, val in distances.items()} + # Filter + if distances["filter"] and avgs["filter"] > self.threshold: + msg = "Rejecting filter face: {} > {}".format(round(avgs["filter"], 2), self.threshold) + retval = False + # nFilter no Filter + elif not distances["filter"] and avgs["nfilter"] < self.threshold: + msg = "Rejecting nFilter face: {} < {}".format(round(avgs["nfilter"], 2), + self.threshold) + retval = False + # Filter with nFilter + elif distances["filter"] and distances["nfilter"] and mins["filter"] > mins["nfilter"]: + msg = ("Rejecting face as distance from nfilter sample is smaller: (filter: {}, " + "nfilter: {})".format(round(mins["filter"], 2), round(mins["nfilter"], 2))) + retval = False + elif distances["filter"] and distances["nfilter"] and avgs["filter"] > avgs["nfilter"]: + msg = ("Rejecting face as average distance from nfilter sample is smaller: (filter: " + "{}, nfilter: {})".format(round(mins["filter"], 2), round(mins["nfilter"], 2))) + retval = False + elif distances["filter"] and distances["nfilter"]: + # k-nn classifier + var_k = min(5, min(len(distances["filter"]), len(distances["nfilter"])) + 1) + var_n = sum(list(map(lambda x: x[0], + list(sorted([(1, d) for d in distances["filter"]] + + [(0, d) for d in distances["nfilter"]], + key=lambda x: x[1]))[:var_k]))) + ratio = var_n/var_k + if ratio < 0.5: + msg = ("Rejecting face as k-nearest neighbors classification is less than " + "0.5: {}".format(round(ratio, 2))) + retval = False + else: + msg = None + retval = True + if msg: + logger.verbose(msg) + else: + logger.trace("Accepted face: (similarity: %s, threshold: %s)", + distances, self.threshold) + return retval + + @staticmethod + def find_cosine_similiarity(source_face, test_face): + """ Find the cosine similarity between a source face and a test face """ + var_a = np.matmul(np.transpose(source_face), test_face) + var_b = np.sum(np.multiply(source_face, source_face)) + var_c = np.sum(np.multiply(test_face, test_face)) + return 1 - (var_a / (np.sqrt(var_b) * np.sqrt(var_c))) diff --git a/lib/multithreading.py b/lib/multithreading.py index c8e7f20c75..c4885e66ef 100644 --- a/lib/multithreading.py +++ b/lib/multithreading.py @@ -371,7 +371,8 @@ def join(self, timeout=None): """ Add logging to join function """ logger.debug("Joining Process: (name: '%s', PID: %s)", self._name, self.pid) super().join(timeout=timeout) - _launched_processes.remove(self) + if self in _launched_processes: + _launched_processes.remove(self) logger.debug("Joined Process: (name: '%s', PID: %s)", self._name, self.pid) diff --git a/lib/utils.py b/lib/utils.py index b22414af4c..2989e1e7cb 100644 --- a/lib/utils.py +++ b/lib/utils.py @@ -225,6 +225,9 @@ class GetModel(): model_filename: The name of the model to be loaded (see notes below) cache_dir: The model cache folder of the current plugin calling this class IE: The folder that holds the model to be loaded. + git_model_id: The second digit in the github tag that identifies this model. + See https://github.com/deepfakes-models/faceswap-models for more + information NB: Models must have a certain naming convention: IE: _v. @@ -236,38 +239,20 @@ class GetModel(): IE: [_v.] EG: [mtcnn_det_v1.1.py, mtcnn_det_v1.2.py, mtcnn_det_v1.3.py] [resnet_ssd_v1.caffemodel, resnet_ssd_v1.prototext] - - Models to be handled by this class must be added to the _model_id property - with their appropriate github identier mapped. - See https://github.com/deepfakes-models/faceswap-models for more information """ - def __init__(self, model_filename, cache_dir): + def __init__(self, model_filename, cache_dir, git_model_id): if not isinstance(model_filename, list): model_filename = [model_filename] self.model_filename = model_filename self.cache_dir = cache_dir + self.git_model_id = git_model_id self.url_base = "https://github.com/deepfakes-models/faceswap-models/releases/download" self.chunk_size = 1024 # Chunk size for downloading and unzipping - + self.retries = 6 self.get() self.model_path = self._model_path - @property - def _model_id(self): - """ Return a mapping of model names to model ids as stored in github """ - ids = { - # EXTRACT (SECTION 1) - "face-alignment-network_2d4": 0, - "cnn-facial-landmark": 1, - "mtcnn_det": 2, - "s3fd": 3, - "resnet_ssd": 4, - # TRAIN (SECTION 2) - # CONVERT (SECTION 3) - } - return ids[self._model_name] - @property def _model_full_name(self): """ Return the model full name from the filename(s) """ @@ -335,11 +320,19 @@ def _url_section(self): @property def _url_download(self): """ Base URL for models """ - tag = "v{}.{}.{}".format(self._url_section, self._model_id, self._model_version) + tag = "v{}.{}.{}".format(self._url_section, self.git_model_id, self._model_version) retval = "{}/{}/{}.zip".format(self.url_base, tag, self._model_full_name) logger.trace("Download url: %s", retval) return retval + @property + def _url_partial_size(self): + """ Return how many bytes have already been downloaded """ + zip_file = self._model_zip_path + retval = os.path.getsize(zip_file) if os.path.exists(zip_file) else 0 + logger.trace(retval) + return retval + def get(self): """ Check the model exists, if not, download and unzip into location """ if self._model_exists: @@ -351,36 +344,46 @@ def get(self): def download_model(self): """ Download model zip to cache dir """ - logger.info("Downloading model: '%s'", self._model_name) - attempts = 3 - for attempt in range(attempts): + logger.info("Downloading model: '%s' from: %s", self._model_name, self._url_download) + for attempt in range(self.retries): try: - response = urllib.request.urlopen(self._url_download, timeout=10) + downloaded_size = self._url_partial_size + req = urllib.request.Request(self._url_download) + if downloaded_size != 0: + req.add_header("Range", "bytes={}-".format(downloaded_size)) + response = urllib.request.urlopen(req, timeout=10) logger.debug("header info: {%s}", response.info()) logger.debug("Return Code: %s", response.getcode()) - self.write_zipfile(response) + self.write_zipfile(response, downloaded_size) break except (socket_error, socket_timeout, urllib.error.HTTPError, urllib.error.URLError) as err: - if attempt + 1 < attempts: + if attempt + 1 < self.retries: logger.warning("Error downloading model (%s). Retrying %s of %s...", - str(err), attempt + 2, attempts) + str(err), attempt + 2, self.retries) else: logger.error("Failed to download model. Exiting. (Error: '%s', URL: '%s')", str(err), self._url_download) - logger.info("You can manually download the model from: %s and unzip the " - "contents to: %s", self._url_download, self.cache_dir) + logger.info("You can try running again to resume the download.") + logger.info("Alternatively, you can manually download the model from: %s and " + "unzip the contents to: %s", self._url_download, self.cache_dir) exit(1) - def write_zipfile(self, response): + def write_zipfile(self, response, downloaded_size): """ Write the model zip file to disk """ - length = int(response.getheader("content-length")) - with open(self._model_zip_path, "wb") as out_file: + length = int(response.getheader("content-length")) + downloaded_size + if length == downloaded_size: + logger.info("Zip already exists. Skipping download") + return + write_type = "wb" if downloaded_size == 0 else "ab" + with open(self._model_zip_path, write_type) as out_file: pbar = tqdm(desc="Downloading", unit="B", total=length, unit_scale=True, unit_divisor=1024) + if downloaded_size != 0: + pbar.update(downloaded_size) while True: buffer = response.read(self.chunk_size) if not buffer: diff --git a/plugins/extract/.cache/.keep b/plugins/extract/.cache/.keep new file mode 100644 index 0000000000..e69de29bb2 diff --git a/plugins/extract/align/_base.py b/plugins/extract/align/_base.py index 713c915b78..ab6193cb49 100644 --- a/plugins/extract/align/_base.py +++ b/plugins/extract/align/_base.py @@ -34,9 +34,11 @@ class Aligner(): """ Landmarks Aligner Object """ - def __init__(self, loglevel, model_filename=None, colorspace="BGR", input_size=256): - logger.debug("Initializing %s: (model_filename: '%s', colorspace: '%s'. input_size: %s)", - self.__class__.__name__, model_filename, colorspace, input_size) + def __init__(self, loglevel, + git_model_id=None, model_filename=None, colorspace="BGR", input_size=256): + logger.debug("Initializing %s: (loglevel: %s, git_model_id: %s, model_filename: '%s', " + "colorspace: '%s'. input_size: %s)", self.__class__.__name__, loglevel, + git_model_id, model_filename, colorspace, input_size) self.loglevel = loglevel self.colorspace = colorspace.upper() self.input_size = input_size @@ -49,7 +51,7 @@ def __init__(self, loglevel, model_filename=None, colorspace="BGR", input_size=2 self.queues = {"in": None, "out": None} # Get model if required - self.model_path = self.get_model(model_filename) + self.model_path = self.get_model(git_model_id, model_filename) # Approximate VRAM required for aligner. Used to calculate # how many parallel processes / batches can be run. @@ -82,13 +84,16 @@ def predict_landmarks(self, feed_dict): # <<< GET MODEL >>> # @staticmethod - def get_model(model_filename): + def get_model(git_model_id, model_filename): """ Check if model is available, if not, download and unzip it """ if model_filename is None: logger.debug("No model_filename specified. Returning None") return None + if git_model_id is None: + logger.debug("No git_model_id specified. Returning None") + return None cache_path = os.path.join(os.path.dirname(__file__), ".cache") - model = GetModel(model_filename, cache_path) + model = GetModel(model_filename, cache_path, git_model_id) return model.model_path # <<< ALIGNMENT WRAPPER >>> # diff --git a/plugins/extract/align/cv2_dnn.py b/plugins/extract/align/cv2_dnn.py index 2498744d13..312634d3d2 100644 --- a/plugins/extract/align/cv2_dnn.py +++ b/plugins/extract/align/cv2_dnn.py @@ -33,8 +33,13 @@ class Align(Aligner): """ Perform transformation to align and get landmarks """ def __init__(self, **kwargs): + git_model_id = 1 model_filename = "cnn-facial-landmark_v1.pb" - super().__init__(model_filename=model_filename, colorspace="RGB", input_size=128, **kwargs) + super().__init__(git_model_id=git_model_id, + model_filename=model_filename, + colorspace="RGB", + input_size=128, + **kwargs) self.vram = 0 # Doesn't use GPU self.model = None diff --git a/plugins/extract/align/fan.py b/plugins/extract/align/fan.py index 3c2a87caeb..3f329312ba 100644 --- a/plugins/extract/align/fan.py +++ b/plugins/extract/align/fan.py @@ -12,8 +12,13 @@ class Align(Aligner): """ Perform transformation to align and get landmarks """ def __init__(self, **kwargs): + git_model_id = 0 model_filename = "face-alignment-network_2d4_v1.pb" - super().__init__(model_filename=model_filename, colorspace="RGB", input_size=256, **kwargs) + super().__init__(git_model_id=git_model_id, + model_filename=model_filename, + colorspace="RGB", + input_size=256, + **kwargs) self.vram = 2240 self.model = None self.reference_scale = 195 diff --git a/plugins/extract/detect/_base.py b/plugins/extract/detect/_base.py index d11a5352d7..b737b6386a 100755 --- a/plugins/extract/detect/_base.py +++ b/plugins/extract/detect/_base.py @@ -33,9 +33,11 @@ def get_config(plugin_name): class Detector(): """ Detector object """ - def __init__(self, loglevel, model_filename=None, rotation=None, min_size=0): - logger.debug("Initializing %s: (model_filename: %s, rotation: %s, min_size: %s)", - self.__class__.__name__, model_filename, rotation, min_size) + def __init__(self, loglevel, + git_model_id=None, model_filename=None, rotation=None, min_size=0): + logger.debug("Initializing %s: (loglevel: %s, git_model_id: %s, model_filename: %s, " + "rotation: %s, min_size: %s)", self.__class__.__name__, loglevel, + git_model_id, model_filename, rotation, min_size) self.config = get_config(".".join(self.__module__.split(".")[-2:])) self.loglevel = loglevel self.rotation = self.get_rotation_angles(rotation) @@ -49,7 +51,7 @@ def __init__(self, loglevel, model_filename=None, rotation=None, min_size=0): self.queues = {"in": None, "out": None} # Path to model if required - self.model_path = self.get_model(model_filename) + self.model_path = self.get_model(git_model_id, model_filename) # Target image size for passing images through the detector # Set to tuple of dimensions (x, y) or int of pixel count @@ -93,13 +95,16 @@ def detect_faces(self, *args, **kwargs): # <<< GET MODEL >>> # @staticmethod - def get_model(model_filename): + def get_model(git_model_id, model_filename): """ Check if model is available, if not, download and unzip it """ if model_filename is None: logger.debug("No model_filename specified. Returning None") return None + if git_model_id is None: + logger.debug("No git_model_id specified. Returning None") + return None cache_path = os.path.join(os.path.dirname(__file__), ".cache") - model = GetModel(model_filename, cache_path) + model = GetModel(model_filename, cache_path, git_model_id) return model.model_path # <<< DETECTION WRAPPER >>> # diff --git a/plugins/extract/detect/cv2_dnn.py b/plugins/extract/detect/cv2_dnn.py index a1496c8897..4538150625 100755 --- a/plugins/extract/detect/cv2_dnn.py +++ b/plugins/extract/detect/cv2_dnn.py @@ -10,8 +10,9 @@ class Detect(Detector): """ CV2 DNN detector for face recognition """ def __init__(self, **kwargs): + git_model_id = 4 model_filename = ["resnet_ssd_v1.caffemodel", "resnet_ssd_v1.prototxt"] - super().__init__(model_filename=model_filename, **kwargs) + super().__init__(git_model_id=git_model_id, model_filename=model_filename, **kwargs) self.parent_is_pool = True self.target = (300, 300) # Doesn't use VRAM self.vram = 0 diff --git a/plugins/extract/detect/mtcnn.py b/plugins/extract/detect/mtcnn.py index 95e5e26264..19e63791d9 100755 --- a/plugins/extract/detect/mtcnn.py +++ b/plugins/extract/detect/mtcnn.py @@ -29,8 +29,9 @@ def import_tensorflow(): class Detect(Detector): """ MTCNN detector for face recognition """ def __init__(self, **kwargs): + git_model_id = 2 model_filename = ["mtcnn_det_v1.1.npy", "mtcnn_det_v1.2.npy", "mtcnn_det_v1.3.npy"] - super().__init__(model_filename=model_filename, **kwargs) + super().__init__(git_model_id=git_model_id, model_filename=model_filename, **kwargs) self.kwargs = self.validate_kwargs() self.name = "mtcnn" self.target = 2073600 # Uses approx 1.30 GB of VRAM diff --git a/plugins/extract/detect/s3fd.py b/plugins/extract/detect/s3fd.py index 61fc33db90..f451b153ca 100644 --- a/plugins/extract/detect/s3fd.py +++ b/plugins/extract/detect/s3fd.py @@ -17,8 +17,9 @@ class Detect(Detector): """ S3FD detector for face recognition """ def __init__(self, **kwargs): + git_model_id = 3 model_filename = "s3fd_v1.pb" - super().__init__(model_filename=model_filename, **kwargs) + super().__init__(git_model_id=git_model_id, model_filename=model_filename, **kwargs) self.name = "s3fd" self.target = (640, 640) # Uses approx 4 GB of VRAM self.vram = 4096 diff --git a/plugins/extract/pipeline.py b/plugins/extract/pipeline.py new file mode 100644 index 0000000000..ae764870b5 --- /dev/null +++ b/plugins/extract/pipeline.py @@ -0,0 +1,246 @@ +#!/usr/bin/env python3 +""" +Return a requested detector/aligner pipeline + +Tensorflow does not like to release GPU VRAM, so these are launched in subprocesses +so that the vram is released on subprocess exit """ + +import logging + +from lib.gpu_stats import GPUStats +from lib.multithreading import PoolProcess, SpawnProcess +from lib.queue_manager import queue_manager, QueueEmpty +from plugins.plugin_loader import PluginLoader + +logger = logging.getLogger(__name__) # pylint:disable=invalid-name + + +class Extractor(): + """ Creates a detect/align pipeline and returns results from a generator + + Input queue is dynamically set depending on the current phase of extraction + and can be accessed from: + Extractor.input_queue + """ + def __init__(self, detector, aligner, loglevel, + multiprocess=False, rotate_images=None, min_size=20): + logger.debug("Initializing %s: (detector: %s, aligner: %s, loglevel: %s, multiprocess: %s" + ", rotate_images: %s, min_size: %s)", self.__class__.__name__, detector, + aligner, loglevel, multiprocess, rotate_images, min_size) + self.phase = "detect" + self.detector = self.load_detector(detector, loglevel, rotate_images, min_size) + self.aligner = self.load_aligner(aligner, loglevel) + self.is_parallel = self.set_parallel_processing(multiprocess) + self.processes = list() + self.queues = self.add_queues() + logger.debug("Initialized %s", self.__class__.__name__) + + @property + def input_queue(self): + """ Return the correct input queue depending on the current phase """ + if self.is_parallel or self.phase == "detect": + qname = "extract_detect_in" + else: + qname = "extract_align_in" + retval = self.queues[qname] + logger.trace("%s: %s", qname, retval) + return retval + + @property + def output_queue(self): + """ Return the correct output queue depending on the current phase """ + qname = "extract_align_out" if self.final_pass else "extract_align_in" + retval = self.queues[qname] + logger.trace("%s: %s", qname, retval) + return retval + + @property + def passes(self): + """ Return the number of passes the extractor needs to make """ + retval = 1 if self.is_parallel else 2 + logger.trace(retval) + return retval + + @property + def final_pass(self): + """ Return true if this is the final extractor pass """ + retval = self.is_parallel or self.phase == "align" + logger.trace(retval) + return retval + + @staticmethod + def load_detector(detector, loglevel, rotation, min_size): + """ Set global arguments and load detector plugin """ + detector_name = detector.replace("-", "_").lower() + logger.debug("Loading Detector: '%s'", detector_name) + detector = PluginLoader.get_detector(detector_name)(loglevel=loglevel, + rotation=rotation, + min_size=min_size) + return detector + + @staticmethod + def load_aligner(aligner, loglevel): + """ Set global arguments and load aligner plugin """ + aligner_name = aligner.replace("-", "_").lower() + logger.debug("Loading Aligner: '%s'", aligner_name) + aligner = PluginLoader.get_aligner(aligner_name)(loglevel=loglevel) + return aligner + + def set_parallel_processing(self, multiprocess): + """ Set whether to run detect and align together or separately """ + 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: + logger.debug("At least one of aligner or detector have no VRAM requirement. " + "Enabling parallel processing.") + return True + + if not multiprocess: + logger.info("NB: Parallel processing disabled.You may get faster " + "extraction speeds by enabling it with the -mp switch") + return False + + required_vram = detector_vram + aligner_vram + 320 # 320MB buffer + stats = gpu_stats.get_card_most_free() + free_vram = int(stats["free"]) + logger.verbose("%s - %sMB free of %sMB", + stats["device"], + free_vram, + int(stats["total"])) + if free_vram <= required_vram: + logger.warning("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 """ + queues = dict() + for task in ("extract_detect_in", "extract_align_in", "extract_align_out"): + size = 0 + if task == "extract_detect_in" or (not self.is_parallel + and task == "extract_align_in"): + # Limit queue size on input queues to avoid stacking ram + size = 100 + queue_manager.add_queue(task, maxsize=size) + queues[task] = queue_manager.get_queue(task) + logger.debug("Queues: %s", queues) + return queues + + def launch(self): + """ Launches the plugins + This can be called multiple times depending on the phase/whether multiprocessing + is enabled. + + If multiprocessing: + launches both plugins, but aligner first so that it's VRAM can be allocated + prior to giving the remaining to the detector + If not multiprocessing: + Launches the relevant plugin for the current phase """ + if self.is_parallel: + logger.debug("Launching aligner and detector") + self.launch_aligner() + self.launch_detector() + elif self.phase == "detect": + logger.debug("Launching detector") + self.launch_detector() + else: + logger.debug("Launching aligner") + self.launch_aligner() + + def launch_aligner(self): + """ Launch the face aligner """ + logger.debug("Launching Aligner") + kwargs = {"in_queue": self.queues["extract_align_in"], + "out_queue": self.queues["extract_align_out"]} + + process = SpawnProcess(self.aligner.run, **kwargs) + event = process.event + error = process.error + process.start() + self.processes.append(process) + + # 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 + for mins in reversed(range(5)): + for seconds in range(60): + event.wait(seconds) + if event.is_set(): + break + if error.is_set(): + break + if event.is_set(): + break + if mins == 0 or error.is_set(): + raise ValueError("Error initializing Aligner") + logger.info("Waiting for Aligner... Time out in %s minutes", mins) + + logger.debug("Launched Aligner") + + def launch_detector(self): + """ Launch the face detector """ + logger.debug("Launching Detector") + kwargs = {"in_queue": self.queues["extract_detect_in"], + "out_queue": self.queues["extract_align_in"]} + mp_func = PoolProcess if self.detector.parent_is_pool else SpawnProcess + process = mp_func(self.detector.run, **kwargs) + + event = process.event if hasattr(process, "event") else None + error = process.error if hasattr(process, "error") else None + process.start() + self.processes.append(process) + + if event is None: + logger.debug("Launched Detector") + return + + for mins in reversed(range(5)): + for seconds in range(60): + event.wait(seconds) + if event.is_set(): + break + if error and error.is_set(): + break + if event.is_set(): + break + if mins == 0 or (error and error.is_set()): + raise ValueError("Error initializing Detector") + logger.info("Waiting for Detector... Time out in %s minutes", mins) + + logger.debug("Launched Detector") + + def detected_faces(self): + """ Detect faces from in an image """ + logger.debug("Running Detection. Phase: '%s'", self.phase) + # If not multiprocessing, intercept the align in queue for + # detection phase + out_queue = self.output_queue + while True: + try: + faces = out_queue.get(True, 1) + if faces == "EOF": + break + if isinstance(faces, dict) and faces.get("exception"): + pid = faces["exception"][0] + t_back = faces["exception"][1].getvalue() + err = "Error in child process {}. {}".format(pid, t_back) + raise Exception(err) + except QueueEmpty: + continue + + yield faces + for process in self.processes: + logger.trace("Joining process: %s", process) + process.join() + del process + if self.final_pass: + # Cleanup queues + for q_name in self.queues.keys(): + queue_manager.del_queue(q_name) + logger.debug("Detection Complete") + else: + logger.debug("Switching to align phase") + self.phase = "align" diff --git a/scripts/convert.py b/scripts/convert.py index a1c6849e18..bdd180c226 100644 --- a/scripts/convert.py +++ b/scripts/convert.py @@ -18,10 +18,9 @@ from lib.multithreading import MultiThread, PoolProcess, total_cpus from lib.queue_manager import queue_manager, QueueEmpty from lib.utils import get_folder, get_image_paths, hash_image_file +from plugins.extract.pipeline import Extractor from plugins.plugin_loader import PluginLoader -from .extract import Plugins as Extractor - logger = logging.getLogger(__name__) # pylint: disable=invalid-name @@ -41,7 +40,6 @@ def __init__(self, arguments): self.add_queues() self.disk_io = DiskIO(self.alignments, self.images, arguments) - self.extractor = None self.predictor = Predict(self.disk_io.load_queue, self.queue_size, arguments) self.converter = Converter(get_folder(self.args.output_dir), self.predictor.output_size, @@ -171,9 +169,7 @@ def __init__(self, alignments, images, arguments): self.imageidxre = re.compile(r"(\d+)(?!.*\d\.)(?=\.\w+$)") # Extractor for on the fly detection - self.extractor = None - if not self.alignments.have_alignments_file: - self.load_extractor() + self.extractor = self.load_extractor() self.load_queue = None self.save_queue = None @@ -236,18 +232,23 @@ def get_frame_ranges(self): def load_extractor(self): """ Set on the fly extraction """ + if self.alignments.have_alignments_file: + return None + logger.debug("Loading extractor") logger.warning("No Alignments file found. Extracting on the fly.") logger.warning("NB: This will use the inferior cv2-dnn for extraction " "and landmarks. It is recommended to perfom Extract first for " "superior results") - extract_args = {"detector": "cv2_dnn", - "aligner": "cv2_dnn", - "loglevel": self.args.loglevel} - self.extractor = Extractor(None, converter_args=extract_args) - self.extractor.launch_detector() - self.extractor.launch_aligner() + extractor = Extractor(detector="cv2-dnn", + aligner="cv2-dnn", + loglevel=self.args.loglevel, + multiprocess=False, + rotate_images=None, + min_size=20) + extractor.launch() logger.debug("Loaded extractor") + return extractor def init_threads(self): """ Initialize queues and threads """ @@ -283,7 +284,6 @@ def start_thread(self, task): def load(self, *args): # pylint: disable=unused-argument """ Load the images with detected_faces""" logger.debug("Load Images: Start") - extract_queue = queue_manager.get_queue("extract_in") if self.extractor else None idx = 0 for filename, image in self.images.load(): idx += 1 @@ -302,7 +302,7 @@ def load(self, *args): # pylint: disable=unused-argument logger.trace("Discarding frame: '%s'", filename) continue - detected_faces = self.get_detected_faces(filename, image, extract_queue) + detected_faces = self.get_detected_faces(filename, image) item = dict(filename=filename, image=image, detected_faces=detected_faces) self.pre_process.do_actions(item) self.load_queue.put(item) @@ -323,13 +323,13 @@ def check_skipframe(self, filename): skipframe = not any(map(lambda b: b[0] <= idx <= b[1], self.frame_ranges)) return skipframe - def get_detected_faces(self, filename, image, extract_queue): + def get_detected_faces(self, filename, image): """ Return detected faces from alignments or detector """ logger.trace("Getting faces for: '%s'", filename) if not self.extractor: detected_faces = self.alignments_faces(os.path.basename(filename), image) else: - detected_faces = self.detect_faces(extract_queue, filename, image) + detected_faces = self.detect_faces(filename, image) logger.trace("Got %s faces for: '%s'", len(detected_faces), filename) return detected_faces @@ -355,12 +355,12 @@ def check_alignments(self, frame): "skipping".format(frame)) return have_alignments - def detect_faces(self, load_queue, filename, image): + def detect_faces(self, 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()) + self.extractor.input_queue.put(inp) + faces = next(self.extractor.detected_faces()) landmarks = faces["landmarks"] detected_faces = faces["detected_faces"] diff --git a/scripts/extract.py b/scripts/extract.py index e6cbd4b97f..4f81ef7f12 100644 --- a/scripts/extract.py +++ b/scripts/extract.py @@ -9,11 +9,10 @@ from tqdm import tqdm from lib.faces_detect import DetectedFace -from lib.gpu_stats import GPUStats -from lib.multithreading import MultiThread, PoolProcess, SpawnProcess -from lib.queue_manager import queue_manager, QueueEmpty +from lib.multithreading import MultiThread +from lib.queue_manager import queue_manager from lib.utils import get_folder, hash_encode_image -from plugins.plugin_loader import PluginLoader +from plugins.extract.pipeline import Extractor from scripts.fsmedia import Alignments, Images, PostProcess, Utils tqdm.monitor_interval = 0 # workaround for TqdmSynchronisationWarning @@ -25,14 +24,20 @@ class Extract(): def __init__(self, arguments): logger.debug("Initializing %s: (args: %s", self.__class__.__name__, arguments) self.args = arguments + Utils.set_verbosity(self.args.loglevel) self.output_dir = get_folder(self.args.output_dir) logger.info("Output Directory: %s", self.args.output_dir) self.images = Images(self.args) self.alignments = Alignments(self.args, True, self.images.is_video) - self.plugins = Plugins(self.args) - self.post_process = PostProcess(arguments) - + self.extractor = Extractor(self.args.detector, + self.args.aligner, + self.args.loglevel, + self.args.multiprocess, + self.args.rotate_images, + self.args.min_size) + + self.save_queue = queue_manager.get_queue("extract_save") self.verify_output = False self.save_interval = None if hasattr(self.args, "save_interval"): @@ -47,8 +52,7 @@ def skip_num(self): def process(self): """ Perform the extraction process """ logger.info('Starting, this may take a while...') - Utils.set_verbosity(self.args.loglevel) - # queue_manager.debug_monitor(1) + # queue_manager.debug_monitor(3) self.threaded_io("load") save_thread = self.threaded_io("save") self.run_extraction() @@ -75,7 +79,7 @@ def threaded_io(self, task, io_args=None): def load_images(self): """ Load the images """ logger.debug("Load Images: Start") - load_queue = queue_manager.get_queue("extract_in") + load_queue = self.extractor.input_queue idx = 0 for filename, image in self.images.load(): idx += 1 @@ -102,7 +106,7 @@ def load_images(self): def reload_images(self, detected_faces): """ Reload the images and pair to detected face """ logger.debug("Reload Images: Start. Detected Faces Count: %s", len(detected_faces)) - load_queue = queue_manager.get_queue("detect") + load_queue = self.extractor.input_queue for filename, image in self.images.load(): if load_queue.shutdown.is_set(): logger.debug("Reload Queue: Stop signal received. Terminating") @@ -117,16 +121,15 @@ def reload_images(self, detected_faces): load_queue.put("EOF") logger.debug("Reload Images: Complete") - @staticmethod - def save_faces(): + def save_faces(self): """ Save the generated faces """ logger.debug("Save Faces: Start") - save_queue = queue_manager.get_queue("extract_out") while True: - if save_queue.shutdown.is_set(): + if self.save_queue.shutdown.is_set(): logger.debug("Save Queue: Stop signal received. Terminating") break - item = save_queue.get() + item = self.save_queue.get() + logger.trace(item) if item == "EOF": break filename, face = item @@ -140,50 +143,6 @@ def save_faces(): continue logger.debug("Save Faces: Complete") - def run_extraction(self): - """ Run Face Detection """ - save_queue = queue_manager.get_queue("extract_out") - to_process = self.process_item_count() - frame_no = 0 - size = self.args.size if hasattr(self.args, "size") else 256 - align_eyes = self.args.align_eyes if hasattr(self.args, "align_eyes") else False - - if self.plugins.is_parallel: - logger.debug("Using parallel processing") - self.plugins.launch_aligner() - self.plugins.launch_detector() - if not self.plugins.is_parallel: - logger.debug("Using serial processing") - self.run_detection(to_process) - self.plugins.launch_aligner() - - for faces in tqdm(self.plugins.detect_faces(extract_pass="align"), - total=to_process, - file=sys.stdout, - desc="Extracting faces"): - - filename = faces["filename"] - - self.align_face(faces, align_eyes, size, filename) - self.post_process.do_actions(faces) - - faces_count = len(faces["detected_faces"]) - if faces_count == 0: - logger.verbose("No faces were detected in image: %s", - os.path.basename(filename)) - - if not self.verify_output and faces_count > 1: - self.verify_output = True - - self.output_faces(filename, faces, save_queue) - - frame_no += 1 - if frame_no == self.save_interval: - self.alignments.save() - frame_no = 0 - - save_queue.put("EOF") - def process_item_count(self): """ Return the number of items to be processedd """ processed = sum(os.path.basename(frame) in self.alignments.data.keys() @@ -203,24 +162,60 @@ def process_item_count(self): exit(0) return to_process - def run_detection(self, to_process): - """ Run detection only """ - self.plugins.launch_detector() - detected_faces = dict() - for detected in tqdm(self.plugins.detect_faces(extract_pass="detect"), - total=to_process, - file=sys.stdout, - desc="Detecting faces"): - exception = detected.get("exception", False) + def run_extraction(self): + """ Run Face Detection """ + to_process = self.process_item_count() + size = self.args.size if hasattr(self.args, "size") else 256 + align_eyes = self.args.align_eyes if hasattr(self.args, "align_eyes") else False + exception = False + + for phase in range(self.extractor.passes): if exception: break + is_final = self.extractor.final_pass + detected_faces = dict() + self.extractor.launch() + for idx, faces in enumerate(tqdm(self.extractor.detected_faces(), + total=to_process, + file=sys.stdout, + desc="Running pass {} of {}: {}".format( + phase + 1, + self.extractor.passes, + self.extractor.phase.title()))): + + exception = faces.get("exception", False) + if exception: + break + filename = faces["filename"] + + if self.extractor.final_pass: + self.output_processing(faces, align_eyes, size, filename) + self.output_faces(filename, faces) + if self.save_interval and idx + 1 % self.save_interval == 0: + self.alignments.save() + else: + del faces["image"] + detected_faces[filename] = faces + + if is_final: + logger.debug("Putting EOF to save") + self.save_queue.put("EOF") + else: + logger.debug("Reloading images") + self.threaded_io("reload", detected_faces) - del detected["image"] - filename = detected["filename"] + def output_processing(self, faces, align_eyes, size, filename): + """ Prepare faces for output """ + self.align_face(faces, align_eyes, size, filename) + self.post_process.do_actions(faces) - detected_faces[filename] = detected + faces_count = len(faces["detected_faces"]) + if faces_count == 0: + logger.verbose("No faces were detected in image: %s", + os.path.basename(filename)) - self.threaded_io("reload", detected_faces) + if not self.verify_output and faces_count > 1: + self.verify_output = True def align_face(self, faces, align_eyes, size, filename): """ Align the detected face and add the destination file path """ @@ -237,7 +232,7 @@ def align_face(self, faces, align_eyes, size, filename): "face": detected_face}) faces["detected_faces"] = final_faces - def output_faces(self, filename, faces, save_queue): + def output_faces(self, filename, faces): """ Output faces to save thread """ final_faces = list() for idx, detected_face in enumerate(faces["detected_faces"]): @@ -249,194 +244,6 @@ def output_faces(self, filename, faces, save_queue): resized_face = face.aligned_face face.hash, img = hash_encode_image(resized_face, extension) - save_queue.put((out_filename, img)) + self.save_queue.put((out_filename, img)) final_faces.append(face.to_alignment()) self.alignments.data[os.path.basename(filename)] = final_faces - - -class Plugins(): - """ Detector and Aligner Plugins and queues """ - def __init__(self, arguments, converter_args=None): - logger.debug("Initializing %s: (converter_args: %s)", - self.__class__.__name__, converter_args) - self.args = arguments - self.converter_args = converter_args # Arguments from converter for on the fly extract - if converter_args is not None: - self.loglevel = converter_args["loglevel"] - else: - self.loglevel = self.args.loglevel - - 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() - logger.debug("Initialized %s", self.__class__.__name__) - - def set_parallel_processing(self): - """ Set whether to run detect and align together or separately """ - 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): - logger.debug("At least one of aligner or detector have no VRAM requirement. " - "Enabling parallel processing.") - return True - - if hasattr(self.args, "multiprocess") and not self.args.multiprocess: - logger.info("NB: Parallel processing disabled.You may get faster " - "extraction speeds by enabling it with the -mp switch") - return False - - required_vram = detector_vram + aligner_vram + 320 # 320MB buffer - stats = gpu_stats.get_card_most_free() - free_vram = int(stats["free"]) - logger.verbose("%s - %sMB free of %sMB", - stats["device"], - free_vram, - int(stats["total"])) - if free_vram <= required_vram: - logger.warning("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 - if task == "load": - q_name = "extract_in" - elif task == "save": - q_name = "extract_out" - else: - q_name = task - queue_manager.add_queue(q_name, maxsize=size) - - def load_detector(self): - """ Set global arguments and load detector plugin """ - if not self.converter_args: - detector_name = self.args.detector.replace("-", "_").lower() - else: - detector_name = self.converter_args["detector"] - logger.debug("Loading Detector: '%s'", detector_name) - # Rotation - rotation = self.args.rotate_images if hasattr(self.args, "rotate_images") else None - # Min acceptable face size: - min_size = self.args.min_size if hasattr(self.args, "min_size") else 0 - - detector = PluginLoader.get_detector(detector_name)( - loglevel=self.loglevel, - rotation=rotation, - min_size=min_size) - - return detector - - def load_aligner(self): - """ Set global arguments and load aligner plugin """ - if not self.converter_args: - aligner_name = self.args.aligner.replace("-", "_").lower() - else: - aligner_name = self.converter_args["aligner"] - - logger.debug("Loading Aligner: '%s'", aligner_name) - - aligner = PluginLoader.get_aligner(aligner_name)( - loglevel=self.loglevel) - - return aligner - - def launch_aligner(self): - """ Launch the face aligner """ - logger.debug("Launching Aligner") - out_queue = queue_manager.get_queue("align") - kwargs = {"in_queue": queue_manager.get_queue("detect"), - "out_queue": out_queue} - - self.process_align = SpawnProcess(self.aligner.run, **kwargs) - event = self.process_align.event - error = self.process_align.error - self.process_align.start() - - # 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 - for mins in reversed(range(5)): - for seconds in range(60): - event.wait(seconds) - if event.is_set(): - break - if error.is_set(): - break - if event.is_set(): - break - if mins == 0 or error.is_set(): - raise ValueError("Error initializing Aligner") - logger.info("Waiting for Aligner... Time out in %s minutes", mins) - - logger.debug("Launched Aligner") - - def launch_detector(self): - """ Launch the face detector """ - logger.debug("Launching Detector") - out_queue = queue_manager.get_queue("detect") - kwargs = {"in_queue": queue_manager.get_queue("extract_in"), - "out_queue": out_queue} - if self.converter_args: - kwargs["processes"] = 1 - mp_func = PoolProcess if self.detector.parent_is_pool else SpawnProcess - self.process_detect = mp_func(self.detector.run, **kwargs) - - event = self.process_detect.event if hasattr(self.process_detect, "event") else None - error = self.process_detect.error if hasattr(self.process_detect, "error") else None - self.process_detect.start() - - if event is None: - logger.debug("Launched Detector") - return - - for mins in reversed(range(5)): - for seconds in range(60): - event.wait(seconds) - if event.is_set(): - break - if error and error.is_set(): - break - if event.is_set(): - break - if mins == 0 or (error and error.is_set()): - raise ValueError("Error initializing Detector") - logger.info("Waiting for Detector... Time out in %s minutes", mins) - - logger.debug("Launched Detector") - - def detect_faces(self, extract_pass="detect"): - """ Detect faces from in an image """ - logger.debug("Running Detection. Pass: '%s'", extract_pass) - 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 - if isinstance(faces, dict) and faces.get("exception"): - pid = faces["exception"][0] - t_back = faces["exception"][1].getvalue() - err = "Error in child process {}. {}".format(pid, t_back) - raise Exception(err) - except QueueEmpty: - continue - - yield faces - logger.debug("Detection Complete") diff --git a/scripts/fsmedia.py b/scripts/fsmedia.py index 71cfea74f2..6c420882f5 100644 --- a/scripts/fsmedia.py +++ b/scripts/fsmedia.py @@ -262,7 +262,10 @@ def get_items(self): if ((hasattr(self.args, "filter") and self.args.filter is not None) or (hasattr(self.args, "nfilter") and self.args.nfilter is not None)): - face_filter = dict() + face_filter = dict(detector=self.args.detector.replace("-", "_").lower(), + aligner=self.args.aligner.replace("-", "_").lower(), + loglevel=self.args.loglevel, + multiprocess=self.args.multiprocess) filter_lists = dict() if hasattr(self.args, "ref_threshold"): face_filter["ref_threshold"] = self.args.ref_threshold @@ -378,12 +381,12 @@ class FaceFilter(PostProcessAction): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - filter_lists = kwargs["filter_lists"] - ref_threshold = kwargs.get("ref_threshold", 0.6) - self.filter = self.load_face_filter(filter_lists, ref_threshold) + logger.info("Extracting and aligning face for Face Filter...") + self.filter = self.load_face_filter(**kwargs) logger.debug("Initialized %s", self.__class__.__name__) - def load_face_filter(self, filter_lists, ref_threshold): + def load_face_filter(self, filter_lists, ref_threshold, aligner, detector, loglevel, + multiprocess): """ Load faces to filter out of images """ if not any(val for val in filter_lists.values()): return None @@ -394,8 +397,12 @@ def load_face_filter(self, filter_lists, ref_threshold): if any(filters for filters in filter_files): facefilter = FilterFunc(filter_files[0], filter_files[1], + detector, + aligner, + loglevel, + multiprocess, ref_threshold) - logger.debug("Face filter: %s", facefilter) + logger.debug("Face filter: %s", facefilter) return facefilter @staticmethod @@ -418,7 +425,7 @@ def process(self, output_item): for idx, detect_face in enumerate(output_item["detected_faces"]): check_item = detect_face["face"] if isinstance(detect_face, dict) else detect_face if not self.filter.check(check_item): - logger.verbose("Skipping not recognized face! Frame: %s Face %s", + logger.verbose("Skipping not recognized face: (Frame: %s Face %s)", output_item["filename"], idx) continue logger.trace("Accepting recognised face. Frame: %s. Face: %s",