mirror of
https://github.com/deepfakes/faceswap.git
synced 2026-10-18 22:45:59 +08:00
torzdf
975 lines
36 KiB
Python
975 lines
36 KiB
Python
#!/usr/bin/env python3
|
|
"""Sorting methods for the sorting tool.
|
|
|
|
All sorting methods inherit from :class:`SortMethod` and control functions for sorting one item,
|
|
sorting a full list of scores and binning based on those sorted scores.
|
|
"""
|
|
from __future__ import annotations
|
|
import logging
|
|
import operator
|
|
import sys
|
|
import typing as T
|
|
from queue import Queue
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from tqdm import tqdm
|
|
|
|
from lib.align import DetectedFace, LandmarkType
|
|
from lib.multithreading import FSThread
|
|
from lib.utils import get_module_objects, FaceswapError
|
|
from lib.infer.identity import Cluster, Identity
|
|
|
|
from .info_loader import InfoLoader
|
|
|
|
if T.TYPE_CHECKING:
|
|
from argparse import Namespace
|
|
import numpy.typing as npt
|
|
from lib.align.objects import PNGAlignments
|
|
from lib.infer.runner import ExtractRunner
|
|
from lib.infer.handler import ExtractHandlerFace
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class SortMethod():
|
|
"""Parent class for sort methods. All sort methods should inherit from this class
|
|
|
|
Parameters
|
|
----------
|
|
arguments
|
|
The command line arguments passed to the sort process
|
|
loader_type
|
|
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
|
|
Set to ``True`` if this class is going to be called exclusively for binning.
|
|
Default: ``False``
|
|
"""
|
|
_log_mask_once = False
|
|
|
|
def __init__(self,
|
|
arguments: Namespace,
|
|
loader_type: T.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: T.Literal["face", "meta", "all"] = loader_type
|
|
self._iterator = self._get_file_iterator(arguments.input_dir)
|
|
|
|
self._result: list[tuple[str, float | np.ndarray]] = []
|
|
self._binned: list[list[str]] = []
|
|
logger.debug("Initialized %s", self.__class__.__name__)
|
|
|
|
@property
|
|
def loader_type(self) -> T.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 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 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]:
|
|
"""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
|
|
Full path to containing folder of faces to be supported
|
|
|
|
Returns
|
|
-------
|
|
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
|
|
-------
|
|
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
|
|
The list of floating point threshold numbers being used as boundary points
|
|
|
|
Returns
|
|
-------
|
|
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
|
|
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
|
|
The amount to multiply the contents in position 1 of :attr:`_results` for displaying in
|
|
the bin folder name
|
|
|
|
Returns
|
|
-------
|
|
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 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: np.ndarray | None,
|
|
alignments: PNGAlignments | None) -> None:
|
|
"""Override for sort method's specific logic. This method should be executed to get a
|
|
single score from a single image and add the result to :attr:`_result`
|
|
|
|
Parameters
|
|
----------
|
|
filename
|
|
The filename of the currently processing image
|
|
image
|
|
A face image loaded from disk or ``None``
|
|
alignments
|
|
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: PNGAlignments) -> np.ndarray:
|
|
"""Function for applying the mask to an aligned face if both the face image and alignment
|
|
data are available.
|
|
|
|
Parameters
|
|
----------
|
|
image
|
|
The aligned face image loaded from disk
|
|
alignments
|
|
The alignments data corresponding to the loaded image
|
|
|
|
Returns
|
|
-------
|
|
The original image with the mask applied
|
|
"""
|
|
det_face = DetectedFace().from_png_meta(alignments)
|
|
det_face.load_aligned(image,
|
|
size=256,
|
|
centering="legacy",
|
|
is_aligned=True)
|
|
aln_face = det_face.aligned
|
|
if aln_face.landmark_type != LandmarkType.LM_2D_68:
|
|
mask = None
|
|
else:
|
|
mask = det_face.get_landmark_mask("face")
|
|
|
|
if mask is None and not cls._log_mask_once:
|
|
logger.warning("Masks cannot be generated for the available landmark types. Results "
|
|
"are likely to be sub-standard")
|
|
cls._log_mask_once = True
|
|
|
|
assert aln_face.face is not None
|
|
if mask is None:
|
|
return aln_face.face
|
|
|
|
return np.minimum(aln_face.face, mask)
|
|
|
|
|
|
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
|
|
The command line arguments passed to the sort process
|
|
sort_method
|
|
A sort method object for sorting the images
|
|
group_method
|
|
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
|
|
Full path to containing folder of faces to be supported
|
|
|
|
Returns
|
|
-------
|
|
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: np.ndarray | None,
|
|
alignments: PNGAlignments | None) -> None:
|
|
"""Score a single image for sort method and add the result to :attr:`_result`
|
|
|
|
Parameters
|
|
----------
|
|
filename
|
|
The filename of the currently processing image
|
|
image
|
|
A face image loaded from disk or ``None``
|
|
alignments
|
|
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 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 # pyright:ignore[reportArgumentType]
|
|
output.append([indices[idx] for idx in sorted(indices)])
|
|
return output
|
|
|
|
|
|
class SortBlur(SortMethod):
|
|
"""Sort images by blur or blur-fft amount
|
|
|
|
Parameters
|
|
----------
|
|
arguments
|
|
The command line arguments passed to the sort process
|
|
is_group
|
|
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
|
|
The face image to calculate blur for
|
|
alignments
|
|
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
|
|
-------
|
|
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 = T.cast(np.ndarray, 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: PNGAlignments | None = None) -> float:
|
|
"""Estimate the amount of blur a fft filtered image has.
|
|
|
|
Parameters
|
|
----------
|
|
image
|
|
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
|
|
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
|
|
-------
|
|
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 = float(np.mean(magnitude))
|
|
|
|
return score
|
|
|
|
def score_image(self,
|
|
filename: str,
|
|
image: np.ndarray | None,
|
|
alignments: PNGAlignments | None) -> None:
|
|
"""Score a single image for blur or blur-fft and add the result to :attr:`_result`
|
|
|
|
Parameters
|
|
----------
|
|
filename
|
|
The filename of the currently processing image
|
|
image
|
|
A face image loaded from disk
|
|
alignments
|
|
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%s image 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 of bins of filenames
|
|
"""
|
|
return self._binning_linear_threshold(multiplier=100)
|
|
|
|
|
|
class SortColor(SortMethod):
|
|
"""Score by channel average intensity or black pixels.
|
|
|
|
Parameters
|
|
----------
|
|
arguments
|
|
The command line arguments passed to the sort process
|
|
is_group
|
|
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
|
|
The original image to convert color space for
|
|
|
|
Returns
|
|
-------
|
|
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
|
|
The range of data to separate into bins
|
|
|
|
Returns
|
|
-------
|
|
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: np.ndarray | None,
|
|
alignments: PNGAlignments | None) -> None:
|
|
"""Score a single image for color
|
|
|
|
Parameters
|
|
----------
|
|
filename
|
|
The filename of the currently processing image
|
|
image
|
|
A face image loaded from disk
|
|
alignments
|
|
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 an Identity plugin
|
|
|
|
Parameters
|
|
----------
|
|
arguments
|
|
The command line arguments passed to the sort process
|
|
is_group
|
|
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._plugin = Identity(arguments.identity, config_file=arguments.config_file)
|
|
self._runner: ExtractRunner[ExtractHandlerFace] | None = None
|
|
self._storage_name = self._plugin.storage_name
|
|
threshold = arguments.threshold
|
|
self._threshold: float | None = 0.25 if threshold < 0 else threshold
|
|
|
|
self._count_seen = 0
|
|
self._plugin_thread = FSThread(self._score_from_plugin)
|
|
self._from_plugin: list[tuple[str, npt.NDArray[np.float32]]] = []
|
|
self._alignment_queue: Queue[tuple[str, PNGAlignments]] = Queue(
|
|
maxsize=self._plugin.batch_size * 3)
|
|
|
|
def _score_from_header(self,
|
|
filename: str,
|
|
alignments: PNGAlignments) -> bool:
|
|
"""Reads header information from the PNG file to look for the identity embedding
|
|
|
|
Parameters
|
|
----------
|
|
filename
|
|
The filename of the currently processing image
|
|
alignments
|
|
The alignments dictionary for the aligned face or ``None``
|
|
|
|
Returns
|
|
-------
|
|
``True`` if embedding information was read from the PNG header, otherwise ``False``
|
|
"""
|
|
embedding = alignments.identity.get(self._storage_name)
|
|
if embedding is None:
|
|
return False
|
|
self._result.append((filename, embedding))
|
|
return True
|
|
|
|
def _score_from_plugin(self):
|
|
"""Obtain the embedding from the identity plugin"""
|
|
logger.info("%s Embeddings are being written to the image header. "
|
|
"Sorting by this method should be quicker next time",
|
|
self._storage_name.title())
|
|
|
|
assert self._runner is not None
|
|
for media in self._runner:
|
|
if self._plugin_thread.error_state.has_error:
|
|
self._plugin_thread.error_state.re_raise()
|
|
embedding = media.detected_faces[0].identity[self._storage_name]
|
|
self._from_plugin.append((media.filename, embedding))
|
|
filename, alignments = self._alignment_queue.get()
|
|
assert filename == media.filename
|
|
alignments.identity[self._storage_name] = embedding
|
|
self._iterator.update_png_header(filename, alignments)
|
|
|
|
def _handle_plugin(self) -> None:
|
|
"""Check if we have seen all input and shutdown the plugin if so"""
|
|
if self._count_seen != self._iterator.filelist_count:
|
|
return
|
|
|
|
if not self._plugin_thread.is_alive():
|
|
return
|
|
assert self._runner is not None
|
|
logger.debug("Shutting down Identity plugin")
|
|
self._runner.stop()
|
|
self._plugin_thread.join()
|
|
self._result += self._from_plugin
|
|
|
|
def score_image(self,
|
|
filename: str,
|
|
image: np.ndarray | None,
|
|
alignments: PNGAlignments | None) -> None:
|
|
"""Score a single image for sort method and add the result to :attr:`_result`. Attempts
|
|
to pull identity information from the PNG metadata. If not available, pulls the information
|
|
from the Identity plugin and stores in the PNG header for future use
|
|
|
|
Parameters
|
|
----------
|
|
filename
|
|
The filename of the currently processing image
|
|
image
|
|
A face image loaded from disk or ``None``
|
|
alignments
|
|
The alignments dictionary for the aligned face or ``None``
|
|
"""
|
|
# pylint:disable=duplicate-code
|
|
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 face set, then you should re-extract the faces from your source "
|
|
"alignments file to generate this data.")
|
|
raise FaceswapError(msg)
|
|
|
|
if self._plugin_thread.error_state.has_error:
|
|
self._plugin_thread.error_state.re_raise()
|
|
|
|
self._count_seen += 1
|
|
if self._score_from_header(filename, alignments):
|
|
self._handle_plugin()
|
|
return
|
|
|
|
if not self._plugin_thread.is_alive():
|
|
logger.debug("Starting Identity plugin")
|
|
self._runner = self._plugin()
|
|
self._plugin_thread.start()
|
|
|
|
self._alignment_queue.put((filename, alignments))
|
|
|
|
face = DetectedFace(left=alignments.x, # Only include required items
|
|
width=alignments.w,
|
|
top=alignments.y,
|
|
height=alignments.h,
|
|
landmarks_xy=alignments.landmarks_xy)
|
|
assert self._runner is not None
|
|
try:
|
|
self._runner.put(filename,
|
|
T.cast("npt.NDArray[np.uint8]", image),
|
|
[face],
|
|
is_aligned=True,
|
|
frame_metadata=self._iterator.cached_source_data[filename])
|
|
except Exception:
|
|
self._plugin_thread.error_state.set(sys.exc_info())
|
|
raise
|
|
self._handle_plugin()
|
|
|
|
def sort(self) -> None:
|
|
"""Sort by dendrogram.
|
|
|
|
Parameters
|
|
----------
|
|
matched_list
|
|
The list of tuples with filename in first position and face encoding in the 2nd
|
|
|
|
Returns
|
|
-------
|
|
The original list, sorted for this metric
|
|
"""
|
|
logger.info("Sorting by ward linkage. This may take some time...")
|
|
results = np.array([item[1] for item in self._result])
|
|
indices = Cluster(np.array(results), "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 binning
|
|
|
|
Returns
|
|
-------
|
|
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
|
|
The command line arguments passed to the sort process
|
|
is_group
|
|
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: PNGAlignments | None) -> 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"""
|
|
result = T.cast(list[tuple[str, np.ndarray]], self._result)
|
|
img_list_len = len(result)
|
|
for i in tqdm(range(0, img_list_len),
|
|
desc="Comparing histograms",
|
|
file=sys.stdout,
|
|
leave=False):
|
|
score_total = 0.0
|
|
for j in range(0, img_list_len):
|
|
if i == j:
|
|
continue
|
|
score_total += cv2.compareHist(result[i][1],
|
|
result[j][1],
|
|
cv2.HISTCMP_BHATTACHARYYA)
|
|
result[i][2] = score_total # pyright:ignore
|
|
|
|
self._result = sorted(result, key=operator.itemgetter(2), reverse=True)
|
|
|
|
def _sort_sim(self) -> None:
|
|
"""Sort histograms by similarity"""
|
|
result = T.cast(list[tuple[str, np.ndarray]], self._result)
|
|
img_list_len = len(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(result[i][1],
|
|
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]) = (result[j_min_score], 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
|
|
The image to test
|
|
references: list
|
|
List of reference images to test the original image against
|
|
|
|
Returns
|
|
-------
|
|
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"""
|
|
# pylint:disable=duplicate-code
|
|
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] = [T.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")
|
|
img = T.cast(np.ndarray, self._result[i][1])
|
|
for key, value in reference_groups.items():
|
|
score = self._get_avg_score(img, value)
|
|
if score < current_score:
|
|
current_key, current_score = key, score
|
|
|
|
if current_score < threshold:
|
|
reference_groups[T.cast(int, current_key)].append(img)
|
|
bins[current_key].append(self._result[i][0])
|
|
else:
|
|
reference_groups[len(reference_groups)] = [img]
|
|
bins.append([self._result[i][0]])
|
|
|
|
return bins
|
|
|
|
def score_image(self,
|
|
filename: str,
|
|
image: np.ndarray | None,
|
|
alignments: PNGAlignments | None) -> None:
|
|
"""Collect the histogram for the given face
|
|
|
|
Parameters
|
|
----------
|
|
filename: str
|
|
The filename of the currently processing image
|
|
image
|
|
A face image loaded from disk
|
|
alignments
|
|
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()
|
|
|
|
|
|
__all__ = get_module_objects(__name__)
|