1
0
Fork 0
mirror of https://github.com/deepfakes/faceswap synced 2025-06-07 10:43:27 -04:00
faceswap/tools/manual/manual.py
2024-04-03 14:03:54 +01:00

942 lines
39 KiB
Python

#!/usr/bin/env python3
""" The Manual Tool is a tkinter driven GUI app for editing alignments files with visual tools.
This module is the main entry point into the Manual Tool. """
from __future__ import annotations
import logging
import os
import sys
import typing as T
import tkinter as tk
from tkinter import ttk
from time import sleep
import cv2
import numpy as np
from lib.gui.control_helper import ControlPanel
from lib.gui.utils import get_images, get_config, initialize_config, initialize_images
from lib.image import SingleFrameLoader, read_image_meta
from lib.multithreading import MultiThread
from lib.utils import _video_extensions
from plugins.extract.pipeline import Extractor, ExtractMedia
from .detected_faces import DetectedFaces
from .faceviewer.frame import FacesFrame
from .frameviewer.frame import DisplayFrame
from .thumbnails import ThumbsCreator
if T.TYPE_CHECKING:
from lib.align import DetectedFace
from lib.align.detected_face import Mask
from lib.queue_manager import EventQueue
logger = logging.getLogger(__name__)
TypeManualExtractor = T.Literal["FAN", "cv2-dnn", "mask"]
class Manual(tk.Tk):
""" The main entry point for Faceswap's Manual Editor Tool. This tool is part of the Faceswap
Tools suite and should be called from ``python tools.py manual`` command.
Allows for visual interaction with frames, faces and alignments file to perform various
adjustments to the alignments file.
Parameters
----------
arguments: :class:`argparse.Namespace`
The :mod:`argparse` arguments as passed in from :mod:`tools.py`
"""
def __init__(self, arguments):
logger.debug("Initializing %s: (arguments: '%s')", self.__class__.__name__, arguments)
super().__init__()
self._validate_non_faces(arguments.frames)
self._initialize_tkinter()
self._globals = TkGlobals(arguments.frames)
extractor = Aligner(self._globals, arguments.exclude_gpus)
self._detected_faces = DetectedFaces(self._globals,
arguments.alignments_path,
arguments.frames,
extractor)
video_meta_data = self._detected_faces.video_meta_data
valid_meta = all(val is not None for val in video_meta_data.values())
loader = FrameLoader(self._globals, arguments.frames, video_meta_data)
if valid_meta: # Load the faces whilst other threads complete if we have valid meta data
self._detected_faces.load_faces()
self._containers = self._create_containers()
self._wait_for_threads(extractor, loader, valid_meta)
if not valid_meta:
# Load the faces after other threads complete if meta data required updating
self._detected_faces.load_faces()
self._generate_thumbs(arguments.frames, arguments.thumb_regen, arguments.single_process)
self._display = DisplayFrame(self._containers["top"],
self._globals,
self._detected_faces)
_Options(self._containers["top"], self._globals, self._display)
self._faces_frame = FacesFrame(self._containers["bottom"],
self._globals,
self._detected_faces,
self._display)
self._display.tk_selected_action.set("View")
self.bind("<Key>", self._handle_key_press)
self._set_initial_layout()
logger.debug("Initialized %s", self.__class__.__name__)
@classmethod
def _validate_non_faces(cls, frames_folder):
""" Quick check on the input to make sure that a folder of extracted faces is not being
passed in. """
if not os.path.isdir(frames_folder):
logger.debug("Input '%s' is not a folder", frames_folder)
return
test_file = next((fname
for fname in os.listdir(frames_folder)
if os.path.splitext(fname)[-1].lower() == ".png"),
None)
if not test_file:
logger.debug("Input '%s' does not contain any .pngs", frames_folder)
return
test_file = os.path.join(frames_folder, test_file)
meta = read_image_meta(test_file)
logger.debug("Test file: (filename: %s, metadata: %s)", test_file, meta)
if "itxt" in meta and "alignments" in meta["itxt"]:
logger.error("The input folder '%s' contains extracted faces.", frames_folder)
logger.error("The Manual Tool works with source frames or a video file, not extracted "
"faces. Please update your input.")
sys.exit(1)
logger.debug("Test input file '%s' does not contain Faceswap header data", test_file)
def _wait_for_threads(self, extractor, loader, valid_meta):
""" The :class:`Aligner` and :class:`FramesLoader` are launched in background threads.
Wait for them to be initialized prior to proceeding.
Parameters
----------
extractor: :class:`Aligner`
The extraction pipeline for the Manual Tool
loader: :class:`FramesLoader`
The frames loader for the Manual Tool
valid_meta: bool
Whether the input video had valid meta-data on import, or if it had to be created.
``True`` if valid meta data existed previously, ``False`` if it needed to be created
Notes
-----
Because some of the initialize checks perform extra work once their threads are complete,
they should only return ``True`` once, and should not be queried again.
"""
extractor_init = False
frames_init = False
while True:
extractor_init = extractor_init if extractor_init else extractor.is_initialized
frames_init = frames_init if frames_init else loader.is_initialized
if extractor_init and frames_init:
logger.debug("Threads inialized")
break
logger.debug("Threads not initialized. Waiting...")
sleep(1)
extractor.link_faces(self._detected_faces)
if not valid_meta:
logger.debug("Saving video meta data to alignments file")
self._detected_faces.save_video_meta_data(**loader.video_meta_data)
def _generate_thumbs(self, input_location, force, single_process):
""" Check whether thumbnails are stored in the alignments file and if not generate them.
Parameters
----------
input_location: str
The input video or folder of images
force: bool
``True`` if the thumbnails should be regenerated even if they exist, otherwise
``False``
single_process: bool
``True`` will extract thumbs from a video in a single process, ``False`` will run
parallel threads
"""
thumbs = ThumbsCreator(self._detected_faces, input_location, single_process)
if thumbs.has_thumbs and not force:
return
logger.debug("Generating thumbnails cache")
thumbs.generate_cache()
logger.debug("Generated thumbnails cache")
def _initialize_tkinter(self):
""" Initialize a standalone tkinter instance. """
logger.debug("Initializing tkinter")
for widget in ("TButton", "TCheckbutton", "TRadiobutton"):
self.unbind_class(widget, "<Key-space>")
initialize_config(self, None, None)
initialize_images()
get_config().set_geometry(940, 600, fullscreen=True)
self.title("Faceswap.py - Visual Alignments")
logger.debug("Initialized tkinter")
def _create_containers(self):
""" Create the paned window containers for various GUI elements
Returns
-------
dict:
The main containers of the manual tool.
"""
logger.debug("Creating containers")
main = ttk.PanedWindow(self,
orient=tk.VERTICAL,
name="pw_main")
main.pack(fill=tk.BOTH, expand=True)
top = ttk.Frame(main, name="frame_top")
main.add(top)
bottom = ttk.Frame(main, name="frame_bottom")
main.add(bottom)
retval = {"main": main, "top": top, "bottom": bottom}
logger.debug("Created containers: %s", retval)
return retval
def _handle_key_press(self, event):
""" Keyboard shortcuts
Parameters
----------
event: :class:`tkinter.Event()`
The tkinter key press event
Notes
-----
The following keys are reserved for the :mod:`tools.lib_manual.editor` classes
* Delete - Used for deleting faces
* [] - decrease / increase brush size
* B, D, E, M - Optional Actions (Brush, Drag, Erase, Zoom)
"""
# Alt modifier appears to be broken in Windows so don't use it.
modifiers = {0x0001: 'shift',
0x0004: 'ctrl'}
tk_pos = self._globals.tk_frame_index
bindings = {
"z": self._display.navigation.decrement_frame,
"x": self._display.navigation.increment_frame,
"space": self._display.navigation.handle_play_button,
"home": self._display.navigation.goto_first_frame,
"end": self._display.navigation.goto_last_frame,
"down": lambda d="down": self._faces_frame.canvas_scroll(d),
"up": lambda d="up": self._faces_frame.canvas_scroll(d),
"next": lambda d="page-down": self._faces_frame.canvas_scroll(d),
"prior": lambda d="page-up": self._faces_frame.canvas_scroll(d),
"f": self._display.cycle_filter_mode,
"f1": lambda k=event.keysym: self._display.set_action(k),
"f2": lambda k=event.keysym: self._display.set_action(k),
"f3": lambda k=event.keysym: self._display.set_action(k),
"f4": lambda k=event.keysym: self._display.set_action(k),
"f5": lambda k=event.keysym: self._display.set_action(k),
"f9": lambda k=event.keysym: self._faces_frame.set_annotation_display(k),
"f10": lambda k=event.keysym: self._faces_frame.set_annotation_display(k),
"c": lambda f=tk_pos.get(), d="prev": self._detected_faces.update.copy(f, d),
"v": lambda f=tk_pos.get(), d="next": self._detected_faces.update.copy(f, d),
"ctrl_s": self._detected_faces.save,
"r": lambda f=tk_pos.get(): self._detected_faces.revert_to_saved(f)}
# Allow keypad keys to be used for numbers
press = event.keysym.replace("KP_", "") if event.keysym.startswith("KP_") else event.keysym
modifier = "_".join(val for key, val in modifiers.items() if event.state & key != 0)
key_press = "_".join([modifier, press]) if modifier else press
if key_press.lower() in bindings:
logger.trace("key press: %s, action: %s", key_press, bindings[key_press.lower()])
self.focus_set()
bindings[key_press.lower()]()
def _set_initial_layout(self):
""" Set the favicon and the bottom frame position to correct location to display full
frame window.
Notes
-----
The favicon pops the tkinter GUI (without loaded elements) as soon as it is called, so
this is set last.
"""
logger.debug("Setting initial layout")
self.tk.call("wm",
"iconphoto",
self._w, get_images().icons["favicon"]) # pylint:disable=protected-access
location = int(self.winfo_screenheight() // 1.5)
self._containers["main"].sashpos(0, location)
self.update_idletasks()
def process(self):
""" The entry point for the Visual Alignments tool from :mod:`lib.tools.manual.cli`.
Launch the tkinter Visual Alignments Window and run main loop.
"""
logger.debug("Launching mainloop")
self.mainloop()
class _Options(ttk.Frame): # pylint:disable=too-many-ancestors
""" Control panel options for currently displayed Editor. This is the right hand panel of the
GUI that holds editor specific settings and annotation display settings.
parent: :class:`tkinter.ttk.Frame`
The parent frame for the control panel options
tk_globals: :class:`~tools.manual.manual.TkGlobals`
The tkinter variables that apply to the whole of the GUI
display_frame: :class:`DisplayFrame`
The frame that holds the editors
"""
def __init__(self, parent, tk_globals, display_frame):
logger.debug("Initializing %s: (parent: %s, tk_globals: %s, display_frame: %s)",
self.__class__.__name__, parent, tk_globals, display_frame)
super().__init__(parent)
self._globals = tk_globals
self._display_frame = display_frame
self._control_panels = self._initialize()
self._set_tk_callbacks()
self._update_options()
self.pack(side=tk.RIGHT, fill=tk.Y)
logger.debug("Initialized %s", self.__class__.__name__)
def _initialize(self):
""" Initialize all of the control panels, then display the default panel.
Adds the control panel to :attr:`_control_panels` and sets the traceback to update
display when a panel option has been changed.
Notes
-----
All panels must be initialized at the beginning so that the global format options are not
reset to default when the editor is first selected.
The Traceback must be set after the panel has first been packed as otherwise it interferes
with the loading of the faces pane.
"""
self._initialize_face_options()
frame = ttk.Frame(self)
frame.pack(side=tk.TOP, fill=tk.BOTH, expand=True)
panels = {}
for name, editor in self._display_frame.editors.items():
logger.debug("Initializing control panel for '%s' editor", name)
controls = editor.controls
panel = ControlPanel(frame, controls["controls"],
option_columns=2,
columns=1,
max_columns=1,
header_text=controls["header"],
blank_nones=False,
label_width=12,
style="CPanel",
scrollbar=False)
panel.pack_forget()
panels[name] = panel
return panels
def _initialize_face_options(self):
""" Set the Face Viewer options panel, beneath the standard control options. """
frame = ttk.Frame(self)
frame.pack(side=tk.BOTTOM, fill=tk.X, padx=5, pady=5)
size_frame = ttk.Frame(frame)
size_frame.pack(side=tk.RIGHT)
lbl = ttk.Label(size_frame, text="Face Size:")
lbl.pack(side=tk.LEFT)
cmb = ttk.Combobox(size_frame,
value=["Tiny", "Small", "Medium", "Large", "Extra Large"],
state="readonly",
textvariable=self._globals.tk_faces_size)
self._globals.tk_faces_size.set("Medium")
cmb.pack(side=tk.RIGHT, padx=5)
def _set_tk_callbacks(self):
""" Sets the callback to change to the relevant control panel options when the selected
editor is changed, and the display update on panel option change."""
self._display_frame.tk_selected_action.trace("w", self._update_options)
seen_controls = set()
for name, editor in self._display_frame.editors.items():
for ctl in editor.controls["controls"]:
if ctl in seen_controls:
# Some controls are re-used (annotation format), so skip if trace has already
# been set
continue
logger.debug("Adding control update callback: (editor: %s, control: %s)",
name, ctl.title)
seen_controls.add(ctl)
ctl.tk_var.trace("w", lambda *e: self._globals.tk_update.set(True))
def _update_options(self, *args): # pylint:disable=unused-argument
""" Update the control panel display for the current editor.
If the options have not already been set, then adds the control panel to
:attr:`_control_panels`. Displays the current editor's control panel
Parameters
----------
args: tuple
Unused but required for tkinter variable callback
"""
self._clear_options_frame()
editor = self._display_frame.tk_selected_action.get()
logger.debug("Displaying control panel for editor: '%s'", editor)
self._control_panels[editor].pack(expand=True, fill=tk.BOTH)
def _clear_options_frame(self):
""" Hides the currently displayed control panel """
for editor, panel in self._control_panels.items():
if panel.winfo_ismapped():
logger.debug("Hiding control panel for: %s", editor)
panel.pack_forget()
class TkGlobals():
""" Holds Tkinter Variables and other frame information that need to be accessible from all
areas of the GUI.
Parameters
----------
input_location: str
The location of the input folder of frames or video file
"""
def __init__(self, input_location):
logger.debug("Initializing %s: (input_location: %s)",
self.__class__.__name__, input_location)
self._tk_vars = self._get_tk_vars()
self._is_video = self._check_input(input_location)
self._frame_count = 0 # set by FrameLoader
self._frame_display_dims = (int(round(896 * get_config().scaling_factor)),
int(round(504 * get_config().scaling_factor)))
self._current_frame = {"image": None,
"scale": None,
"interpolation": None,
"display_dims": None,
"filename": None}
logger.debug("Initialized %s", self.__class__.__name__)
@classmethod
def _get_tk_vars(cls):
""" Create and initialize the tkinter variables.
Returns
-------
dict
The variable name as key, the variable as value
"""
retval = {}
for name in ("frame_index", "transport_index", "face_index", "filter_distance"):
var = tk.IntVar()
var.set(10 if name == "filter_distance" else 0)
retval[name] = var
for name in ("update", "update_active_viewport", "is_zoomed"):
var = tk.BooleanVar()
var.set(False)
retval[name] = var
for name in ("filter_mode", "faces_size"):
retval[name] = tk.StringVar()
return retval
@property
def current_frame(self):
""" dict: The currently displayed frame in the frame viewer with it's meta information. Key
and Values are as follows:
**image** (:class:`numpy.ndarry`): The currently displayed frame in original dimensions
**scale** (`float`): The scaling factor to use to resize the image to the display
window
**interpolation** (`int`): The opencv interpolator ID to use for resizing the image to
the display window
**display_dims** (`tuple`): The size of the currently displayed frame, sized for the
display window
**filename** (`str`): The filename of the currently displayed frame
"""
return self._current_frame
@property
def frame_count(self):
""" int: The total number of frames for the input location """
return self._frame_count
@property
def tk_face_index(self):
""" :class:`tkinter.IntVar`: The variable that holds the face index of the selected face
within the current frame when in zoomed mode. """
return self._tk_vars["face_index"]
@property
def tk_update_active_viewport(self):
""" :class:`tkinter.BooleanVar`: Boolean Variable that is traced by the viewport's active
frame to update.. """
return self._tk_vars["update_active_viewport"]
@property
def face_index(self):
""" int: The currently displayed face index when in zoomed mode. """
return self._tk_vars["face_index"].get()
@property
def frame_display_dims(self):
""" tuple: The (`width`, `height`) of the video display frame in pixels. """
return self._frame_display_dims
@property
def frame_index(self):
""" int: The currently displayed frame index. NB This returns -1 if there are no frames
that meet the currently selected filter criteria. """
return self._tk_vars["frame_index"].get()
@property
def tk_frame_index(self):
""" :class:`tkinter.IntVar`: The variable holding the current frame index. """
return self._tk_vars["frame_index"]
@property
def filter_mode(self):
""" str: The currently selected navigation mode. """
return self._tk_vars["filter_mode"].get()
@property
def tk_filter_mode(self):
""" :class:`tkinter.StringVar`: The variable holding the currently selected navigation
filter mode. """
return self._tk_vars["filter_mode"]
@property
def tk_filter_distance(self):
""" :class:`tkinter.DoubleVar`: The variable holding the currently selected threshold
distance for misaligned filter mode. """
return self._tk_vars["filter_distance"]
@property
def tk_faces_size(self):
""" :class:`tkinter.StringVar`: The variable holding the currently selected Faces Viewer
thumbnail size. """
return self._tk_vars["faces_size"]
@property
def is_video(self):
""" bool: ``True`` if the input is a video file, ``False`` if it is a folder of images. """
return self._is_video
@property
def tk_is_zoomed(self):
""" :class:`tkinter.BooleanVar`: The variable holding the value indicating whether the
frame viewer is zoomed into a face or zoomed out to the full frame. """
return self._tk_vars["is_zoomed"]
@property
def is_zoomed(self):
""" bool: ``True`` if the frame viewer is zoomed into a face, ``False`` if the frame viewer
is displaying a full frame. """
return self._tk_vars["is_zoomed"].get()
@property
def tk_transport_index(self):
""" :class:`tkinter.IntVar`: The current index of the display frame's transport slider. """
return self._tk_vars["transport_index"]
@property
def tk_update(self):
""" :class:`tkinter.BooleanVar`: The variable holding the trigger that indicates that a
full update needs to occur. """
return self._tk_vars["update"]
@staticmethod
def _check_input(frames_location):
""" Check whether the input is a video
Parameters
----------
frames_location: str
The input location for video or images
Returns
-------
bool: 'True' if input is a video 'False' if it is a folder.
"""
if os.path.isdir(frames_location):
retval = False
elif os.path.splitext(frames_location)[1].lower() in _video_extensions:
retval = True
else:
logger.error("The input location '%s' is not valid", frames_location)
sys.exit(1)
logger.debug("Input '%s' is_video: %s", frames_location, retval)
return retval
def set_frame_count(self, count):
""" Set the count of total number of frames to :attr:`frame_count` when the
:class:`FramesLoader` has completed loading.
Parameters
----------
count: int
The number of frames that exist for this session
"""
logger.debug("Setting frame_count to : %s", count)
self._frame_count = count
def set_current_frame(self, image, filename):
""" Set the frame and meta information for the currently displayed frame. Populates the
attribute :attr:`current_frame`
Parameters
----------
image: :class:`numpy.ndarray`
The image used to display in the Frame Viewer
filename: str
The filename of the current frame
"""
scale = min(self.frame_display_dims[0] / image.shape[1],
self.frame_display_dims[1] / image.shape[0])
self._current_frame["image"] = image
self._current_frame["filename"] = filename
self._current_frame["scale"] = scale
self._current_frame["interpolation"] = cv2.INTER_CUBIC if scale > 1.0 else cv2.INTER_AREA
self._current_frame["display_dims"] = (int(round(image.shape[1] * scale)),
int(round(image.shape[0] * scale)))
logger.trace({k: v.shape if isinstance(v, np.ndarray) else v
for k, v in self._current_frame.items()})
def set_frame_display_dims(self, width, height):
""" Set the size, in pixels, of the video frame display window and resize the displayed
frame.
Used on a frame resize callback, sets the :attr:frame_display_dims`.
Parameters
----------
width: int
The width of the frame holding the video canvas in pixels
height: int
The height of the frame holding the video canvas in pixels
"""
self._frame_display_dims = (int(width), int(height))
image = self._current_frame["image"]
scale = min(self.frame_display_dims[0] / image.shape[1],
self.frame_display_dims[1] / image.shape[0])
self._current_frame["scale"] = scale
self._current_frame["interpolation"] = cv2.INTER_CUBIC if scale > 1.0 else cv2.INTER_AREA
self._current_frame["display_dims"] = (int(round(image.shape[1] * scale)),
int(round(image.shape[0] * scale)))
logger.trace({k: v.shape if isinstance(v, np.ndarray) else v
for k, v in self._current_frame.items()})
class Aligner():
""" The :class:`Aligner` class sets up an extraction pipeline for each of the current Faceswap
Aligners, along with the Landmarks based Maskers. When new landmarks are required, the bounding
boxes from the GUI are passed to this class for pushing through the pipeline. The resulting
Landmarks and Masks are then returned.
Parameters
----------
tk_globals: :class:`~tools.manual.manual.TkGlobals`
The tkinter variables that apply to the whole of the GUI
exclude_gpus: list or ``None``
A list of indices correlating to connected GPUs that Tensorflow should not use. Pass
``None`` to not exclude any GPUs.
"""
def __init__(self, tk_globals: TkGlobals, exclude_gpus: list[int] | None) -> None:
logger.debug("Initializing: %s (tk_globals: %s, exclude_gpus: %s)",
self.__class__.__name__, tk_globals, exclude_gpus)
self._globals = tk_globals
self._exclude_gpus = exclude_gpus
self._detected_faces: DetectedFaces | None = None
self._frame_index: int | None = None
self._face_index: int | None = None
self._aligners: dict[TypeManualExtractor, Extractor | None] = {"cv2-dnn": None,
"FAN": None,
"mask": None}
self._aligner: TypeManualExtractor = "FAN"
self._init_thread = self._background_init_aligner()
logger.debug("Initialized: %s", self.__class__.__name__)
@property
def _in_queue(self) -> EventQueue:
""" :class:`queue.Queue` - The input queue to the extraction pipeline. """
aligner = self._aligners[self._aligner]
assert aligner is not None
return aligner.input_queue
@property
def _feed_face(self) -> ExtractMedia:
""" :class:`plugins.extract.pipeline.ExtractMedia`: The current face for feeding into the
aligner, formatted for the pipeline """
assert self._frame_index is not None
assert self._face_index is not None
assert self._detected_faces is not None
face = self._detected_faces.current_faces[self._frame_index][self._face_index]
return ExtractMedia(
self._globals.current_frame["filename"],
self._globals.current_frame["image"],
detected_faces=[face])
@property
def is_initialized(self) -> bool:
""" bool: The Aligners are initialized in a background thread so that other tasks can be
performed whilst we wait for initialization. ``True`` is returned if the aligner has
completed initialization otherwise ``False``."""
thread_is_alive = self._init_thread.is_alive()
if thread_is_alive:
logger.trace("Aligner not yet initialized") # type:ignore[attr-defined]
self._init_thread.check_and_raise_error()
else:
logger.trace("Aligner initialized") # type:ignore[attr-defined]
self._init_thread.join()
return not thread_is_alive
def _background_init_aligner(self) -> MultiThread:
""" Launch the aligner in a background thread so we can run other tasks whilst
waiting for initialization
Returns
-------
:class:`lib.multithreading.MultiThread
The background aligner loader thread
"""
logger.debug("Launching aligner initialization thread")
thread = MultiThread(self._init_aligner,
thread_count=1,
name=f"{self.__class__.__name__}.init_aligner")
thread.start()
logger.debug("Launched aligner initialization thread")
return thread
def _init_aligner(self) -> None:
""" Initialize Aligner in a background thread, and set it to :attr:`_aligner`. """
logger.debug("Initialize Aligner")
# Make sure non-GPU aligner is allocated first
for model in T.get_args(TypeManualExtractor):
logger.debug("Initializing aligner: %s", model)
plugin = None if model == "mask" else model
exclude_gpus = self._exclude_gpus if model == "FAN" else None
aligner = Extractor(None,
plugin,
["components", "extended"],
exclude_gpus=exclude_gpus,
multiprocess=True,
normalize_method="hist",
disable_filter=True)
if plugin:
aligner.set_batchsize("align", 1) # Set the batchsize to 1
aligner.launch()
logger.debug("Initialized %s Extractor", model)
self._aligners[model] = aligner
def link_faces(self, detected_faces: DetectedFaces) -> None:
""" As the Aligner has the potential to take the longest to initialize, it is kicked off
as early as possible. At this time :class:`~tools.manual.detected_faces.DetectedFaces` is
not yet available.
Once the Aligner has initialized, this function is called to add the
:class:`~tools.manual.detected_faces.DetectedFaces` class as a property of the Aligner.
Parameters
----------
detected_faces: :class:`~tools.manual.detected_faces.DetectedFaces`
The class that holds the :class:`~lib.align.DetectedFace` objects for the
current Manual session
"""
logger.debug("Linking detected_faces: %s", detected_faces)
self._detected_faces = detected_faces
def get_landmarks(self, frame_index: int, face_index: int, aligner: TypeManualExtractor
) -> np.ndarray:
""" Feed the detected face into the alignment pipeline and retrieve the landmarks.
The face to feed into the aligner is generated from the given frame and face indices.
Parameters
----------
frame_index: int
The frame index to extract the aligned face for
face_index: int
The face index within the current frame to extract the face for
aligner: Literal["FAN", "cv2-dnn"]
The aligner to use to extract the face
Returns
-------
:class:`numpy.ndarray`
The 68 point landmark alignments
"""
logger.trace("frame_index: %s, face_index: %s, aligner: %s", # type:ignore[attr-defined]
frame_index, face_index, aligner)
self._frame_index = frame_index
self._face_index = face_index
self._aligner = aligner
self._in_queue.put(self._feed_face)
extractor = self._aligners[aligner]
assert extractor is not None
detected_face = next(extractor.detected_faces()).detected_faces[0]
logger.trace("landmarks: %s", detected_face.landmarks_xy) # type:ignore[attr-defined]
return detected_face.landmarks_xy
def _remove_nn_masks(self, detected_face: DetectedFace) -> None:
""" Remove any non-landmarks based masks on a landmark edit
Parameters
----------
detected_face:
The detected face object to remove masks from
"""
del_masks = {m for m in detected_face.mask if m not in ("components", "extended")}
logger.debug("Removing masks after landmark update: %s", del_masks)
for mask in del_masks:
del detected_face.mask[mask]
def get_masks(self, frame_index: int, face_index: int) -> dict[str, Mask]:
""" Feed the aligned face into the mask pipeline and retrieve the updated masks.
The face to feed into the aligner is generated from the given frame and face indices.
This is to be called when a manual update is done on the landmarks, and new masks need
generating.
Parameters
----------
frame_index: int
The frame index to extract the aligned face for
face_index: int
The face index within the current frame to extract the face for
Returns
-------
dict[str, :class:`~lib.align.detected_face.Mask`]
The updated masks
"""
logger.trace("frame_index: %s, face_index: %s", # type:ignore[attr-defined]
frame_index, face_index)
self._frame_index = frame_index
self._face_index = face_index
self._aligner = "mask"
self._in_queue.put(self._feed_face)
assert self._aligners["mask"] is not None
detected_face = next(self._aligners["mask"].detected_faces()).detected_faces[0]
self._remove_nn_masks(detected_face)
logger.debug("mask: %s", detected_face.mask)
return detected_face.mask
def set_normalization_method(self, method: T.Literal["none", "clahe", "hist", "mean"]) -> None:
""" Change the normalization method for faces fed into the aligner.
The normalization method is user adjustable from the GUI. When this method is triggered
the method is updated for all aligner pipelines.
Parameters
----------
method: Literal["none", "clahe", "hist", "mean"]
The normalization method to use
"""
logger.debug("Setting normalization method to: '%s'", method)
for plugin, aligner in self._aligners.items():
assert aligner is not None
if plugin == "mask":
continue
logger.debug("Setting to: '%s'", method)
aligner.aligner.set_normalize_method(method)
class FrameLoader():
""" Loads the frames, sets the frame count to :attr:`TkGlobals.frame_count` and handles the
return of the correct frame for the GUI.
Parameters
----------
tk_globals: :class:`~tools.manual.manual.TkGlobals`
The tkinter variables that apply to the whole of the GUI
frames_location: str
The path to the input frames
video_meta_data: dict
The meta data held within the alignments file, if it exists and the input is a video
"""
def __init__(self, tk_globals, frames_location, video_meta_data):
logger.debug("Initializing %s: (tk_globals: %s, frames_location: '%s', "
"video_meta_data: %s)", self.__class__.__name__, tk_globals, frames_location,
video_meta_data)
self._globals = tk_globals
self._loader = None
self._current_idx = 0
self._init_thread = self._background_init_frames(frames_location, video_meta_data)
self._globals.tk_frame_index.trace("w", self._set_frame)
logger.debug("Initialized %s", self.__class__.__name__)
@property
def is_initialized(self):
""" bool: ``True`` if the Frame Loader has completed initialization otherwise
``False``. """
thread_is_alive = self._init_thread.is_alive()
if thread_is_alive:
self._init_thread.check_and_raise_error()
else:
self._init_thread.join()
# Setting the initial frame cannot be done in the thread, so set when queried from main
self._set_frame(initialize=True)
return not thread_is_alive
@property
def video_meta_data(self):
""" dict: The pts_time and key frames for the loader. """
return self._loader.video_meta_data
def _background_init_frames(self, frames_location, video_meta_data):
""" Launch the images loader in a background thread so we can run other tasks whilst
waiting for initialization. """
thread = MultiThread(self._load_images,
frames_location,
video_meta_data,
thread_count=1,
name=f"{self.__class__.__name__}.init_frames")
thread.start()
return thread
def _load_images(self, frames_location, video_meta_data):
""" Load the images in a background thread. """
self._loader = SingleFrameLoader(frames_location, video_meta_data=video_meta_data)
self._globals.set_frame_count(self._loader.count)
def _set_frame(self, *args, initialize=False): # pylint:disable=unused-argument
""" Set the currently loaded frame to :attr:`_current_frame` and trigger a full GUI update.
If the loader has not been initialized, or the navigation position is the same as the
current position and the face is not zoomed in, then this returns having done nothing.
Parameters
----------
args: tuple
:class:`tkinter.Event` arguments. Required but not used.
initialize: bool, optional
``True`` if initializing for the first frame to be displayed otherwise ``False``.
Default: ``False``
"""
position = self._globals.frame_index
if not initialize and (position == self._current_idx and not self._globals.is_zoomed):
logger.trace("Update criteria not met. Not updating: (initialize: %s, position: %s, "
"current_idx: %s, is_zoomed: %s)", initialize, position,
self._current_idx, self._globals.is_zoomed)
return
if position == -1:
filename = "No Frame"
frame = np.ones(self._globals.frame_display_dims + (3, ), dtype="uint8")
else:
filename, frame = self._loader.image_from_index(position)
logger.trace("filename: %s, frame: %s, position: %s", filename, frame.shape, position)
self._globals.set_current_frame(frame, filename)
self._current_idx = position
self._globals.tk_update.set(True)
self._globals.tk_update_active_viewport.set(True)