mirror of
https://github.com/deepfakes/faceswap
synced 2025-06-07 10:43:27 -04:00
936 lines
39 KiB
Python
936 lines
39 KiB
Python
#!/usr/bin/env python3
|
|
""" Alignments handling for Faceswap's Manual Adjustments tool. Handles the conversion of
|
|
alignments data to :class:`~lib.align.DetectedFace` objects, and the update of these faces
|
|
when edits are made in the GUI. """
|
|
from __future__ import annotations
|
|
import logging
|
|
import os
|
|
import sys
|
|
import tkinter as tk
|
|
import typing as T
|
|
from copy import deepcopy
|
|
from queue import Queue, Empty
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from lib.align import Alignments, AlignedFace, DetectedFace
|
|
from lib.gui.custom_widgets import PopupProgress
|
|
from lib.gui.utils import FileHandler
|
|
from lib.image import ImagesLoader, ImagesSaver, encode_image, generate_thumbnail
|
|
from lib.multithreading import MultiThread
|
|
from lib.utils import get_folder
|
|
|
|
if T.TYPE_CHECKING:
|
|
from . import manual
|
|
from lib.align.alignments import AlignmentFileDict, PNGHeaderDict
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class DetectedFaces():
|
|
""" Handles the manipulation of :class:`~lib.align.DetectedFace` objects stored
|
|
in the alignments file. Acts as a parent class for the IO operations (saving and loading from
|
|
an alignments file), the face update operations (when changes are made to alignments in the
|
|
GUI) and the face filters (when a user changes the filter navigation mode.)
|
|
|
|
Parameters
|
|
----------
|
|
tk_globals: :class:`~tools.manual.manual.TkGlobals`
|
|
The tkinter variables that apply to the whole of the GUI
|
|
alignments_path: str
|
|
The full path to the alignments file
|
|
input_location: str
|
|
The location of the input folder of frames or video file
|
|
extractor: :class:`~tools.manual.manual.Aligner`
|
|
The pipeline for passing faces through the aligner and retrieving results
|
|
"""
|
|
def __init__(self,
|
|
tk_globals: manual.TkGlobals,
|
|
alignments_path: str,
|
|
input_location: str,
|
|
extractor: manual.Aligner) -> None:
|
|
logger.debug("Initializing %s: (tk_globals: %s. alignments_path: %s, input_location: %s "
|
|
"extractor: %s)", self.__class__.__name__, tk_globals, alignments_path,
|
|
input_location, extractor)
|
|
self._globals = tk_globals
|
|
self._frame_faces: list[list[DetectedFace]] = []
|
|
self._updated_frame_indices: set[int] = set()
|
|
|
|
self._alignments: Alignments = self._get_alignments(alignments_path, input_location)
|
|
self._alignments.update_legacy_has_source(os.path.basename(input_location))
|
|
|
|
self._extractor = extractor
|
|
self._tk_vars = self._set_tk_vars()
|
|
|
|
self._io = _DiskIO(self, input_location)
|
|
self._update = FaceUpdate(self)
|
|
self._filter = Filter(self)
|
|
logger.debug("Initialized %s", self.__class__.__name__)
|
|
|
|
# <<<< PUBLIC PROPERTIES >>>> #
|
|
@property
|
|
def extractor(self) -> manual.Aligner:
|
|
""" :class:`~tools.manual.manual.Aligner`: The pipeline for passing faces through the
|
|
aligner and retrieving results. """
|
|
return self._extractor
|
|
|
|
@property
|
|
def filter(self) -> Filter:
|
|
""" :class:`Filter`: Handles returning of faces and stats based on the current user set
|
|
navigation mode filter. """
|
|
return self._filter
|
|
|
|
@property
|
|
def update(self) -> FaceUpdate:
|
|
""" :class:`FaceUpdate`: Handles the adding, removing and updating of
|
|
:class:`~lib.align.DetectedFace` stored within the alignments file. """
|
|
return self._update
|
|
|
|
# << TKINTER VARIABLES >> #
|
|
@property
|
|
def tk_unsaved(self) -> tk.BooleanVar:
|
|
""" :class:`tkinter.BooleanVar`: The variable indicating whether the alignments have been
|
|
updated since the last save. """
|
|
return self._tk_vars["unsaved"]
|
|
|
|
@property
|
|
def tk_edited(self) -> tk.BooleanVar:
|
|
""" :class:`tkinter.BooleanVar`: The variable indicating whether an edit has occurred
|
|
meaning a GUI redraw needs to be triggered. """
|
|
return self._tk_vars["edited"]
|
|
|
|
@property
|
|
def tk_face_count_changed(self) -> tk.BooleanVar:
|
|
""" :class:`tkinter.BooleanVar`: The variable indicating whether a face has been added or
|
|
removed meaning the :class:`FaceViewer` grid redraw needs to be triggered. """
|
|
return self._tk_vars["face_count_changed"]
|
|
|
|
# << STATISTICS >> #
|
|
@property
|
|
def frame_list(self) -> list[str]:
|
|
""" list[str]: The list of all frame names that appear in the alignments file """
|
|
return list(self._alignments.data)
|
|
|
|
@property
|
|
def available_masks(self) -> dict[str, int]:
|
|
""" dict[str, int]: The mask type names stored in the alignments; type as key with the
|
|
number of faces which possess the mask type as value. """
|
|
return self._alignments.mask_summary
|
|
|
|
@property
|
|
def current_faces(self) -> list[list[DetectedFace]]:
|
|
""" list[list[:class:`~lib.align.DetectedFace`]]: The most up to date full list of detected
|
|
face objects. """
|
|
return self._frame_faces
|
|
|
|
@property
|
|
def video_meta_data(self) -> dict[str, list[int] | list[float] | None]:
|
|
""" dict[str, list[int] | list[float] | None]: The frame meta data stored in the alignments
|
|
file. If data does not exist in the alignments file then ``None`` is returned for each
|
|
Key """
|
|
return self._alignments.video_meta_data
|
|
|
|
@property
|
|
def face_count_per_index(self) -> list[int]:
|
|
""" list[int]: Count of faces for each frame. List is in frame index order.
|
|
|
|
The list needs to be calculated on the fly as the number of faces in a frame
|
|
can change based on user actions. """
|
|
return [len(faces) for faces in self._frame_faces]
|
|
|
|
# <<<< PUBLIC METHODS >>>> #
|
|
def is_frame_updated(self, frame_index: int) -> bool:
|
|
""" Check whether the given frame index has been updated
|
|
|
|
Parameters
|
|
----------
|
|
frame_index: int
|
|
The frame index to check
|
|
|
|
Returns
|
|
-------
|
|
bool:
|
|
``True`` if the given frame index has updated faces within it otherwise ``False``
|
|
"""
|
|
return frame_index in self._updated_frame_indices
|
|
|
|
def load_faces(self) -> None:
|
|
""" Load the faces as :class:`~lib.align.DetectedFace` objects from the alignments
|
|
file. """
|
|
self._io.load()
|
|
|
|
def save(self) -> None:
|
|
""" Save the alignments file with the latest edits. """
|
|
self._io.save()
|
|
|
|
def revert_to_saved(self, frame_index):
|
|
""" Revert the frame's alignments to their saved version for the given frame index.
|
|
|
|
Parameters
|
|
----------
|
|
frame_index: int
|
|
The frame that should have their faces reverted to their saved version
|
|
"""
|
|
self._io.revert_to_saved(frame_index)
|
|
|
|
def extract(self) -> None:
|
|
""" Extract the faces in the current video to a user supplied folder. """
|
|
self._io.extract()
|
|
|
|
def save_video_meta_data(self, pts_time: list[float], keyframes: list[int]) -> None:
|
|
""" Save video meta data to the alignments file. This is executed if the video meta data
|
|
does not already exist in the alignments file, so the video does not need to be scanned
|
|
on every use of the Manual Tool.
|
|
|
|
Parameters
|
|
----------
|
|
pts_time: list[float]
|
|
A list of presentation timestamps in frame index order for every frame in the input
|
|
video
|
|
keyframes: list[int]
|
|
A list of frame indices corresponding to the key frames in the input video.
|
|
"""
|
|
if self._globals.is_video:
|
|
self._alignments.save_video_meta_data(pts_time, keyframes)
|
|
|
|
# <<<< PRIVATE METHODS >>> #
|
|
# << INIT >> #
|
|
@staticmethod
|
|
def _set_tk_vars() -> dict[T.Literal["unsaved", "edited", "face_count_changed"],
|
|
tk.BooleanVar]:
|
|
""" Set the required tkinter variables.
|
|
|
|
The alignments specific `unsaved` and `edited` are set here.
|
|
The global variables are added into the dictionary with `None` as value, so the
|
|
objects exist. Their actual variables are populated during :func:`load_faces`.
|
|
|
|
Returns
|
|
-------
|
|
dict
|
|
The internal variable name as key with the tkinter variable as value
|
|
"""
|
|
retval = {}
|
|
for name in T.get_args(T.Literal["unsaved", "edited", "face_count_changed"]):
|
|
var = tk.BooleanVar()
|
|
var.set(False)
|
|
retval[name] = var
|
|
logger.debug(retval)
|
|
return retval
|
|
|
|
def _get_alignments(self, alignments_path: str, input_location: str) -> Alignments:
|
|
""" Get the :class:`~lib.align.Alignments` object for the given location.
|
|
|
|
Parameters
|
|
----------
|
|
alignments_path: str
|
|
Full path to the alignments file. If empty string is passed then location is calculated
|
|
from the source folder
|
|
input_location: str
|
|
The location of the input folder of frames or video file
|
|
|
|
Returns
|
|
-------
|
|
:class:`~lib.align.Alignments`
|
|
The alignments object for the given input location
|
|
"""
|
|
logger.debug("alignments_path: %s, input_location: %s", alignments_path, input_location)
|
|
if alignments_path:
|
|
folder, filename = os.path.split(alignments_path)
|
|
else:
|
|
filename = "alignments.fsa"
|
|
if self._globals.is_video:
|
|
folder, vid = os.path.split(os.path.splitext(input_location)[0])
|
|
filename = f"{vid}_{filename}"
|
|
else:
|
|
folder = input_location
|
|
retval = Alignments(folder, filename)
|
|
if retval.version == 1.0:
|
|
logger.error("The Manual Tool is not compatible with legacy Alignments files.")
|
|
logger.info("You can update legacy Alignments files by using the Extract job in the "
|
|
"Alignments tool to re-extract the faces in full-head format.")
|
|
sys.exit(0)
|
|
logger.debug("folder: %s, filename: %s, alignments: %s", folder, filename, retval)
|
|
return retval
|
|
|
|
|
|
class _DiskIO():
|
|
""" Handles the loading of :class:`~lib.align.DetectedFaces` from the alignments file
|
|
into :class:`DetectedFaces` and the saving of this data (in the opposite direction) to an
|
|
alignments file.
|
|
|
|
Parameters
|
|
----------
|
|
detected_faces: :class:`DetectedFaces`
|
|
The parent :class:`DetectedFaces` object
|
|
input_location: str
|
|
The location of the input folder of frames or video file
|
|
"""
|
|
def __init__(self, detected_faces: DetectedFaces, input_location: str) -> None:
|
|
logger.debug("Initializing %s: (detected_faces: %s, input_location: %s)",
|
|
self.__class__.__name__, detected_faces, input_location)
|
|
self._input_location = input_location
|
|
self._alignments = detected_faces._alignments
|
|
self._frame_faces = detected_faces._frame_faces
|
|
self._updated_frame_indices = detected_faces._updated_frame_indices
|
|
self._tk_unsaved = detected_faces.tk_unsaved
|
|
self._tk_edited = detected_faces.tk_edited
|
|
self._tk_face_count_changed = detected_faces.tk_face_count_changed
|
|
self._globals = detected_faces._globals
|
|
|
|
# Must be populated after loading faces as video_meta_data may have increased frame count
|
|
self._sorted_frame_names: list[str] = []
|
|
logger.debug("Initialized %s", self.__class__.__name__)
|
|
|
|
def load(self) -> None:
|
|
""" Load the faces from the alignments file, convert to
|
|
:class:`~lib.align.DetectedFace`. objects and add to :attr:`_frame_faces`. """
|
|
for key in sorted(self._alignments.data):
|
|
this_frame_faces: list[DetectedFace] = []
|
|
for item in self._alignments.data[key]["faces"]:
|
|
face = DetectedFace()
|
|
face.from_alignment(item, with_thumb=True)
|
|
face.load_aligned(None)
|
|
_ = face.aligned.average_distance # cache the distances
|
|
this_frame_faces.append(face)
|
|
self._frame_faces.append(this_frame_faces)
|
|
self._sorted_frame_names = sorted(self._alignments.data)
|
|
|
|
def save(self) -> None:
|
|
""" Convert updated :class:`~lib.align.DetectedFace` objects to alignments format
|
|
and save the alignments file. """
|
|
if not self._tk_unsaved.get():
|
|
logger.debug("Alignments not updated. Returning")
|
|
return
|
|
frames = list(self._updated_frame_indices)
|
|
logger.verbose("Saving alignments for %s updated frames", # type:ignore[attr-defined]
|
|
len(frames))
|
|
|
|
for idx, faces in zip(frames,
|
|
np.array(self._frame_faces, dtype="object")[np.array(frames)]):
|
|
frame = self._sorted_frame_names[idx]
|
|
self._alignments.data[frame]["faces"] = [face.to_alignment() for face in faces]
|
|
|
|
self._alignments.backup()
|
|
self._alignments.save()
|
|
self._updated_frame_indices.clear()
|
|
self._tk_unsaved.set(False)
|
|
|
|
def revert_to_saved(self, frame_index: int) -> None:
|
|
""" Revert the frame's alignments to their saved version for the given frame index.
|
|
|
|
Parameters
|
|
----------
|
|
frame_index: int
|
|
The frame that should have their faces reverted to their saved version
|
|
"""
|
|
if frame_index not in self._updated_frame_indices:
|
|
logger.debug("Alignments not amended. Returning")
|
|
return
|
|
logger.verbose("Reverting alignments for frame_index %s", # type:ignore[attr-defined]
|
|
frame_index)
|
|
alignments = self._alignments.data[self._sorted_frame_names[frame_index]]["faces"]
|
|
faces = self._frame_faces[frame_index]
|
|
|
|
reset_grid = self._add_remove_faces(alignments, faces)
|
|
|
|
for detected_face, face in zip(faces, alignments):
|
|
detected_face.from_alignment(face, with_thumb=True)
|
|
detected_face.load_aligned(None, force=True)
|
|
_ = detected_face.aligned.average_distance # cache the distances
|
|
|
|
self._updated_frame_indices.remove(frame_index)
|
|
if not self._updated_frame_indices:
|
|
self._tk_unsaved.set(False)
|
|
|
|
if reset_grid:
|
|
self._tk_face_count_changed.set(True)
|
|
else:
|
|
self._tk_edited.set(True)
|
|
self._globals.var_full_update.set(True)
|
|
|
|
@classmethod
|
|
def _add_remove_faces(cls,
|
|
alignments: list[AlignmentFileDict],
|
|
faces: list[DetectedFace]) -> bool:
|
|
""" On a revert, ensure that the alignments and detected face object counts for each frame
|
|
are in sync.
|
|
|
|
Parameters
|
|
----------
|
|
alignments: list[:class:`~lib.align.alignments.AlignmentFileDict`]
|
|
Alignments stored for a frame
|
|
|
|
faces: list[:class:`~lib.align.DetectedFace`]
|
|
List of detected faces for a frame
|
|
|
|
Returns
|
|
-------
|
|
bool
|
|
``True`` if a face was added or removed otherwise ``False``
|
|
"""
|
|
num_alignments = len(alignments)
|
|
num_faces = len(faces)
|
|
if num_alignments == num_faces:
|
|
retval = False
|
|
elif num_alignments > num_faces:
|
|
faces.extend([DetectedFace() for _ in range(num_faces, num_alignments)])
|
|
retval = True
|
|
else:
|
|
del faces[num_alignments:]
|
|
retval = True
|
|
return retval
|
|
|
|
def extract(self) -> None:
|
|
""" Extract the current faces to a folder.
|
|
|
|
To stop the GUI becoming completely unresponsive (particularly in Windows) the extract is
|
|
done in a background thread, with the process count passed back in a queue to the main
|
|
thread to update the progress bar.
|
|
"""
|
|
dirname = FileHandler("dir", None,
|
|
initial_folder=os.path.dirname(self._input_location),
|
|
title="Select output folder...").return_file
|
|
if not dirname:
|
|
return
|
|
logger.debug(dirname)
|
|
|
|
queue: Queue = Queue()
|
|
pbar = PopupProgress("Extracting Faces...", self._alignments.frames_count + 1)
|
|
thread = MultiThread(self._background_extract, dirname, queue)
|
|
thread.start()
|
|
self._monitor_extract(thread, queue, pbar)
|
|
|
|
def _monitor_extract(self,
|
|
thread: MultiThread,
|
|
queue: Queue,
|
|
progress_bar: PopupProgress) -> None:
|
|
""" Monitor the extraction thread, and update the progress bar.
|
|
|
|
On completion, save alignments and clear progress bar.
|
|
|
|
Parameters
|
|
----------
|
|
thread: :class:`~lib.multithreading.MultiThread`
|
|
The thread that is performing the extraction task
|
|
queue: :class:`queue.Queue`
|
|
The queue that the worker thread is putting it's incremental counts to
|
|
progress_bar: :class:`~lib.gui.custom_widget.PopupProgress`
|
|
The popped up progress bar
|
|
"""
|
|
thread.check_and_raise_error()
|
|
if not thread.is_alive():
|
|
thread.join()
|
|
progress_bar.stop()
|
|
return
|
|
|
|
while True:
|
|
try:
|
|
progress_bar.step(queue.get(False, 0))
|
|
except Empty:
|
|
break
|
|
progress_bar.after(100, self._monitor_extract, thread, queue, progress_bar)
|
|
|
|
def _background_extract(self, output_folder: str, progress_queue: Queue) -> None:
|
|
""" Perform the background extraction in a thread so GUI doesn't become unresponsive.
|
|
|
|
Parameters
|
|
----------
|
|
output_folder: str
|
|
The location to save the output faces to
|
|
progress_queue: :class:`queue.Queue`
|
|
The queue to place incremental counts to for updating the GUI's progress bar
|
|
"""
|
|
saver = ImagesSaver(get_folder(output_folder), as_bytes=True)
|
|
loader = ImagesLoader(self._input_location, count=self._alignments.frames_count)
|
|
for frame_idx, (filename, image) in enumerate(loader.load()):
|
|
logger.trace("Outputting frame: %s: %s", # type:ignore[attr-defined]
|
|
frame_idx, filename)
|
|
src_filename = os.path.basename(filename)
|
|
progress_queue.put(1)
|
|
|
|
for face_idx, face in enumerate(self._frame_faces[frame_idx]):
|
|
output = f"{os.path.splitext(src_filename)[0]}_{face_idx}.png"
|
|
aligned = AlignedFace(face.landmarks_xy,
|
|
image=image,
|
|
centering="head",
|
|
size=512) # TODO user selectable size
|
|
meta: PNGHeaderDict = {"alignments": face.to_png_meta(),
|
|
"source": {"alignments_version": self._alignments.version,
|
|
"original_filename": output,
|
|
"face_index": face_idx,
|
|
"source_filename": src_filename,
|
|
"source_is_video": self._globals.is_video,
|
|
"source_frame_dims": image.shape[:2]}}
|
|
|
|
assert aligned.face is not None
|
|
b_image = encode_image(aligned.face, ".png", metadata=meta)
|
|
saver.save(output, b_image)
|
|
saver.close()
|
|
|
|
|
|
class Filter():
|
|
""" Returns stats and frames for filtered frames based on the user selected navigation mode
|
|
filter.
|
|
|
|
Parameters
|
|
----------
|
|
detected_faces: :class:`DetectedFaces`
|
|
The parent :class:`DetectedFaces` object
|
|
"""
|
|
def __init__(self, detected_faces: DetectedFaces) -> None:
|
|
logger.debug("Initializing %s: (detected_faces: %s)",
|
|
self.__class__.__name__, detected_faces)
|
|
self._globals = detected_faces._globals
|
|
self._detected_faces = detected_faces
|
|
logger.debug("Initialized %s", self.__class__.__name__)
|
|
|
|
@property
|
|
def frame_meets_criteria(self) -> bool:
|
|
""" bool: ``True`` if the current frame meets the selected filter criteria otherwise
|
|
``False`` """
|
|
filter_mode = self._globals.var_filter_mode.get()
|
|
frame_faces = self._detected_faces.current_faces[self._globals.frame_index]
|
|
distance = self._filter_distance
|
|
|
|
retval = (
|
|
filter_mode == "All Frames" or
|
|
(filter_mode == "No Faces" and not frame_faces) or
|
|
(filter_mode == "Has Face(s)" and len(frame_faces) > 0) or
|
|
(filter_mode == "Multiple Faces" and len(frame_faces) > 1) or
|
|
(filter_mode == "Misaligned Faces" and any(face.aligned.average_distance > distance
|
|
for face in frame_faces)))
|
|
assert isinstance(retval, bool)
|
|
logger.trace("filter_mode: %s, frame meets criteria: %s", # type:ignore[attr-defined]
|
|
filter_mode, retval)
|
|
return retval
|
|
|
|
@property
|
|
def _filter_distance(self) -> float:
|
|
""" float: The currently selected distance when Misaligned Faces filter is selected. """
|
|
try:
|
|
retval = self._globals.var_filter_distance.get()
|
|
except tk.TclError:
|
|
# Suppress error when distance box is empty
|
|
retval = 0
|
|
return retval / 100.
|
|
|
|
@property
|
|
def count(self) -> int:
|
|
""" int: The number of frames that meet the filter criteria returned by
|
|
:attr:`~tools.manual.manual.TkGlobals.var_filter_mode.get()`. """
|
|
face_count_per_index = self._detected_faces.face_count_per_index
|
|
if self._globals.var_filter_mode.get() == "No Faces":
|
|
retval = sum(1 for fcount in face_count_per_index if fcount == 0)
|
|
elif self._globals.var_filter_mode.get() == "Has Face(s)":
|
|
retval = sum(1 for fcount in face_count_per_index if fcount != 0)
|
|
elif self._globals.var_filter_mode.get() == "Multiple Faces":
|
|
retval = sum(1 for fcount in face_count_per_index if fcount > 1)
|
|
elif self._globals.var_filter_mode.get() == "Misaligned Faces":
|
|
distance = self._filter_distance
|
|
retval = sum(1 for frame in self._detected_faces.current_faces
|
|
if any(face.aligned.average_distance > distance for face in frame))
|
|
else:
|
|
retval = len(face_count_per_index)
|
|
logger.trace("filter mode: %s, frame count: %s", # type:ignore[attr-defined]
|
|
self._globals.var_filter_mode.get(), retval)
|
|
return retval
|
|
|
|
@property
|
|
def raw_indices(self) -> dict[T.Literal["frame", "face"], list[int]]:
|
|
""" dict[str, int]: The frame and face indices that meet the current filter criteria for
|
|
each displayed face. """
|
|
frame_indices: list[int] = []
|
|
face_indices: list[int] = []
|
|
face_counts = self._detected_faces.face_count_per_index # Copy to avoid recalculations
|
|
|
|
for frame_idx in self.frames_list:
|
|
for face_idx in range(face_counts[frame_idx]):
|
|
frame_indices.append(frame_idx)
|
|
face_indices.append(face_idx)
|
|
|
|
retval: dict[T.Literal["frame", "face"], list[int]] = {"frame": frame_indices,
|
|
"face": face_indices}
|
|
logger.trace("frame_indices: %s, face_indices: %s", # type:ignore[attr-defined]
|
|
frame_indices, face_indices)
|
|
return retval
|
|
|
|
@property
|
|
def frames_list(self) -> list[int]:
|
|
""" list[int]: The list of frame indices that meet the filter criteria returned by
|
|
:attr:`~tools.manual.manual.TkGlobals.var_filter_mode.get()`. """
|
|
face_count_per_index = self._detected_faces.face_count_per_index
|
|
if self._globals.var_filter_mode.get() == "No Faces":
|
|
retval = [idx for idx, count in enumerate(face_count_per_index) if count == 0]
|
|
elif self._globals.var_filter_mode.get() == "Multiple Faces":
|
|
retval = [idx for idx, count in enumerate(face_count_per_index) if count > 1]
|
|
elif self._globals.var_filter_mode.get() == "Has Face(s)":
|
|
retval = [idx for idx, count in enumerate(face_count_per_index) if count != 0]
|
|
elif self._globals.var_filter_mode.get() == "Misaligned Faces":
|
|
distance = self._filter_distance
|
|
retval = [idx for idx, frame in enumerate(self._detected_faces.current_faces)
|
|
if any(face.aligned.average_distance > distance for face in frame)]
|
|
else:
|
|
retval = list(range(len(face_count_per_index)))
|
|
logger.trace("filter mode: %s, number_frames: %s", # type:ignore[attr-defined]
|
|
self._globals.var_filter_mode.get(), len(retval))
|
|
return retval
|
|
|
|
|
|
class FaceUpdate():
|
|
""" Perform updates on :class:`~lib.align.DetectedFace` objects stored in
|
|
:class:`DetectedFaces` when changes are made within the GUI.
|
|
|
|
Parameters
|
|
----------
|
|
detected_faces: :class:`DetectedFaces`
|
|
The parent :class:`DetectedFaces` object
|
|
"""
|
|
def __init__(self, detected_faces: DetectedFaces) -> None:
|
|
logger.debug("Initializing %s: (detected_faces: %s)",
|
|
self.__class__.__name__, detected_faces)
|
|
self._detected_faces = detected_faces
|
|
self._globals = detected_faces._globals
|
|
self._frame_faces = detected_faces._frame_faces
|
|
self._updated_frame_indices = detected_faces._updated_frame_indices
|
|
self._tk_unsaved = detected_faces.tk_unsaved
|
|
self._extractor = detected_faces.extractor
|
|
logger.debug("Initialized %s", self.__class__.__name__)
|
|
|
|
@property
|
|
def _tk_edited(self) -> tk.BooleanVar:
|
|
""" :class:`tkinter.BooleanVar`: The variable indicating whether an edit has occurred
|
|
meaning a GUI redraw needs to be triggered.
|
|
|
|
Notes
|
|
-----
|
|
The variable is still a ``None`` when this class is initialized, so referenced explicitly.
|
|
"""
|
|
return self._detected_faces.tk_edited
|
|
|
|
@property
|
|
def _tk_face_count_changed(self) -> tk.BooleanVar:
|
|
""" :class:`tkinter.BooleanVar`: The variable indicating whether an edit has occurred
|
|
meaning a GUI redraw needs to be triggered.
|
|
|
|
Notes
|
|
-----
|
|
The variable is still a ``None`` when this class is initialized, so referenced explicitly.
|
|
"""
|
|
return self._detected_faces.tk_face_count_changed
|
|
|
|
def _faces_at_frame_index(self, frame_index: int) -> list[DetectedFace]:
|
|
""" Checks whether the frame has already been added to :attr:`_updated_frame_indices` and
|
|
adds it. Triggers the unsaved variable if this is the first edited frame. Returns the
|
|
detected face objects for the given frame.
|
|
|
|
Parameters
|
|
----------
|
|
frame_index: int
|
|
The frame index to check whether there are updated alignments available
|
|
|
|
Returns
|
|
-------
|
|
list
|
|
The :class:`~lib.align.DetectedFace` objects for the requested frame
|
|
"""
|
|
if not self._updated_frame_indices and not self._tk_unsaved.get():
|
|
self._tk_unsaved.set(True)
|
|
self._updated_frame_indices.add(frame_index)
|
|
retval = self._frame_faces[frame_index]
|
|
return retval
|
|
|
|
def add(self, frame_index: int, pnt_x: int, width: int, pnt_y: int, height: int) -> None:
|
|
""" Add a :class:`~lib.align.DetectedFace` object to the current frame with the
|
|
given dimensions.
|
|
|
|
Parameters
|
|
----------
|
|
frame_index: int
|
|
The frame that the face is being set for
|
|
pnt_x: int
|
|
The left point of the bounding box
|
|
width: int
|
|
The width of the bounding box
|
|
pnt_y: int
|
|
The top point of the bounding box
|
|
height: int
|
|
The height of the bounding box
|
|
"""
|
|
face = DetectedFace()
|
|
faces = self._faces_at_frame_index(frame_index)
|
|
faces.append(face)
|
|
face_index = len(faces) - 1
|
|
|
|
self.bounding_box(frame_index, face_index, pnt_x, width, pnt_y, height, aligner="cv2-dnn")
|
|
face.load_aligned(None)
|
|
self._tk_face_count_changed.set(True)
|
|
|
|
def delete(self, frame_index: int, face_index: int) -> None:
|
|
""" Delete the :class:`~lib.align.DetectedFace` object for the given frame and face
|
|
indices.
|
|
|
|
Parameters
|
|
----------
|
|
frame_index: int
|
|
The frame that the face is being set for
|
|
face_index: int
|
|
The face index within the frame
|
|
"""
|
|
logger.debug("Deleting face at frame index: %s face index: %s", frame_index, face_index)
|
|
faces = self._faces_at_frame_index(frame_index)
|
|
del faces[face_index]
|
|
self._tk_face_count_changed.set(True)
|
|
self._globals.var_full_update.set(True)
|
|
|
|
def bounding_box(self,
|
|
frame_index: int,
|
|
face_index: int,
|
|
pnt_x: int,
|
|
width: int,
|
|
pnt_y: int,
|
|
height: int,
|
|
aligner: manual.TypeManualExtractor = "FAN") -> None:
|
|
""" Update the bounding box for the :class:`~lib.align.DetectedFace` object at the
|
|
given frame and face indices, with the given dimensions and update the 68 point landmarks
|
|
from the :class:`~tools.manual.manual.Aligner` for the updated bounding box.
|
|
|
|
Parameters
|
|
----------
|
|
frame_index: int
|
|
The frame that the face is being set for
|
|
face_index: int
|
|
The face index within the frame
|
|
pnt_x: int
|
|
The left point of the bounding box
|
|
width: int
|
|
The width of the bounding box
|
|
pnt_y: int
|
|
The top point of the bounding box
|
|
height: int
|
|
The height of the bounding box
|
|
aligner: ["cv2-dnn", "FAN"], optional
|
|
The aligner to use to generate the landmarks. Default: "FAN"
|
|
"""
|
|
logger.trace("frame_index: %s, face_index %s, pnt_x %s, " # type:ignore[attr-defined]
|
|
"width %s, pnt_y %s, height %s, aligner: %s",
|
|
frame_index, face_index, pnt_x, width, pnt_y, height, aligner)
|
|
face = self._faces_at_frame_index(frame_index)[face_index]
|
|
face.left = pnt_x
|
|
face.width = width
|
|
face.top = pnt_y
|
|
face.height = height
|
|
face.add_landmarks_xy(self._extractor.get_landmarks(frame_index, face_index, aligner))
|
|
self._globals.var_full_update.set(True)
|
|
|
|
def landmark(self,
|
|
frame_index: int, face_index: int,
|
|
landmark_index: int,
|
|
shift_x: int,
|
|
shift_y: int,
|
|
is_zoomed: bool) -> None:
|
|
""" Shift a single landmark point for the :class:`~lib.align.DetectedFace` object
|
|
at the given frame and face indices by the given x and y values.
|
|
|
|
Parameters
|
|
----------
|
|
frame_index: int
|
|
The frame that the face is being set for
|
|
face_index: int
|
|
The face index within the frame
|
|
landmark_index: int or list
|
|
The landmark index to shift. If a list is provided, this should be a list of landmark
|
|
indices to be shifted
|
|
shift_x: int
|
|
The amount to shift the landmark by along the x axis
|
|
shift_y: int
|
|
The amount to shift the landmark by along the y axis
|
|
is_zoomed: bool
|
|
``True`` if landmarks are being adjusted on a zoomed image otherwise ``False``
|
|
"""
|
|
face = self._faces_at_frame_index(frame_index)[face_index]
|
|
if is_zoomed:
|
|
aligned = AlignedFace(face.landmarks_xy,
|
|
centering="face",
|
|
size=min(self._globals.frame_display_dims))
|
|
landmark = aligned.landmarks[landmark_index]
|
|
landmark += (shift_x, shift_y)
|
|
matrix = aligned.adjusted_matrix
|
|
matrix = cv2.invertAffineTransform(matrix)
|
|
if landmark.ndim == 1:
|
|
landmark = np.reshape(landmark, (1, 1, 2))
|
|
landmark = cv2.transform(landmark, matrix, landmark.shape).squeeze()
|
|
face.landmarks_xy[landmark_index] = landmark
|
|
else:
|
|
for lmk, idx in zip(landmark, landmark_index): # type:ignore[call-overload]
|
|
lmk = np.reshape(lmk, (1, 1, 2))
|
|
lmk = cv2.transform(lmk, matrix, lmk.shape).squeeze()
|
|
face.landmarks_xy[idx] = lmk
|
|
else:
|
|
face.landmarks_xy[landmark_index] += (shift_x, shift_y)
|
|
self._globals.var_full_update.set(True)
|
|
|
|
def landmarks(self, frame_index: int, face_index: int, shift_x: int, shift_y: int) -> None:
|
|
""" Shift all of the landmarks and bounding box for the
|
|
:class:`~lib.align.DetectedFace` object at the given frame and face indices by the
|
|
given x and y values and update the masks.
|
|
|
|
Parameters
|
|
----------
|
|
frame_index: int
|
|
The frame that the face is being set for
|
|
face_index: int
|
|
The face index within the frame
|
|
shift_x: int
|
|
The amount to shift the landmarks by along the x axis
|
|
shift_y: int
|
|
The amount to shift the landmarks by along the y axis
|
|
|
|
Notes
|
|
-----
|
|
Whilst the bounding box does not need to be shifted, it is anyway, to ensure that it is
|
|
aligned with the newly adjusted landmarks.
|
|
"""
|
|
face = self._faces_at_frame_index(frame_index)[face_index]
|
|
assert face.left is not None and face.top is not None
|
|
face.left += shift_x
|
|
face.top += shift_y
|
|
face.add_landmarks_xy(face.landmarks_xy + (shift_x, shift_y))
|
|
self._globals.var_full_update.set(True)
|
|
|
|
def landmarks_rotate(self,
|
|
frame_index: int,
|
|
face_index: int,
|
|
angle: float,
|
|
center: np.ndarray) -> None:
|
|
""" Rotate the landmarks on an Extract Box rotate for the
|
|
:class:`~lib.align.DetectedFace` object at the given frame and face indices for the
|
|
given angle from the given center point.
|
|
|
|
Parameters
|
|
----------
|
|
frame_index: int
|
|
The frame that the face is being set for
|
|
face_index: int
|
|
The face index within the frame
|
|
angle: float
|
|
The angle, in radians to rotate the points by
|
|
center: :class:`numpy.ndarray`
|
|
The center point of the Landmark's Extract Box
|
|
"""
|
|
face = self._faces_at_frame_index(frame_index)[face_index]
|
|
rot_mat = cv2.getRotationMatrix2D(tuple(center.astype("float32")), angle, 1.)
|
|
face.add_landmarks_xy(cv2.transform(np.expand_dims(face.landmarks_xy, axis=0),
|
|
rot_mat).squeeze())
|
|
self._globals.var_full_update.set(True)
|
|
|
|
def landmarks_scale(self,
|
|
frame_index: int,
|
|
face_index: int,
|
|
scale: np.ndarray,
|
|
center: np.ndarray) -> None:
|
|
""" Scale the landmarks on an Extract Box resize for the
|
|
:class:`~lib.align.DetectedFace` object at the given frame and face indices from the
|
|
given center point.
|
|
|
|
Parameters
|
|
----------
|
|
frame_index: int
|
|
The frame that the face is being set for
|
|
face_index: int
|
|
The face index within the frame
|
|
scale: float
|
|
The amount to scale the landmarks by
|
|
center: :class:`numpy.ndarray`
|
|
The center point of the Landmark's Extract Box
|
|
"""
|
|
face = self._faces_at_frame_index(frame_index)[face_index]
|
|
face.add_landmarks_xy(((face.landmarks_xy - center) * scale) + center)
|
|
self._globals.var_full_update.set(True)
|
|
|
|
def mask(self, frame_index: int, face_index: int, mask: np.ndarray, mask_type: str) -> None:
|
|
""" Update the mask on an edit for the :class:`~lib.align.DetectedFace` object at
|
|
the given frame and face indices, for the given mask and mask type.
|
|
|
|
Parameters
|
|
----------
|
|
frame_index: int
|
|
The frame that the face is being set for
|
|
face_index: int
|
|
The face index within the frame
|
|
mask: class:`numpy.ndarray`:
|
|
The mask to replace
|
|
mask_type: str
|
|
The name of the mask that is to be replaced
|
|
"""
|
|
face = self._faces_at_frame_index(frame_index)[face_index]
|
|
face.mask[mask_type].replace_mask(mask)
|
|
self._tk_edited.set(True)
|
|
self._globals.var_full_update.set(True)
|
|
|
|
def copy(self, frame_index: int, direction: T.Literal["prev", "next"]) -> None:
|
|
""" Copy the alignments from the previous or next frame that has alignments
|
|
to the current frame.
|
|
|
|
Parameters
|
|
----------
|
|
frame_index: int
|
|
The frame that the needs to have alignments copied to it
|
|
direction: ["prev", "next"]
|
|
Whether to copy alignments from the previous frame with alignments, or the next
|
|
frame with alignments
|
|
"""
|
|
logger.debug("frame: %s, direction: %s", frame_index, direction)
|
|
faces = self._faces_at_frame_index(frame_index)
|
|
frames_with_faces = [idx for idx, faces in enumerate(self._detected_faces.current_faces)
|
|
if len(faces) > 0]
|
|
if direction == "prev":
|
|
idx = next((idx for idx in reversed(frames_with_faces)
|
|
if idx < frame_index), None)
|
|
else:
|
|
idx = next((idx for idx in frames_with_faces
|
|
if idx > frame_index), None)
|
|
if idx is None:
|
|
# No previous/next frame available
|
|
return
|
|
logger.debug("Copying alignments from frame %s to frame: %s", idx, frame_index)
|
|
|
|
# aligned_face cannot be deep copied, so remove and recreate
|
|
to_copy = self._faces_at_frame_index(idx)
|
|
for face in to_copy:
|
|
face._aligned = None # pylint:disable=protected-access
|
|
copied = deepcopy(to_copy)
|
|
|
|
for old_face, new_face in zip(to_copy, copied):
|
|
old_face.load_aligned(None)
|
|
new_face.load_aligned(None)
|
|
|
|
faces.extend(copied)
|
|
self._tk_face_count_changed.set(True)
|
|
self._globals.var_full_update.set(True)
|
|
|
|
def post_edit_trigger(self, frame_index: int, face_index: int) -> None:
|
|
""" Update the jpg thumbnail, the viewport thumbnail, the landmark masks and the aligned
|
|
face on a face edit.
|
|
|
|
Parameters
|
|
----------
|
|
frame_index: int
|
|
The frame that the face is being set for
|
|
face_index: int
|
|
The face index within the frame
|
|
"""
|
|
face = self._frame_faces[frame_index][face_index]
|
|
face.load_aligned(None, force=True) # Update average distance
|
|
face.mask = self._extractor.get_masks(frame_index, face_index)
|
|
face.clear_all_identities()
|
|
|
|
aligned = AlignedFace(face.landmarks_xy,
|
|
image=self._globals.current_frame.image,
|
|
centering="head",
|
|
size=96)
|
|
assert aligned.face is not None
|
|
face.thumbnail = generate_thumbnail(aligned.face, size=96)
|
|
if self._globals.var_filter_mode.get() == "Misaligned Faces":
|
|
self._detected_faces.tk_face_count_changed.set(True)
|
|
self._tk_edited.set(True)
|