#!/usr/bin/env python3 """ Sorting methods for the sorting tool. All sorting methods inherit from :class:`SortMethod` and control functions for scorting one item, sorting a full list of scores and binning based on those sorted scores. """ import logging import operator import sys from typing import Any, cast, Dict, Generator, List, Optional, Tuple, TYPE_CHECKING, Union import cv2 import numpy as np from tqdm import tqdm from lib.align import AlignedFace, DetectedFace from lib.image import FacesLoader, ImagesLoader, read_image_meta_batch, update_existing_metadata from lib.utils import FaceswapError from plugins.extract.recognition.vgg_face2 import Cluster, Recognition as VGGFace if sys.version_info < (3, 8): from typing_extensions import Literal else: from typing import Literal if TYPE_CHECKING: from argparse import Namespace from lib.align.alignments import PNGHeaderAlignmentsDict, PNGHeaderSourceDict logger = logging.getLogger(__name__) ImgMetaType = Generator[Tuple[str, Optional[np.ndarray], Optional["PNGHeaderAlignmentsDict"]], None, None] class InfoLoader(): """ Loads aligned faces and/or face metadata Parameters ---------- input_dir: str Full path to containing folder of faces to be supported loader_type: ["face", "meta", "all"] Dictates the type of iterator that will be used. "face" just loads the image with the filename, "meta" just loads the image alignment data with the filename. "all" loads the image and the alignment data with the filename """ def __init__(self, input_dir: str, info_type: Literal["face", "meta", "all"]) -> None: logger.debug("Initializing: %s (input_dir: %s, info_type: %s)", self.__class__.__name__, input_dir, info_type) self._info_type = info_type self._iterator = None self._description = "Reading image statistics..." self._loader = ImagesLoader(input_dir) if info_type == "face" else FacesLoader(input_dir) self._cached_source_data: Dict[str, "PNGHeaderSourceDict"] = {} if self._loader.count == 0: logger.error("No images to process in location: '%s'", input_dir) sys.exit(1) logger.debug("Initialized: %s", self.__class__.__name__) @property def filelist_count(self) -> int: """ int: The number of files to be processed """ return len(self._loader.file_list) def _get_iterator(self) -> ImgMetaType: """ Obtain the iterator for the selected :attr:`info_type`. Returns ------- generator The correct generator for the given info_type """ if self._info_type == "all": return self._full_data_reader() if self._info_type == "meta": return self._metadata_reader() return self._image_data_reader() def __call__(self) -> ImgMetaType: """ Return the selected iterator The resulting generator: Yields ------ filename: str The filename that has been read image: :class:`numpy.ndarray or ``None`` The aligned face image loaded from disk for 'face' and 'all' info_types otherwise ``None`` alignments: dict or ``None`` The alignments dict for 'all' and 'meta' infor_types otherwise ``None`` """ iterator = self._get_iterator() return iterator def _get_alignments(self, filename: str, metadata: Dict[str, Any]) -> Optional["PNGHeaderAlignmentsDict"]: """ Obtain the alignments from a PNG Header. The other image metadata is cached locally in case a sort method needs to write back to the PNG header Parameters ---------- filename: str Full path to the image PNG file metadata: dict The header data from a PNG file Returns ------- dict or ``None`` The alignments dictionary from the PNG header, if it exists, otherwise ``None`` """ if not metadata or not metadata.get("alignments") or not metadata.get("source"): return None self._cached_source_data[filename] = metadata["source"] return metadata["alignments"] def _metadata_reader(self) -> ImgMetaType: """ Load metadata from saved aligned faces Yields ------ filename: str The filename that has been read image: None This will always be ``None`` with the metadata reader alignments: dict or ``None`` The alignment data for the given face or ``None`` if no alignments found """ for filename, metadata in tqdm(read_image_meta_batch(self._loader.file_list), total=self._loader.count, desc=self._description, leave=False): alignments = self._get_alignments(filename, metadata.get("itxt", {})) yield filename, None, alignments def _full_data_reader(self) -> ImgMetaType: """ Load the image and metadata from a folder of aligned faces Yields ------ filename: str The filename that has been read image: :class:`numpy.ndarray The aligned face image loaded from disk alignments: dict or ``None`` The alignment data for the given face or ``None`` if no alignments found """ for filename, image, metadata in tqdm(self._loader.load(), desc=self._description, total=self._loader.count, leave=False): alignments = self._get_alignments(filename, metadata) yield filename, image, alignments def _image_data_reader(self) -> ImgMetaType: """ Just loads the images with their filenames Yields ------ filename: str The filename that has been read image: :class:`numpy.ndarray The aligned face image loaded from disk alignments: ``None`` Alignments will always be ``None`` with the image data reader """ for filename, image in tqdm(self._loader.load(), desc=self._description, total=self._loader.count, leave=False): yield filename, image, None def update_png_header(self, filename: str, alignments: "PNGHeaderAlignmentsDict") -> None: """ Update the PNG header of the given file with the given alignments. NB: Header information can only be updated if the face is already on at least alignment version 2.2. If below this version, then the header is not updated Parameters ---------- filename: str Full path to the PNG file to update alignments: dict The alignments to update into the PNG header """ vers = self._cached_source_data[filename]["alignments_version"] if vers < 2.2: return self._cached_source_data[filename]["alignments_version"] = 2.3 if vers == 2.2 else vers header = dict(alignments=alignments, source=self._cached_source_data[filename]) update_existing_metadata(filename, header) class SortMethod(): """ Parent class for sort methods. All sort methods should inherit from this class Parameters ---------- arguments: :class:`argparse.Namespace` The command line arguments passed to the sort process loader_type: ["face", "meta", "all"] The type of image loader to use. "face" just loads the image with the filename, "meta" just loads the image alignment data with the filename. "all" loads the image and the alignment data with the filename is_group: bool, optional Set to ``True`` if this class is going to be called exclusively for binning. Default: ``False`` """ def __init__(self, arguments: "Namespace", loader_type: Literal["face", "meta", "all"] = "meta", is_group: bool = False) -> None: logger.debug("Initializing %s: loader_type: '%s' is_group: %s, arguments: %s", self.__class__.__name__, loader_type, is_group, arguments) self._is_group = is_group self._log_once = True self._method = arguments.group_method if self._is_group else arguments.sort_method self._num_bins: int = arguments.num_bins self._bin_names: List[str] = [] self._loader_type = loader_type self._iterator = self._get_file_iterator(arguments.input_dir) self._result: List[Tuple[str, Union[float, np.ndarray]]] = [] self._binned: List[List[str]] = [] logger.debug("Initialized %s", self.__class__.__name__) @property def loader_type(self) -> Literal["face", "meta", "all"]: """ ["face", "meta", "all"]: The loader that this sorter uses """ return self._loader_type @property def binned(self) -> List[List[str]]: """ list: List of bins (list) containing the filenames belonging to the bin. The binning process is called when this property is first accessed""" if not self._binned: self._binned = self._binning() logger.debug({f"bin_{idx}": len(bin_) for idx, bin_ in enumerate(self._binned)}) return self._binned @property def sorted_filelist(self) -> List[str]: """ list: List of sorted filenames for given sorter in a single list. The sort process is called when this property is first accessed """ if not self._result: self._sort_filelist() retval = [item[0] for item in self._result] logger.debug(retval) else: retval = [item[0] for item in self._result] return retval @property def bin_names(self) -> List[str]: """ list: The name of each created bin, if they exist, otherwise an empty list """ return self._bin_names def _get_file_iterator(self, input_dir: str) -> InfoLoader: """ Override for method specific iterators. Parameters ---------- input_dir: str Full path to containing folder of faces to be supported Returns ------- :class:`InfoLoader` The correct InfoLoader iterator for the current sort method """ return InfoLoader(input_dir, self.loader_type) def _sort_filelist(self) -> None: """ Call the sort method's logic to populate the :attr:`_results` attribute. Put logic for scoring an individual frame in in :attr:`score_image` of the child Returns ------- list The sorted file. A list of tuples with the filename in the first position and score in the second position """ for filename, image, alignments in self._iterator(): self.score_image(filename, image, alignments) self.sort() logger.debug("sorted list: %s", [r[0] if isinstance(r, (tuple, list)) else r for r in self._result]) @classmethod def _get_unique_labels(cls, numbers: np.ndarray) -> List[str]: """ For a list of threshold values for displaying in the bin name, get the lowest number of decimal figures (down to int) required to have a unique set of folder names and return the formatted numbers. Parameters ---------- numbers: :class:`numpy.ndarray` The list of floating point threshold numbers being used as boundary points Returns ------- list[str] The string formatted numbers at the lowest precision possible to represent them uniquely """ i = 0 while True: rounded = [round(n, i) for n in numbers] if len(set(rounded)) == len(numbers): break i += 1 if i == 0: retval = [str(int(n)) for n in rounded] else: pre, post = zip(*[str(r).split(".") for r in rounded]) rpad = max(len(x) for x in post) retval = [f"{str(int(left))}.{str(int(right)).ljust(rpad, '0')}" for left, right in zip(pre, post)] logger.debug("rounded values: %s, formatted labels: %s", rounded, retval) return retval def _binning_linear_threshold(self, units: str = "", multiplier: int = 1) -> List[List[str]]: """ Standard linear binning method for binning by threshold. The minimum and maximum result from :attr:`_result` are taken, A range is created between these min and max values and is divided to get the number of bins to hold the data Parameters ---------- units, str, optional The units to use for the bin name for displaying the threshold values. This this should correspond the value in position 1 of :attr:`_result`. Default: "" (no units) multiplier: int, optional The amount to multiply the contents in position 1 of :attr:`_results` for displaying in the bin folder name Returns ------- list List of bins of filenames """ sizes = np.array([i[1] for i in self._result]) thresholds = np.linspace(sizes.min(), sizes.max(), self._num_bins + 1) labels = self._get_unique_labels(thresholds * multiplier) self._bin_names = [f"{self._method}_{idx:03d}_" f"{labels[idx]}{units}_to_{labels[idx + 1]}{units}" for idx in range(self._num_bins)] bins: List[List[str]] = [[] for _ in range(self._num_bins)] for filename, result in self._result: bin_idx = next(bin_id for bin_id, thresh in enumerate(thresholds) if result <= thresh) - 1 bins[bin_idx].append(filename) return bins def _binning(self) -> List[List[str]]: """ Called when :attr:`binning` is first accessed. Checks if sorting has been done, if not triggers it, then does binning Returns ------- list List of bins of filenames """ if not self._result: self._sort_filelist() retval = self.binning() if not self._bin_names: self._bin_names = [f"{self._method}_{i:03d}" for i in range(len(retval))] logger.debug({bin_name: len(bin_) for bin_name, bin_ in zip(self._bin_names, retval)}) return retval def sort(self) -> None: """ Override for method specific logic for sorting the loaded statistics The scored list :attr:`_result` should be sorted in place """ raise NotImplementedError() def score_image(self, filename: str, image: Optional[np.ndarray], alignments: Optional["PNGHeaderAlignmentsDict"]) -> None: """ Override for sort method's specificic logic. This method should be executed to get a single score from a single image and add the result to :attr:`_result` Parameters ---------- filename: str The filename of the currently processing image image: :class:`np.ndarray` or ``None`` A face image loaded from disk or ``None`` alignments: dict or ``None`` The alignments dictionary for the aligned face or ``None`` """ raise NotImplementedError() def binning(self) -> List[List[str]]: """ Group into bins by their sorted score. Override for method specific binning techniques. Binning takes the results from :attr:`_result` compiled during :func:`_sort_filelist` and organizes into bins for output. Returns ------- list List of bins of filenames """ raise NotImplementedError() @classmethod def _mask_face(cls, image: np.ndarray, alignments: "PNGHeaderAlignmentsDict") -> np.ndarray: """ Function for applying the mask to an aligned face if both the face image and alignment data are available. Parameters ---------- image: :class:`numpy.ndarray` The aligned face image loaded from disk alignments: Dict The alignments data corresponding to the loaded image Returns ------- :class:`numpy.ndarray` The original image with the mask applied """ det_face = DetectedFace() det_face.from_png_meta(alignments) aln_face = AlignedFace(np.array(alignments["landmarks_xy"], dtype="float32"), image=image, centering="legacy", size=256, is_aligned=True) mask = det_face.mask["components"] mask.set_sub_crop(aln_face.pose.offset[mask.stored_centering], aln_face.pose.offset["legacy"], centering="legacy") nmask = cv2.resize(mask.mask, (256, 256), interpolation=cv2.INTER_CUBIC)[..., None] assert aln_face.face is not None return np.minimum(aln_face.face, nmask) class SortMultiMethod(SortMethod): """ A Parent sort method that runs 2 different underlying methods (one for sorting one for binning) in instances where grouping has been requested, but the sort method is different from the group method Parameters ---------- arguments: :class:`argparse.Namespace` The command line arguments passed to the sort process sort_method: :class:`SortMethod` A sort method object for sorting the images group_method: :class:`SortMethod` A sort method object used for sorting and binning the images """ def __init__(self, arguments: "Namespace", sort_method: SortMethod, group_method: SortMethod) -> None: self._sorter = sort_method self._grouper = group_method self._is_built = False super().__init__(arguments) def _get_file_iterator(self, input_dir: str) -> InfoLoader: """ Override to get a group specific iterator. If the sorter and grouper use the same kind of iterator, use that. Otherwise return the 'all' iterator, as which ever way it is cut all outputs will be required. Monkey patch the actual loader used into the children in case of any callbacks. Parameters ---------- input_dir: str Full path to containing folder of faces to be supported Returns ------- :class:`InfoLoader` The correct InfoLoader iterator for the current sort method """ if self._sorter.loader_type == self._grouper.loader_type: retval = InfoLoader(input_dir, self._sorter.loader_type) else: retval = InfoLoader(input_dir, "all") self._sorter._iterator = retval # pylint: disable=protected-access self._grouper._iterator = retval # pylint: disable=protected-access return retval def score_image(self, filename: str, image: Optional[np.ndarray], alignments: Optional["PNGHeaderAlignmentsDict"]) -> None: """ Score a single image for sort method: "distance", "yaw" "pitch" or "size" and add the result to :attr:`_result` Parameters ---------- filename: str The filename of the currently processing image image: :class:`np.ndarray` or ``None`` A face image loaded from disk or ``None`` alignments: dict or ``None`` The alignments dictionary for the aligned face or ``None`` """ self._sorter.score_image(filename, image, alignments) self._grouper.score_image(filename, image, alignments) def sort(self) -> None: """ Sort the sorter and grouper methods """ logger.debug("Sorting") self._sorter.sort() self._result = self._sorter.sorted_filelist # type:ignore self._grouper.sort() self._binned = self._grouper.binned self._bin_names = self._grouper.bin_names logger.debug("Sorted") def binning(self) -> List[List[str]]: """ Override standard binning, to bin by the group-by method and sort by the sorting method. Go through the grouped binned results, and reorder each bin contents based on the sorted list Returns ------- list List of bins of filenames """ sorted_ = self._result output: List[List[str]] = [] for bin_ in tqdm(self._binned, desc="Binning and sorting", file=sys.stdout, leave=False): indices: Dict[int, str] = {} for filename in bin_: indices[sorted_.index(filename)] = filename output.append([indices[idx] for idx in sorted(indices)]) return output class SortBlur(SortMethod): """ Sort images by blur or blur-fft amount Parameters ---------- arguments: :class:`argparse.Namespace` The command line arguments passed to the sort process is_group: bool, optional Set to ``True`` if this class is going to be called exclusively for binning. Default: ``False`` """ def __init__(self, arguments: "Namespace", is_group: bool = False) -> None: super().__init__(arguments, loader_type="all", is_group=is_group) method = arguments.group_method if self._is_group else arguments.sort_method self._use_fft = method == "blur_fft" def estimate_blur(self, image: np.ndarray, alignments=None) -> float: """ Estimate the amount of blur an image has with the variance of the Laplacian. Normalize by pixel number to offset the effect of image size on pixel gradients & variance. Parameters ---------- image: :class:`numpy.ndarray` The face image to calculate blur for alignments: dict, optional The metadata for the face image or ``None`` if no metadata is available. If metadata is provided the face will be masked by the "components" mask prior to calculating blur. Default:``None`` Returns ------- float The estimated blur score for the face """ if alignments is not None: image = self._mask_face(image, alignments) if image.ndim == 3: image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) blur_map = cv2.Laplacian(image, cv2.CV_32F) score = np.var(blur_map) / np.sqrt(image.shape[0] * image.shape[1]) return score def estimate_blur_fft(self, image: np.ndarray, alignments: Optional["PNGHeaderAlignmentsDict"] = None) -> float: """ Estimate the amount of blur a fft filtered image has. Parameters ---------- image: :class:`numpy.ndarray` Use Fourier Transform to analyze the frequency characteristics of the masked face using 2D Discrete Fourier Transform (DFT) filter to find the frequency domain. A mean value is assigned to the magnitude spectrum and returns a blur score. Adapted from https://www.pyimagesearch.com/2020/06/15/ opencv-fast-fourier-transform-fft-for-blur-detection-in-images-and-video-streams/ alignments: dict, optional The metadata for the face image or ``None`` if no metadata is available. If metadata is provided the face will be masked by the "components" mask prior to calculating blur. Default:``None`` Returns ------- float The estimated fft blur score for the face """ if alignments is not None: image = self._mask_face(image, alignments) if image.ndim == 3: image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) height, width = image.shape c_height, c_width = (int(height / 2.0), int(width / 2.0)) fft = np.fft.fft2(image) fft_shift = np.fft.fftshift(fft) fft_shift[c_height - 75:c_height + 75, c_width - 75:c_width + 75] = 0 ifft_shift = np.fft.ifftshift(fft_shift) shift_back = np.fft.ifft2(ifft_shift) magnitude = np.log(np.abs(shift_back)) score = np.mean(magnitude) return score def score_image(self, filename: str, image: Optional[np.ndarray], alignments: Optional["PNGHeaderAlignmentsDict"]) -> None: """ Score a single image for blur or blur-fft and add the result to :attr:`_result` Parameters ---------- filename: str The filename of the currently processing image image: :class:`np.ndarray` A face image loaded from disk alignments: dict or ``None`` The alignments dictionary for the aligned face or ``None`` """ assert image is not None if self._log_once: msg = "Grouping" if self._is_group else "Sorting" inf = "fft_filtered " if self._use_fft else " " logger.info("%s by estimated %simage blur...", msg, inf) self._log_once = False estimator = self.estimate_blur_fft if self._use_fft else self.estimate_blur self._result.append((filename, estimator(image, alignments))) def sort(self) -> None: """ Sort by metric score. Order in reverse for distance sort. """ logger.info("Sorting...") self._result = sorted(self._result, key=operator.itemgetter(1), reverse=True) def binning(self) -> List[List[str]]: """ Create bins to split linearly from the lowest to the highest sample value Returns ------- list List of bins of filenames """ return self._binning_linear_threshold(multiplier=100) class SortColor(SortMethod): """ Score by channel average intensity or black pixels. Parameters ---------- arguments: :class:`argparse.Namespace` The command line arguments passed to the sort process is_group: bool, optional Set to ``True`` if this class is going to be called exclusively for binning. Default: ``False`` """ def __init__(self, arguments: "Namespace", is_group: bool = False) -> None: super().__init__(arguments, loader_type="face", is_group=is_group) self._desired_channel = {'gray': 0, 'luma': 0, 'orange': 1, 'green': 2} method = arguments.group_method if self._is_group else arguments.sort_method self._method = method.replace("color_", "") def _convert_color(self, image: np.ndarray) -> np.ndarray: """ Helper function to convert color spaces Parameters ---------- image: :class:`numpy.ndarray` The original image to convert color space for Returns ------- :class:`numpy.ndarray` The color converted image """ if self._method == 'gray': conversion = np.array([[0.0722], [0.7152], [0.2126]]) else: conversion = np.array([[0.25, 0.5, 0.25], [-0.5, 0.0, 0.5], [-0.25, 0.5, -0.25]]) operation = 'ijk, kl -> ijl' if self._method == "gray" else 'ijl, kl -> ijk' path = np.einsum_path(operation, image[..., :3], conversion, optimize='optimal')[0] return np.einsum(operation, image[..., :3], conversion, optimize=path).astype('float32') def _near_split(self, bin_range: int) -> List[int]: """ Obtain the split for the given number of bins for the given range Parameters ---------- bin_range: int The range of data to separate into bins Returns ------- list The split dividers for the given number of bins for the given range """ quotient, remainder = divmod(bin_range, self._num_bins) seps = [quotient + 1] * remainder + [quotient] * (self._num_bins - remainder) uplimit = 0 bins = [0] for sep in seps: bins.append(uplimit + sep) uplimit += sep return bins def binning(self) -> List[List[str]]: """ Group into bins by percentage of black pixels """ # TODO. Only grouped by black pixels. Check color logger.info("Grouping by percentage of %s...", self._method) # Starting the binning process bins: List[List[str]] = [[] for _ in range(self._num_bins)] # Get edges of bins from 0 to 100 bins_edges = self._near_split(100) # Get the proper bin number for each img order img_bins = np.digitize([float(x[1]) for x in self._result], bins_edges, right=True) # Place imgs in bins for idx, _bin in enumerate(img_bins): bins[_bin].append(self._result[idx][0]) retval = [b for b in bins if b] return retval def score_image(self, filename: str, image: Optional[np.ndarray], alignments: Optional["PNGHeaderAlignmentsDict"]) -> None: """ Score a single image for color Parameters ---------- filename: str The filename of the currently processing image image: :class:`np.ndarray` A face image loaded from disk alignments: dict or ``None`` The alignments dictionary for the aligned face or ``None`` """ if self._log_once: msg = "Grouping" if self._is_group else "Sorting" if self._method == "black": logger.info("%s by percentage of black pixels...", msg) else: logger.info("%s by channel average intensity...", msg) self._log_once = False assert image is not None if self._method == "black": score = np.ndarray.all(image == [0, 0, 0], axis=2).sum()/image.size*100*3 else: channel_to_sort = self._desired_channel[self._method] score = np.average(self._convert_color(image), axis=(0, 1))[channel_to_sort] self._result.append((filename, score)) def sort(self) -> None: """ Sort by metric score. Order in reverse for distance sort. """ if self._method == "black": self._sort_black_pixels() return self._result = sorted(self._result, key=operator.itemgetter(1), reverse=True) def _sort_black_pixels(self) -> None: """ Sort by percentage of black pixels Calculates the sum of black pixels, gets the percentage X 3 channels """ img_list_len = len(self._result) for i in tqdm(range(0, img_list_len - 1), desc="Comparing black pixels", file=sys.stdout, leave=False): for j in range(0, img_list_len-i-1): if self._result[j][1] > self._result[j+1][1]: temp = self._result[j] self._result[j] = self._result[j+1] self._result[j+1] = temp class SortFace(SortMethod): """ Sort by identity similarity using VGG Face 2 Parameters ---------- arguments: :class:`argparse.Namespace` The command line arguments passed to the sort process is_group: bool, optional Set to ``True`` if this class is going to be called exclusively for binning. Default: ``False`` """ def __init__(self, arguments: "Namespace", is_group: bool = False) -> None: super().__init__(arguments, loader_type="all", is_group=is_group) self._vgg_face = VGGFace(exclude_gpus=arguments.exclude_gpus) self._vgg_face.init_model() threshold = arguments.threshold self._output_update_info = True self._threshold: Optional[float] = 0.25 if threshold < 0 else threshold def score_image(self, filename: str, image: Optional[np.ndarray], alignments: Optional["PNGHeaderAlignmentsDict"]) -> None: """ Processing logic for sort by face method. Reads header information from the PNG file to look for VGGFace2 embedding. If it does not exist, the embedding is obtained and added back into the PNG Header. Parameters ---------- filename: str The filename of the currently processing image image: :class:`np.ndarray` A face image loaded from disk alignments: dict or ``None`` The alignments dictionary for the aligned face or ``None`` """ if not alignments: msg = ("The images to be sorted do not contain alignment data. Images must have " "been generated by Faceswap's Extract process.\nIf you are sorting an " "older faceset, then you should re-extract the faces from your source " "alignments file to generate this data.") raise FaceswapError(msg) if self._log_once: msg = "Grouping" if self._is_group else "Sorting" logger.info("%s by identity similarity...", msg) self._log_once = False if alignments.get("identity", {}).get("vggface2"): embedding = np.array(alignments["identity"]["vggface2"], dtype="float32") self._result.append((filename, embedding)) return if self._output_update_info: logger.info("VGG Face2 Embeddings are being written to the image header. " "Sorting by this method will be quicker next time") self._output_update_info = False face = AlignedFace(np.array(alignments["landmarks_xy"], dtype="float32"), image=image, centering="legacy", size=self._vgg_face.input_size, is_aligned=True).face assert face is not None embedding = self._vgg_face.predict(face[None, ...])[0] alignments.setdefault("identity", {})["vggface2"] = embedding.tolist() self._iterator.update_png_header(filename, alignments) self._result.append((filename, embedding)) def sort(self) -> None: """ Sort by dendogram. Parameters ---------- matched_list: list The list of tuples with filename in first position and face encoding in the 2nd Returns ------- list The original list, sorted for this metric """ logger.info("Sorting by ward linkage. This may take some time...") preds = np.array([item[1] for item in self._result]) indices = Cluster(np.array(preds), "ward", threshold=self._threshold)() self._result = [(self._result[idx][0], float(score)) for idx, score in indices] def binning(self) -> List[List[str]]: """ Group into bins by their sorted score The bin ID has been output in the 2nd column of :attr:`_result` so use that for binnin Returns ------- list List of bins of filenames """ num_bins = len(set(int(i[1]) for i in self._result)) logger.info("Grouping by %s...", self.__class__.__name__.replace("Sort", "")) bins: List[List[str]] = [[] for _ in range(num_bins)] for filename, bin_id in self._result: bins[int(bin_id)].append(filename) return bins class SortHistogram(SortMethod): """ Sort by image histogram similarity or dissimilarity Parameters ---------- arguments: :class:`argparse.Namespace` The command line arguments passed to the sort process is_group: bool, optional Set to ``True`` if this class is going to be called exclusively for binning. Default: ``False`` """ def __init__(self, arguments: "Namespace", is_group: bool = False) -> None: super().__init__(arguments, loader_type="all", is_group=is_group) method = arguments.group_method if self._is_group else arguments.sort_method self._is_dissim = method == "hist-dissim" self._threshold: float = 0.3 if arguments.threshold < 0.0 else arguments.threshold def _calc_histogram(self, image: np.ndarray, alignments: Optional["PNGHeaderAlignmentsDict"]) -> np.ndarray: if alignments: image = self._mask_face(image, alignments) return cv2.calcHist([image], [0], None, [256], [0, 256]) def _sort_dissim(self) -> None: """ Sort histograms by dissimilarity """ img_list_len = len(self._result) for i in tqdm(range(0, img_list_len), desc="Comparing histograms", file=sys.stdout, leave=False): score_total = 0 for j in range(0, img_list_len): if i == j: continue score_total += cv2.compareHist(self._result[i][1], self._result[j][1], cv2.HISTCMP_BHATTACHARYYA) self._result[i][2] = score_total self._result = sorted(self._result, key=operator.itemgetter(2), reverse=True) def _sort_sim(self) -> None: """ Sort histograms by similarity """ img_list_len = len(self._result) for i in tqdm(range(0, img_list_len - 1), desc="Comparing histograms", file=sys.stdout, leave=False): min_score = float("inf") j_min_score = i + 1 for j in range(i + 1, img_list_len): score = cv2.compareHist(self._result[i][1], self._result[j][1], cv2.HISTCMP_BHATTACHARYYA) if score < min_score: min_score = score j_min_score = j (self._result[i + 1], self._result[j_min_score]) = (self._result[j_min_score], self._result[i + 1]) @classmethod def _get_avg_score(cls, image: np.ndarray, references: List[np.ndarray]) -> float: """ Return the average histogram score between a face and reference images Parameters ---------- image: :class:`numpy.ndarray` The image to test references: list List of reference images to test the original image against Returns ------- float The average score between the histograms """ scores = [] for img2 in references: score = cv2.compareHist(image, img2, cv2.HISTCMP_BHATTACHARYYA) scores.append(score) return sum(scores) / len(scores) def binning(self) -> List[List[str]]: """ Group into bins by histogram """ msg = "dissimilarity" if self._is_dissim else "similarity" logger.info("Grouping by %s...", msg) # Groups are of the form: group_num -> reference histogram reference_groups: Dict[int, List[np.ndarray]] = {} # Bins array, where index is the group number and value is # an array containing the file paths to the images in that group bins: List[List[str]] = [] threshold = self._threshold img_list_len = len(self._result) reference_groups[0] = [cast(np.ndarray, self._result[0][1])] bins.append([self._result[0][0]]) for i in tqdm(range(1, img_list_len), desc="Grouping", file=sys.stdout, leave=False): current_key = -1 current_score = float("inf") for key, value in reference_groups.items(): score = self._get_avg_score(self._result[i][1], value) if score < current_score: current_key, current_score = key, score if current_score < threshold: reference_groups[cast(int, current_key)].append(self._result[i][1]) bins[current_key].append(self._result[i][0]) else: reference_groups[len(reference_groups)] = [self._result[i][1]] bins.append([self._result[i][0]]) return bins def score_image(self, filename: str, image: Optional[np.ndarray], alignments: Optional["PNGHeaderAlignmentsDict"]) -> None: """ Collect the histogram for the given face Parameters ---------- filename: str The filename of the currently processing image image: :class:`np.ndarray` A face image loaded from disk alignments: dict or ``None`` The alignments dictionary for the aligned face or ``None`` """ if self._log_once: msg = "Grouping" if self._is_group else "Sorting" logger.info("%s by histogram similarity...", msg) self._log_once = False assert image is not None self._result.append((filename, self._calc_histogram(image, alignments))) def sort(self) -> None: """ Sort by histogram. """ logger.info("Comparing histograms and sorting...") if self._is_dissim: self._sort_dissim() return self._sort_sim()