From 3fd26b51a6745e5081f91e21e436b32bb255ccf6 Mon Sep 17 00:00:00 2001 From: torzdf <36920800+torzdf@users.noreply.github.com> Date: Sat, 25 Jul 2020 11:05:29 +0100 Subject: [PATCH] Manual Tool (#1038) Initial Commit --- .gitignore | 2 + docs/full/modules.rst | 2 +- docs/full/tools/manual.faceviewer.rst | 50 + docs/full/tools/manual.frameviewer.rst | 109 ++ docs/full/tools/manual.rst | 57 + docs/full/{ => tools}/tools.rst | 18 +- tools/alignments/alignments.py | 39 +- tools/manual/__init__.py | 0 tools/manual/cli.py | 56 + tools/manual/detected_faces.py | 1040 +++++++++++++++++ tools/manual/faceviewer/__init__.py | 0 tools/manual/faceviewer/frame.py | 742 ++++++++++++ tools/manual/faceviewer/viewport.py | 1011 ++++++++++++++++ tools/manual/frameviewer/__init__.py | 0 tools/manual/frameviewer/control.py | 289 +++++ tools/manual/frameviewer/editor/__init__.py | 8 + tools/manual/frameviewer/editor/_base.py | 623 ++++++++++ .../manual/frameviewer/editor/bounding_box.py | 403 +++++++ .../manual/frameviewer/editor/extract_box.py | 401 +++++++ tools/manual/frameviewer/editor/landmarks.py | 457 ++++++++ tools/manual/frameviewer/editor/mask.py | 544 +++++++++ tools/manual/frameviewer/frame.py | 741 ++++++++++++ tools/manual/manual.py | 845 ++++++++++++++ 23 files changed, 7426 insertions(+), 11 deletions(-) create mode 100644 docs/full/tools/manual.faceviewer.rst create mode 100644 docs/full/tools/manual.frameviewer.rst create mode 100644 docs/full/tools/manual.rst rename docs/full/{ => tools}/tools.rst (79%) create mode 100644 tools/manual/__init__.py create mode 100644 tools/manual/cli.py create mode 100644 tools/manual/detected_faces.py create mode 100644 tools/manual/faceviewer/__init__.py create mode 100644 tools/manual/faceviewer/frame.py create mode 100644 tools/manual/faceviewer/viewport.py create mode 100644 tools/manual/frameviewer/__init__.py create mode 100644 tools/manual/frameviewer/control.py create mode 100644 tools/manual/frameviewer/editor/__init__.py create mode 100644 tools/manual/frameviewer/editor/_base.py create mode 100644 tools/manual/frameviewer/editor/bounding_box.py create mode 100644 tools/manual/frameviewer/editor/extract_box.py create mode 100644 tools/manual/frameviewer/editor/landmarks.py create mode 100644 tools/manual/frameviewer/editor/mask.py create mode 100644 tools/manual/frameviewer/frame.py create mode 100644 tools/manual/manual.py diff --git a/.gitignore b/.gitignore index 8eaaea6b..40e94ba2 100644 --- a/.gitignore +++ b/.gitignore @@ -38,6 +38,8 @@ !tests/*/* !tools !tools/* +!tools/*/* +!tools/*/*/* !_travis !_travis/* !.travis.yml diff --git a/docs/full/modules.rst b/docs/full/modules.rst index d5eefe90..8dd05810 100644 --- a/docs/full/modules.rst +++ b/docs/full/modules.rst @@ -7,4 +7,4 @@ faceswap lib/lib plugins/plugins scripts - tools + tools/tools diff --git a/docs/full/tools/manual.faceviewer.rst b/docs/full/tools/manual.faceviewer.rst new file mode 100644 index 00000000..ed3209f4 --- /dev/null +++ b/docs/full/tools/manual.faceviewer.rst @@ -0,0 +1,50 @@ +****************** +faceviewer package +****************** + +Handles the display of faces in the Face Viewer section of Faceswap's Manual Tool. + +.. contents:: Contents + :local: + +frame module +============ + +.. rubric:: Module Summary + +.. autosummary:: + :nosignatures: + + ~tools.manual.faceviewer.frame.ContextMenu + ~tools.manual.faceviewer.frame.FacesActionsFrame + ~tools.manual.faceviewer.frame.FacesFrame + ~tools.manual.faceviewer.frame.FacesViewer + ~tools.manual.faceviewer.frame.Grid + +.. rubric:: Module + +.. automodule:: tools.manual.faceviewer.frame + :members: + :undoc-members: + :show-inheritance: + +viewport module +=============== + +.. rubric:: Module Summary + +.. autosummary:: + :nosignatures: + + ~tools.manual.faceviewer.viewport.ActiveFrame + ~tools.manual.faceviewer.viewport.HoverBox + ~tools.manual.faceviewer.viewport.TKFace + ~tools.manual.faceviewer.viewport.Viewport + ~tools.manual.faceviewer.viewport.VisibleObjects + +.. rubric:: Module + +.. automodule:: tools.manual.faceviewer.viewport + :members: + :undoc-members: + :show-inheritance: \ No newline at end of file diff --git a/docs/full/tools/manual.frameviewer.rst b/docs/full/tools/manual.frameviewer.rst new file mode 100644 index 00000000..ee6f084a --- /dev/null +++ b/docs/full/tools/manual.frameviewer.rst @@ -0,0 +1,109 @@ +****************** +frameviewer module +****************** + +Handles the display of frames in the Frame Viewer section of Faceswap's Manual Tool. + +.. contents:: Contents + :local: + +frame module +============ + +.. rubric:: Module Summary + +.. autosummary:: + :nosignatures: + + ~tools.manual.frameviewer.frame.ActionsFrame + ~tools.manual.frameviewer.frame.BackgroundImage + ~tools.manual.frameviewer.frame.DisplayFrame + ~tools.manual.frameviewer.frame.FrameViewer + ~tools.manual.frameviewer.frame.Navigation + +.. rubric:: Module + +.. automodule:: tools.manual.frameviewer.frame + :members: + :undoc-members: + :show-inheritance: + +control module +============== + +.. rubric:: Module Summary + +.. autosummary:: + :nosignatures: + + ~tools.manual.frameviewer.control.BackgroundImage + ~tools.manual.frameviewer.control.Navigation + +.. rubric:: Module + +.. automodule:: tools.manual.frameviewer.control + :members: + :undoc-members: + :show-inheritance: + +editor package +============== +.. contents:: Contents + :local: + +_base module +------------ + +.. rubric:: Module Summary + +.. autosummary:: + :nosignatures: + + ~tools.manual.frameviewer.editor._base.Editor + ~tools.manual.frameviewer.editor._base.View + +.. rubric:: Module + +.. automodule:: tools.manual.frameviewer.editor._base + :members: + :undoc-members: + :show-inheritance: + +bounding_box module +------------------- +.. automodule:: tools.manual.frameviewer.editor.bounding_box + :members: + :undoc-members: + :show-inheritance: + +extract_box module +------------------ +.. automodule:: tools.manual.frameviewer.editor.extract_box + :members: + :undoc-members: + :show-inheritance: + +landmarks module +---------------- + +.. rubric:: Module Summary + +.. autosummary:: + :nosignatures: + + ~tools.manual.frameviewer.editor.landmarks.Landmarks + ~tools.manual.frameviewer.editor.landmarks.Mesh + +.. rubric:: Module + +.. automodule:: tools.manual.frameviewer.editor.landmarks + :members: + :undoc-members: + :show-inheritance: + +mask module +----------- +.. automodule:: tools.manual.frameviewer.editor.mask + :members: + :undoc-members: + :show-inheritance: \ No newline at end of file diff --git a/docs/full/tools/manual.rst b/docs/full/tools/manual.rst new file mode 100644 index 00000000..f2428395 --- /dev/null +++ b/docs/full/tools/manual.rst @@ -0,0 +1,57 @@ +************** +manual package +************** + +.. contents:: Contents + :local: + +Subpackages +=========== +The following subpackages handle the main two display areas of the Manual Tool's GUI. + +.. toctree:: + :maxdepth: 4 + + manual.faceviewer + manual.frameviewer + +manual module +============= +The Manual Module is the main entry point into the Manual Editor Tool. + +.. rubric:: Module Summary + +.. autosummary:: + :nosignatures: + + ~tools.manual.manual.Aligner + ~tools.manual.manual.FrameLoader + ~tools.manual.manual.Manual + ~tools.manual.manual.TkGlobals + +.. rubric:: Module + +.. automodule:: tools.manual.manual + :members: + :undoc-members: + :show-inheritance: + +detected_faces module +===================== + +.. rubric:: Module Summary + +.. autosummary:: + :nosignatures: + + ~tools.manual.detected_faces.DetectedFaces + ~tools.manual.detected_faces.FaceUpdate + ~tools.manual.detected_faces.Filter + ~tools.manual.detected_faces.ThumbsCreator + +.. rubric:: Module + +.. automodule:: tools.manual.detected_faces + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/full/tools.rst b/docs/full/tools/tools.rst similarity index 79% rename from docs/full/tools.rst rename to docs/full/tools/tools.rst index 9c892843..07468d3c 100644 --- a/docs/full/tools.rst +++ b/docs/full/tools/tools.rst @@ -7,6 +7,21 @@ The Tools Package provides various tools for working with Faceswap outside of th .. contents:: Contents :local: +Subpackages +=========== + +.. toctree:: + :maxdepth: 1 + + manual + +alignments module +================= +.. automodule:: tools.alignments.alignments + :members: + :undoc-members: + :show-inheritance: + mask module =========== @@ -16,11 +31,10 @@ mask module :show-inheritance: preview module -============== +=============== .. rubric:: Module Summary - .. autosummary:: :nosignatures: diff --git a/tools/alignments/alignments.py b/tools/alignments/alignments.py index fb103c33..8fd277df 100644 --- a/tools/alignments/alignments.py +++ b/tools/alignments/alignments.py @@ -11,16 +11,33 @@ from .jobs_manual import Manual # noqa pylint: disable=unused-import logger = logging.getLogger(__name__) # pylint: disable=invalid-name -class Alignments(): - """ Perform tasks relating to alignments file """ +class Alignments(): # pylint:disable=too-few-public-methods + """ The main entry point for Faceswap's Alignments Tool. This tool is part of the Faceswap + Tools suite and should be called from the ``python tools.py alignments`` command. + + The tool allows for manipulation, and working with Faceswap alignments files. + + 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) self.args = arguments - self.alignments = self.load_alignments() + self.alignments = self._load_alignments() logger.debug("Initialized %s", self.__class__.__name__) - def load_alignments(self): - """ Loading alignments """ + def _load_alignments(self): + """ Loads the given alignments file(s) prior to running the selected job. + + Returns + ------- + :class:`~tools.alignments.media.AlignmentData` or list + The alignments data formatted for use by the alignments tool. If multiple alignments + files have been selected, then this will be a list of + :class:`~tools.alignments.media.AlignmentData` objects + """ logger.debug("Loading alignments") if len(self.args.alignments_file) > 1 and self.args.job != "merge": logger.error("Multiple alignments files are only permitted for merging") @@ -37,13 +54,19 @@ class Alignments(): return retval def process(self): - """ Main processing function of the Align tool """ + """ The entry point for the Alignments tool from :mod:`lib.tools.alignments.cli`. + + Launches the selected alignments job. + """ + if self.args.job == "manual": + logger.warning("The 'manual' job is deprecated and will be removed from a future " + "update. Please use the new 'manual' tool.") if self.args.job == "update-hashes": job = UpdateHashes elif self.args.job.startswith("remove-"): job = RemoveAlignments - elif self.args.job in("missing-alignments", "missing-frames", - "multi-faces", "leftover-faces", "no-faces"): + elif self.args.job in ("missing-alignments", "missing-frames", + "multi-faces", "leftover-faces", "no-faces"): job = Check else: job = globals()[self.args.job.title()] diff --git a/tools/manual/__init__.py b/tools/manual/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tools/manual/cli.py b/tools/manual/cli.py new file mode 100644 index 00000000..db9278eb --- /dev/null +++ b/tools/manual/cli.py @@ -0,0 +1,56 @@ +#!/usr/bin/env python3 +""" The Command Line Arguments for the Manual Editor tool. """ +from lib.cli.args import FaceSwapArgs, DirOrFileFullPaths, FileFullPaths + +_HELPTEXT = ("This command lets you perform various actions on frames, " + "faces and alignments files using visual tools.") + + +class ManualArgs(FaceSwapArgs): + """ Generate the command line options for the Manual Editor Tool.""" + + @staticmethod + def get_info(): + """ Obtain the information about what the Manual Tool does. """ + return ("A tool to perform various actions on frames, faces and alignments files using " + "visual tools") + + @staticmethod + def get_argument_list(): + """ Generate the command line argument list for the Manual Tool. """ + argument_list = list() + argument_list.append(dict( + opts=("-al", "--alignments"), + action=FileFullPaths, + filetypes="alignments", + type=str, + group="data", + dest="alignments_path", + help="Path to the alignments file for the input, if not at the default location")) + argument_list.append(dict( + opts=("-fr", "--frames"), + action=DirOrFileFullPaths, + filetypes="video", + required=True, + group="data", + help="Video file or directory containing source frames that faces were extracted " + "from.")) + argument_list.append(dict( + opts=("-t", "--thumb-regen"), + action="store_true", + dest="thumb_regen", + default=False, + group="options", + help="Force regeneration of the low resolution jpg thumbnails in the alignments " + "file.")) + argument_list.append(dict( + opts=("-s", "--single-process"), + action="store_true", + dest="single_process", + default=False, + group="options", + help="The process attempts to speed up generation of thumbnails by extracting from " + "the video in parallel threads. For some videos, this causes the caching " + "process to hang. If this happens, then set this option to generate the " + "thumbnails in a slower, but more stable single thread.")) + return argument_list diff --git a/tools/manual/detected_faces.py b/tools/manual/detected_faces.py new file mode 100644 index 00000000..27ddbd31 --- /dev/null +++ b/tools/manual/detected_faces.py @@ -0,0 +1,1040 @@ +#!/usr/bin/env python3 +""" Alignments handling for Faceswap's Manual Adjustments tool. Handles the conversion of +alignments data to :class:`~lib.faces_detect.DetectedFace` objects, and the update of these faces +when edits are made in the GUI. """ + +import logging +import os +import tkinter as tk +from copy import deepcopy +from queue import Queue, Empty +from time import sleep +from threading import Lock + +import cv2 +import imageio +import numpy as np +from tqdm import tqdm + +from lib.aligner import Extract as AlignerExtract +from lib.alignments import Alignments +from lib.faces_detect import DetectedFace +from lib.gui.custom_widgets import PopupProgress +from lib.gui.utils import FileHandler +from lib.image import SingleFrameLoader, ImagesLoader, ImagesSaver, encode_image_with_hash +from lib.multithreading import MultiThread +from lib.utils import get_folder + + +logger = logging.getLogger(__name__) # pylint: disable=invalid-name + + +class DetectedFaces(): + """ Handles the manipulation of :class:`~lib.faces_detect.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, alignments_path, input_location, extractor): + 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 = [] + self._updated_frame_indices = set() + + self._alignments = self._get_alignments(alignments_path, input_location) + self._extractor = extractor + self._tk_vars = self._set_tk_vars() + self._children = dict(io=_DiskIO(self, input_location), + update=FaceUpdate(self), + filter=Filter(self)) + logger.debug("Initialized %s", self.__class__.__name__) + + # <<<< PUBLIC PROPERTIES >>>> # + # << SUBCLASSES >> # + @property + def extractor(self): + """ :class:`~tools.manual.manual.Aligner`: The pipeline for passing faces through the + aligner and retrieving results. """ + return self._extractor + + @property + def filter(self): + """ :class:`Filter`: Handles returning of faces and stats based on the current user set + navigation mode filter. """ + return self._children["filter"] + + @property + def update(self): + """ :class:`FaceUpdate`: Handles the adding, removing and updating of + :class:`~lib.faces_detect.DetectedFace` stored within the alignments file. """ + return self._children["update"] + + # << TKINTER VARIABLES >> # + @property + def tk_unsaved(self): + """ :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): + """ :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): + """ :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 available_masks(self): + """ dict: 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: The most up to date full list of :class:`~lib.faces_detect.DetectedFace` + objects. """ + return self._frame_faces + + @property + def video_meta_data(self): + """ dict: 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: 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): + """ 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): + """ Load the faces as :class:`~lib.faces_detect.DetectedFace` objects from the alignments + file. """ + self._children["io"].load() + + def save(self): + """ Save the alignments file with the latest edits. """ + self._children["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._children["io"].revert_to_saved(frame_index) + + def extract(self): + """ Extract the faces in the current video to a user supplied folder. """ + self._children["io"].extract() + + def save_video_meta_data(self, pts_time, keyframes): + """ 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 + A list of presentation timestamps (`float`) in frame index order for every frame in + the input video + keyframes: list + 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(): + """ 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 = dict() + for name in ("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, input_location): + """ Get the :class:`~lib.alignments.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.alignments.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 = "{}_{}".format(vid, filename) + else: + folder = input_location + retval = Alignments(folder, filename) + logger.debug("folder: %s, filename: %s, alignments: %s", folder, filename, retval) + return retval + + +class _DiskIO(): # pylint:disable=too-few-public-methods + """ Handles the loading of :class:`~lib.faces_detect.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, input_location): + 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 + self._sorted_frame_names = sorted(self._alignments.data) + logger.debug("Initialized %s", self.__class__.__name__) + + def load(self): + """ Load the faces from the alignments file, convert to + :class:`~lib.faces_detect.DetectedFace`. objects and add to :attr:`_frame_faces`. """ + for key in sorted(self._alignments.data): + this_frame_faces = [] + for item in self._alignments.data[key]["faces"]: + face = DetectedFace() + face.from_alignment(item, with_thumb=True) + this_frame_faces.append(face) + self._frame_faces.append(this_frame_faces) + + def save(self): + """ Convert updated :class:`~lib.faces_detect.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", len(frames)) + + for idx, faces in zip(frames, np.array(self._frame_faces)[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): + """ 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", 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) + + 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.tk_update.set(True) + + @classmethod + def _add_remove_faces(cls, alignments, faces): + """ On a revert, ensure that the alignments and detected face object counts for each frame + are in sync. """ + 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): + """ 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...").retfile + if not dirname: + return + logger.debug(dirname) + + 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, queue, pbar): + """ 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 + pbar: :class:`lib.gui.custom_widget.PopupProgress` + The popped up progress bar + """ + thread.check_and_raise_error() + if not thread.is_alive(): + thread.join() + # Update hashes in alignments file. + pbar.update_title("Saving Alignments...") + self._alignments.backup() + self._alignments.save() + self._updated_frame_indices.clear() + self._tk_unsaved.set(False) + pbar.stop() + return + + while True: + try: + pbar.step(queue.get(False, 0)) + except Empty: + break + pbar.after(100, self._monitor_extract, thread, queue, pbar) + + def _background_extract(self, output_folder, progress_queue): + """ 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 incrememental counts to for updating the GUI's progress bar + """ + saver = ImagesSaver(str(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", frame_idx, filename) + frame_name, extension = os.path.splitext(filename) + final_faces = [] + progress_queue.put(1) + for face_idx, face in enumerate(self._frame_faces[frame_idx]): + output = "{}_{}{}".format(frame_name, str(face_idx), extension) + face.load_aligned(image, size=256, force=True) # TODO user selectable size + face.hash, b_image = encode_image_with_hash(face.aligned_face, extension) + saver.save(output, b_image) + final_faces.append(face.to_alignment()) + face.aligned = dict() + self._alignments.data[filename]["faces"] = final_faces + 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): + 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 count(self): + """ int: The number of frames that meet the filter criteria returned by + :attr:`~tools.manual.manual.TkGlobals.filter_mode`. """ + face_count_per_index = self._detected_faces.face_count_per_index + if self._globals.filter_mode == "No Faces": + retval = sum(1 for fcount in face_count_per_index if fcount == 0) + elif self._globals.filter_mode == "Has Face(s)": + retval = sum(1 for fcount in face_count_per_index if fcount != 0) + elif self._globals.filter_mode == "Multiple Faces": + retval = sum(1 for fcount in face_count_per_index if fcount > 1) + else: + retval = len(face_count_per_index) + logger.trace("filter mode: %s, frame count: %s", self._globals.filter_mode, retval) + return retval + + @property + def raw_indices(self): + """ dict: The frame and face indices that meet the current filter criteria for each + displayed face. """ + frame_indices = [] + face_indices = [] + if self._globals.filter_mode != "No Faces": + for frame_idx, face_count in enumerate(self._detected_faces.face_count_per_index): + if face_count <= 1 and self._globals.filter_mode == "Multiple Faces": + continue + for face_idx in range(face_count): + frame_indices.append(frame_idx) + face_indices.append(face_idx) + logger.trace("frame_indices: %s, face_indices: %s", frame_indices, face_indices) + retval = dict(frame=frame_indices, face=face_indices) + return retval + + @property + def frames_list(self): + """ list: The list of frame indices that meet the filter criteria returned by + :attr:`~tools.manual.manual.TkGlobals.filter_mode`. """ + face_count_per_index = self._detected_faces.face_count_per_index + if self._globals.filter_mode == "No Faces": + retval = [idx for idx, count in enumerate(face_count_per_index) if count == 0] + elif self._globals.filter_mode == "Multiple Faces": + retval = [idx for idx, count in enumerate(face_count_per_index) if count > 1] + elif self._globals.filter_mode == "Has Face(s)": + retval = [idx for idx, count in enumerate(face_count_per_index) if count != 0] + else: + retval = range(len(face_count_per_index)) + logger.trace("filter mode: %s, number_frames: %s", self._globals.filter_mode, len(retval)) + return retval + + +class FaceUpdate(): + """ Perform updates on :class:`~lib.faces_detect.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): + 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): + """ :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): + """ :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): + """ 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.faces_detect.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 _generate_thumbnail(self, face): + """ Generate the jpg thumbnail from the currently active frame for the detected face and + assign to it's `thumbnail` attribute. + + Parameters + ---------- + face: class:`~lib.faces_detect.DetectedFace` + The detected face object to generate the thumbnail for + """ + face.load_aligned(self._globals.current_frame["image"], 80, force=True) + jpg = cv2.imencode(".jpg", face.aligned_face, [cv2.IMWRITE_JPEG_QUALITY, 60])[1] + face.thumbnail = jpg + face.aligned = dict() + + def add(self, frame_index, pnt_x, width, pnt_y, height): + """ Add a :class:`~lib.faces_detect.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") + self._tk_face_count_changed.set(True) + + def delete(self, frame_index, face_index): + """ Delete the :class:`~lib.faces_detect.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.tk_update.set(True) + + def bounding_box(self, frame_index, face_index, pnt_x, width, pnt_y, height, aligner="FAN"): + """ Update the bounding box for the :class:`~lib.faces_detect.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, 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.x = pnt_x + face.w = width + face.y = pnt_y + face.h = height + face.landmarks_xy = self._extractor.get_landmarks(frame_index, face_index, aligner) + self._globals.tk_update.set(True) + + def landmark(self, frame_index, face_index, landmark_index, shift_x, shift_y, is_zoomed): + """ Shift a single landmark point for the :class:`~lib.faces_detect.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: + if not np.any(face.aligned_landmarks): # This will be None on a resize + face.load_aligned(None, size=min(self._globals.frame_display_dims)) + landmark = face.aligned_landmarks[landmark_index] + landmark += (shift_x, shift_y) + matrix = AlignerExtract.transform_matrix(face.aligned["matrix"], + face.aligned["size"], + face.aligned["padding"]) + 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): + 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) + face.mask = self._extractor.get_masks(frame_index, face_index) + self._globals.tk_update.set(True) + + def landmarks(self, frame_index, face_index, shift_x, shift_y): + """ Shift all of the landmarks and bounding box for the + :class:`~lib.faces_detect.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] + face.x += shift_x + face.y += shift_y + face.landmarks_xy += (shift_x, shift_y) + face.mask = self._extractor.get_masks(frame_index, face_index) + self._globals.tk_update.set(True) + + def landmarks_rotate(self, frame_index, face_index, angle, center): + """ Rotate the landmarks on an Extract Box rotate for the + :class:`~lib.faces_detect.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: :class:`numpy.ndarray` + 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), angle, 1.) + face.landmarks_xy = cv2.transform(np.expand_dims(face.landmarks_xy, axis=0), + rot_mat).squeeze() + face.mask = self._extractor.get_masks(frame_index, face_index) + self._globals.tk_update.set(True) + + def landmarks_scale(self, frame_index, face_index, scale, center): + """ Scale the landmarks on an Extract Box resize for the + :class:`~lib.faces_detect.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.landmarks_xy = ((face.landmarks_xy - center) * scale) + center + face.mask = self._extractor.get_masks(frame_index, face_index) + self._globals.tk_update.set(True) + + def mask(self, frame_index, face_index, mask, mask_type): + """ Update the mask on an edit for the :class:`~lib.faces_detect.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.tk_update.set(True) + + def copy(self, frame_index, direction): + """ 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) + faces.extend(deepcopy(self._faces_at_frame_index(idx))) + self._tk_face_count_changed.set(True) + self._globals.tk_update.set(True) + + def post_edit_trigger(self, frame_index, face_index): + """ Update the jpg thumbnail and the viewport thumbnail 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(self._globals.current_frame["image"], 80, force=True) + jpg = cv2.imencode(".jpg", face.aligned_face, [cv2.IMWRITE_JPEG_QUALITY, 60])[1] + face.thumbnail = jpg + face.aligned = dict() + self._tk_edited.set(True) + + +class ThumbsCreator(): + """ Background loader to generate thumbnails for the alignments file. Generates low resolution + thumbnails in parallel threads for faster processing. + + Parameters + ---------- + detected_faces: :class:`~tool.manual.faces.DetectedFaces` + The :class:`~lib.faces_detect.DetectedFace` objects for this video + input_location: str + The location of the input folder of frames or video file + """ + def __init__(self, detected_faces, input_location, single_process): + logger.debug("Initializing %s: (detected_faces: %s, input_location: %s, " + "single_process: %s)", self.__class__.__name__, detected_faces, + input_location, single_process) + self._size = 80 + self._jpeg_quality = 60 + self._pbar = dict(pbar=None, lock=Lock()) + self._meta = dict(key_frames=detected_faces.video_meta_data.get("keyframes", None), + pts_times=detected_faces.video_meta_data.get("pts_time", None)) + self._location = input_location + self._alignments = detected_faces._alignments + self._frame_faces = detected_faces._frame_faces + + self._is_video = all(val is not None for val in self._meta.values()) + self._num_threads = os.cpu_count() - 2 + if self._is_video and single_process: + self._num_threads = 1 + elif self._is_video and not single_process: + self._num_threads = min(self._num_threads, len(self._meta["key_frames"])) + else: + self._num_threads = max(self._num_threads, 32) + self._threads = [] + logger.debug("Initialized %s", self.__class__.__name__) + + @property + def has_thumbs(self): + """ bool: ``True`` if the underlying alignments file holds thumbnail images + otherwise ``False``. """ + return self._alignments.thumbnails.has_thumbnails + + def generate_cache(self): + """ Extract the face thumbnails from a video or folder of images into the + alignments file. """ + self._pbar["pbar"] = tqdm(desc="Caching Thumbnails", + leave=False, + total=len(self._frame_faces)) + if self._is_video: + self._launch_video() + else: + self._launch_folder() + while True: + self._check_and_raise_error() + if all(not thread.is_alive() for thread in self._threads): + break + sleep(1) + self._join_threads() + self._pbar["pbar"].close() + self._alignments.save() + + # << PRIVATE METHODS >> # + def _check_and_raise_error(self): + """ Monitor the loading threads for errors and raise if any occur. """ + for thread in self._threads: + thread.check_and_raise_error() + + def _join_threads(self): + """ Join the loading threads """ + logger.debug("Joining face viewer loading threads") + for thread in self._threads: + thread.join() + + def _launch_video(self): + """ Launch multiple :class:`lib.multithreading.MultiThread` objects to load faces from + a video file. + + Splits the video into segments and passes each of these segments to separate background + threads for some speed up. + """ + key_frame_split = len(self._meta["key_frames"]) // self._num_threads + key_frames = self._meta["key_frames"] + pts_times = self._meta["pts_times"] + for idx in range(self._num_threads): + is_final = idx == self._num_threads - 1 + start_idx = idx * key_frame_split + keyframe_idx = len(key_frames) - 1 if is_final else start_idx + key_frame_split + end_idx = key_frames[keyframe_idx] + start_pts = pts_times[key_frames[start_idx]] + end_pts = False if idx + 1 == self._num_threads else pts_times[end_idx] + starting_index = pts_times.index(start_pts) + if end_pts: + segment_count = len(pts_times[key_frames[start_idx]:end_idx]) + else: + segment_count = len(pts_times[key_frames[start_idx]:]) + logger.debug("thread index: %s, start_idx: %s, end_idx: %s, start_pts: %s, " + "end_pts: %s, starting_index: %s, segment_count: %s", idx, start_idx, + end_idx, start_pts, end_pts, starting_index, segment_count) + thread = MultiThread(self._load_from_video, + start_pts, + end_pts, + starting_index, + segment_count) + thread.start() + self._threads.append(thread) + + def _launch_folder(self): + """ Launch :class:`lib.multithreading.MultiThread` to retrieve faces from a + folder of images. + + Goes through the file list one at a time, passing each file to a separate background + thread for some speed up. + """ + reader = SingleFrameLoader(self._location) + num_threads = min(reader.count, self._num_threads) + frame_split = reader.count // self._num_threads + logger.debug("total images: %s, num_threads: %s, frames_per_thread: %s", + reader.count, num_threads, frame_split) + for idx in range(num_threads): + is_final = idx == num_threads - 1 + start_idx = idx * frame_split + end_idx = reader.count if is_final else start_idx + frame_split + thread = MultiThread(self._load_from_folder, reader, start_idx, end_idx) + thread.start() + self._threads.append(thread) + + def _load_from_video(self, pts_start, pts_end, start_index, segment_count): + """ Loads faces from video for the given segment of the source video. + + Each segment of the video is extracted from in a different background thread. + + Parameters + ---------- + pts_start: float + The start time to cut the segment out of the video + pts_end: float + The end time to cut the segment out of the video + start_index: int + The frame index that this segment starts from. Used for calculating the actual frame + index of each frame extracted + segment_count: int + The number of frames that appear in this segment. Used for ending early in case more + frames come out of the segment than should appear (sometimes more frames are picked up + at the end of the segment, so these are discarded) + """ + logger.debug("pts_start: %s, pts_end: %s, start_index: %s, segment_count: %s", + pts_start, pts_end, start_index, segment_count) + reader = self._get_reader(pts_start, pts_end) + idx = 0 + sample_filename = next(fname for fname in self._alignments.data) + vidname = sample_filename[:sample_filename.rfind("_")] + for idx, frame in enumerate(reader): + frame_idx = idx + start_index + filename = "{}_{:06d}.png".format(vidname, frame_idx + 1) + self._set_thumbail(filename, frame[..., ::-1], frame_idx) + if idx == segment_count - 1: + # Sometimes extra frames are picked up at the end of a segment, so stop + # processing when segment frame count has been hit. + break + reader.close() + logger.debug("Segment complete: (starting_frame_index: %s, processed_count: %s)", + start_index, idx) + + def _get_reader(self, pts_start, pts_end): + """ Get an imageio iterator for this thread's segment. + + Parameters + ---------- + pts_start: float + The start time to cut the segment out of the video + pts_end: float + The end time to cut the segment out of the video + + Returns + ------- + :class:`imageio.Reader` + A reader iterator for the requested segment of video + """ + input_params = ["-ss", str(pts_start)] + if pts_end: + input_params.extend(["-to", str(pts_end)]) + logger.debug("pts_start: %s, pts_end: %s, input_params: %s", + pts_start, pts_end, input_params) + return imageio.get_reader(self._location, "ffmpeg", input_params=input_params) + + def _load_from_folder(self, reader, start_index, end_index): + """ Loads faces from the given range of frame indices from a folder of images. + + Each frame range is extracted in a different background thread. + + Parameters + ---------- + reader: :class:`lib.image.SingleFrameLoader` + The reader that is used to retrieve the requested frame + start_index: int + The starting frame index for the images to extract faces from + end_index: int + The end frame index for the images to extract faces from + """ + logger.debug("reader: %s, start_index: %s, end_index: %s", + reader, start_index, end_index) + for frame_index in range(start_index, end_index): + filename, frame = reader.image_from_index(frame_index) + self._set_thumbail(filename, frame, frame_index) + logger.debug("Segment complete: (start_index: %s, processed_count: %s)", + start_index, end_index - start_index) + + def _set_thumbail(self, filename, frame, frame_index): + """ Extracts the faces from the frame and adds to alignments file + + Parameters + ---------- + filename: str + The filename of the frame within the alignments file + frame: :class:`numpy.ndarray` + The frame that contains the faces + frame_index: int + The frame index of this frame in the :attr:`_frame_faces` + """ + for face_idx, face in enumerate(self._frame_faces[frame_index]): + face.load_aligned(frame, size=self._size, force=True) + jpg = cv2.imencode(".jpg", + face.aligned_face, + [cv2.IMWRITE_JPEG_QUALITY, self._jpeg_quality])[1] + face.thumbnail = jpg + self._alignments.thumbnails.add_thumbnail(filename, face_idx, jpg) + face.aligned["face"] = None + with self._pbar["lock"]: + self._pbar["pbar"].update(1) diff --git a/tools/manual/faceviewer/__init__.py b/tools/manual/faceviewer/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tools/manual/faceviewer/frame.py b/tools/manual/faceviewer/frame.py new file mode 100644 index 00000000..169a87d8 --- /dev/null +++ b/tools/manual/faceviewer/frame.py @@ -0,0 +1,742 @@ +#!/usr/bin/env python3 +""" The Faces Viewer Frame and Canvas for Faceswap's Manual Tool. """ +import colorsys +import logging +import platform +import tkinter as tk +from tkinter import ttk +from math import floor, ceil +from threading import Thread, Event + +import numpy as np + +from lib.gui.custom_widgets import RightClickMenu, Tooltip +from lib.gui.utils import get_config, get_images +from lib.image import hex_to_rgb, rgb_to_hex + +from .viewport import Viewport + +logger = logging.getLogger(__name__) # pylint: disable=invalid-name + + +class FacesFrame(ttk.Frame): # pylint:disable=too-many-ancestors + """ The faces display frame (bottom section of GUI). This frame holds the faces viewport and + the tkinter objects. + + Parameters + ---------- + parent: :class:`tkinter.PanedWindow` + The paned window that the faces frame resides in + tk_globals: :class:`~tools.manual.manual.TkGlobals` + The tkinter variables that apply to the whole of the GUI + detected_faces: :class:`~tools.manual.detected_faces.DetectedFaces` + The :class:`~lib.faces_detect.DetectedFace` objects for this video + display_frame: :class:`~tools.manual.frameviewer.frame.DisplayFrame` + The section of the Manual Tool that holds the frames viewer + """ + def __init__(self, parent, tk_globals, detected_faces, display_frame): + logger.debug("Initializing %s: (parent: %s, tk_globals: %s, detected_faces: %s, " + "display_frame: %s)", self.__class__.__name__, parent, tk_globals, + detected_faces, display_frame) + super().__init__(parent) + self.pack(side=tk.TOP, fill=tk.BOTH, expand=True) + self._actions_frame = FacesActionsFrame(self) + + self._faces_frame = ttk.Frame(self) + self._faces_frame.pack_propagate(0) + self._faces_frame.pack(side=tk.LEFT, fill=tk.BOTH, expand=True) + self._event = Event() + self._canvas = FacesViewer(self._faces_frame, + tk_globals, + self._actions_frame._tk_vars, + detected_faces, + display_frame, + self._event) + self._add_scrollbar() + logger.debug("Initialized %s", self.__class__.__name__) + + def _add_scrollbar(self): + """ Add a scrollbar to the faces frame """ + logger.debug("Add Faces Viewer Scrollbar") + scrollbar = ttk.Scrollbar(self._faces_frame, command=self._on_scroll) + scrollbar.pack(side=tk.RIGHT, fill=tk.Y) + self._canvas.config(yscrollcommand=scrollbar.set) + self.bind("", self._update_viewport) + logger.debug("Added Faces Viewer Scrollbar") + self.update_idletasks() # Update so scrollbar width is correct + return scrollbar.winfo_width() + + def _on_scroll(self, *event): + """ Callback on scrollbar scroll. Updates the canvas location and displays/hides + thumbnail images. + + Parameters + ---------- + event :class:`tkinter.Event` + The scrollbar callback event + """ + self._canvas.yview(*event) + self._canvas.viewport.update() + + def _update_viewport(self, event): # pylint: disable=unused-argument + """ Update the faces viewport and scrollbar. + + Parameters + ---------- + event: :class:`tkinter.Event` + Unused but required + """ + self._canvas.viewport.update() + self._canvas.configure(scrollregion=self._canvas.bbox("backdrop")) + + def canvas_scroll(self, direction): + """ Scroll the canvas on an up/down or page-up/page-down key press. + + Notes + ----- + To protect against a held down key press stacking tasks and locking up the GUI + a background thread is launched and discards subsequent key presses whilst the + previous update occurs. + + Parameters + ---------- + direction: ["up", "down", "page-up", "page-down"] + The request page scroll direction and amount. + """ + + if self._event.is_set(): + logger.trace("Update already running. Aborting repeated keypress") + return + logger.trace("Running update on received key press: %s", direction) + + amount = 1 if direction.endswith("down") else -1 + units = "pages" if direction.startswith("page") else "units" + self._event.set() + thread = Thread(target=self._canvas.canvas_scroll, + args=(amount, units, self._event)) + thread.start() + + def set_annotation_display(self, key): + """ Set the optional annotation overlay based on keyboard shortcut. + + Parameters + ---------- + key: str + The pressed key + """ + self._actions_frame.on_click(self._actions_frame.key_bindings[key]) + + +class FacesActionsFrame(ttk.Frame): # pylint:disable=too-many-ancestors + """ The left hand action frame holding the optional annotation buttons. + + Parameters + ---------- + parent: :class:`FacesFrame` + The Faces frame that this actions frame reside in + """ + def __init__(self, parent): + logger.debug("Initializing %s: (parent: %s)", + self.__class__.__name__, parent) + super().__init__(parent) + self.pack(side=tk.LEFT, fill=tk.Y, padx=(2, 4), pady=2) + self._tk_vars = dict() + self._configure_styles() + self._buttons = self._add_buttons() + logger.debug("Initialized %s", self.__class__.__name__) + + @property + def key_bindings(self): + """ dict: The mapping of key presses to optional annotations to display. Keyboard shortcuts + utilize the function keys. """ + return {"F{}".format(idx + 9): display for idx, display in enumerate(("mesh", "mask"))} + + @property + def _helptext(self): + """ dict: `button key`: `button helptext`. The help text to display for each button. """ + inverse_keybindings = {val: key for key, val in self.key_bindings.items()} + retval = dict(mesh="Display the landmarks mesh", + mask="Display the mask") + for item in retval: + retval[item] += " ({})".format(inverse_keybindings[item]) + return retval + + def _configure_styles(self): + """ Configure the background color for button frame and the button styles. """ + style = ttk.Style() + style.configure("display.TFrame", background='#d3d3d3') + style.configure("display_selected.TButton", relief="flat", background="#bedaf1") + style.configure("display_deselected.TButton", relief="flat") + self.config(style="display.TFrame") + + def _add_buttons(self): + """ Add the display buttons to the Faces window. + + Returns + ------- + dict + The display name and its associated button. + """ + frame = ttk.Frame(self) + frame.pack(side=tk.TOP, fill=tk.Y) + buttons = dict() + for display in self.key_bindings.values(): + var = tk.BooleanVar() + var.set(False) + self._tk_vars[display] = var + + lookup = "landmarks" if display == "mesh" else display + button = ttk.Button(frame, + image=get_images().icons[lookup], + command=lambda t=display: self.on_click(t), + style="display_deselected.TButton") + button.state(["!pressed", "!focus"]) + button.pack() + Tooltip(button, text=self._helptext[display]) + buttons[display] = button + return buttons + + def on_click(self, display): + """ Click event for the optional annotation buttons. Loads and unloads the annotations from + the faces viewer. + + Parameters + ---------- + display: str + The display name for the button that has called this event as exists in + :attr:`_buttons` + """ + is_pressed = not self._tk_vars[display].get() + style = "display_selected.TButton" if is_pressed else "display_deselected.TButton" + state = ["pressed", "focus"] if is_pressed else ["!pressed", "!focus"] + btn = self._buttons[display] + btn.configure(style=style) + btn.state(state) + self._tk_vars[display].set(is_pressed) + + +class FacesViewer(tk.Canvas): # pylint:disable=too-many-ancestors + """ The :class:`tkinter.Canvas` that holds the faces viewer section of the Manual Tool. + + Parameters + ---------- + parent: :class:`tkinter.ttk.Frame` + The parent frame for the canvas + tk_globals: :class:`~tools.manual.manual.TkGlobals` + The tkinter variables that apply to the whole of the GUI + tk_action_vars: dict + The :class:`tkinter.BooleanVar` objects for selectable optional annotations + as set by the buttons in the :class:`FacesActionsFrame` + detected_faces: :class:`~tools.manual.detected_faces.DetectedFaces` + The :class:`~lib.faces_detect.DetectedFace` objects for this video + display_frame: :class:`~tools.manual.frameviewer.frame.DisplayFrame` + The section of the Manual Tool that holds the frames viewer + event: :class:`threading.Event` + The threading event object for repeated key press protection + """ + def __init__(self, parent, tk_globals, tk_action_vars, detected_faces, display_frame, event): + logger.debug("Initializing %s: (parent: %s, tk_globals: %s, tk_action_vars: %s, " + "detected_faces: %s, display_frame: %s, event: %s)", self.__class__.__name__, + parent, tk_globals, tk_action_vars, detected_faces, display_frame, event) + super().__init__(parent, bd=0, highlightthickness=0, bg="#bcbcbc") + self.pack(side=tk.LEFT, fill=tk.BOTH, expand=True, anchor=tk.E) + self._sizes = dict(tiny=32, small=64, medium=96, large=128, extralarge=192) + + self._globals = tk_globals + self._tk_optional_annotations = tk_action_vars + self._event = event + self._display_frame = display_frame + self._grid = Grid(self, detected_faces) + self._view = Viewport(self, detected_faces.tk_edited) + self._annotation_colors = dict(mesh=self.get_muted_color("Mesh"), + box=self.control_colors["ExtractBox"]) + + ContextMenu(self, detected_faces) + self._bind_mouse_wheel_scrolling() + self._set_tk_callbacks(detected_faces) + logger.debug("Initialized %s", self.__class__.__name__) + + @property + def face_size(self): + """ int: The currently selected thumbnail size in pixels """ + scaling = get_config().scaling_factor + size = self._sizes[self._globals.tk_faces_size.get().lower().replace(" ", "")] + return int(round(size * scaling)) + + @property + def viewport(self): + """ :class:`~tools.manual.faceviewer.viewport.Viewport`: The viewport area of the + faces viewer. """ + return self._view + + @property + def grid(self): + """ :class:`Grid`: The grid for the current :class:`FacesViewer`. """ + return self._grid + + @property + def optional_annotations(self): + """ dict: The values currently set for the selectable optional annotations. """ + return {opt: val.get() for opt, val in self._tk_optional_annotations.items()} + + @property + def selected_mask(self): + """ str: The currently selected mask from the display frame control panel. """ + return self._display_frame.tk_selected_mask.get().lower() + + @property + def control_colors(self): + """ :dict: The frame Editor name as key with the current user selected hex code as + value. """ + return ({key: val.get() for key, val in self._display_frame.tk_control_colors.items()}) + + # << CALLBACK FUNCTIONS >> # + def _set_tk_callbacks(self, detected_faces): + """ Set the tkinter variable call backs. + + Redraw the grid on a face size change, a filter change or on add/remove faces. + Updates the annotation colors when user amends a color drop down. + Updates the mask type when the user changes the selected mask types + Toggles the face viewer annotations on an optional annotation button press. + """ + for var in (self._globals.tk_faces_size, self._globals.tk_filter_mode): + var.trace("w", lambda *e, v=var: self.refresh_grid(v)) + var = detected_faces.tk_face_count_changed + var.trace("w", lambda *e, v=var: self.refresh_grid(v, retain_position=True)) + + self._display_frame.tk_control_colors["Mesh"].trace( + "w", lambda *e: self._update_mesh_color()) + self._display_frame.tk_control_colors["ExtractBox"].trace( + "w", lambda *e: self._update_box_color()) + self._display_frame.tk_selected_mask.trace("w", lambda *e: self._update_mask_type()) + + for opt, var in self._tk_optional_annotations.items(): + var.trace("w", lambda *e, o=opt: self._toggle_annotations(o)) + + self.bind("", lambda *e: self._view.update()) + + def refresh_grid(self, trigger_var, retain_position=False): + """ Recalculate the full grid and redraw. Used when the active filter pull down is used, a + face has been added or removed, or the face thumbnail size has changed. + + Parameters + ---------- + trigger_var: :class:`tkinter.BooleanVar` + The tkinter variable that has triggered the grid update. Will either be the variable + indicating that the face size have been changed, or the variable indicating that the + selected filter mode has been changed. + retain_position: bool, optional + ``True`` if the grid should be set back to the position it was at after the update has + been processed, otherwise ``False``. Default: ``False``. + """ + if not trigger_var.get(): + return + size_change = isinstance(trigger_var, tk.StringVar) + move_to = self.yview()[0] if retain_position else 0.0 + self._grid.update() + if move_to != 0.0: + self.yview_moveto(move_to) + if size_change: + self._view.reset() + self._view.update() + if not size_change: + trigger_var.set(False) + + def _update_mask_type(self): + """ Update the displayed mask in the :class:`FacesViewer` canvas when the user changes + the mask type. """ + state = "normal" if self.optional_annotations["mask"] else "hidden" + logger.debug("Updating mask type: (mask_type: %s. state: %s)", self.selected_mask, state) + self._view.toggle_mask(state, self.selected_mask) + + # << MOUSE HANDLING >> + def _bind_mouse_wheel_scrolling(self): + """ Bind mouse wheel to scroll the :class:`FacesViewer` canvas. """ + if platform.system() == "Linux": + self.bind("", self._scroll) + self.bind("", self._scroll) + else: + self.bind("", self._scroll) + + def _scroll(self, event): + """ Handle mouse wheel scrolling over the :class:`FacesViewer` canvas. + + Update is run in a thread to avoid repeated scroll actions stacking and locking up the GUI. + + Parameters + ---------- + event: :class:`tkinter.Event` + The event fired by the mouse scrolling + """ + if self._event.is_set(): + logger.trace("Update already running. Aborting repeated mousewheel") + return + if platform.system() == "Darwin": + adjust = event.delta + elif platform.system() == "Windows": + adjust = event.delta / 120 + elif event.num == 5: + adjust = -1 + else: + adjust = 1 + self._event.set() + thread = Thread(target=self.canvas_scroll, args=(-1 * adjust, "units", self._event)) + thread.start() + + def canvas_scroll(self, amount, units, event): + """ Scroll the canvas on an up/down or page-up/page-down key press. + + Parameters + ---------- + amount: int + The number of units to scroll the canvas + units: ["page", "units"] + The unit type to scroll by + event: :class:`threading.Event` + event to indicate to the calling process whether the scroll is still updating + """ + self.yview_scroll(int(amount), units) + self._view.update() + self._view.hover_box.on_hover(None) + event.clear() + + # << OPTIONAL ANNOTATION METHODS >> # + def _update_mesh_color(self): + """ Update the mesh color when user updates the control panel. """ + color = self.get_muted_color("Mesh") + if self._annotation_colors["mesh"] == color: + return + highlight_color = self.control_colors["Mesh"] + + self.itemconfig("viewport_polygon", outline=color) + self.itemconfig("viewport_line", fill=color) + self.itemconfig("active_mesh_polygon", outline=highlight_color) + self.itemconfig("active_mesh_line", fill=highlight_color) + self._annotation_colors["mesh"] = color + + def _update_box_color(self): + """ Update the active box color when user updates the control panel. """ + color = self.control_colors["ExtractBox"] + + if self._annotation_colors["box"] == color: + return + self.itemconfig("active_highlighter", outline=color) + self._annotation_colors["box"] = color + + def get_muted_color(self, color_key): + """ Creates a muted version of the given annotation color for non-active faces. + + Parameters + ---------- + color_key: str + The annotation key to obtain the color for from :attr:`control_colors` + """ + scale = 0.65 + hls = np.array(colorsys.rgb_to_hls(*hex_to_rgb(self.control_colors[color_key]))) + scale = (1 - scale) + 1 if hls[1] < 120 else scale + hls[1] = max(0., min(256., scale * hls[1])) + rgb = np.clip(np.rint(colorsys.hls_to_rgb(*hls)).astype("uint8"), 0, 255) + retval = rgb_to_hex(rgb) + return retval + + def _toggle_annotations(self, annotation): + """ Toggle optional annotations on or off after the user depresses an optional button. + + Parameters + ---------- + annotation: ["mesh", "mask"] + The optional annotation to toggle on or off + """ + state = "normal" if self.optional_annotations[annotation] else "hidden" + logger.debug("Toggle annotation: (annotation: %s, state: %s)", annotation, state) + if annotation == "mesh": + self._view.toggle_mesh(state) + if annotation == "mask": + self._view.toggle_mask(state, self.selected_mask) + + +class Grid(): + """ Holds information on the current filtered grid layout. + + The grid keeps information on frame indices, face indices, x and y positions and detected face + objects laid out in a numpy array to reflect the current full layout of faces within the face + viewer based on the currently selected filter and face thumbnail size. + + Parameters + ---------- + canvas: :class:`tkinter.Canvas` + The :class:`~tools.manual.faceviewer.frame.FacesViewer` canvas + detected_faces: :class:`~tools.manual.detected_faces.DetectedFaces` + The :class:`~lib.faces_detect.DetectedFace` objects for this video + """ + def __init__(self, canvas, detected_faces): + logger.debug("Initializing %s: (detected_faces: %s)", + self.__class__.__name__, detected_faces) + self._canvas = canvas + self._detected_faces = detected_faces + self._raw_indices = detected_faces.filter.raw_indices + self._frames_list = detected_faces.filter.frames_list + + self._is_valid = False + self._face_size = None + self._grid = None + self._display_faces = None + + self._canvas.update_idletasks() + self._canvas.create_rectangle(0, 0, 0, 0, tags=["backdrop"]) + self.update() + logger.debug("Initialized %s", self.__class__.__name__) + + @property + def face_size(self): + """ int: The pixel size of each thumbnail within the face viewer. """ + return self._face_size + + @property + def is_valid(self): + """ bool: ``True`` if the current filter means that the grid holds faces. ``False`` if + there are no faces displayed in the grid. """ + return self._is_valid + + @property + def columns_rows(self): + """ tuple: the (`columns`, `rows`) required to hold all display images. """ + retval = tuple(reversed(self._grid.shape[1:])) if self._is_valid else (0, 0) + return retval + + @property + def dimensions(self): + """ tuple: The (`width`, `height`) required to hold all display images. """ + if self._is_valid: + retval = tuple(dim * self._face_size for dim in reversed(self._grid.shape[1:])) + else: + retval = (0, 0) + return retval + + @property + def _visible_row_indices(self): + """tuple: A 1 dimensional array of the (`top_row_index`, `bottom_row_index`) of the grid + currently in the viewable area. + """ + height = self.dimensions[1] + visible = (max(0, floor(height * self._canvas.yview()[0]) - self._face_size), + ceil(height * self._canvas.yview()[1])) + logger.trace("visible: %s", visible) + y_points = self._grid[3, :, 1] + top = np.searchsorted(y_points, visible[0], side="left") + bottom = np.searchsorted(y_points, visible[1], side="right") + return top, bottom + + @property + def visible_area(self): + """:class:`numpy.ndarray`: A numpy array of shape (`4`, `rows`, `columns`) corresponding + to the viewable area of the display grid. 1st dimension contains frame indices, 2nd + dimension face indices. The 3rd and 4th dimension contain the x and y position of the top + left corner of the face respectively. + + Any locations that are not populated by a face will have a frame and face index of -1 + """ + if not self._is_valid: + retval = None, None + else: + top, bottom = self._visible_row_indices + retval = self._grid[:, top:bottom, :], self._display_faces[top:bottom, :] + logger.trace([r if r is None else r.shape for r in retval]) + return retval + + def y_coord_from_frame(self, frame_index): + """ Return the y coordinate for the first face that appears in the given frame. + + Parameters + ---------- + frame_index: int + The frame index to locate in the grid + + Returns + ------- + int + The y coordinate of the first face for the given frame + """ + return min(self._grid[3][np.where(self._grid[0] == frame_index)]) + + def frame_has_faces(self, frame_index): + """ Check whether the given frame index contains any faces. + + Parameters + ---------- + frame_index: int + The frame index to locate in the grid + + Returns + ------- + bool + ``True`` if there are faces in the given frame otherwise ``False`` + """ + return self._is_valid and np.any(self._grid[0] == frame_index) + + def update(self): + """ Update the underlying grid. + + Called on initialization, on a filter change or on add/remove faces. Recalculates the + underlying grid for the current filter view and updates the attributes :attr:`_grid`, + :attr:`_display_faces`, :attr:`_raw_indices`, :attr:`_frames_list` and :attr:`is_valid` + """ + self._face_size = self._canvas.face_size + self._raw_indices = self._detected_faces.filter.raw_indices + self._frames_list = self._detected_faces.filter.frames_list + self._get_grid() + self._get_display_faces() + self._canvas.coords("backdrop", 0, 0, *self.dimensions) + self._canvas.configure(scrollregion=(self._canvas.bbox("backdrop"))) + self._canvas.yview_moveto(0.0) + + def _get_grid(self): + """ Get the grid information for faces currently displayed in the :class:`FacesViewer`. + + Returns + :class:`numpy.ndarray` + A numpy array of shape (`4`, `rows`, `columns`) corresponding to the display grid. + 1st dimension contains frame indices, 2nd dimension face indices. The 3rd and 4th + dimension contain the x and y position of the top left corner of the face respectively. + + Any locations that are not populated by a face will have a frame and face index of -1 + """ + labels = self._get_labels() + if not self._is_valid: + logger.debug("Setting grid to None for no faces.") + self._grid = None + return + x_coords = np.linspace(0, + labels.shape[2] * self._face_size, + num=labels.shape[2], + endpoint=False, + dtype="int") + y_coords = np.linspace(0, + labels.shape[1] * self._face_size, + num=labels.shape[1], + endpoint=False, + dtype="int") + self._grid = np.array((*labels, *np.meshgrid(x_coords, y_coords)), dtype="int") + logger.debug(self._grid.shape) + + def _get_labels(self): + """ Get the frame and face index for each grid position for the current filter. + + Returns + ------- + :class:`numpy.ndarray` + Array of dimensions (2, rows, columns) corresponding to the display grid, with frame + index as the first dimension and face index within the frame as the 2nd dimension. + + Any remaining placeholders at the end of the grid which are not populated with a face + are given the index -1 + """ + face_count = len(self._raw_indices["frame"]) + self._is_valid = face_count != 0 + if not self._is_valid: + return None + columns = self._canvas.winfo_width() // self._face_size + rows = ceil(face_count / columns) + remainder = face_count % columns + padding = [] if remainder == 0 else [-1 for _ in range(columns - remainder)] + labels = np.array((self._raw_indices["frame"] + padding, + self._raw_indices["face"] + padding), + dtype="int").reshape((2, rows, columns)) + logger.debug(labels.shape) + return labels + + def _get_display_faces(self): + """ Get the detected faces for the current filter and arrange to grid. + + Returns + ------- + :class:`numpy.ndarray` + Array of dimensions (rows, columns) corresponding to the display grid, containing the + corresponding :class:`lib.faces_detect.DetectFace` object + + Any remaining placeholders at the end of the grid which are not populated with a face + are replaced with ``None`` + """ + if not self._is_valid: + logger.debug("Setting display_faces to None for no faces.") + self._display_faces = None + return + current_faces = self._detected_faces.current_faces + columns, rows = self.columns_rows + face_count = len(self._raw_indices["frame"]) + padding = [None for _ in range(face_count, columns * rows)] + self._display_faces = np.array([None if idx is None else current_faces[idx][face_idx] + for idx, face_idx + in zip(self._raw_indices["frame"] + padding, + self._raw_indices["face"] + padding)], + dtype="object").reshape(rows, columns) + logger.debug("faces: (shape: %s, dtype: %s)", + self._display_faces.shape, self._display_faces.dtype) + + def transport_index_from_frame(self, frame_index): + """ Return the main frame's transport index for the given frame index based on the current + filter criteria. + + Parameters + ---------- + frame_index: int + The absolute index for the frame within the full frames list + + Returns + ------- + int + The index of the requested frame within the filtered frames view. + """ + retval = self._frames_list.index(frame_index) if frame_index in self._frames_list else None + logger.trace("frame_index: %s, transport_index: %s", frame_index, retval) + return retval + + +class ContextMenu(): # pylint:disable=too-few-public-methods + """ Enables a right click context menu for the + :class:`~tools.manual.faceviewer.frame.FacesViewer`. + + Parameters + ---------- + canvas: :class:`tkinter.Canvas` + The :class:`FacesViewer` canvas + detected_faces: :class:`~tools.manual.detected_faces` + The manual tool's detected faces class + """ + def __init__(self, canvas, detected_faces): + logger.debug("Initializing: %s (canvas: %s, detected_faces: %s)", + self.__class__.__name__, canvas, detected_faces) + self._canvas = canvas + self._detected_faces = detected_faces + self._menu = RightClickMenu(["Delete Face"], [self._delete_face]) + self._frame_index = None + self._face_index = None + self._canvas.bind("" if platform.system() == "Darwin" else "", + self._pop_menu) + logger.debug("Initialized: %s", self.__class__.__name__) + + def _pop_menu(self, event): + """ Pop up the context menu on a right click mouse event. + + Parameters + ---------- + event: :class:`tkinter.Event` + The mouse event that has triggered the pop up menu + """ + frame_idx, face_idx = self._canvas.viewport.face_from_point( + self._canvas.canvasx(event.x), self._canvas.canvasy(event.y))[:2] + if frame_idx == -1: + logger.trace("No valid item under mouse") + self._frame_index = self._face_index = None + return + self._frame_index = frame_idx + self._face_index = face_idx + logger.trace("Popping right click menu") + self._menu.popup(event) + + def _delete_face(self): + """ Delete the selected face on a right click mouse delete action. """ + logger.trace("Right click delete received. frame_id: %s, face_id: %s", + self._frame_index, self._face_index) + self._detected_faces.update.delete(self._frame_index, self._face_index) + self._frame_index = self._face_index = None diff --git a/tools/manual/faceviewer/viewport.py b/tools/manual/faceviewer/viewport.py new file mode 100644 index 00000000..a467f967 --- /dev/null +++ b/tools/manual/faceviewer/viewport.py @@ -0,0 +1,1011 @@ +#!/usr/bin/env python3 +""" Handles the visible area of the :class:`~tools.manual.faceviewer.frame.FacesViewer` canvas. """ + +import logging +import tkinter as tk + +import cv2 +import numpy as np +from PIL import Image, ImageTk + + +logger = logging.getLogger(__name__) # pylint: disable=invalid-name + + +class Viewport(): + """ Handles the display of faces and annotations in the currently viewable area of the canvas. + + Parameters + ---------- + canvas: :class:`tkinter.Canvas` + The :class:`~tools.manual.faceviewer.frame.FacesViewer` canvas + tk_edited_variable: :class:`tkinter.BooleanVar` + The variable that indicates that a face has been edited + """ + def __init__(self, canvas, tk_edited_variable): + logger.debug("Initializing: %s: (canvas: %s, tk_edited_variable: %s)", + self.__class__.__name__, canvas, tk_edited_variable) + self._canvas = canvas + self._grid = canvas.grid + self._tk_selected_editor = canvas._display_frame.tk_selected_action + self._landmark_mapping = dict(mouth_inner=(60, 68), + mouth_outer=(48, 60), + right_eyebrow=(17, 22), + left_eyebrow=(22, 27), + right_eye=(36, 42), + left_eye=(42, 48), + nose=(27, 36), + jaw=(0, 17), + chin=(8, 11)) + self._landmarks = dict() + self._tk_faces = dict() + self._objects = VisibleObjects(self) + self._hoverbox = HoverBox(self) + self._active_frame = ActiveFrame(self, tk_edited_variable) + self._tk_selected_editor.trace( + "w", lambda *e: self._active_frame.reload_annotations()) + + @property + def face_size(self): + """ int: The pixel size of each thumbnail """ + return self._grid.face_size + + @property + def mesh_kwargs(self): + """ dict: The color and state keyword arguments for the objects that make up a single + face's mesh annotation based on the current user selected options. Key is the object + type (`polygon` or `line`), value are the keyword arguments for that type. """ + state = "normal" if self._canvas.optional_annotations["mesh"] else "hidden" + color = self._canvas.get_muted_color("Mesh") + kwargs = dict(polygon=dict(fill="", outline=color, state=state), + line=dict(fill=color, state=state)) + return kwargs + + @property + def hover_box(self): + """ :class:`HoverBox`: The hover box for the viewport. """ + return self._hoverbox + + @property + def selected_editor(self): + """ str: The currently selected editor. """ + return self._tk_selected_editor.get().lower() + + def toggle_mesh(self, state): + """ Toggles the mesh optional annotations on and off. + + Parameters + ---------- + state: ["hidden", "normal"] + The state to set the mesh annotations to + """ + logger.debug("Toggling mesh annotations to: %s", state) + self._canvas.itemconfig("viewport_mesh", state=state) + self.update() + + def toggle_mask(self, state, mask_type): + """ Toggles the mask optional annotation on and off. + + Parameters + ---------- + state: ["hidden", "normal"] + Whether the mask should be displayed or hidden + mask_type: str + The type of mask to overlay onto the face + """ + logger.debug("Toggling mask annotations to: %s. mask_type: %s", state, mask_type) + for (frame_idx, face_idx), det_faces in zip( + self._objects.visible_grid[:2].transpose(1, 2, 0).reshape(-1, 2), + self._objects.visible_faces.flatten()): + if frame_idx == -1: + continue + key = "_".join([str(frame_idx), str(face_idx)]) + mask = None if state == "hidden" else det_faces.mask.get(mask_type, None) + mask = mask if mask is None else mask.mask.squeeze() + self._tk_faces[key].update_mask(mask) + self.update() + + def reset(self): + """ Reset all the cached objects on a face size change. """ + self._landmarks = dict() + self._tk_faces = dict() + + def update(self): + """ Update the viewport. + + Obtains the objects that are currently visible. Updates the visible area of the canvas + and reloads the active frame's annotations. """ + self._objects.update() + self._update_viewport() + self._active_frame.reload_annotations() + + def _update_viewport(self): + """ Update the viewport + + Clear out cached objects that are not currently in view. Populate the cache for any + faces that are now in view. Populate the correct face image and annotations for each + object in the viewport based on current location. If optional mesh annotations are + enabled, then calculates newly displayed meshes. """ + if not self._grid.is_valid: + return + self._discard_tk_faces() + + if self._canvas.optional_annotations["mesh"]: # Display any hidden end of row meshes + self._canvas.itemconfig("viewport_mesh", state="normal") + + for collection in zip(self._objects.visible_grid.transpose(1, 2, 0), + self._objects.images, + self._objects.meshes, + self._objects.visible_faces): + for (frame_idx, face_idx, pnt_x, pnt_y), image_id, mesh_ids, face in zip(*collection): + top_left = np.array((pnt_x, pnt_y)) + if frame_idx == self._active_frame.frame_index: + logger.trace("Skipping active frame: %s", frame_idx) + continue + if frame_idx == -1: + logger.debug("Blanking non-existant face") + self._canvas.itemconfig(image_id, image="") + for area in mesh_ids.values(): + for mesh_id in area: + self._canvas.itemconfig(mesh_id, state="hidden") + continue + + tk_face = self.get_tk_face(frame_idx, face_idx, face) + self._canvas.itemconfig(image_id, image=tk_face.photo) + if (self._canvas.optional_annotations["mesh"] + or frame_idx == self._active_frame.frame_index): + landmarks = self.get_landmarks(frame_idx, face_idx, face, top_left) + self._locate_mesh(mesh_ids, landmarks) + + def _discard_tk_faces(self): + """ Remove any :class:`TKFace` objects from the cache that are not currently displayed. """ + keys = ["{}_{}".format(pnt_x, pnt_y) + for pnt_x, pnt_y in self._objects.visible_grid[:2].T.reshape(-1, 2)] + for key in list(self._tk_faces): + if key not in keys: + del self._tk_faces[key] + logger.trace("keys: %s allocated_faces: %s", keys, len(self._tk_faces)) + + def get_tk_face(self, frame_index, face_index, face): + """ Obtain the :class:`TKFace` object for the given face from the cache. If the face does + not exist in the cache, then it is generated and added prior to returning. + + Parameters + ---------- + frame_index: int + The frame index to obtain the face for + face_index: int + The face index of the face within the requested frame + face: :class:`~lib.faces_detect.DetectedFace` + The detected face object, containing the thumbnail jpg + + Returns + ------- + :class:`TKFace` + An object for displaying in the faces viewer canvas populated with the aligned mesh + landmarks and face thumbnail + """ + is_active = frame_index == self._active_frame.frame_index + key = "_".join([str(frame_index), str(face_index)]) + if key not in self._tk_faces or is_active: + logger.trace("creating new tk_face: (key: %s, is_active: %s)", key, is_active) + if is_active: + face.load_aligned(self._active_frame.current_frame, + size=self.face_size, + force=True) + image = face.aligned_face + face.aligned = dict() + else: + image = face.thumbnail + tk_face = self._get_tk_face_object(face, image, is_active) + self._tk_faces[key] = tk_face + else: + logger.trace("tk_face exists: %s", key) + tk_face = self._tk_faces[key] + return tk_face + + def _get_tk_face_object(self, face, image, is_active): + """ Obtain an existing unallocated, or a newly created :class:`TKFace` and populate it with + face information from the requested frame and face index. + + If the face is currently active, then the face is generated from the currently displayed + frame, otherwise it is generated from the jpg thumbnail. + + Parameters + ---------- + face: :class:`lib.faces_detect.DetectedFace` + A detected face object to create the :class:`TKFace` from + image: :class:`numpy.ndarray` + The jpg thumbnail or the 3 channel image for the face + is_active: bool + ``True`` if the face in the currently active frame otherwise ``False`` + + Returns + ------- + :class:`TKFace` + An object for displaying in the faces viewer canvas populated with the aligned face + image with a mask applied, if required. + """ + get_mask = (self._canvas.optional_annotations["mask"] or + (is_active and self.selected_editor == "mask")) + mask = face.mask.get(self._canvas.selected_mask, None) if get_mask else None + mask = mask if mask is None else mask.mask.squeeze() + tk_face = TKFace(image, size=self.face_size, mask=mask) + logger.trace("face: %s, tk_face: %s", face, tk_face) + return tk_face + + def get_landmarks(self, frame_index, face_index, face, top_left, refresh=False): + """ Obtain the landmark points for each mesh annotation. + + First tries to obtain the aligned landmarks from the cache. If the landmarks do not exist + in the cache, or a refresh has been requested, then the landmarks are calculated from the + detected face object. + + Parameters + ---------- + frame_index: int + The frame index to obtain the face for + face_index: int + The face index of the face within the requested frame + top_left: tuple + The top left (x, y) points of the face's bounding box within the viewport + refresh: bool, optional + Whether to force a reload of the face's aligned landmarks, even if they already exist + within the cache. Default: ``False`` + + Returns + ------- + dict + The key is the tkinter canvas object type for each part of the mesh annotation + (`polygon`, `line`). The value is a list containing the (x, y) coordinates of each + part of the mesh annotation, from the top left corner location. + """ + key = "{}_{}".format(frame_index, face_index) + landmarks = self._landmarks.get(key, None) + if not landmarks or refresh: + face.load_aligned(None, size=self.face_size, force=True) + landmarks = dict(polygon=[], line=[]) + for area, val in self._landmark_mapping.items(): + points = face.aligned_landmarks[val[0]:val[1]] + top_left + shape = "polygon" if area.endswith("eye") or area.startswith("mouth") else "line" + landmarks[shape].append(points) + self._landmarks[key] = landmarks + return landmarks + + def _locate_mesh(self, mesh_ids, landmarks): + """ Place the mesh annotation canvas objects in the correct location. + + Parameters + ---------- + mesh_ids: list + The list of mesh id objects to set coordinates for + landmarks: dict + The mesh point groupings and whether each group should be a line or a polygon + """ + for key, area in landmarks.items(): + for coords, mesh_id in zip(area, mesh_ids[key]): + self._canvas.coords(mesh_id, *coords.flatten()) + + def face_from_point(self, point_x, point_y): + """ Given an (x, y) point on the :class:`Viewport`, obtain the face information at that + location. + + Parameters + ---------- + point_x: int + The x position on the canvas of the point to retrieve the face for + point_y: int + The y position on the canvas of the point to retrieve the face for + + Returns + ------- + :class:`numpy.ndarray` + Array of shape (4, ) containing the (`frame index`, `face index`, `x_point of top left + corner`, `y point of top left corner`) of the face at the given coordinates. + + If the given coordinates are not over a face, then the frame and face indices will be + -1 + """ + if point_x > self._grid.dimensions[0]: + retval = np.array((-1, -1, -1, -1)) + else: + x_idx = np.searchsorted(self._objects.visible_grid[2, 0, :], point_x, side="left") - 1 + y_idx = np.searchsorted(self._objects.visible_grid[3, :, 0], point_y, side="left") - 1 + if x_idx < 0 or y_idx < 0: + retval = np.array((-1, -1, -1, -1)) + else: + retval = self._objects.visible_grid[:, y_idx, x_idx] + logger.trace(retval) + return retval + + def move_active_to_top(self): + """ Check whether the active frame is going off the bottom of the viewport, if so: move it + to the top of the viewport. """ + self._active_frame.move_to_top() + + +class VisibleObjects(): + """ Holds the objects from the :class:`~tools.manual.faceviewer.frame.Grid` that appear in the + viewable area of the :class:`Viewport`. + + Parameters + ---------- + viewport: :class:`Viewport` + The viewport object for the :class:`~tools.manual.faceviewer.frame.FacesViewer` canvas + """ + def __init__(self, viewport): + self._viewport = viewport + self._canvas = viewport._canvas + self._grid = viewport._grid + self._size = viewport.face_size + + self._visible_grid = None + self._visible_faces = None + self._images = [] + self._meshes = [] + self._recycled = dict(images=[], meshes=[]) + + @property + def visible_grid(self): + """ :class:`numpy.ndarray`: The currently visible section of the + :class:`~tools.manual.faceviewer.frame.Grid` + + A numpy array of shape (`4`, `rows`, `columns`) corresponding to the viewable area of the + display grid. 1st dimension contains frame indices, 2nd dimension face indices. The 3rd and + 4th dimension contain the x and y position of the top left corner of the face respectively. + + Any locations that are not populated by a face will have a frame and face index of -1. """ + return self._visible_grid + + @property + def visible_faces(self): + """ :class:`numpy.ndarray`: The currently visible :class:`~lib.faces_detect.DetectedFace` + objects. + + A numpy array of shape (`rows`, `columns`) corresponding to the viewable area of the + display grid and containing the detected faces at their currently viewable position. + + Any locations that are not populated by a face will have ``None`` in it's place. """ + return self._visible_faces + + @property + def images(self): + """ :class:`numpy.ndarray`: The viewport's tkinter canvas image objects. + + A numpy array of shape (`rows`, `columns`) corresponding to the viewable area of the + display grid and containing the tkinter canvas image object for the face at the + corresponding location. """ + return self._images + + @property + def meshes(self): + """ :class:`numpy.ndarray`: The viewport's tkinter canvas mesh annotation objects. + + A numpy array of shape (`rows`, `columns`) corresponding to the viewable area of the + display grid and containing a dictionary of the corresponding tkinter polygon and line + objects required to build a face's mesh annotation for the face at the corresponding + location. """ + return self._meshes + + @property + def _top_left(self): + """ :class:`numpy.ndarray`: The canvas (`x`, `y`) position of the face currently in the + viewable area's top left position. """ + return np.array(self._canvas.coords(self._images[0][0]), dtype="int") + + def update(self): + """ Load and unload thumbnails in the visible area of the faces viewer. """ + self._visible_grid, self._visible_faces = self._grid.visible_area + if (isinstance(self._images, np.ndarray) and + self._visible_grid.shape[-1] != self._images.shape[-1]): + self._recycle_objects() + + required_rows = self._visible_grid.shape[1] if self._grid.is_valid else 0 + existing_rows = len(self._images) + logger.trace("existing_rows: %s. required_rows: %s", existing_rows, required_rows) + + if existing_rows > required_rows: + for image_id in self._images[required_rows: existing_rows].flatten(): + logger.trace("Hiding image id: %s", image_id) + self._canvas.itemconfig(image_id, image="") + + if existing_rows < required_rows: + self._add_rows(existing_rows, required_rows) + + self._shift() + + def _recycle_objects(self): + """ On a column count change, place all existing objects into the recycle bin so that + they can be used for the new grid shape and reset the objects size to the new size. """ + self._size = self._viewport.face_size + images = self._images.flatten().tolist() + meshes = self._meshes.flatten().tolist() + + for image_id in images: + self._canvas.itemconfig(image_id, image="") + self._canvas.coords(image_id, 0, 0) + for mesh in meshes: + for key, mesh_ids in mesh.items(): + coords = (0, 0, 0, 0) if key == "line" else (0, 0) + for mesh_id in mesh_ids: + self._canvas.coords(mesh_id, *coords) + + self._recycled["images"].extend(images) + self._recycled["meshes"].extend(meshes) + logger.trace("Recycled objects: %s", self._recycled) + + self._images = [] + self._meshes = [] + + def _add_rows(self, existing_rows, required_rows): + """ Add rows to the viewport. + + Parameters + ---------- + existing_rows: int + The number of existing rows within the viewport + required_rows: int + The number of rows required by the viewport + """ + columns = self._grid.columns_rows[0] + if not isinstance(self._images, np.ndarray): + base_coords = [(col * self._size, 0) for col in range(columns)] + else: + base_coords = [self._canvas.coords(item_id) for item_id in self._images[0]] + logger.debug("existing rows: %s, required_rows: %s, base_coords: %s", + existing_rows, required_rows, base_coords) + images = [] + meshes = [] + for row in range(existing_rows, required_rows): + y_coord = base_coords[0][1] + (row * self._size) + images.append(np.array([self._get_image((coords[0], y_coord)) + for coords in base_coords])) + meshes.append(np.array([self._get_mesh() for _ in range(columns)])) + images = np.array(images) + meshes = np.array(meshes) + + if not isinstance(self._images, np.ndarray): + logger.debug("Adding initial viewport objects: (image shapes: %s, mesh shapes: %s)", + images.shape, meshes.shape) + self._images = images + self._meshes = meshes + else: + logger.debug("Adding new viewport objects: (image shapes: %s, mesh shapes: %s)", + images.shape, meshes.shape) + self._images = np.concatenate((self._images, images)) + self._meshes = np.concatenate((self._meshes, meshes)) + logger.debug("self._images: %s, self._meshes: %s", self._images.shape, self._meshes.shape) + + def _get_image(self, coordinates): + """ Create or recycle a tkinter canvas image object with the given coordinates. + + Parameters + ---------- + coordinates: tuple + The (`x`, `y`) coordinates for the top left corner of the image + + Returns + ------- + int + The canvas object id for the created image + """ + if self._recycled["images"]: + image_id = self._recycled["images"].pop() + self._canvas.coords(image_id, *coordinates) + logger.trace("Recycled image: %s", image_id) + else: + image_id = self._canvas.create_image(*coordinates, + anchor=tk.NW, + tags=["viewport", "viewport_image"]) + logger.trace("Created new image: %s", image_id) + return image_id + + def _get_mesh(self): + """ Get the mesh annotation for the landmarks. This is made up of a series of polygons + or lines, depending on which part of the face is being annotated. Creates a new series of + objects, or pulls existing objects from the recycled objects pool if they are available. + + Returns + ------- + dict + The dictionary of line and polygon tkinter canvas object ids for the mesh annotation + """ + kwargs = self._viewport.mesh_kwargs + logger.trace("self.mesh_kwargs: %s", kwargs) + if self._recycled["meshes"]: + mesh = self._recycled["meshes"].pop() + for key, mesh_ids in mesh.items(): + for mesh_id in mesh_ids: + self._canvas.itemconfig(mesh_id, **kwargs[key]) + logger.trace("Recycled mesh: %s", mesh) + else: + tags = ["viewport", "viewport_mesh"] + mesh = dict(polygon=[self._canvas.create_polygon(0, 0, + width=1, + tags=tags + ["viewport_polygon"], + **kwargs["polygon"]) + for _ in range(4)], + line=[self._canvas.create_line(0, 0, 0, 0, + width=1, + tags=tags + ["viewport_line"], + **kwargs["line"]) + for _ in range(5)]) + logger.trace("Created new mesh: %s", mesh) + return mesh + + def _shift(self): + """ Shift the viewport in the y direction if required + + Returns + ------- + bool + ``True`` if the viewport was shifted otherwise ``False`` + """ + current_y = self._top_left[1] + required_y = self._visible_grid[3, 0, 0] if self._grid.is_valid else 0 + logger.trace("current_y: %s, required_y: %s", current_y, required_y) + if current_y == required_y: + logger.trace("No move required") + return False + shift_amount = required_y - current_y + logger.trace("Shifting viewport: %s", shift_amount) + self._canvas.move("viewport", 0, shift_amount) + return True + + +class HoverBox(): # pylint:disable=too-few-public-methods + """ Handle the current mouse location when over the :class:`Viewport`. + + Highlights the face currently underneath the cursor and handles actions when clicking + on a face. + + Parameters + ---------- + viewport: :class:`Viewport` + The viewport object for the :class:`~tools.manual.faceviewer.frame.FacesViewer` canvas + """ + def __init__(self, viewport): + logger.debug("Initializing: %s (viewport: %s)", self.__class__.__name__, viewport) + self._viewport = viewport + self._canvas = viewport._canvas + self._grid = viewport._canvas.grid + self._globals = viewport._canvas._globals + self._navigation = viewport._canvas._display_frame.navigation + self._box = self._canvas.create_rectangle(0, 0, self._size, self._size, + outline="#0000ff", + width=2, + state="hidden", + fill="#0000ff", + stipple="gray12", + tags="hover_box") + self._current_frame_index = None + self._current_face_index = None + self._canvas.bind("", lambda e: self._clear()) + self._canvas.bind("", self.on_hover) + self._canvas.bind("", lambda e: self._select_frame()) + logger.debug("Initialized: %s", self.__class__.__name__) + + @property + def _size(self): + """ int: the currently set viewport face size in pixels. """ + return self._viewport.face_size + + def on_hover(self, event): + """ Highlight the face and set the mouse cursor for the mouse's current location. + + Parameters + ---------- + event: :class:`tkinter.Event` or ``None`` + The tkinter mouse event. Provides the current location of the mouse cursor. If ``None`` + is passed as the event (for example when this function is being called outside of a + mouse event) then the location of the cursor will be calculated + """ + if event is None: + pnts = np.array((self._canvas.winfo_pointerx(), self._canvas.winfo_pointery())) + pnts -= np.array((self._canvas.winfo_rootx(), self._canvas.winfo_rooty())) + else: + pnts = (event.x, event.y) + + coords = (int(self._canvas.canvasx(pnts[0])), int(self._canvas.canvasy(pnts[1]))) + face = self._viewport.face_from_point(*coords) + frame_idx, face_idx = face[:2] + is_zoomed = self._globals.is_zoomed + + if (-1 in face or (frame_idx == self._globals.frame_index + and (not is_zoomed or + (is_zoomed and face_idx == self._globals.tk_face_index.get())))): + self._clear() + self._canvas.config(cursor="") + self._current_frame_index = None + self._current_face_index = None + return + + self._canvas.config(cursor="hand2") + self._highlight(face[2:]) + self._current_frame_index = frame_idx + self._current_face_index = face_idx + + def _clear(self): + """ Hide the hover box when the mouse is not over a face. """ + if self._canvas.itemcget(self._box, "state") != "hidden": + self._canvas.itemconfig(self._box, state="hidden") + + def _highlight(self, top_left): + """ Display the hover box around the face that the mouse is currently over. + + Parameters + ---------- + top_left: tuple + The top left point of the highlight box location + """ + coords = (*top_left, *top_left + self._size) + self._canvas.coords(self._box, *coords) + self._canvas.itemconfig(self._box, state="normal") + self._canvas.tag_raise(self._box) + + def _select_frame(self): + """ Select the face and the subsequent frame (in the editor view) when a face is clicked + on in the :class:`Viewport`. + """ + frame_id = self._current_frame_index + is_zoomed = self._globals.is_zoomed + if frame_id is None or (frame_id == self._globals.frame_index and not is_zoomed): + return + face_idx = self._current_face_index if is_zoomed else 0 + self._globals.tk_face_index.set(face_idx) + transport_id = self._grid.transport_index_from_frame(frame_id) + logger.trace("frame_index: %s, transport_id: %s, face_idx: %s", + frame_id, transport_id, face_idx) + if transport_id is None: + return + self._navigation.stop_playback() + self._globals.tk_transport_index.set(transport_id) + self._viewport.move_active_to_top() + self.on_hover(None) + + +class ActiveFrame(): + """ Handles the display of faces and annotations for the currently active frame. + + Parameters + ---------- + canvas: :class:`tkinter.Canvas` + The :class:`~tools.manual.faceviewer.frame.FacesViewer` canvas + tk_edited_variable: :class:`tkinter.BooleanVar` + The tkinter callback variable indicating that a face has been edited + """ + def __init__(self, viewport, tk_edited_variable): + logger.debug("Initializing: %s (viewport: %s, tk_edited_variable: %s)", + self.__class__.__name__, viewport, tk_edited_variable) + self._objects = viewport._objects + self._viewport = viewport + self._grid = viewport._grid + self._tk_faces = viewport._tk_faces + self._canvas = viewport._canvas + self._globals = viewport._canvas._globals + self._navigation = viewport._canvas._display_frame.navigation + self._last_execution = dict(frame_index=-1, size=viewport.face_size) + self._tk_vars = dict(selected_editor=self._canvas._display_frame.tk_selected_action, + edited=tk_edited_variable) + self._assets = dict(images=[], meshes=[], faces=[], boxes=[]) + + self._globals.tk_update_active_viewport.trace("w", lambda *e: self._reload_callback()) + tk_edited_variable.trace("w", lambda *e: self._update_on_edit()) + logger.debug("Initialized: %s", self.__class__.__name__) + + @property + def frame_index(self): + """ int: The frame index of the currently displayed frame. """ + return self._globals.frame_index + + @property + def current_frame(self): + """ :class:`numpy.ndarray`: A BGR version of the frame currently being displayed. """ + return self._globals.current_frame["image"] + + @property + def _size(self): + """ int: The size of the thumbnails displayed in the viewport, in pixels. """ + return self._viewport.face_size + + @property + def _optional_annotations(self): + """ dict: The currently selected optional annotations """ + return self._canvas.optional_annotations + + def _reload_callback(self): + """ If a frame has changed, triggering the variable, then update the active frame. Return + having done nothing if the variable is resetting. """ + if self._globals.tk_update_active_viewport.get(): + self.reload_annotations() + + def reload_annotations(self): + """ Handles the reloading of annotations for the currently active faces. + + Highlights the faces within the viewport of those faces that exist in the currently + displaying frame. Applies annotations based on the optional annotations and current + editor selections. + """ + logger.trace("Reloading annotations") + if np.any(self._assets["images"]): + self._clear_previous() + + self._set_active_objects() + self._check_active_in_view() + + if not np.any(self._assets["images"]): + logger.trace("No active faces. Returning") + self._last_execution["frame_index"] = self.frame_index + return + + if self._last_execution["frame_index"] != self.frame_index: + self.move_to_top() + self._create_new_boxes() + + self._update_face() + self._canvas.tag_raise("active_highlighter") + self._globals.tk_update_active_viewport.set(False) + self._last_execution["frame_index"] = self.frame_index + + def _clear_previous(self): + """ Reverts the previously selected annotations to their default state. """ + logger.trace("Clearing previous active frame") + self._canvas.itemconfig("active_highlighter", state="hidden") + + for key in ("polygon", "line"): + tag = "active_mesh_{}".format(key) + self._canvas.itemconfig(tag, **self._viewport.mesh_kwargs[key]) + self._canvas.dtag(tag) + + if self._viewport.selected_editor == "mask" and not self._optional_annotations["mask"]: + for key, tk_face in self._tk_faces.items(): + if key.startswith("{}_".format(self._last_execution["frame_index"])): + tk_face.update_mask(None) + + def _set_active_objects(self): + """ Collect the objects that exist in the currently active frame from the main grid. """ + if self._grid.is_valid: + rows, cols = np.where(self._objects.visible_grid[0] == self.frame_index) + logger.trace("Setting active objects: (rows: %s, columns: %s)", rows, cols) + self._assets["images"] = self._objects.images[rows, cols] + self._assets["meshes"] = self._objects.meshes[rows, cols] + self._assets["faces"] = self._objects.visible_faces[rows, cols] + else: + logger.trace("No valid grid. Clearing active objects") + self._assets["images"] = [] + self._assets["meshes"] = [] + self._assets["faces"] = [] + + def _check_active_in_view(self): + """ If the frame has changed, there are faces in the frame, but they don't appear in the + viewport, then bring the active faces to the top of the viewport. """ + if (not np.any(self._assets["images"]) and + self._last_execution["frame_index"] != self.frame_index and + self._grid.frame_has_faces(self.frame_index)): + y_coord = self._grid.y_coord_from_frame(self.frame_index) + logger.trace("Active not in view. Moving to: %s", y_coord) + self._canvas.yview_moveto(y_coord / self._canvas.bbox("backdrop")[3]) + self._viewport.update() + + def move_to_top(self): + """ Move the currently selected frame's faces to the top of the viewport if they are moving + off the bottom of the viewer. """ + height = self._canvas.bbox("backdrop")[3] + bot = int(self._canvas.coords(self._assets["images"][-1])[1] + self._size) + + y_top, y_bot = (int(round(pnt * height)) for pnt in self._canvas.yview()) + + if y_top < bot < y_bot: # bottom face is still in fully visible area + logger.trace("Active faces in frame. Returning") + return + + top = int(self._canvas.coords(self._assets["images"][0])[1]) + if y_top == top: + logger.trace("Top face already on top row. Returning") + return + + if self._canvas.winfo_height() > self._size: + logger.trace("Viewport taller than single face height. Moving Active faces to top: %s", + top) + self._canvas.yview_moveto(top / height) + self._viewport.update() + elif self._canvas.winfo_height() <= self._size and y_top != top: + logger.trace("Viewport shorter than single face height. Moving Active faces to " + "top: %s", top) + self._canvas.yview_moveto(top / height) + self._viewport.update() + + def _create_new_boxes(self): + """ The highlight boxes (border around selected faces) are the only additional annotations + that are required for the highlighter. If more faces are displayed in the current frame + than highlight boxes are available, then new boxes are created to accommodate the + additional faces. """ + new_boxes_count = max(0, len(self._assets["images"]) - len(self._assets["boxes"])) + if new_boxes_count == 0: + return + logger.debug("new_boxes_count: %s", new_boxes_count) + for _ in range(new_boxes_count): + box = self._canvas.create_rectangle(0, + 0, + self._viewport.face_size, self._viewport.face_size, + outline="#00FF00", + width=2, + state="hidden", + tags=["active_highlighter"]) + logger.trace("Created new highlight_box: %s", box) + self._assets["boxes"].append(box) + + def _update_on_edit(self): + """ Update the active faces on a frame edit. """ + if not self._tk_vars["edited"].get(): + return + self._set_active_objects() + self._update_face() + self._tk_vars["edited"].set(False) + + def _update_face(self): + """ Update the highlighted annotations for faces in the currently selected frame. """ + for face_idx, (image_id, mesh_ids, box_id, det_face), in enumerate( + zip(self._assets["images"], + self._assets["meshes"], + self._assets["boxes"], + self._assets["faces"])): + top_left = np.array(self._canvas.coords(image_id)) + coords = (*top_left, *top_left + self._size) + tk_face = self._viewport.get_tk_face(self.frame_index, face_idx, det_face) + self._canvas.itemconfig(image_id, image=tk_face.photo) + self._show_box(box_id, coords) + self._show_mesh(mesh_ids, face_idx, det_face, top_left) + self._last_execution["size"] = self._viewport.face_size + + def _show_box(self, item_id, coordinates): + """ Display the highlight box around the given coordinates. + + Parameters + ---------- + item_id: int + The tkinter canvas object identifier for the highlight box + coordinates: :class:`numpy.ndarray` + The (x, y, x1, y1) coordinates of the top left corner of the box + """ + self._canvas.coords(item_id, *coordinates) + self._canvas.itemconfig(item_id, state="normal") + + def _show_mesh(self, mesh_ids, face_index, detected_face, top_left): + """ Display the mesh annotation for the given face, at the given location. + + Parameters + ---------- + mesh_ids: dict + Dictionary containing the `polygon` and `line` tkinter canvas identifiers that make up + the mesh for the given face + face_index: int + The face index within the frame for the given face + detected_face: :class:`~lib.faces_detect.DetectedFace` + The detected face object that contains the landmarks for generating the mesh + top_left: tuple + The (x, y) top left co-ordinates of the mesh's bounding box + """ + state = "normal" if (self._tk_vars["selected_editor"].get() != "Mask" or + self._optional_annotations["mesh"]) else "hidden" + kwargs = dict(polygon=dict(fill="", outline=self._canvas.control_colors["Mesh"]), + line=dict(fill=self._canvas.control_colors["Mesh"])) + + edited = (self._tk_vars["edited"].get() and + self._tk_vars["selected_editor"].get() not in ("Mask", "View")) + relocate = self._viewport.face_size != self._last_execution["size"] or ( + state == "normal" and not self._optional_annotations["mesh"]) + if relocate or edited: + landmarks = self._viewport.get_landmarks(self.frame_index, + face_index, + detected_face, + top_left, + edited) + for key, kwarg in kwargs.items(): + for idx, mesh_id in enumerate(mesh_ids[key]): + if relocate: + self._canvas.coords(mesh_id, *landmarks[key][idx].flatten()) + self._canvas.itemconfig(mesh_id, state=state, **kwarg) + self._canvas.addtag_withtag("active_mesh_{}".format(key), mesh_id) + + +class TKFace(): + """ An object that holds a single :class:`tkinter.PhotoImage` face, ready for placement in the + :class:`Viewport`, Handles the placement of and removal of masks for the face as well as + updates on any edits. + + Parameters + ---------- + face: :class:`numpy.ndarray` + The face, sized correctly as a 3 channel BGR image or an encoded jpg to create a + :class:`tkinter.PhotoImage` from + size: int, optional + The pixel size of the face image. Default: `128` + mask: :class:`numpy.ndarray` or ``None``, optional + The mask to be applied to the face image. Pass ``None`` if no mask is to be used. + Default ``None`` + """ + def __init__(self, face, size=128, mask=None): + logger.trace("Initializing %s: (face: %s, size: %s, mask: %s)", + self.__class__.__name__, + face if face is None else face.shape, + size, + mask if mask is None else mask.shape) + self._size = size + if face.ndim == 2 and face.shape[1] == 1: + self._face = self._image_from_jpg(face) + else: + self._face = face[..., 2::-1] + self._photo = ImageTk.PhotoImage(self._generate_tk_face_data(mask)) + + logger.trace("Initialized %s", self.__class__.__name__) + + # << PUBLIC PROPERTIES >> # + @property + def photo(self): + """ :class:`tkinter.PhotoImage`: The face in a format that can be placed on the + :class:`~tools.manual.faceviewer.frame.FacesViewer` canvas. """ + return self._photo + + # << PUBLIC METHODS >> # + def update(self, face, mask): + """ Update the :attr:`photo` with the given face and mask. + + Parameters + ---------- + face: :class:`numpy.ndarray` + The face, sized correctly as a 3 channel BGR image + mask: :class:`numpy.ndarray` or ``None`` + The mask to be applied to the face image. Pass ``None`` if no mask is to be used + """ + self._face = face[..., 2::-1] + self._photo.paste(self._generate_tk_face_data(mask)) + + def update_mask(self, mask): + """ Update the mask in the 4th channel of :attr:`photo` to the given mask. + + Parameters + ---------- + mask: :class:`numpy.ndarray` or ``None`` + The mask to be applied to the face image. Pass ``None`` if no mask is to be used + """ + self._photo.paste(self._generate_tk_face_data(mask)) + + # << PRIVATE METHODS >> # + def _image_from_jpg(self, face): + """ Convert an encoded jpg into 3 channel BGR image. + + Parameters + ---------- + face: :class:`numpy.ndarray` + The encoded jpg as a two dimension numpy array + + Returns + ------- + :class:`numpy.ndarray` + The decoded jpg as a 3 channel BGR image + """ + face = cv2.imdecode(face, cv2.IMREAD_UNCHANGED) + interp = cv2.INTER_CUBIC if face.shape[0] < self._size else cv2.INTER_AREA + if face.shape[0] != self._size: + face = cv2.resize(face, (self._size, self._size), interpolation=interp) + return face[..., 2::-1] + + def _generate_tk_face_data(self, mask): + """ Create the :class:`tkinter.PhotoImage` from the currant :attr:`_face`. + + Parameters + ---------- + mask: :class:`numpy.ndarray` or ``None`` + The mask to add to the image. ``None`` if a mask is not being used + + Returns + ------- + :class:`tkinter.PhotoImage` + The face formatted for the :class:`~tools.manual.faceviewer.frame.FacesViewer` canvas. + """ + mask = np.ones(self._face.shape[:2], dtype="uint8") * 255 if mask is None else mask + if mask.shape[0] != self._size: + mask = cv2.resize(mask, self._face.shape[:2], interpolation=cv2.INTER_AREA) + img = np.concatenate((self._face, mask[..., None]), axis=-1) + return Image.fromarray(img) diff --git a/tools/manual/frameviewer/__init__.py b/tools/manual/frameviewer/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tools/manual/frameviewer/control.py b/tools/manual/frameviewer/control.py new file mode 100644 index 00000000..6d76ea82 --- /dev/null +++ b/tools/manual/frameviewer/control.py @@ -0,0 +1,289 @@ +#!/usr/bin/env python3 +""" Handles Navigation and Background Image for the Frame Viewer section of the manual +tool GUI. """ + +import logging +import tkinter as tk + +import cv2 +import numpy as np +from PIL import Image, ImageTk + +logger = logging.getLogger(__name__) # pylint: disable=invalid-name + + +class Navigation(): + """ Handles playback and frame navigation for the Frame Viewer Window. + + Parameters + ---------- + display_frame: :class:`DisplayFrame` + The parent frame viewer window + """ + def __init__(self, display_frame): + logger.debug("Initializing %s", self.__class__.__name__) + self._globals = display_frame._globals + self._det_faces = display_frame._det_faces + self._nav = display_frame._nav + self._tk_is_playing = tk.BooleanVar() + self._tk_is_playing.set(False) + logger.debug("Initialized %s", self.__class__.__name__) + + @property + def _current_nav_frame_count(self): + """ int: The current frame count for the transport slider """ + return self._nav["scale"].cget("to") + 1 + + def nav_scale_callback(self, *args, reset_progress=True): # pylint:disable=unused-argument + """ Adjust transport slider scale for different filters. + + Returns + ------- + bool + ``True`` if the navigation scale has been updated otherwise ``False`` + """ + if reset_progress: + self.stop_playback() + frame_count = self._det_faces.filter.count + if self._current_nav_frame_count == frame_count: + logger.trace("Filtered count has not changed. Returning") + return False + max_frame = max(0, frame_count - 1) + logger.debug("Filtered frame count has changed. Updating from %s to %s", + self._current_nav_frame_count, frame_count) + self._nav["scale"].config(to=max_frame) + self._nav["label"].config(text="/{}".format(max_frame)) + state = "disabled" if max_frame == 0 else "normal" + self._nav["entry"].config(state=state) + if reset_progress: + self._globals.tk_transport_index.set(0) + return True + + @property + def tk_is_playing(self): + """ :class:`tkinter.BooleanVar`: Whether the stream is currently playing. """ + return self._tk_is_playing + + def handle_play_button(self): + """ Handle the play button. + + Switches the :attr:`tk_is_playing` variable. + """ + is_playing = self.tk_is_playing.get() + self.tk_is_playing.set(not is_playing) + + def stop_playback(self): + """ Stop play back if playing """ + if self.tk_is_playing.get(): + logger.trace("Stopping playback") + self.tk_is_playing.set(False) + + def increment_frame(self, frame_count=None, is_playing=False): + """ Update The frame navigation position to the next frame based on filter. """ + if not is_playing: + self.stop_playback() + position = self._globals.tk_transport_index.get() + face_count_change = self._check_face_count_change() + if face_count_change: + position -= 1 + frame_count = self._det_faces.filter.count if frame_count is None else frame_count + if not face_count_change and (frame_count == 0 or position == frame_count - 1): + logger.debug("End of Stream. Not incrementing") + self.stop_playback() + return + self._globals.tk_transport_index.set(min(position + 1, max(0, frame_count - 1))) + + def decrement_frame(self): + """ Update The frame navigation position to the previous frame based on filter. """ + self.stop_playback() + position = self._globals.tk_transport_index.get() + face_count_change = self._check_face_count_change() + if face_count_change: + position += 1 + if not face_count_change and (self._det_faces.filter.count == 0 or position == 0): + logger.debug("End of Stream. Not incrementing") + return + self._globals.tk_transport_index.set(min(max(0, self._det_faces.filter.count - 1), + max(0, position - 1))) + + def _check_face_count_change(self): + """ Check whether the face count for the current filter has changed, and update the + transport scale appropriately. + + Perform additional check on whether the current frame still meets the selected navigation + mode filter criteria. + + Returns + ------- + bool + ``True`` if the currently active frame no longer meets the filter criteria otherwise + ``False`` + """ + filter_mode = self._globals.filter_mode + if filter_mode not in ("No Faces", "Multiple Faces"): + return False + if not self.nav_scale_callback(reset_progress=False): + return False + face_count = len(self._det_faces.current_faces[self._globals.frame_index]) + if (filter_mode == "No Faces" and face_count != 0) or (filter_mode == "Multiple Faces" + and face_count < 2): + return True + return False + + def goto_first_frame(self): + """ Go to the first frame that meets the filter criteria. """ + self.stop_playback() + position = self._globals.tk_transport_index.get() + if position == 0: + return + self._globals.tk_transport_index.set(0) + + def goto_last_frame(self): + """ Go to the last frame that meets the filter criteria. """ + self.stop_playback() + position = self._globals.tk_transport_index.get() + frame_count = self._det_faces.filter.count + if position == frame_count - 1: + return + self._globals.tk_transport_index.set(frame_count - 1) + + +class BackgroundImage(): + """ The background image of the canvas """ + def __init__(self, canvas): + self._canvas = canvas + self._globals = canvas._globals + self._det_faces = canvas._det_faces + placeholder = np.ones((*reversed(self._globals.frame_display_dims), 3), dtype="uint8") + self._tk_frame = ImageTk.PhotoImage(Image.fromarray(placeholder)) + self._tk_face = ImageTk.PhotoImage(Image.fromarray(placeholder)) + self._image = self._canvas.create_image(self._globals.frame_display_dims[0] / 2, + self._globals.frame_display_dims[1] / 2, + image=self._tk_frame, + anchor=tk.CENTER, + tags="main_image") + + @property + def _current_view_mode(self): + """ str: `frame` if global zoom mode variable is set to ``False`` other wise `face`. """ + retval = "face" if self._globals.is_zoomed else "frame" + logger.trace(retval) + return retval + + def refresh(self, view_mode): + """ Update the displayed frame. + + Parameters + ---------- + view_mode: ["frame", "face"] + The currently active editor's selected view mode. + """ + self._switch_image(view_mode) + getattr(self, "_update_tk_{}".format(self._current_view_mode))() + logger.trace("Updating background frame") + + def _switch_image(self, view_mode): + """ Switch the image between the full frame image and the zoomed face image. + + Parameters + ---------- + view_mode: ["frame", "face"] + The currently active editor's selected view mode. + """ + if view_mode == self._current_view_mode: + return + logger.trace("Switching background image from '%s' to '%s'", + self._current_view_mode, view_mode) + img = getattr(self, "_tk_{}".format(view_mode)) + self._canvas.itemconfig(self._image, image=img) + self._globals.tk_is_zoomed.set(view_mode == "face") + self._globals.tk_face_index.set(0) + + def _update_tk_face(self): + """ Update the currently zoomed face. """ + face = self._get_zoomed_face() + padding = self._get_padding((min(self._globals.frame_display_dims), + min(self._globals.frame_display_dims))) + face = cv2.copyMakeBorder(face, *padding, cv2.BORDER_CONSTANT) + if self._tk_frame.height() != face.shape[0]: + self._resize_frame() + + logger.trace("final shape: %s", face.shape) + self._tk_face.paste(Image.fromarray(face)) + + def _get_zoomed_face(self): + """ Get the zoomed face or a blank image if no faces are available. + + Returns + ------- + :class:`numpy.ndarray` + The face sized to the shortest dimensions of the face viewer + """ + frame_idx = self._globals.frame_index + face_idx = self._globals.face_index + faces_in_frame = self._det_faces.face_count_per_index[frame_idx] + size = min(self._globals.frame_display_dims) + + if face_idx + 1 > faces_in_frame: + logger.debug("Resetting face index to 0 for more faces in frame than current index: (" + "faces_in_frame: %s, zoomed_face_index: %s", faces_in_frame, face_idx) + self._globals.tk_face_index.set(0) + + if faces_in_frame == 0: + face = np.ones((size, size, 3), dtype="uint8") + else: + det_face = self._det_faces.current_faces[frame_idx][face_idx] + det_face.load_aligned(self._globals.current_frame["image"], size=size, force=True) + face = det_face.aligned_face.copy() + det_face.aligned["image"] = None + + logger.trace("face shape: %s", face.shape) + return face[..., 2::-1] + + def _update_tk_frame(self): + """ Place the currently held frame into :attr:`_tk_frame`. """ + img = cv2.resize(self._globals.current_frame["image"], + self._globals.current_frame["display_dims"], + interpolation=self._globals.current_frame["interpolation"])[..., 2::-1] + padding = self._get_padding(img.shape[:2]) + if any(padding): + img = cv2.copyMakeBorder(img, *padding, cv2.BORDER_CONSTANT) + logger.trace("final shape: %s", img.shape) + + if self._tk_frame.height() != img.shape[0]: + self._resize_frame() + + self._tk_frame.paste(Image.fromarray(img)) + + def _get_padding(self, size): + """ Obtain the Left, Top, Right, Bottom padding required to place the square face or frame + in to the Photo Image + + Returns + ------- + tuple + The (Left, Top, Right, Bottom) padding to apply to the face image in pixels + """ + pad_lt = ((self._globals.frame_display_dims[1] - size[0]) // 2, + (self._globals.frame_display_dims[0] - size[1]) // 2) + padding = (pad_lt[0], + self._globals.frame_display_dims[1] - size[0] - pad_lt[0], + pad_lt[1], + self._globals.frame_display_dims[0] - size[1] - pad_lt[1]) + logger.debug("Frame dimensions: %s, size: %s, padding: %s", + self._globals.frame_display_dims, size, padding) + return padding + + def _resize_frame(self): + """ Resize the :attr:`_tk_frame`, attr:`_tk_face` photo images, update the canvas to + offset the image correctly. + """ + logger.trace("Resizing video frame on resize event: %s", self._globals.frame_display_dims) + placeholder = np.ones((*reversed(self._globals.frame_display_dims), 3), dtype="uint8") + self._tk_frame = ImageTk.PhotoImage(Image.fromarray(placeholder)) + self._tk_face = ImageTk.PhotoImage(Image.fromarray(placeholder)) + self._canvas.coords(self._image, + self._globals.frame_display_dims[0] / 2, + self._globals.frame_display_dims[1] / 2) + img = self._tk_face if self._current_view_mode == "face" else self._tk_frame + self._canvas.itemconfig(self._image, image=img) diff --git a/tools/manual/frameviewer/editor/__init__.py b/tools/manual/frameviewer/editor/__init__.py new file mode 100644 index 00000000..8a7244ab --- /dev/null +++ b/tools/manual/frameviewer/editor/__init__.py @@ -0,0 +1,8 @@ +#!/usr/bin/env python3 +""" The Frame Viewer for Faceswap's Manual Tool. """ + +from ._base import View # noqa +from .bounding_box import BoundingBox # noqa +from .extract_box import ExtractBox # noqa +from .landmarks import Landmarks, Mesh # noqa +from .mask import Mask # noqa diff --git a/tools/manual/frameviewer/editor/_base.py b/tools/manual/frameviewer/editor/_base.py new file mode 100644 index 00000000..0aa11745 --- /dev/null +++ b/tools/manual/frameviewer/editor/_base.py @@ -0,0 +1,623 @@ +#!/usr/bin/env python3 +""" Editor objects for the manual adjustments tool """ + +import logging +import tkinter as tk + +from collections import OrderedDict + +import numpy as np + +from lib.gui.control_helper import ControlPanelOption + +logger = logging.getLogger(__name__) # pylint: disable=invalid-name + + +class Editor(): + """ Parent Class for Object Editors. + + Editors allow the user to use a variety of tools to manipulate alignments from the main + display frame. + + Parameters + ---------- + canvas: :class:`tkinter.Canvas` + The canvas that holds the image and annotations + detected_faces: :class:`~tools.manual.detected_faces.DetectedFaces` + The _detected_faces data for this manual session + control_text: str + The text that is to be displayed at the top of the Editor's control panel. + """ + def __init__(self, canvas, detected_faces, control_text="", key_bindings=None): + logger.debug("Initializing %s: (canvas: '%s', detected_faces: %s, control_text: %s)", + self.__class__.__name__, canvas, detected_faces, control_text) + self._canvas = canvas + self._globals = canvas._globals + self._det_faces = detected_faces + + self._current_color = dict() + self._actions = OrderedDict() + self._controls = dict(header=control_text, controls=[]) + self._add_key_bindings(key_bindings) + + self._add_actions() + self._add_controls() + self._add_annotation_format_controls() + + self._mouse_location = None + self._drag_data = dict() + self._drag_callback = None + self.bind_mouse_motion() + logger.debug("Initialized %s", self.__class__.__name__) + + @property + def _default_colors(self): + """ dict: The default colors for each annotation """ + return {"BoundingBox": "#0000ff", + "ExtractBox": "#00ff00", + "Landmarks": "#ff00ff", + "Mask": "#ff0000", + "Mesh": "#00ffff"} + + @property + def _is_active(self): + """ bool: ``True`` if this editor is currently active otherwise ``False``. + + Notes + ----- + When initializing, the active_editor parameter will not be set in the parent, + so return ``False`` in this instance + """ + return hasattr(self._canvas, "active_editor") and self._canvas.active_editor == self + + @property + def view_mode(self): + """ ["frame", "face"]: The view mode for the currently selected editor. If the editor does + not have a view mode that can be updated, then `"frame"` will be returned. """ + tk_var = self._actions.get("magnify", dict()).get("tk_var", None) + retval = "frame" if tk_var is None or not tk_var.get() else "face" + return retval + + @property + def _zoomed_roi(self): + """ :class:`numpy.ndarray`: The (`left`, `top`, `right`, `bottom`) roi of the zoomed face + in the display frame. """ + half_size = min(self._globals.frame_display_dims) / 2 + left = self._globals.frame_display_dims[0] / 2 - half_size + top = 0 + right = self._globals.frame_display_dims[0] / 2 + half_size + bottom = self._globals.frame_display_dims[1] + retval = np.rint(np.array((left, top, right, bottom))).astype("int32") + logger.trace("Zoomed ROI: %s", retval) + return retval + + @property + def _zoomed_dims(self): + """ tuple: The (`width`, `height`) of the zoomed ROI. """ + roi = self._zoomed_roi + return (roi[2] - roi[0], roi[3] - roi[1]) + + @property + def _control_vars(self): + """ dict: The tk control panel variables for the currently selected editor. """ + return self._canvas.control_tk_vars.get(self.__class__.__name__, dict()) + + @property + def controls(self): + """ dict: The control panel options and header text for the current editor """ + return self._controls + + @property + def _control_color(self): + """ str: The hex color code set in the control panel for the current editor. """ + annotation = self.__class__.__name__ + return self._annotation_formats[annotation]["color"].get() + + @property + def _annotation_formats(self): + """ dict: The format (color, opacity etc.) of each editor's annotation display. """ + return self._canvas.annotation_formats + + @property + def actions(self): + """ list: The optional action buttons for the actions frame in the GUI for the + current editor """ + return self._actions + + @property + def _face_iterator(self): + """ list: The detected face objects to be iterated. This will either be all faces in the + frame (normal view) or the single zoomed in face (zoom mode). """ + if self._globals.frame_index == -1: + faces = [] + else: + faces = self._det_faces.current_faces[self._globals.frame_index] + faces = ([faces[self._globals.face_index]] + if self._globals.is_zoomed and faces else faces) + return faces + + def _add_key_bindings(self, key_bindings): + """ Add the editor specific key bindings for the currently viewed editor. + + Parameters + ---------- + key_bindings: dict + The key binding to method dictionary for this editor. + """ + if key_bindings is None: + return + for key, method in key_bindings.items(): + logger.debug("Binding key '%s' to method %s for editor '%s'", + key, method, self.__class__.__name__) + self._canvas.key_bindings.setdefault(key, dict())["bound_to"] = None + self._canvas.key_bindings[key][self.__class__.__name__] = method + + @staticmethod + def _get_anchor_points(bounding_box): + """ Retrieve the (x, y) co-ordinates for each of the 4 corners of a bounding box's anchors + for both the displayed anchors and the anchor grab locations. + + Parameters + ---------- + bounding_box: tuple + The (`top-left`, `top-right`, `bottom-right`, `bottom-left`) (x, y) coordinates of the + bounding box + + Returns + display_anchors: tuple + The (`top`, `left`, `bottom`, `right`) co-ordinates for each circle at each point + of the bounding box corners, sized for display + grab_anchors: tuple + The (`top`, `left`, `bottom`, `right`) co-ordinates for each circle at each point + of the bounding box corners, at a larger size for grabbing with a mouse + """ + radius = 3 + grab_radius = radius * 3 + display_anchors = tuple((cnr[0] - radius, cnr[1] - radius, + cnr[0] + radius, cnr[1] + radius) + for cnr in bounding_box) + grab_anchors = tuple((cnr[0] - grab_radius, cnr[1] - grab_radius, + cnr[0] + grab_radius, cnr[1] + grab_radius) + for cnr in bounding_box) + return display_anchors, grab_anchors + + def update_annotation(self): # pylint:disable=no-self-use + """ Update the display annotations for the current objects. + + Override for specific editors. + """ + logger.trace("Default annotations. Not storing Objects") + + def hide_annotation(self, tag=None): + """ Hide annotations for this editor. + + Parameters + ---------- + tag: str, optional + The specific tag to hide annotations for. If ``None`` then all annotations for this + editor are hidden, otherwise only the annotations specified by the given tag are + hidden. Default: ``None`` + """ + tag = self.__class__.__name__ if tag is None else tag + logger.trace("Hiding annotations for tag: %s", tag) + self._canvas.itemconfig(tag, state="hidden") + + def _object_tracker(self, key, object_type, face_index, + coordinates, object_kwargs): + """ Create an annotation object and add it to :attr:`_objects` or update an existing + annotation if it has already been created. + + Parameters + ---------- + key: str + The key for this annotation in :attr:`_objects` + object_type: str + This can be any string that is a natural extension to :class:`tkinter.Canvas.create_` + face_index: int + The index of the face within the current frame + coordinates: tuple or list + The bounding box coordinates for this object + object_kwargs: dict + The keyword arguments for this object + + Returns + ------- + int: + The tkinter canvas item identifier for the created object + """ + object_color_keys = self._get_object_color_keys(key, object_type) + tracking_id = "_".join((key, str(face_index))) + face_tag = "face_{}".format(face_index) + face_objects = set(self._canvas.find_withtag(face_tag)) + annotation_objects = set(self._canvas.find_withtag(key)) + existing_object = tuple(face_objects.intersection(annotation_objects)) + if not existing_object: + item_id = self._add_new_object(key, + object_type, + face_index, + coordinates, + object_kwargs) + update_color = bool(object_color_keys) + else: + item_id = existing_object[0] + update_color = self._update_existing_object( + existing_object[0], + coordinates, + object_kwargs, + tracking_id, + object_color_keys) + if update_color: + self._current_color[tracking_id] = object_kwargs[object_color_keys[0]] + return item_id + + @staticmethod + def _get_object_color_keys(key, object_type): + """ The canvas object's parameter that needs to be adjusted for color varies based on + the type of object that is being used. Returns the correct parameter based on object. + + Parameters + ---------- + key: str + The key for this annotation's tag creation + object_type: str + This can be any string that is a natural extension to :class:`tkinter.Canvas.create_` + + Returns + ------- + list: + The list of keyword arguments for this objects color parameter(s) or an empty list + if it is not relevant for this object + """ + if object_type in ("line", "text"): + retval = ["fill"] + elif object_type == "image": + retval = [] + elif object_type == "oval" and key.startswith("lm_dsp_"): + retval = ["fill", "outline"] + else: + retval = ["outline"] + logger.trace("returning %s for key: %s, object_type: %s", retval, key, object_type) + return retval + + def _add_new_object(self, key, object_type, face_index, coordinates, object_kwargs): + """ Add a new object to the canvas. + + Parameters + ---------- + key: str + The key for this annotation's tag creation + object_type: str + This can be any string that is a natural extension to :class:`tkinter.Canvas.create_` + face_index: int + The index of the face within the current frame + coordinates: tuple or list + The bounding box coordinates for this object + object_kwargs: dict + The keyword arguments for this object + + Returns + ------- + int: + The tkinter canvas item identifier for the created object + """ + logger.debug("Adding object: (key: '%s', object_type: '%s', face_index: %s, " + "coordinates: %s, object_kwargs: %s)", key, object_type, face_index, + coordinates, object_kwargs) + object_kwargs["tags"] = self._set_object_tags(face_index, key) + item_id = getattr(self._canvas, + "create_{}".format(object_type))(*coordinates, **object_kwargs) + return item_id + + def _set_object_tags(self, face_index, key): + """ Create the tkinter object tags for the incoming object. + + Parameters + ---------- + face_index: int + The face index within the current frame for the face that tags are being created for + key: str + The base tag for this object, for which additional tags will be generated + + Returns + ------- + list + The generated tags for the current object + """ + tags = ["face_{}".format(face_index), + self.__class__.__name__, + "{}_face_{}".format(self.__class__.__name__, face_index), + key, + "{}_face_{}".format(key, face_index)] + if "_" in key: + split_key = key.split("_") + if split_key[-1].isdigit(): + base_tag = "_".join(split_key[:-1]) + tags.append(base_tag) + tags.append("{}_face_{}".format(base_tag, face_index)) + return tags + + def _update_existing_object(self, item_id, coordinates, object_kwargs, + tracking_id, object_color_keys): + """ Update an existing tracked object. + + Parameters + ---------- + item_id: int + The canvas object item_id to be updated + coordinates: tuple or list + The bounding box coordinates for this object + object_kwargs: dict + The keyword arguments for this object + tracking_id: str + The tracking identifier for this object's color + object_color_keys: list + The list of keyword arguments for this object to update for color + + Returns + ------- + bool + ``True`` if :attr:`_current_color` should be updated otherwise ``False`` + """ + update_color = (object_color_keys and + object_kwargs[object_color_keys[0]] != self._current_color[tracking_id]) + update_kwargs = dict(state=object_kwargs.get("state", "normal")) + if update_color: + for key in object_color_keys: + update_kwargs[key] = object_kwargs[object_color_keys[0]] + if self._canvas.type(item_id) == "image" and "image" in object_kwargs: + update_kwargs["image"] = object_kwargs["image"] + logger.trace("Updating coordinates: (item_id: '%s', object_kwargs: %s, " + "coordinates: %s, update_kwargs: %s", item_id, object_kwargs, + coordinates, update_kwargs) + self._canvas.itemconfig(item_id, **update_kwargs) + self._canvas.coords(item_id, *coordinates) + return update_color + + # << MOUSE CALLBACKS >> + # Mouse cursor display + def bind_mouse_motion(self): + """ Binds the mouse motion for the current editor's mouse event to the editor's + :func:`_update_cursor` function. + + Called on initialization and active editor update. + """ + self._canvas.bind("", self._update_cursor) + + def _update_cursor(self, event): # pylint: disable=unused-argument + """ The mouse cursor display as bound to the mouse's event.. + + The default is to always return a standard cursor, so this method should be overridden for + editor specific cursor update. + + Parameters + ---------- + event: :class:`tkinter.Event` + The tkinter mouse event. Unused for default tracking, but available for specific editor + tracking. + """ + self._canvas.config(cursor="") + + # Mouse click and drag actions + def set_mouse_click_actions(self): + """ Add the bindings for left mouse button click and drag actions. + + This binds the mouse to the :func:`_drag_start`, :func:`_drag` and :func:`_drag_stop` + methods. + + By default these methods do nothing (except for :func:`_drag_stop` which resets + :attr:`_drag_data`. + + This bindings should be added for all editors. To add additional bindings, + `super().set_mouse_click_actions` should be called prior to adding them.. + """ + logger.debug("Setting mouse bindings") + self._canvas.bind("", self._drag_start) + self._canvas.bind("", self._drag_stop) + self._canvas.bind("", self._drag) + + def _drag_start(self, event): # pylint:disable=unused-argument + """ The action to perform when the user starts clicking and dragging the mouse. + + The default does nothing except reset the attr:`drag_data` and attr:`drag_callback`. + Override for Editor specific click and drag start actions. + + Parameters + ---------- + event: :class:`tkinter.Event` + The tkinter mouse event. Unused but for default action, but available for editor + specific actions + """ + self._drag_data = dict() + self._drag_callback = None + + def _drag(self, event): + """ The default callback for the drag part of a mouse click and drag action. + + :attr:`_drag_callback` should be set in :func:`self._drag_start`. This callback will then + be executed on a mouse drag event. + + Parameters + ---------- + event: :class:`tkinter.Event` + The tkinter mouse event. + """ + if self._drag_callback is None: + return + self._drag_callback(event) + + def _drag_stop(self, event): # pylint:disable=unused-argument + """ The action to perform when the user stops clicking and dragging the mouse. + + Default is to set :attr:`_drag_data` to `dict`. Override for Editor specific stop actions. + + Parameters + ---------- + event: :class:`tkinter.Event` + The tkinter mouse event. Unused but required + """ + self._drag_data = dict() + + def _scale_to_display(self, points): + """ Scale and offset the given points to the current display scale and offset values. + + Parameters + ---------- + points: :class:`numpy.ndarray` + Array of x, y co-ordinates to adjust + + Returns + ------- + :class:`numpy.ndarray` + The adjusted x, y co-ordinates for display purposes rounded to the nearest integer + """ + retval = np.rint((points * self._globals.current_frame["scale"]) + + self._canvas.offset).astype("int32") + logger.trace("Original points: %s, scaled points: %s", points, retval) + return retval + + def scale_from_display(self, points, do_offset=True): + """ Scale and offset the given points from the current display to the correct original + values. + + Parameters + ---------- + points: :class:`numpy.ndarray` + Array of x, y co-ordinates to adjust + offset: bool, optional + ``True`` if the offset should be calculated otherwise ``False``. Default: ``True`` + + Returns + ------- + :class:`numpy.ndarray` + The adjusted x, y co-ordinates to the original frame location rounded to the nearest + integer + """ + offset = self._canvas.offset if do_offset else (0, 0) + retval = np.rint((points - offset) / self._globals.current_frame["scale"]).astype("int32") + logger.trace("Original points: %s, scaled points: %s", points, retval) + return retval + + # << ACTION CONTROL PANEL OPTIONS >> + def _add_actions(self): + """ Add the Action buttons for this editor's optional left hand side action sections. + + The default does nothing. Override for editor specific actions. + """ + self._actions = self._actions + + def _add_action(self, title, icon, helptext, group=None, hotkey=None): + """ Add an action dictionary to :attr:`_actions`. This will create a button in the optional + actions frame to the left hand side of the frames viewer. + + Parameters + ---------- + title: str + The title of the action to be generated + icon: str + The name of the icon that is used to display this action's button + helptext: str + The tooltip text to display for this action + group: str, optional + If a group is passed in, then any buttons belonging to that group will be linked (i.e. + only one button can be active at a time.). If ``None`` is passed in then the button + will act independently. Default: ``None`` + hotkey: str, optional + The hotkey binding for this action. Set to ``None`` if there is no hotkey binding. + Default: ``None`` + """ + var = tk.BooleanVar() + action = dict(icon=icon, helptext=helptext, group=group, tk_var=var, hotkey=hotkey) + logger.debug("Adding action: %s", action) + self._actions[title] = action + + def _add_controls(self): + """ Add the controls for this editor's control panel. + + The default does nothing. Override for editor specific controls. + """ + self._controls = self._controls + + def _add_control(self, option, global_control=False): + """ Add a control panel control to :attr:`_controls` and add a trace to the variable + to update display. + + Parameters + ---------- + option: :class:`lib.gui.control_helper.ControlPanelOption' + The control panel option to add to this editor's control + global_control: bool, optional + Whether the given control is a global control (i.e. annotation formatting). + Default: ``False`` + """ + self._controls["controls"].append(option) + if global_control: + logger.debug("Added global control: '%s' for editor: '%s'", + option.title, self.__class__.__name__) + return + logger.debug("Added local control: '%s' for editor: '%s'", + option.title, self.__class__.__name__) + editor_key = self.__class__.__name__ + group_key = option.group.replace(" ", "").lower() + group_key = "none" if group_key == "_master" else group_key + annotation_key = option.title.replace(" ", "") + self._canvas.control_tk_vars.setdefault( + editor_key, dict()).setdefault(group_key, dict())[annotation_key] = option.tk_var + + def _add_annotation_format_controls(self): + """ Add the annotation display (color/size) controls to :attr:`_annotation_formats`. + + These should be universal and available for all editors. + """ + editors = ("Bounding Box", "Extract Box", "Landmarks", "Mask", "Mesh") + if not self._annotation_formats: + opacity = ControlPanelOption("Mask Opacity", + int, + group="Color", + min_max=(0, 100), + default=40, + rounding=1, + helptext="Set the mask opacity") + for editor in editors: + annotation_key = editor.replace(" ", "") + logger.debug("Adding to global format controls: '%s'", editor) + colors = ControlPanelOption(editor, + str, + group="Color", + subgroup="colors", + choices="colorchooser", + default=self._default_colors[annotation_key], + helptext="Set the annotation color") + colors.set(self._default_colors[annotation_key]) + self._annotation_formats.setdefault(annotation_key, dict())["color"] = colors + self._annotation_formats[annotation_key]["mask_opacity"] = opacity + + for editor in editors: + annotation_key = editor.replace(" ", "") + for group, ctl in self._annotation_formats[annotation_key].items(): + logger.debug("Adding global format control to editor: (editor:'%s', group: '%s')", + editor, group) + self._add_control(ctl, global_control=True) + + +class View(Editor): + """ The view Editor. + + Does not allow any editing, just used for previewing annotations. + + This is the default start-up editor. + + Parameters + ---------- + canvas: :class:`tkinter.Canvas` + The canvas that holds the image and annotations + detected_faces: :class:`~tools.manual.detected_faces.DetectedFaces` + The _detected_faces data for this manual session + """ + def __init__(self, canvas, detected_faces): + control_text = "Viewer\nPreview the frame's annotations." + super().__init__(canvas, detected_faces, control_text) + + def _add_actions(self): + """ Add the optional action buttons to the viewer. Current actions are Zoom. """ + self._add_action("magnify", "zoom", "Magnify/Demagnify the View", group=None, hotkey="M") + self._actions["magnify"]["tk_var"].trace("w", lambda *e: self._globals.tk_update.set(True)) diff --git a/tools/manual/frameviewer/editor/bounding_box.py b/tools/manual/frameviewer/editor/bounding_box.py new file mode 100644 index 00000000..58a7b30c --- /dev/null +++ b/tools/manual/frameviewer/editor/bounding_box.py @@ -0,0 +1,403 @@ +#!/usr/bin/env python3 +""" Bounding Box Editor for the manual adjustments tool """ + +import platform +from functools import partial + +import numpy as np + +from lib.gui.custom_widgets import RightClickMenu +from ._base import ControlPanelOption, Editor, logger + + +class BoundingBox(Editor): + """ The Bounding Box Editor. + + Adjusting the bounding box feeds the aligner to generate new 68 point landmarks. + + Parameters + ---------- + canvas: :class:`tkinter.Canvas` + The canvas that holds the image and annotations + detected_faces: :class:`~tools.manual.detected_faces.DetectedFaces` + The _detected_faces data for this manual session + """ + def __init__(self, canvas, detected_faces): + self._tk_aligner = None + self._right_click_menu = RightClickMenu(["Delete Face"], + [self._delete_current_face], + ["Del"]) + control_text = ("Bounding Box Editor\nEdit the bounding box being fed into the aligner " + "to recalculate the landmarks.\n\n" + " - Grab the corner anchors to resize the bounding box.\n" + " - Click and drag the bounding box to relocate.\n" + " - Click in empty space to create a new bounding box.\n" + " - Right click a bounding box to delete a face.") + key_bindings = {"": self._delete_current_face} + super().__init__(canvas, detected_faces, + control_text=control_text, key_bindings=key_bindings) + + @property + def _corner_order(self): + """ dict: The position index of bounding box corners """ + return {0: ("top", "left"), + 1: ("top", "right"), + 2: ("bottom", "right"), + 3: ("bottom", "left")} + + @property + def _bounding_boxes(self): + """ list: The :func:`tkinter.Canvas.coords` for all displayed bounding boxes. """ + item_ids = self._canvas.find_withtag("bb_box") + return [self._canvas.coords(item_id) for item_id in item_ids + if self._canvas.itemcget(item_id, "state") != "hidden"] + + def _add_controls(self): + """ Controls for feeding the Aligner. Exposes Normalization Method as a parameter. """ + align_ctl = ControlPanelOption( + "Aligner", + str, + group="Aligner", + choices=["cv2-dnn", "FAN"], + default="FAN", + is_radio=True, + helptext="Aligner to use. FAN will obtain better alignments, but cv2-dnn can be " + "useful if FAN cannot get decent alignments and you want to set a base to " + "edit from.") + self._tk_aligner = align_ctl.tk_var + self._add_control(align_ctl) + + norm_ctl = ControlPanelOption( + "Normalization method", + str, + group="Aligner", + choices=["none", "clahe", "hist", "mean"], + default="hist", + is_radio=True, + helptext="Normalization method to use for feeding faces to the aligner. This can help " + "the aligner better align faces with difficult lighting conditions. " + "Different methods will yield different results on different sets. NB: This " + "does not impact the output face, just the input to the aligner." + "\n\tnone: Don't perform normalization on the face." + "\n\tclahe: Perform Contrast Limited Adaptive Histogram Equalization on the " + "face." + "\n\thist: Equalize the histograms on the RGB channels." + "\n\tmean: Normalize the face colors to the mean.") + var = norm_ctl.tk_var + var.trace("w", + lambda *e, v=var: self._det_faces.extractor.set_normalization_method(v.get())) + self._add_control(norm_ctl) + + def update_annotation(self): + """ Get the latest bounding box data from alignments and update. """ + if self._globals.is_zoomed: + logger.trace("Image is zoomed. Hiding Bounding Box.") + self.hide_annotation() + return + key = "bb_box" + color = self._control_color + for idx, face in enumerate(self._face_iterator): + box = np.array([(face.left, face.top), (face.right, face.bottom)]) + box = self._scale_to_display(box).astype("int32").flatten() + kwargs = dict(outline=color, width=1) + logger.trace("frame_index: %s, face_index: %s, box: %s, kwargs: %s", + self._globals.frame_index, idx, box, kwargs) + self._object_tracker(key, "rectangle", idx, box, kwargs) + self._update_anchor_annotation(idx, box, color) + logger.trace("Updated bounding box annotations") + + def _update_anchor_annotation(self, face_index, bounding_box, color): + """ Update the anchor annotations for each corner of the bounding box. + + The anchors only display when the bounding box editor is active. + + Parameters + ---------- + face_index: int + The index of the face being annotated + bounding_box: :class:`numpy.ndarray` + The scaled bounding box to get the corner anchors for + color: str + The hex color of the bounding box line + """ + if not self._is_active: + self.hide_annotation("bb_anc_dsp") + self.hide_annotation("bb_anc_grb") + return + fill_color = "gray" + activefill_color = "white" if self._is_active else "" + anchor_points = self._get_anchor_points(((bounding_box[0], bounding_box[1]), + (bounding_box[2], bounding_box[1]), + (bounding_box[2], bounding_box[3]), + (bounding_box[0], bounding_box[3]))) + for idx, (anc_dsp, anc_grb) in enumerate(zip(*anchor_points)): + dsp_kwargs = dict(outline=color, fill=fill_color, width=1) + grb_kwargs = dict(outline="", fill="", width=1, activefill=activefill_color) + dsp_key = "bb_anc_dsp_{}".format(idx) + grb_key = "bb_anc_grb_{}".format(idx) + self._object_tracker(dsp_key, "oval", face_index, anc_dsp, dsp_kwargs) + self._object_tracker(grb_key, "oval", face_index, anc_grb, grb_kwargs) + logger.trace("Updated bounding box anchor annotations") + + # << MOUSE HANDLING >> + # Mouse cursor display + def _update_cursor(self, event): + """ Set the cursor action. + + Update :attr:`_mouse_location` with the current cursor position and display appropriate + icon. + + If the cursor is over a corner anchor, then pop resize icon. + If the cursor is over a bounding box, then pop move icon. + If the cursor is over the image, then pop add icon. + + Parameters + ---------- + event: :class:`tkinter.Event` + The current tkinter mouse event + """ + if self._check_cursor_anchors(): + return + if self._check_cursor_bounding_box(event): + return + if self._check_cursor_image(event): + return + + self._canvas.config(cursor="") + self._mouse_location = None + + def _check_cursor_anchors(self): + """ Check whether the cursor is over a corner anchor. + + If it is, set the appropriate cursor type and set :attr:`_mouse_location` to + ("anchor", (`face index`, `anchor index`) + + Returns + ------- + bool + ``True`` if cursor is over an anchor point otherwise ``False`` + """ + anchors = set(self._canvas.find_withtag("bb_anc_grb")) + item_ids = set(self._canvas.find_withtag("current")).intersection(anchors) + if not item_ids: + return False + item_id = list(item_ids)[0] + tags = self._canvas.gettags(item_id) + face_idx = int(next(tag for tag in tags if tag.startswith("face_")).split("_")[-1]) + corner_idx = int(next(tag for tag in tags + if tag.startswith("bb_anc_grb_") + and "face_" not in tag).split("_")[-1]) + self._canvas.config(cursor="{}_{}_corner".format(*self._corner_order[corner_idx])) + self._mouse_location = ("anchor", "{}_{}".format(face_idx, corner_idx)) + return True + + def _check_cursor_bounding_box(self, event): + """ Check whether the cursor is over a bounding box. + + If it is, set the appropriate cursor type and set :attr:`_mouse_location` to: + ("box", `face index`) + + Parameters + ---------- + event: :class:`tkinter.Event` + The tkinter mouse event + + Returns + ------- + bool + ``True`` if cursor is over a bounding box otherwise ``False`` + + Notes + ----- + We can't use tags on unfilled rectangles as the interior of the rectangle is not tagged. + """ + for face_idx, bbox in enumerate(self._bounding_boxes): + if bbox[0] <= event.x <= bbox[2] and bbox[1] <= event.y <= bbox[3]: + self._canvas.config(cursor="fleur") + self._mouse_location = ("box", str(face_idx)) + return True + return False + + def _check_cursor_image(self, event): + """ Check whether the cursor is over the image. + + If it is, set the appropriate cursor type and set :attr:`_mouse_location` to: + ("image", ) + + Parameters + ---------- + event: :class:`tkinter.Event` + The tkinter mouse event + + Returns + ------- + bool + ``True`` if cursor is over a bounding box otherwise ``False`` + """ + if self._globals.frame_index == -1: + return False + display_dims = self._globals.current_frame["display_dims"] + if (self._canvas.offset[0] <= event.x <= display_dims[0] + self._canvas.offset[0] and + self._canvas.offset[1] <= event.y <= display_dims[1] + self._canvas.offset[1]): + self._canvas.config(cursor="plus") + self._mouse_location = ("image", ) + return True + return False + + # Mouse Actions + def set_mouse_click_actions(self): + """ Add context menu to OS specific right click action. """ + super().set_mouse_click_actions() + self._canvas.bind("" if platform.system() == "Darwin" else "", + self._context_menu) + + def _drag_start(self, event): + """ The action to perform when the user starts clicking and dragging the mouse. + + If :attr:`_mouse_location` indicates a corner anchor, then the bounding box is resized + based on the adjusted corner, and the alignments re-generated. + + If :attr:`_mouse_location` indicates a bounding box, then the bounding box is moved, and + the alignments re-generated. + + If :attr:`_mouse_location` indicates being over the main image, then a new bounding box is + created, and alignments generated. + + Parameters + ---------- + event: :class:`tkinter.Event` + The tkinter mouse event. + """ + if self._mouse_location is None: + self._drag_data = dict() + self._drag_callback = None + return + if self._mouse_location[0] == "anchor": + corner_idx = int(self._mouse_location[1].split("_")[-1]) + self._drag_data["corner"] = self._corner_order[corner_idx] + self._drag_callback = self._resize + elif self._mouse_location[0] == "box": + self._drag_data["current_location"] = (event.x, event.y) + self._drag_callback = self._move + elif self._mouse_location[0] == "image": + self._create_new_bounding_box(event) + # Refresh cursor and _mouse_location for new bounding box and reset _drag_start + self._update_cursor(event) + self._drag_start(event) + + def _drag_stop(self, event): # pylint: disable=unused-argument + """ Trigger a viewport thumbnail update on click + drag release + + Parameters + ---------- + event: :class:`tkinter.Event` + The tkinter mouse event. Required but unused. + """ + if self._mouse_location is None: + return + face_idx = int(self._mouse_location[1].split("_")[0]) + self._det_faces.update.post_edit_trigger(self._globals.frame_index, face_idx) + + def _create_new_bounding_box(self, event): + """ Create a new bounding box when user clicks on image, outside of existing boxes. + + The bounding box is created as a square located around the click location, with dimensions + 1 quarter the size of the frame's shortest side + + Parameters + ---------- + event: :class:`tkinter.Event` + The tkinter mouse event + """ + size = min(self._globals.current_frame["display_dims"]) // 8 + box = (event.x - size, event.y - size, event.x + size, event.y + size) + logger.debug("Creating new bounding box: %s ", box) + self._det_faces.update.add(self._globals.frame_index, *self._coords_to_bounding_box(box)) + + def _resize(self, event): + """ Resizes a bounding box on a corner anchor drag event. + + Parameters + ---------- + event: :class:`tkinter.Event` + The tkinter mouse event. + """ + face_idx = int(self._mouse_location[1].split("_")[0]) + face_tag = "bb_box_face_{}".format(face_idx) + box = self._canvas.coords(face_tag) + logger.trace("Face Index: %s, Corner Index: %s. Original ROI: %s", + face_idx, self._drag_data["corner"], box) + # Switch top/bottom and left/right and set partial so indices match and we don't + # need branching logic for min/max. + limits = (partial(min, box[2] - 20), + partial(min, box[3] - 20), + partial(max, box[0] + 20), + partial(max, box[1] + 20)) + rect_xy_indices = [("left", "top", "right", "bottom").index(pnt) + for pnt in self._drag_data["corner"]] + box[rect_xy_indices[1]] = limits[rect_xy_indices[1]](event.x) + box[rect_xy_indices[0]] = limits[rect_xy_indices[0]](event.y) + logger.trace("New ROI: %s", box) + self._det_faces.update.bounding_box(self._globals.frame_index, + face_idx, + *self._coords_to_bounding_box(box), + aligner=self._tk_aligner.get()) + + def _move(self, event): + """ Moves the bounding box on a bounding box drag event. + + Parameters + ---------- + event: :class:`tkinter.Event` + The tkinter mouse event. + """ + logger.trace("event: %s, mouse_location: %s", event, self._mouse_location) + face_idx = int(self._mouse_location[1]) + shift = (event.x - self._drag_data["current_location"][0], + event.y - self._drag_data["current_location"][1]) + face_tag = "bb_box_face_{}".format(face_idx) + coords = np.array(self._canvas.coords(face_tag)) + (*shift, *shift) + logger.trace("face_tag: %s, shift: %s, new co-ords: %s", face_tag, shift, coords) + self._det_faces.update.bounding_box(self._globals.frame_index, + face_idx, + *self._coords_to_bounding_box(coords), + aligner=self._tk_aligner.get()) + self._drag_data["current_location"] = (event.x, event.y) + + def _coords_to_bounding_box(self, coords): + """ Converts tkinter coordinates to :class:`lib.faces_detect.DetectedFace` bounding + box format, scaled up and offset for feeding the model. + + Returns + ------- + tuple + The (`x`, `width`, `y`, `height`) integer points of the bounding box. + """ + logger.trace("in: %s", coords) + coords = self.scale_from_display( + np.array(coords).reshape((2, 2))).flatten().astype("int32") + logger.trace("out: %s", coords) + return (coords[0], coords[2] - coords[0], coords[1], coords[3] - coords[1]) + + def _context_menu(self, event): + """ Create a right click context menu to delete the alignment that is being + hovered over. """ + if self._mouse_location is None or self._mouse_location[0] != "box": + return + self._right_click_menu.popup(event) + + def _delete_current_face(self, *args): # pylint:disable=unused-argument + """ Called by the right click delete event. Deletes the face that the mouse is currently + over. + + Parameters + ---------- + args: tuple (unused) + The event parameter is passed in by the hot key binding, so args is required + """ + if self._mouse_location is None or self._mouse_location[0] != "box": + logger.debug("Delete called without valid location. _mouse_location: %s", + self._mouse_location) + return + logger.debug("Deleting face. _mouse_location: %s", self._mouse_location) + self._det_faces.update.delete(self._globals.frame_index, int(self._mouse_location[1])) diff --git a/tools/manual/frameviewer/editor/extract_box.py b/tools/manual/frameviewer/editor/extract_box.py new file mode 100644 index 00000000..d781221c --- /dev/null +++ b/tools/manual/frameviewer/editor/extract_box.py @@ -0,0 +1,401 @@ +#!/usr/bin/env python3 +""" Extract Box Editor for the manual adjustments tool """ + +import platform + +import numpy as np + +from lib.gui.custom_widgets import RightClickMenu +from lib.gui.utils import get_config +from ._base import Editor, logger + + +class ExtractBox(Editor): + """ The Extract Box Editor. + + Adjust the calculated Extract Box to shift all of the 68 point landmarks in place. + + Parameters + ---------- + canvas: :class:`tkinter.Canvas` + The canvas that holds the image and annotations + detected_faces: :class:`~tools.manual.detected_faces.DetectedFaces` + The _detected_faces data for this manual session + """ + def __init__(self, canvas, detected_faces): + self._right_click_menu = RightClickMenu(["Delete Face"], + [self._delete_current_face], + ["Del"]) + control_text = ("Extract Box Editor\nMove the extract box that has been generated by the " + "aligner. Click and drag:\n\n" + " - Inside the bounding box to relocate the landmarks.\n" + " - The corner anchors to resize the landmarks.\n" + " - Outside of the corners to rotate the landmarks.") + key_bindings = {"": self._delete_current_face} + super().__init__(canvas, detected_faces, + control_text=control_text, key_bindings=key_bindings) + + @property + def _corner_order(self): + """ dict: The position index of bounding box corners """ + return {0: ("top", "left"), + 3: ("top", "right"), + 2: ("bottom", "right"), + 1: ("bottom", "left")} + + def update_annotation(self): + """ Draw the latest Extract Boxes around the faces. """ + color = self._control_color + roi = self._zoomed_roi + for idx, face in enumerate(self._face_iterator): + logger.trace("Drawing Extract Box: (idx: %s, roi: %s)", idx, face.original_roi) + if self._globals.is_zoomed: + box = np.array((roi[0], roi[1], roi[2], roi[1], roi[2], roi[3], roi[0], roi[3])) + else: + face.load_aligned(None, force=True) + box = self._scale_to_display(face.original_roi).flatten() + top_left = box[:2] - 10 + kwargs = dict(fill=color, font=("Default", 20, "bold"), text=str(idx)) + self._object_tracker("eb_text", "text", idx, top_left, kwargs) + kwargs = dict(fill="", outline=color, width=1) + self._object_tracker("eb_box", "polygon", idx, box, kwargs) + self._update_anchor_annotation(idx, box, color) + logger.trace("Updated extract box annotations") + + def _update_anchor_annotation(self, face_index, extract_box, color): + """ Update the anchor annotations for each corner of the extract box. + + The anchors only display when the extract box editor is active. + + Parameters + ---------- + face_index: int + The index of the face being annotated + extract_box: :class:`numpy.ndarray` + The scaled extract box to get the corner anchors for + color: str + The hex color of the extract box line + """ + if not self._is_active or self._globals.is_zoomed: + self.hide_annotation("eb_anc_dsp") + self.hide_annotation("eb_anc_grb") + return + fill_color = "gray" + activefill_color = "white" if self._is_active else "" + anchor_points = self._get_anchor_points((extract_box[:2], + extract_box[2:4], + extract_box[4:6], + extract_box[6:])) + for idx, (anc_dsp, anc_grb) in enumerate(zip(*anchor_points)): + dsp_kwargs = dict(outline=color, fill=fill_color, width=1) + grb_kwargs = dict(outline="", fill="", width=1, activefill=activefill_color) + dsp_key = "eb_anc_dsp_{}".format(idx) + grb_key = "eb_anc_grb_{}".format(idx) + self._object_tracker(dsp_key, "oval", face_index, anc_dsp, dsp_kwargs) + self._object_tracker(grb_key, "oval", face_index, anc_grb, grb_kwargs) + logger.trace("Updated extract box anchor annotations") + + # << MOUSE HANDLING >> + # Mouse cursor display + def _update_cursor(self, event): + """ Update the cursor when it is hovering over an extract box and update + :attr:`_mouse_location` with the current cursor position. + + Parameters + ---------- + event: :class:`tkinter.Event` + The current tkinter mouse event + """ + if self._check_cursor_anchors(): + return + if self._check_cursor_box(): + return + if self._check_cursor_rotate(event): + return + self._canvas.config(cursor="") + self._mouse_location = None + + def _check_cursor_anchors(self): + """ Check whether the cursor is over a corner anchor. + + If it is, set the appropriate cursor type and set :attr:`_mouse_location` to + ("anchor", `face index`, `corner_index`) + + Returns + ------- + bool + ``True`` if cursor is over an anchor point otherwise ``False`` + """ + anchors = set(self._canvas.find_withtag("eb_anc_grb")) + item_ids = set(self._canvas.find_withtag("current")).intersection(anchors) + if not item_ids: + return False + item_id = list(item_ids)[0] + tags = self._canvas.gettags(item_id) + face_idx = int(next(tag for tag in tags if tag.startswith("face_")).split("_")[-1]) + corner_idx = int(next(tag for tag in tags + if tag.startswith("eb_anc_grb_") + and "face_" not in tag).split("_")[-1]) + + self._canvas.config(cursor="{}_{}_corner".format(*self._corner_order[corner_idx])) + self._mouse_location = ("anchor", face_idx, corner_idx) + return True + + def _check_cursor_box(self): + """ Check whether the cursor is inside an extract box. + + If it is, set the appropriate cursor type and set :attr:`_mouse_location` to + ("box", `face index`) + + Returns + ------- + bool + ``True`` if cursor is over a rotate point otherwise ``False`` + """ + extract_boxes = set(self._canvas.find_withtag("eb_box")) + item_ids = set(self._canvas.find_withtag("current")).intersection(extract_boxes) + if not item_ids: + return False + item_id = list(item_ids)[0] + self._canvas.config(cursor="fleur") + self._mouse_location = ("box", next(int(tag.split("_")[-1]) + for tag in self._canvas.gettags(item_id) + if tag.startswith("face_"))) + return True + + def _check_cursor_rotate(self, event): + """ Check whether the cursor is in an area to rotate the extract box. + + If it is, set the appropriate cursor type and set :attr:`_mouse_location` to + ("rotate", `face index`) + + Notes + ----- + This code is executed after the check has been completed to see if the mouse is inside + the extract box. For this reason, we don't bother running a check to see if the mouse + is inside the box, as this code will never run if that is the case. + + Parameters + ---------- + event: :class:`tkinter.Event` + The current tkinter mouse event + + Returns + ------- + bool + ``True`` if cursor is over a rotate point otherwise ``False`` + """ + distance = 30 + boxes = np.array([np.array(self._canvas.coords(item_id)).reshape(4, 2) + for item_id in self._canvas.find_withtag("eb_box") + if self._canvas.itemcget(item_id, "state") != "hidden"]) + position = np.array((event.x, event.y)).astype("float32") + for face_idx, points in enumerate(boxes): + if any(np.all(position > point - distance) and np.all(position < point + distance) + for point in points): + self._canvas.config(cursor="exchange") + self._mouse_location = ("rotate", face_idx) + return True + return False + + # Mouse click actions + def set_mouse_click_actions(self): + """ Add context menu to OS specific right click action. """ + super().set_mouse_click_actions() + self._canvas.bind("" if platform.system() == "Darwin" else "", + self._context_menu) + + def _drag_start(self, event): + """ The action to perform when the user starts clicking and dragging the mouse. + + Selects the correct extract box action based on the initial cursor position. + + Parameters + ---------- + event: :class:`tkinter.Event` + The tkinter mouse event. + """ + if self._mouse_location is None: + self._drag_data = dict() + self._drag_callback = None + return + self._drag_data["current_location"] = np.array((event.x, event.y)) + callback = dict(anchor=self._resize, rotate=self._rotate, box=self._move) + self._drag_callback = callback[self._mouse_location[0]] + + def _drag_stop(self, event): # pylint: disable=unused-argument + """ Trigger a viewport thumbnail update on click + drag release + + Parameters + ---------- + event: :class:`tkinter.Event` + The tkinter mouse event. Required but unused. + """ + if self._mouse_location is None: + return + self._det_faces.update.post_edit_trigger(self._globals.frame_index, + self._mouse_location[1]) + + def _move(self, event): + """ Updates the underlying detected faces landmarks based on mouse dragging delta, + which moves the Extract box on a drag event. + + Parameters + ---------- + event: :class:`tkinter.Event` + The tkinter mouse event. + """ + if not self._drag_data: + return + shift_x = event.x - self._drag_data["current_location"][0] + shift_y = event.y - self._drag_data["current_location"][1] + scaled_shift = self.scale_from_display(np.array((shift_x, shift_y)), do_offset=False) + self._det_faces.update.landmarks(self._globals.frame_index, + self._mouse_location[1], + *scaled_shift) + self._drag_data["current_location"] = (event.x, event.y) + + def _resize(self, event): + """ Resizes the landmarks contained within an extract box on a corner anchor drag event. + + Parameters + ---------- + event: :class:`tkinter.Event` + The tkinter mouse event. + """ + face_idx = self._mouse_location[1] + face_tag = "eb_box_face_{}".format(face_idx) + position = np.array((event.x, event.y)) + box = np.array(self._canvas.coords(face_tag)) + center = np.array((sum(box[0::2]) / 4, sum(box[1::2]) / 4)) + if not self._check_in_bounds(center, box, position): + logger.trace("Drag out of bounds. Not updating") + self._drag_data["current_location"] = position + return + + start = self._drag_data["current_location"] + distance = ((np.linalg.norm(center - start) - np.linalg.norm(center - position)) + * get_config().scaling_factor) + size = ((box[2] - box[0]) ** 2 + (box[3] - box[1]) ** 2) ** 0.5 + scale = 1 - (distance / size) + logger.trace("face_index: %s, center: %s, start: %s, position: %s, distance: %s, " + "size: %s, scale: %s", face_idx, center, start, position, distance, size, + scale) + if size * scale < 20: + # Don't over shrink the box + logger.trace("Box would size to less than 20px. Not updating") + self._drag_data["current_location"] = position + return + + self._det_faces.update.landmarks_scale(self._globals.frame_index, + face_idx, + scale, + self.scale_from_display(center)) + self._drag_data["current_location"] = position + + def _check_in_bounds(self, center, box, position): + """ Ensure that a resize drag does is not going to cross the center point from it's initial + corner location. + + Parameters + ---------- + center: :class:`numpy.ndarray` + The (`x`, `y`) center point of the face extract box + box: :class:`numpy.ndarray` + The canvas coordinates of the extract box polygon's corners + position: : class:`numpy.ndarray` + The current (`x`, `y`) position of the mouse cursor + + Returns + ------- + bool + ``True`` if the drag operation does not cross the center point otherwise ``False`` + """ + # Generate lines that span the full frame (x and y) along the center point + center_x = np.array(((center[0], 0), (center[0], self._globals.frame_display_dims[1]))) + center_y = np.array(((0, center[1]), (self._globals.frame_display_dims[0], center[1]))) + + # Generate a line coming from the current corner location to the current cursor position + full_line = np.array((box[self._mouse_location[2] * 2:self._mouse_location[2] * 2 + 2], + position)) + logger.trace("center: %s, center_x_line: %s, center_y_line: %s, full_line: %s", + center, center_x, center_y, full_line) + + # Check whether any of the generated lines intersect + for line in (center_x, center_y): + if (self._is_ccw(full_line[0], *line) != self._is_ccw(full_line[1], *line) and + self._is_ccw(*full_line, line[0]) != self._is_ccw(*full_line, line[1])): + logger.trace("line: %s crosses center: %s", full_line, center) + return False + return True + + @staticmethod + def _is_ccw(point_a, point_b, point_c): + """ Check whether 3 points are counter clockwise from each other. + + Parameters + ---------- + point_a: :class:`numpy.ndarray` + The first (`x`, `y`) point to check for counter clockwise ordering + point_b: :class:`numpy.ndarray` + The second (`x`, `y`) point to check for counter clockwise ordering + point_c: :class:`numpy.ndarray` + The third (`x`, `y`) point to check for counter clockwise ordering + + Returns + ------- + bool + ``True`` if the 3 points are provided in counter clockwise order otherwise ``False`` + """ + return ((point_c[1] - point_a[1]) * (point_b[0] - point_a[0]) > + (point_b[1] - point_a[1]) * (point_c[0] - point_a[0])) + + def _rotate(self, event): + """ Rotates the landmarks contained within an extract box on a corner rotate drag event. + + Parameters + ---------- + event: :class:`tkinter.Event` + The tkinter mouse event. + """ + face_idx = self._mouse_location[1] + face_tag = "eb_box_face_{}".format(face_idx) + box = np.array(self._canvas.coords(face_tag)) + position = np.array((event.x, event.y)) + + center = np.array((sum(box[0::2]) / 4, sum(box[1::2]) / 4)) + init_to_center = self._drag_data["current_location"] - center + new_to_center = position - center + angle = np.rad2deg(np.arctan2(*new_to_center) - np.arctan2(*init_to_center)) + logger.trace("face_index: %s, box: %s, center: %s, init_to_center: %s, new_to_center: %s" + "center: %s, angle: %s", face_idx, box, center, init_to_center, new_to_center, + center, angle) + + self._det_faces.update.landmarks_rotate(self._globals.frame_index, + face_idx, + angle, + self.scale_from_display(center)) + self._drag_data["current_location"] = position + + def _get_scale(self): + """ Obtain the scaling for the extract box resize """ + + def _context_menu(self, event): + """ Create a right click context menu to delete the alignment that is being + hovered over. """ + if self._mouse_location is None or self._mouse_location[0] != "box": + return + self._right_click_menu.popup(event) + + def _delete_current_face(self, *args): # pylint:disable=unused-argument + """ Called by the right click delete event. Deletes the face that the mouse is currently + over. + + Parameters + ---------- + args: tuple (unused) + The event parameter is passed in by the hot key binding, so args is required + """ + if self._mouse_location is None or self._mouse_location[0] != "box": + return + self._det_faces.update.delete(self._globals.frame_index, self._mouse_location[1]) diff --git a/tools/manual/frameviewer/editor/landmarks.py b/tools/manual/frameviewer/editor/landmarks.py new file mode 100644 index 00000000..c2d2b346 --- /dev/null +++ b/tools/manual/frameviewer/editor/landmarks.py @@ -0,0 +1,457 @@ +#!/usr/bin/env python3 +""" Landmarks Editor and Landmarks Mesh viewer for the manual adjustments tool """ +import numpy as np + +from ._base import Editor, logger + + +class Landmarks(Editor): + """ The Landmarks Editor. + + Adjust individual landmark points and re-generate Extract Box. + + Parameters + ---------- + canvas: :class:`tkinter.Canvas` + The canvas that holds the image and annotations + detected_faces: :class:`~tools.manual.detected_faces.DetectedFaces` + The _detected_faces data for this manual session + """ + def __init__(self, canvas, detected_faces): + control_text = ("Landmark Point Editor\nEdit the individual landmark points.\n\n" + " - Click and drag individual points to relocate.\n" + " - Draw a box to select multiple points to relocate.") + self._selection_box = canvas.create_rectangle(0, 0, 0, 0, + dash=(2, 4), + state="hidden", + outline="gray", + fill="blue", + stipple="gray12") + super().__init__(canvas, detected_faces, control_text) + # Clear selection box on an editor or frame change + self._canvas._tk_action_var.trace("w", lambda *e: self._reset_selection()) + self._globals.tk_frame_index.trace("w", lambda *e: self._reset_selection()) + + def _add_actions(self): + """ Add the optional action buttons to the viewer. Current actions are Point, Select + and Zoom. """ + self._add_action("magnify", "zoom", "Magnify/Demagnify the View", group=None, hotkey="M") + self._actions["magnify"]["tk_var"].trace("w", self._toggle_zoom) + + # CALLBACKS + def _toggle_zoom(self, *args): # pylint:disable=unused-argument + """ Clear any selections when switching mode and perform an update. + + Parameters + ---------- + args: tuple + tkinter callback arguments. Required but unused. + """ + self._reset_selection() + self._globals.tk_update.set(True) + + def _reset_selection(self, event=None): # pylint:disable=unused-argument + """ Reset the selection box and the selected landmark annotations. """ + self._canvas.itemconfig("lm_selected", outline=self._control_color) + self._canvas.dtag("lm_selected") + self._canvas.itemconfig(self._selection_box, + stipple="gray12", + fill="blue", + outline="gray", + state="hidden") + self._canvas.coords(self._selection_box, 0, 0, 0, 0) + self._drag_data = dict() + if event is not None: + self._drag_start(event) + + def update_annotation(self): + """ Get the latest Landmarks points and update. """ + zoomed_offset = self._zoomed_roi[:2] + for face_idx, face in enumerate(self._face_iterator): + face_index = self._globals.face_index if self._globals.is_zoomed else face_idx + if self._globals.is_zoomed: + landmarks = face.aligned_landmarks + zoomed_offset + # Hide all landmarks and only display selected + self._canvas.itemconfig("lm_dsp", state="hidden") + self._canvas.itemconfig("lm_dsp_face_{}".format(face_index), state="normal") + else: + landmarks = self._scale_to_display(face.landmarks_xy) + for lm_idx, landmark in enumerate(landmarks): + self._display_landmark(landmark, face_index, lm_idx) + self._label_landmark(landmark, face_index, lm_idx) + self._grab_landmark(landmark, face_index, lm_idx) + logger.trace("Updated landmark annotations") + + def _display_landmark(self, bounding_box, face_index, landmark_index): + """ Add an individual landmark display annotation to the canvas. + + Parameters + ---------- + bounding_box: :class:`numpy.ndarray` + The (left, top), (right, bottom) (x, y) coordinates of the oval bounding box for this + landmark + face_index: int + The index of the face within the current frame + landmark_index: int + The index point of this landmark + """ + radius = 1 + color = self._control_color + bbox = (bounding_box[0] - radius, bounding_box[1] - radius, + bounding_box[0] + radius, bounding_box[1] + radius) + key = "lm_dsp_{}".format(landmark_index) + kwargs = dict(outline=color, fill=color, width=radius) + self._object_tracker(key, "oval", face_index, bbox, kwargs) + + def _label_landmark(self, bounding_box, face_index, landmark_index): + """ Add a text label for a landmark to the canvas. + + Parameters + ---------- + bounding_box: :class:`numpy.ndarray` + The (left, top), (right, bottom) (x, y) coordinates of the oval bounding box for this + landmark + face_index: int + The index of the face within the current frame + landmark_index: int + The index point of this landmark + """ + if not self._is_active: + return + top_left = np.array(bounding_box[:2]) - 20 + # NB The text must be visible to be able to get the bounding box, so set to hidden + # after the bounding box has been retrieved + + keys = ["lm_lbl_{}".format(landmark_index), "lm_lbl_bg_{}".format(landmark_index)] + text_kwargs = dict(fill="black", font=("Default", 10), text=str(landmark_index + 1)) + bg_kwargs = dict(fill="#ffffea", outline="black") + + text_id = self._object_tracker(keys[0], "text", face_index, top_left, text_kwargs) + bbox = self._canvas.bbox(text_id) + bbox = [bbox[0] - 2, bbox[1] - 2, bbox[2] + 2, bbox[3] + 2] + bg_id = self._object_tracker(keys[1], "rectangle", face_index, bbox, bg_kwargs) + self._canvas.tag_lower(bg_id, text_id) + self._canvas.itemconfig(text_id, state="hidden") + self._canvas.itemconfig(bg_id, state="hidden") + + def _grab_landmark(self, bounding_box, face_index, landmark_index): + """ Add an individual landmark grab anchor to the canvas. + + Parameters + ---------- + bounding_box: :class:`numpy.ndarray` + The (left, top), (right, bottom) (x, y) coordinates of the oval bounding box for this + landmark + face_index: int + The index of the face within the current frame + landmark_index: int + The index point of this landmark + """ + if not self._is_active: + return + radius = 7 + bbox = (bounding_box[0] - radius, bounding_box[1] - radius, + bounding_box[0] + radius, bounding_box[1] + radius) + key = "lm_grb_{}".format(landmark_index) + kwargs = dict(outline="", + fill="", + width=1, + dash=(2, 4)) + self._object_tracker(key, "oval", face_index, bbox, kwargs) + + # << MOUSE HANDLING >> + # Mouse cursor display + def _update_cursor(self, event): + """ Set the cursor action. + + Launch the cursor update action for the currently selected edit mode. + + Parameters + ---------- + event: :class:`tkinter.Event` + The current tkinter mouse event + """ + self._hide_labels() + if self._drag_data: + self._update_cursor_select_mode(event) + else: + objs = self._canvas.find_withtag("lm_grb_face_{}".format(self._globals.face_index) + if self._globals.is_zoomed else "lm_grb") + item_ids = set(self._canvas.find_overlapping(event.x - 6, + event.y - 6, + event.x + 6, + event.y + 6)).intersection(objs) + bboxes = [self._canvas.bbox(idx) for idx in item_ids] + item_id = next((idx for idx, bbox in zip(item_ids, bboxes) + if bbox[0] <= event.x <= bbox[2] and bbox[1] <= event.y <= bbox[3]), + None) + if item_id: + self._update_cursor_point_mode(item_id) + else: + self._canvas.config(cursor="") + self._mouse_location = None + return + + def _hide_labels(self): + """ Clear all landmark text labels from display """ + self._canvas.itemconfig("lm_lbl", state="hidden") + self._canvas.itemconfig("lm_lbl_bg", state="hidden") + self._canvas.itemconfig("lm_grb", fill="", outline="") + + def _update_cursor_point_mode(self, item_id): + """ Update the cursor when the mouse is over an individual landmark's grab anchor. Displays + the landmark label for the landmark under the cursor. Updates :attr:`_mouse_location` with + the current cursor position. + + Parameters + ---------- + item_id: int + The tkinter canvas object id for the landmark point that the cursor is over + """ + self._canvas.itemconfig(item_id, outline="yellow") + tags = self._canvas.gettags(item_id) + face_idx = int(next(tag for tag in tags if tag.startswith("face_")).split("_")[-1]) + lm_idx = int(next(tag for tag in tags if tag.startswith("lm_grb_")).split("_")[-1]) + obj_idx = (face_idx, lm_idx) + + self._canvas.config(cursor="none") + for prefix in ("lm_lbl_", "lm_lbl_bg_"): + tag = "{}{}_face_{}".format(prefix, lm_idx, face_idx) + logger.trace("Displaying: %s tag: %s", self._canvas.type(tag), tag) + self._canvas.itemconfig(tag, state="normal") + self._mouse_location = obj_idx + + def _update_cursor_select_mode(self, event): + """ Update the mouse cursor when in select mode. + + Standard cursor returned when creating a new selection box. Move cursor returned when over + an existing selection box + + Parameters + ---------- + event: :class:`tkinter.Event` + The current tkinter mouse event + """ + bbox = self._canvas.coords(self._selection_box) + if bbox[0] <= event.x <= bbox[2] and bbox[1] <= event.y <= bbox[3]: + self._canvas.config(cursor="fleur") + else: + self._canvas.config(cursor="") + + # Mouse actions + def _drag_start(self, event): + """ The action to perform when the user starts clicking and dragging the mouse. + + The underlying Detected Face's landmark is updated for the point being edited. + + Parameters + ---------- + event: :class:`tkinter.Event` + The tkinter mouse event. + """ + sel_box = self._canvas.coords(self._selection_box) + if self._mouse_location is not None: # Point edit mode + self._drag_data["start_location"] = (event.x, event.y) + self._drag_callback = self._move_point + elif not self._drag_data: # Initial point selection box + self._drag_data["start_location"] = (event.x, event.y) + self._drag_callback = self._select + elif sel_box[0] <= event.x <= sel_box[2] and sel_box[1] <= event.y <= sel_box[3]: + # Move point selection box + self._drag_data["start_location"] = (event.x, event.y) + self._drag_callback = self._move_selection + else: # Reset + self._drag_data = dict() + self._drag_callback = None + self._reset_selection(event) + + def _drag_stop(self, event): # pylint: disable=unused-argument + """ In select mode, call the select mode callback. + + In point mode: trigger a viewport thumbnail update on click + drag release + + If there is drag data, and there are selected points in the drag data then + trigger the selected points stop code. + + Otherwise reset the selection box and return + + Parameters + ---------- + event: :class:`tkinter.Event` + The tkinter mouse event. Required but unused. + """ + if self._mouse_location is not None: # Point edit mode + self._det_faces.update.post_edit_trigger(self._globals.frame_index, + self._mouse_location[0]) + self._mouse_location = None + self._drag_data = dict() + elif self._drag_data and self._drag_data.get("selected", False): + self._drag_stop_selected() + else: + logger.debug("No selected data. Clearing. drag_data: %s", self._drag_data) + self._reset_selection() + + def _drag_stop_selected(self): + """ Action to perform when mouse drag is stopped in selected points editor mode. + + If there is already a selection, update the viewport thumbnail + + If this is a new selection, then obtain the selected points and track + """ + if "face_index" in self._drag_data: # Selected data has been moved + self._det_faces.update.post_edit_trigger(self._globals.frame_index, + self._drag_data["face_index"]) + return + + # This is a new selection + face_idx = set() + landmark_indices = [] + + for item_id in self._canvas.find_withtag("lm_selected"): + tags = self._canvas.gettags(item_id) + face_idx.add(next(int(tag.split("_")[-1]) + for tag in tags if tag.startswith("face_"))) + landmark_indices.append(next(int(tag.split("_")[-1]) + for tag in tags + if tag.startswith("lm_dsp_") and "face" not in tag)) + if len(face_idx) != 1: + logger.trace("Not exactly 1 face in selection. Aborting. Face indices: %s", face_idx) + self._reset_selection() + return + + self._drag_data["face_index"] = face_idx.pop() + self._drag_data["landmarks"] = landmark_indices + self._canvas.itemconfig(self._selection_box, stipple="", fill="", outline="#ffff00") + self._snap_selection_to_points() + + def _snap_selection_to_points(self): + """ Snap the selection box to the selected points. + + As the landmarks are calculated and redrawn, the selection box can drift. This is + particularly true in zoomed mode. The selection box is therefore redrawn to bind just + outside of the selected points. + """ + all_coords = np.array([self._canvas.coords(item_id) + for item_id in self._canvas.find_withtag("lm_selected")]) + mins = np.min(all_coords, axis=0) + maxes = np.max(all_coords, axis=0) + box_coords = [np.min(mins[[0, 2]] - 5), + np.min(mins[[1, 3]] - 5), + np.max(maxes[[0, 2]] + 5), + np.max(maxes[[1, 3]]) + 5] + self._canvas.coords(self._selection_box, *box_coords) + + def _move_point(self, event): + """ Moves the selected landmark point box and updates the underlying landmark on a point + drag event. + + Parameters + ---------- + event: :class:`tkinter.Event` + The tkinter mouse event. + """ + face_idx, lm_idx = self._mouse_location + shift_x = event.x - self._drag_data["start_location"][0] + shift_y = event.y - self._drag_data["start_location"][1] + + if self._globals.is_zoomed: + scaled_shift = np.array((shift_x, shift_y)) + else: + scaled_shift = self.scale_from_display(np.array((shift_x, shift_y)), do_offset=False) + self._det_faces.update.landmark(self._globals.frame_index, + face_idx, + lm_idx, + *scaled_shift, + self._globals.is_zoomed) + self._drag_data["start_location"] = (event.x, event.y) + + def _select(self, event): + """ Create a selection box on mouse drag event when in "select" mode + + Parameters + ---------- + event: :class:`tkinter.Event` + The tkinter mouse event. + """ + if self._canvas.itemcget(self._selection_box, "state") == "hidden": + self._canvas.itemconfig(self._selection_box, state="normal") + coords = (*self._drag_data["start_location"], event.x, event.y) + self._canvas.coords(self._selection_box, *coords) + enclosed = set(self._canvas.find_enclosed(*coords)) + landmarks = set(self._canvas.find_withtag("lm_dsp")) + + for item_id in list(enclosed.intersection(landmarks)): + self._canvas.addtag_withtag("lm_selected", item_id) + self._canvas.itemconfig("lm_selected", outline="#ffff00") + self._drag_data["selected"] = True + + def _move_selection(self, event): + """ Move a selection box and the landmarks contained when in "select" mode and a selection + box has been drawn. """ + shift_x = event.x - self._drag_data["start_location"][0] + shift_y = event.y - self._drag_data["start_location"][1] + if self._globals.is_zoomed: + scaled_shift = np.array((shift_x, shift_y)) + else: + scaled_shift = self.scale_from_display(np.array((shift_x, shift_y)), do_offset=False) + self._canvas.move(self._selection_box, shift_x, shift_y) + + self._det_faces.update.landmark(self._globals.frame_index, + self._drag_data["face_index"], + self._drag_data["landmarks"], + *scaled_shift, + self._globals.is_zoomed) + self._snap_selection_to_points() + self._drag_data["start_location"] = (event.x, event.y) + + +class Mesh(Editor): + """ The Landmarks Mesh Display. + + There are no editing options for Mesh editor. It is purely aesthetic and updated when other + editors are used. + + Parameters + ---------- + canvas: :class:`tkinter.Canvas` + The canvas that holds the image and annotations + detected_faces: :class:`~tools.manual.detected_faces.DetectedFaces` + The _detected_faces data for this manual session + """ + def __init__(self, canvas, detected_faces): + self._landmark_mapping = dict(mouth_inner=(60, 68), + mouth_outer=(48, 60), + right_eyebrow=(17, 22), + left_eyebrow=(22, 27), + right_eye=(36, 42), + left_eye=(42, 48), + nose=(27, 36), + jaw=(0, 17), + chin=(8, 11)) + super().__init__(canvas, detected_faces, None) + + def update_annotation(self): + """ Get the latest Landmarks and update the mesh.""" + key = "mesh" + color = self._control_color + zoomed_offset = self._zoomed_roi[:2] + for face_idx, face in enumerate(self._face_iterator): + face_index = self._globals.face_index if self._globals.is_zoomed else face_idx + if self._globals.is_zoomed: + landmarks = face.aligned_landmarks + zoomed_offset + # Hide all meshes and only display selected + self._canvas.itemconfig("Mesh", state="hidden") + self._canvas.itemconfig("Mesh_face_{}".format(face_index), state="normal") + else: + landmarks = self._scale_to_display(face.landmarks_xy) + logger.trace("Drawing Landmarks Mesh: (landmarks: %s, color: %s)", landmarks, color) + for idx, (segment, val) in enumerate(self._landmark_mapping.items()): + key = "mesh_{}".format(idx) + pts = landmarks[val[0]:val[1]].flatten() + if segment in ("right_eye", "left_eye", "mouth_inner", "mouth_outer"): + kwargs = dict(fill="", outline=color, width=1) + self._object_tracker(key, "polygon", face_index, pts, kwargs) + else: + self._object_tracker(key, "line", face_index, pts, dict(fill=color, width=1)) + # Place mesh as bottom annotation + self._canvas.tag_raise(self.__class__.__name__, "main_image") diff --git a/tools/manual/frameviewer/editor/mask.py b/tools/manual/frameviewer/editor/mask.py new file mode 100644 index 00000000..9a102ae8 --- /dev/null +++ b/tools/manual/frameviewer/editor/mask.py @@ -0,0 +1,544 @@ +#!/usr/bin/env python3 +""" Mask Editor for the manual adjustments tool """ +import tkinter as tk + +import numpy as np +import cv2 +from PIL import Image, ImageTk + +from ._base import ControlPanelOption, Editor, logger + + +class Mask(Editor): + """ The mask Editor. + + Edit a mask in the alignments file. + + Parameters + ---------- + canvas: :class:`tkinter.Canvas` + The canvas that holds the image and annotations + detected_faces: :class:`~tools.manual.detected_faces.DetectedFaces` + The _detected_faces data for this manual session + """ + def __init__(self, canvas, detected_faces): + self._meta = [] + self._tk_faces = [] + self._internal_size = 512 + control_text = ("Mask Editor\nEdit the mask." + "\n - NB: For Landmark based masks (e.g. components/extended) it is " + "better to make sure the landmarks are correct rather than editing the " + "mask directly. Any change to the landmarks after editing the mask will " + "override your manual edits.") + key_bindings = {"[": lambda *e, i=False: self._adjust_brush_radius(increase=i), + "]": lambda *e, i=True: self._adjust_brush_radius(increase=i)} + super().__init__(canvas, detected_faces, + control_text=control_text, key_bindings=key_bindings) + # Bind control click for reverse painting + self._canvas.bind("", self._control_click) + self._mask_type = self._set_tk_mask_change_callback() + self._mouse_location = [ + self._canvas.create_oval(0, 0, 0, 0, outline="black", state="hidden"), False] + + @property + def _opacity(self): + """ float: The mask opacity setting from the control panel from 0.0 - 1.0. """ + annotation = self.__class__.__name__ + return self._annotation_formats[annotation]["mask_opacity"].get() / 100.0 + + @property + def _brush_radius(self): + """ int: The radius of the brush to use as set in control panel options """ + return self._control_vars["brush"]["BrushSize"].get() + + @property + def _edit_mode(self): + """ str: The currently selected edit mode based on optional action button. + One of "draw" or "erase" """ + action = [name for name, option in self._actions.items() + if option["group"] == "paint" and option["tk_var"].get()] + return "draw" if not action else action[0] + + @property + def _cursor_color(self): + """ str: The hex code for the selected cursor color """ + return self._control_vars["brush"]["CursorColor"].get() + + def _add_actions(self): + """ Add the optional action buttons to the viewer. Current actions are Draw, Erase + and Zoom. """ + self._add_action("magnify", "zoom", "Magnify/Demagnify the View", group=None, hotkey="M") + self._add_action("draw", "draw", "Draw Tool", group="paint", hotkey="D") + self._add_action("erase", "erase", "Erase Tool", group="paint", hotkey="E") + self._actions["magnify"]["tk_var"].trace("w", lambda *e: self._globals.tk_update.set(True)) + + def _add_controls(self): + """ Add the mask specific control panel controls. + + Current controls are: + - the mask type to edit + - the size of brush to use + - the cursor display color + """ + masks = sorted(msk.title() for msk in list(self._det_faces.available_masks) + ["None"]) + default = masks[0] if len(masks) == 1 else [mask for mask in masks if mask != "None"][0] + self._add_control(ControlPanelOption("Mask type", + str, + group="Display", + choices=masks, + default=default, + is_radio=True, + helptext="Select which mask to edit")) + self._add_control(ControlPanelOption("Brush Size", + int, + group="Brush", + min_max=(1, 100), + default=10, + rounding=1, + helptext="Set the brush size. ([ - decrease, " + "] - increase)")) + self._add_control(ControlPanelOption("Cursor Color", + str, + group="Brush", + choices="colorchooser", + default="#ffffff", + helptext="Select the brush cursor color.")) + + def _set_tk_mask_change_callback(self): + """ Add a trace to change the displayed mask on a mask type change. """ + var = self._control_vars["display"]["MaskType"] + var.trace("w", lambda *e: self._on_mask_type_change()) + return var.get() + + def _on_mask_type_change(self): + """ Update the displayed mask on a mask type change """ + mask_type = self._control_vars["display"]["MaskType"].get() + if mask_type == self._mask_type: + return + self._meta = dict(position=self._globals.frame_index) + self._mask_type = mask_type + self._globals.tk_update.set(True) + + def hide_annotation(self, tag=None): + """ Clear the mask :attr:`_meta` dict when hiding the annotation. """ + super().hide_annotation() + self._meta = dict() + + def update_annotation(self): + """ Update the mask annotation with the latest mask. """ + position = self._globals.frame_index + if position != self._meta.get("position", -1): + # Reset meta information when moving to a new frame + self._meta = dict(position=position) + key = self.__class__.__name__ + mask_type = self._control_vars["display"]["MaskType"].get().lower() + color = self._control_color[1:] + rgb_color = np.array(tuple(int(color[i:i + 2], 16) for i in (0, 2, 4))) + roi_color = self._annotation_formats["ExtractBox"]["color"].get() + opacity = self._opacity + for idx, face in enumerate(self._face_iterator): + face_idx = self._globals.face_index if self._globals.is_zoomed else idx + mask = face.mask.get(mask_type, None) + if mask is None: + continue + self._set_face_meta_data(mask, face_idx) + self._update_mask_image(key.lower(), face_idx, rgb_color, opacity) + self._update_roi_box(mask, face_idx, roi_color) + + self._canvas.tag_raise(self._mouse_location[0]) # Always keep brush cursor on top + logger.trace("Updated mask annotation") + + def _set_face_meta_data(self, mask, face_index): + """ Set the metadata for the current face if it has changed or is new. + + Parameters + ---------- + mask: :class:`numpy.ndarray` + The one channel mask cropped to the ROI + face_index: int + The index pertaining to the current face + """ + masks = self._meta.get("mask", None) + if masks is not None and len(masks) - 1 == face_index: + logger.trace("Meta information already defined for face: %s", face_index) + return + + logger.debug("Defining meta information for face: %s", face_index) + scale = self._internal_size / mask.mask.shape[0] + self._set_full_frame_meta(mask, scale) + dims = (self._internal_size, self._internal_size) + self._meta.setdefault("mask", []).append(cv2.resize(mask.mask, + dims, + interpolation=cv2.INTER_CUBIC)) + + def _set_full_frame_meta(self, mask, mask_scale): + """ Sets the meta information for displaying the mask in full frame mode. + + Parameters + ---------- + mask: :class:`lib.faces_detect.Mask` + The mask object + mask_scale: float + The scaling factor from the stored mask size to the internal mask size + + Sets the following parameters to :attr:`_meta`: + - roi_mask: the rectangular ROI box from the full frame that contains the original ROI + for the full frame mask + - top_left: The location that the roi_mask should be placed in the display frame + - affine_matrix: The matrix for transposing the mask to a full frame + - interpolator: The cv2 interpolation method to use for transposing mask to a + full frame + - slices: The (`x`, `y`) slice objects required to extract the mask ROI + from the full frame + """ + frame_dims = self._globals.current_frame["display_dims"] + scaled_mask_roi = np.rint(mask.original_roi * + self._globals.current_frame["scale"]).astype("int32") + + # Scale and clip the ROI to fit within display frame boundaries + clipped_roi = scaled_mask_roi.clip(min=(0, 0), max=frame_dims) + + # Obtain min and max points to get ROI as a rectangle + min_max = dict(min=clipped_roi.min(axis=0), max=clipped_roi.max(axis=0)) + + # Create a bounding box rectangle ROI + roi_dims = np.rint((min_max["max"][1] - min_max["min"][1], + min_max["max"][0] - min_max["min"][0])).astype("uint16") + roi = dict(mask=np.zeros(roi_dims, dtype="uint8")[..., None], + corners=np.expand_dims(scaled_mask_roi - min_max["min"], axis=0)) + # Block out areas outside of the actual mask ROI polygon + cv2.fillPoly(roi["mask"], roi["corners"], 255) + logger.trace("Setting Full Frame mask ROI. shape: %s", roi["mask"].shape) + + # obtain the slices for cropping mask from full frame + xy_slices = (slice(int(round(min_max["min"][1])), int(round(min_max["max"][1]))), + slice(int(round(min_max["min"][0])), int(round(min_max["max"][0])))) + + # Adjust affine matrix for internal mask size and display dimensions + adjustments = (np.array([[mask_scale, 0., 0.], [0., mask_scale, 0.]]), + np.array([[1 / self._globals.current_frame["scale"], 0., 0.], + [0., 1 / self._globals.current_frame["scale"], 0.], + [0., 0., 1.]])) + in_matrix = np.dot(adjustments[0], + np.concatenate((mask.affine_matrix, np.array([[0., 0., 1.]])))) + affine_matrix = np.dot(in_matrix, adjustments[1]) + + # Get the size of the mask roi box in the frame + side_sizes = (scaled_mask_roi[1][0] - scaled_mask_roi[0][0], + scaled_mask_roi[1][1] - scaled_mask_roi[0][1]) + mask_roi_size = (side_sizes[0] ** 2 + side_sizes[1] ** 2) ** 0.5 + + self._meta.setdefault("roi_mask", []).append(roi["mask"]) + self._meta.setdefault("affine_matrix", []).append(affine_matrix) + self._meta.setdefault("interpolator", []).append(mask.interpolator) + self._meta.setdefault("slices", []).append(xy_slices) + self._meta.setdefault("top_left", []).append(min_max["min"] + self._canvas.offset) + self._meta.setdefault("mask_roi_size", []).append(mask_roi_size) + + def _update_mask_image(self, key, face_index, rgb_color, opacity): + """ Obtain a mask, overlay over image and add to canvas or update. + + Parameters + ---------- + key: str + The base annotation name for creating tags + face_index: int + The index of the face within the current frame + rgb_color: tuple + The color that the mask should be displayed as + opacity: float + The opacity to apply to the mask + """ + mask = (self._meta["mask"][face_index] * opacity).astype("uint8") + if self._globals.is_zoomed: + display_image = self._update_mask_image_zoomed(mask, rgb_color) + top_left = self._zoomed_roi[:2] + # Hide all masks and only display selected + self._canvas.itemconfig("Mask", state="hidden") + self._canvas.itemconfig("Mask_face_{}".format(face_index), state="normal") + else: + display_image = self._update_mask_image_full_frame(mask, rgb_color, face_index) + top_left = self._meta["top_left"][face_index] + + if len(self._tk_faces) < face_index + 1: + logger.trace("Adding new Photo Image for face index: %s", face_index) + self._tk_faces.append(ImageTk.PhotoImage(display_image)) + elif self._tk_faces[face_index].width() != display_image.width: + logger.trace("Replacing existing Photo Image on width change for face index: %s", + face_index) + self._tk_faces[face_index] = ImageTk.PhotoImage(display_image) + else: + logger.trace("Updating existing image") + self._tk_faces[face_index].paste(display_image) + + self._object_tracker(key, + "image", + face_index, + top_left, + dict(image=self._tk_faces[face_index], anchor=tk.NW)) + + def _update_mask_image_zoomed(self, mask, rgb_color): + """ Update the mask image when zoomed in. + + Parameters + ---------- + mask: :class:`numpy.ndarray` + The raw mask + rgb_color: tuple + The rgb color selected for the mask + + Returns + ------- + :class: `PIL.Image` + The zoomed mask image formatted for display + """ + rgb = np.tile(rgb_color, self._zoomed_dims + (1, )).astype("uint8") + mask = cv2.resize(mask, + tuple(reversed(self._zoomed_dims)), + interpolation=cv2.INTER_CUBIC)[..., None] + rgba = np.concatenate((rgb, mask), axis=2) + return Image.fromarray(rgba) + + def _update_mask_image_full_frame(self, mask, rgb_color, face_index): + """ Update the mask image when in full frame view. + + Parameters + ---------- + mask: :class:`numpy.ndarray` + The raw mask + rgb_color: tuple + The rgb color selected for the mask + face_index: int + The index of the face being displayed + + Returns + ------- + :class: `PIL.Image` + The full frame mask image formatted for display + """ + frame_dims = self._globals.current_frame["display_dims"] + frame = np.zeros(frame_dims + (1, ), dtype="uint8") + interpolator = self._meta["interpolator"][face_index] + slices = self._meta["slices"][face_index] + mask = cv2.warpAffine(mask, + self._meta["affine_matrix"][face_index], + frame_dims, + frame, + flags=cv2.WARP_INVERSE_MAP | interpolator, + borderMode=cv2.BORDER_CONSTANT)[slices[0], slices[1]][..., None] + rgb = np.tile(rgb_color, mask.shape).astype("uint8") + rgba = np.concatenate((rgb, np.minimum(mask, self._meta["roi_mask"][face_index])), axis=2) + return Image.fromarray(rgba) + + def _update_roi_box(self, mask, face_index, color): + """ Update the region of interest box for the current mask. + + mask: :class:`~lib.faces_detect.Mask` + The current mask object to create an ROI box for + face_index: int + The index of the face within the current frame + color: str + The hex color code that the mask should be displayed as + """ + if self._globals.is_zoomed: + roi = self._zoomed_roi + box = np.array((roi[0], roi[1], roi[2], roi[1], roi[2], roi[3], roi[0], roi[3])) + else: + box = self._scale_to_display(mask.original_roi).flatten() + top_left = box[:2] - 10 + kwargs = dict(fill=color, font=("Default", 20, "bold"), text=str(face_index)) + self._object_tracker("mask_text", "text", face_index, top_left, kwargs) + kwargs = dict(fill="", outline=color, width=1) + self._object_tracker("mask_roi", "polygon", face_index, box, kwargs) + if self._globals.is_zoomed: + # Raise box above zoomed image + self._canvas.tag_raise("mask_roi_face_{}".format(face_index)) + + # << MOUSE HANDLING >> + # Mouse cursor display + def _update_cursor(self, event): + """ Set the cursor action. + + Update :attr:`_mouse_location` with the current cursor position and display appropriate + icon. + + Checks whether the mouse is over a mask ROI box and pops the paint icon. + + Parameters + ---------- + event: :class:`tkinter.Event` + The current tkinter mouse event + """ + roi_boxes = self._canvas.find_withtag("mask_roi") + item_ids = set(self._canvas.find_withtag("current")).intersection(roi_boxes) + if not item_ids: + self._canvas.config(cursor="") + self._canvas.itemconfig(self._mouse_location[0], state="hidden") + self._mouse_location[1] = None + return + item_id = list(item_ids)[0] + tags = self._canvas.gettags(item_id) + face_idx = int(next(tag for tag in tags if tag.startswith("face_")).split("_")[-1]) + + radius = self._brush_radius + coords = (event.x - radius, event.y - radius, event.x + radius, event.y + radius) + self._canvas.config(cursor="none") + self._canvas.coords(self._mouse_location[0], *coords) + self._canvas.itemconfig(self._mouse_location[0], + state="normal", + outline=self._cursor_color) + self._mouse_location[1] = face_idx + self._canvas.update_idletasks() + + def _control_click(self, event): + """ The action to perform when the user starts clicking and dragging the mouse whilst + pressing the control button. + + For editing the mask this will activate the opposite action than what is currently selected + (e.g. it will erase if draw is set and it will draw if erase is set) + + Parameters + ---------- + event: :class:`tkinter.Event` + The tkinter mouse event. + """ + self._drag_start(event, control_click=True) + + def _drag_start(self, event, control_click=False): # pylint:disable=arguments-differ + """ The action to perform when the user starts clicking and dragging the mouse. + + Paints on the mask with the appropriate draw or erase action. + + Parameters + ---------- + event: :class:`tkinter.Event` + The tkinter mouse event. + control_click: bool, optional + Indicates whether the control button is depressed when drag has commenced. If ``True`` + then the opposite of the selected action is performed. Default: ``False`` + """ + face_idx = self._mouse_location[1] + if face_idx is None: + self._drag_data = dict() + self._drag_callback = None + else: + self._drag_data["starting_location"] = np.array((event.x, event.y)) + self._drag_data["control_click"] = control_click + self._drag_data["color"] = np.array(tuple(int(self._control_color[1:][i:i + 2], 16) + for i in (0, 2, 4))) + self._drag_data["opacity"] = self._opacity + self._drag_callback = self._paint + + def _paint(self, event): + """ Paint or erase from Mask and update cursor on click and drag. + + Parameters + ---------- + event: :class:`tkinter.Event` + The tkinter mouse event. + """ + face_idx = self._mouse_location[1] + line = np.array((self._drag_data["starting_location"], (event.x, event.y))) + line, scale = self._transform_points(face_idx, line) + brush_radius = int(round(self._brush_radius * scale)) + color = 0 if self._edit_mode == "erase" else 255 + # Reverse action on control click + color = abs(color - 255) if self._drag_data["control_click"] else color + cv2.line(self._meta["mask"][face_idx], + tuple(line[0]), + tuple(line[1]), + color, + brush_radius * 2) + self._update_mask_image("mask", + face_idx, + self._drag_data["color"], + self._drag_data["opacity"]) + self._drag_data["starting_location"] = np.array((event.x, event.y)) + self._update_cursor(event) + + def _transform_points(self, face_index, points): + """ Transform the edit points from a full frame or zoomed view back to the mask. + + Parameters + ---------- + face_index: int + The index of the face within the current frame + points: :class:`numpy.ndarray` + The points that are to be translated from the viewer to the underlying + Detected Face + """ + if self._globals.is_zoomed: + offset = self._zoomed_roi[:2] + scale = self._internal_size / self._zoomed_dims[0] + t_points = np.rint((points - offset) * scale).astype("int32").squeeze() + else: + scale = self._internal_size / self._meta["mask_roi_size"][face_index] + t_points = np.expand_dims(points - self._canvas.offset, axis=0) + t_points = cv2.transform(t_points, self._meta["affine_matrix"][face_index]).squeeze() + t_points = np.rint(t_points).astype("int32") + logger.trace("original points: %s, transformed points: %s, scale: %s", + points, t_points, scale) + return t_points, scale + + def _drag_stop(self, event): + """ The action to perform when the user stops clicking and dragging the mouse. + + If a line hasn't been drawn then draw a circle. Update alignments. + + Parameters + ---------- + event: :class:`tkinter.Event` + The tkinter mouse event. Unused but required + """ + if not self._drag_data: + return + face_idx = self._mouse_location[1] + location = np.array(((event.x, event.y), )) + color = 0 if self._edit_mode == "erase" else 255 + # Reverse action on control click + color = abs(color - 255) if self._drag_data["control_click"] else color + if np.array_equal(self._drag_data["starting_location"], location[0]): + points, scale = self._transform_points(face_idx, location) + brush_radius = int(round(self._brush_radius * scale)) + cv2.circle(self._meta["mask"][face_idx], tuple(points), brush_radius, color, + thickness=-1) + self._mask_to_alignments(face_idx) + self._drag_data = dict() + self._update_cursor(event) + + def _mask_to_alignments(self, face_index): + """ Update the annotated mask to alignments. + + Parameters + ---------- + face_index: int + The index of the face in the current frame + """ + mask_type = self._control_vars["display"]["MaskType"].get().lower() + mask = self._meta["mask"][face_index].astype("float32") / 255.0 + self._det_faces.update.mask(self._globals.frame_index, face_index, mask, mask_type) + + def _adjust_brush_radius(self, increase=True): # pylint:disable=unused-argument + """ Adjust the brush radius up or down by 2px. + + Sets the control panel option for brush radius to 2 less or 2 more than its current value + + Parameters + ---------- + increase: bool, optional + ``True`` to increment brush radius, ``False`` to decrement. Default: ``True`` + """ + radius_var = self._control_vars["brush"]["BrushSize"] + current_val = radius_var.get() + new_val = min(100, current_val + 2) if increase else max(1, current_val - 2) + logger.trace("Adjusting brush radius from %s to %s", current_val, new_val) + radius_var.set(new_val) + + delta = new_val - current_val + if delta == 0: + return + current_coords = self._canvas.coords(self._mouse_location[0]) + new_coords = tuple(coord - delta if idx < 2 else coord + delta + for idx, coord in enumerate(current_coords)) + logger.trace("Adjusting brush coordinates from %s to %s", current_coords, new_coords) + self._canvas.coords(self._mouse_location[0], new_coords) diff --git a/tools/manual/frameviewer/frame.py b/tools/manual/frameviewer/frame.py new file mode 100644 index 00000000..e42e581a --- /dev/null +++ b/tools/manual/frameviewer/frame.py @@ -0,0 +1,741 @@ +#!/usr/bin/env python3 +""" The frame viewer section of the manual tool GUI """ +import logging +import tkinter as tk +from tkinter import ttk, TclError + +from functools import partial +from time import time + +from lib.gui.control_helper import set_slider_rounding +from lib.gui.custom_widgets import Tooltip +from lib.gui.utils import get_images + +from .control import Navigation, BackgroundImage +from .editor import (BoundingBox, ExtractBox, Landmarks, Mask, # noqa pylint:disable=unused-import + Mesh, View) + +logger = logging.getLogger(__name__) # pylint: disable=invalid-name + + +class DisplayFrame(ttk.Frame): # pylint:disable=too-many-ancestors + """ The main video display frame (top left section of GUI). + + Parameters + ---------- + parent: :class:`tkinter.PanedWindow` + The paned window that the display frame resides in + tk_globals: :class:`~tools.manual.manual.TkGlobals` + The tkinter variables that apply to the whole of the GUI + detected_faces: :class:`tools.manual.detected_faces.DetectedFaces` + The detected faces stored in the alignments file + """ + def __init__(self, parent, tk_globals, detected_faces): + logger.debug("Initializing %s: (parent: %s, tk_globals: %s, detected_faces: %s)", + self.__class__.__name__, parent, tk_globals, detected_faces) + super().__init__(parent) + + self._globals = tk_globals + self._det_faces = detected_faces + + self._actions_frame = ActionsFrame(self) + main_frame = ttk.Frame(self) + + self._transport_frame = ttk.Frame(main_frame) + self._nav = self._add_nav() + self._navigation = Navigation(self) + self._buttons = self._add_transport() + self._add_transport_tk_trace() + + video_frame = ttk.Frame(main_frame) + video_frame.bind("", self._resize) + + self._canvas = FrameViewer(video_frame, + self._globals, + self._det_faces, + self._actions_frame.actions, + self._actions_frame.tk_selected_action) + + self._actions_frame.add_optional_buttons(self.editors) + + self._transport_frame.pack(side=tk.BOTTOM, padx=5, fill=tk.X) + video_frame.pack(side=tk.TOP, expand=True, fill=tk.BOTH) + main_frame.pack(side=tk.RIGHT, expand=True, fill=tk.BOTH) + self.pack(side=tk.LEFT, anchor=tk.NW, expand=True, fill=tk.BOTH) + + logger.debug("Initialized %s", self.__class__.__name__) + + @property + def _helptext(self): + """ dict: {`name`: `help text`} Helptext lookup for navigation buttons """ + return dict( + play="Play/Pause (SPACE)", + beginning="Go to First Frame (HOME)", + prev="Go to Previous Frame (Z)", + next="Go to Next Frame (X)", + end="Go to Last Frame (END)", + extract="Extract the faces to a folder... (Ctrl+E)", + save="Save the Alignments file (Ctrl+S)", + mode="Filter Frames to only those Containing the Selected Item (F)") + + @property + def _btn_action(self): + """ dict: {`name`: `action`} Command lookup for navigation buttons """ + actions = dict(play=self._navigation.handle_play_button, + beginning=self._navigation.goto_first_frame, + prev=self._navigation.decrement_frame, + next=self._navigation.increment_frame, + end=self._navigation.goto_last_frame, + extract=self._det_faces.extract, + save=self._det_faces.save) + return actions + + @property + def tk_selected_action(self): + """ :class:`tkinter.StringVar`: The variable holding the currently selected action """ + return self._actions_frame.tk_selected_action + + @property + def active_editor(self): + """ :class:`Editor`: The current editor in use based on :attr:`selected_action`. """ + return self._canvas.active_editor + + @property + def editors(self): + """ dict: All of the :class:`Editor` that the canvas holds """ + return self._canvas.editors + + @property + def navigation(self): + """ :class:`~tools.manual.frameviewer.control.Navigation`: Class that handles frame + Navigation and transport. """ + return self._navigation + + @property + def tk_control_colors(self): + """ :dict: Editor key with :class:`tkinter.StringVar` containing the selected color hex + code for each annotation """ + return {key: val["color"].tk_var for key, val in self._canvas.annotation_formats.items()} + + @property + def tk_selected_mask(self): + """ :dict: Editor key with :class:`tkinter.StringVar` containing the selected color hex + code for each annotation """ + return self._canvas.control_tk_vars["Mask"]["display"]["MaskType"] + + @property + def _filter_modes(self): + """ list: The filter modes combo box values """ + return ["All Frames", "Has Face(s)", "No Faces", "Multiple Faces"] + + def _add_nav(self): + """ Add the slider to navigate through frames """ + self._globals.tk_transport_index.trace("w", self._set_frame_index) + max_frame = self._globals.frame_count - 1 + + frame = ttk.Frame(self._transport_frame) + + frame.pack(side=tk.TOP, fill=tk.X, pady=(0, 5)) + lbl_frame = ttk.Frame(frame) + lbl_frame.pack(side=tk.RIGHT) + tbox = ttk.Entry(lbl_frame, + width=7, + textvariable=self._globals.tk_transport_index, + justify=tk.RIGHT) + tbox.pack(padx=0, side=tk.LEFT) + lbl = ttk.Label(lbl_frame, text="/{}".format(max_frame)) + lbl.pack(side=tk.RIGHT) + + cmd = partial(set_slider_rounding, + var=self._globals.tk_transport_index, + d_type=int, + round_to=1, + min_max=(0, max_frame)) + + nav = ttk.Scale(frame, + variable=self._globals.tk_transport_index, + from_=0, + to=max_frame, + command=cmd) + nav.pack(side=tk.LEFT, fill=tk.BOTH, expand=True) + return dict(entry=tbox, scale=nav, label=lbl) + + def _set_frame_index(self, *args): # pylint:disable=unused-argument + """ Set the actual frame index based on current slider position and filter mode. """ + try: + slider_position = self._globals.tk_transport_index.get() + except TclError: + # don't update the slider when the entry box has been cleared of any value + return + frames = self._det_faces.filter.frames_list + actual_position = max(0, min(len(frames) - 1, slider_position)) + if actual_position != slider_position: + self._globals.tk_transport_index.set(actual_position) + frame_idx = frames[actual_position] if frames else -1 + logger.trace("slider_position: %s, frame_idx: %s", actual_position, frame_idx) + self._globals.tk_frame_index.set(frame_idx) + + def _add_transport(self): + """ Add video transport controls """ + frame = ttk.Frame(self._transport_frame) + frame.pack(side=tk.BOTTOM, fill=tk.X) + icons = get_images().icons + buttons = dict() + for action in ("play", "beginning", "prev", "next", "end", "save", "extract", "mode"): + padx = (0, 6) if action in ("play", "prev", "mode") else (0, 0) + side = tk.RIGHT if action in ("extract", "save", "mode") else tk.LEFT + state = ["!disabled"] if action != "save" else ["disabled"] + if action != "mode": + icon = action if action != "extract" else "folder" + wgt = ttk.Button(frame, image=icons[icon], command=self._btn_action[action]) + wgt.state(state) + else: + wgt = self._add_filter_mode_combo(frame) + wgt.pack(side=side, padx=padx) + Tooltip(wgt, text=self._helptext[action]) + buttons[action] = wgt + logger.debug("Transport buttons: %s", buttons) + return buttons + + def _add_transport_tk_trace(self): + """ Add the tkinter variable traces to buttons """ + self._navigation.tk_is_playing.trace("w", self._play) + self._det_faces.tk_unsaved.trace("w", self._toggle_save_state) + + def _add_filter_mode_combo(self, frame): + """ Add the navigation mode combo box to the transport frame """ + self._globals.tk_filter_mode.set("All Frames") + self._globals.tk_filter_mode.trace("w", self._navigation.nav_scale_callback) + nav_frame = ttk.Frame(frame) + lbl = ttk.Label(nav_frame, text="Filter:") + lbl.pack(side=tk.LEFT, padx=(0, 5)) + combo = ttk.Combobox( + nav_frame, + textvariable=self._globals.tk_filter_mode, + state="readonly", + values=self._filter_modes) + combo.pack(side=tk.RIGHT) + return nav_frame + + def cycle_filter_mode(self): + """ Cycle the navigation mode combo entry """ + current_mode = self._globals.filter_mode + idx = (self._filter_modes.index(current_mode) + 1) % len(self._filter_modes) + self._globals.tk_filter_mode.set(self._filter_modes[idx]) + + def set_action(self, key): + """ Set the current action based on keyboard shortcut + + Parameters + ---------- + key: str + The pressed key + """ + # Allow key pad keys for numeric presses + key = key.replace("KP_", "") if key.startswith("KP_") else key + self._actions_frame.on_click(self._actions_frame.key_bindings[key]) + + def _resize(self, event): + """ Resize the image to fit the frame, maintaining aspect ratio """ + framesize = (event.width, event.height) + logger.trace("Resizing video frame. Framesize: %s", framesize) + self._globals.set_frame_display_dims(*framesize) + self._globals.tk_update.set(True) + + # << TRANSPORT >> # + def _play(self, *args, frame_count=None): # pylint:disable=unused-argument + """ Play the video file. """ + start = time() + is_playing = self._navigation.tk_is_playing.get() + icon = "pause" if is_playing else "play" + self._buttons["play"].config(image=get_images().icons[icon]) + + if not is_playing: + logger.debug("Pause detected. Stopping.") + return + + # Populate the filtered frames count on first frame + frame_count = self._det_faces.filter.count if frame_count is None else frame_count + self._navigation.increment_frame(frame_count=frame_count, is_playing=True) + delay = 16 # Cap speed at approx 60fps max. Unlikely to hit, but just in case + duration = int((time() - start) * 1000) + delay = max(1, delay - duration) + self.after(delay, lambda f=frame_count: self._play(f)) + + def _toggle_save_state(self, *args): # pylint:disable=unused-argument + """ Toggle the state of the save button when alignments are updated. """ + state = ["!disabled"] if self._det_faces.tk_unsaved.get() else ["disabled"] + self._buttons["save"].state(state) + + +class ActionsFrame(ttk.Frame): # pylint:disable=too-many-ancestors + """ The left hand action frame holding the action buttons. + + Parameters + ---------- + parent: :class:`DisplayFrame` + The Display frame that the Actions reside in + """ + def __init__(self, parent): + super().__init__(parent) + self.pack(side=tk.LEFT, fill=tk.Y, padx=(2, 4), pady=2) + self._globals = parent._globals + self._det_faces = parent._det_faces + + self._configure_styles() + self._actions = ("View", "BoundingBox", "ExtractBox", "Landmarks", "Mask") + self._initial_action = "View" + self._buttons = self._add_buttons() + self._static_buttons = self._add_static_buttons() + self._selected_action = self._set_selected_action_tkvar() + self._optional_buttons = dict() # Has to be set from parent after canvas is initialized + + @property + def actions(self): + """ tuple: The available action names as a tuple of strings. """ + return self._actions + + @property + def tk_selected_action(self): + """ :class:`tkinter.StringVar`: The variable holding the currently selected action """ + return self._selected_action + + @property + def key_bindings(self): + """ dict: {`key`: `action`}. The mapping of key presses to actions. Keyboard shortcut is + the first letter of each action. """ + return {"F{}".format(idx + 1): action for idx, action in enumerate(self._actions)} + + @property + def _helptext(self): + """ dict: `button key`: `button helptext`. The help text to display for each button. """ + inverse_keybindings = {val: key for key, val in self.key_bindings.items()} + retval = dict(View="View alignments", + BoundingBox="Bounding box editor", + ExtractBox="Location editor", + Mask="Mask editor", + Landmarks="Landmark point editor") + for item in retval: + retval[item] += " ({})".format(inverse_keybindings[item]) + return retval + + def _configure_styles(self): + """ Configure background color for Actions widget """ + style = ttk.Style() + style.configure("actions.TFrame", background='#d3d3d3') + style.configure("actions_selected.TButton", relief="flat", background="#bedaf1") + style.configure("actions_deselected.TButton", relief="flat") + self.config(style="actions.TFrame") + + def _add_buttons(self): + """ Add the action buttons to the Display window. + + Returns + ------- + dict: + The action name and its associated button. + """ + frame = ttk.Frame(self) + frame.pack(side=tk.TOP, fill=tk.Y) + buttons = dict() + for action in self.key_bindings.values(): + if action == self._initial_action: + btn_style = "actions_selected.TButton" + state = (["pressed", "focus"]) + else: + btn_style = "actions_deselected.TButton" + state = (["!pressed", "!focus"]) + + button = ttk.Button(frame, + image=get_images().icons[action.lower()], + command=lambda t=action: self.on_click(t), + style=btn_style) + button.state(state) + button.pack() + Tooltip(button, text=self._helptext[action]) + buttons[action] = button + return buttons + + def on_click(self, action): + """ Click event for all of the main buttons. + + Parameters + ---------- + action: str + The action name for the button that has called this event as exists in :attr:`_buttons` + """ + for title, button in self._buttons.items(): + if action == title: + button.configure(style="actions_selected.TButton") + button.state(["pressed", "focus"]) + else: + button.configure(style="actions_deselected.TButton") + button.state(["!pressed", "!focus"]) + self._selected_action.set(action) + + def _set_selected_action_tkvar(self): + """ Set the tkinter string variable that holds the currently selected editor action. + Add traceback to display or hide editor specific optional buttons. + + Returns + ------- + :class:`tkinter.StringVar + The variable that holds the currently selected action + """ + var = tk.StringVar() + var.set(self._initial_action) + var.trace("w", self._display_optional_buttons) + return var + + def _add_static_buttons(self): + """ Add the buttons to copy alignments from previous and next frames """ + lookup = dict(copy_prev=("Previous", "C"), copy_next=("Next", "V"), reload=("", "R")) + frame = ttk.Frame(self) + frame.pack(side=tk.TOP, fill=tk.Y) + sep = ttk.Frame(frame, height=2, relief=tk.RIDGE) + sep.pack(fill=tk.X, pady=5, side=tk.TOP) + buttons = dict() + tk_frame_index = self._globals.tk_frame_index + for action in ("copy_prev", "copy_next", "reload"): + if action == "reload": + icon = "reload3" + cmd = lambda f=tk_frame_index: self._det_faces.revert_to_saved(f.get()) # noqa + helptext = "Revert to saved Alignments ({})".format(lookup[action][1]) + else: + icon = action + direction = action.replace("copy_", "") + cmd = lambda f=tk_frame_index, d=direction: self._det_faces.update.copy( # noqa + f.get(), d) + helptext = "Copy {} Alignments ({})".format(*lookup[action]) + state = ["!disabled"] if action == "copy_next" else ["disabled"] + button = ttk.Button(frame, + image=get_images().icons[icon], + command=cmd, + style="actions_deselected.TButton") + button.state(state) + button.pack() + Tooltip(button, text=helptext) + buttons[action] = button + self._globals.tk_frame_index.trace("w", self._disable_enable_copy_buttons) + self._globals.tk_update.trace("w", self._disable_enable_reload_button) + return buttons + + def _disable_enable_copy_buttons(self, *args): # pylint: disable=unused-argument + """ Disable or enable the static buttons """ + position = self._globals.frame_index + face_count_per_index = self._det_faces.face_count_per_index + prev_exists = position != -1 and any(count != 0 + for count in face_count_per_index[:position]) + next_exists = position != -1 and any(count != 0 + for count in face_count_per_index[position + 1:]) + states = dict(prev=["!disabled"] if prev_exists else ["disabled"], + next=["!disabled"] if next_exists else ["disabled"]) + for direction in ("prev", "next"): + self._static_buttons["copy_{}".format(direction)].state(states[direction]) + + def _disable_enable_reload_button(self, *args): # pylint: disable=unused-argument + """ Disable or enable the static buttons """ + position = self._globals.frame_index + state = ["!disabled"] if (position != -1 and + self._det_faces.is_frame_updated(position)) else ["disabled"] + self._static_buttons["reload"].state(state) + + def add_optional_buttons(self, editors): + """ Add the optional editor specific action buttons """ + for name, editor in editors.items(): + actions = editor.actions + if not actions: + self._optional_buttons[name] = None + continue + frame = ttk.Frame(self) + sep = ttk.Frame(frame, height=2, relief=tk.RIDGE) + sep.pack(fill=tk.X, pady=5, side=tk.TOP) + seen_groups = set() + for action in actions.values(): + group = action["group"] + if group is not None and group not in seen_groups: + btn_style = "actions_selected.TButton" + state = (["pressed", "focus"]) + action["tk_var"].set(True) + seen_groups.add(group) + else: + btn_style = "actions_deselected.TButton" + state = (["!pressed", "!focus"]) + action["tk_var"].set(False) + button = ttk.Button(frame, + image=get_images().icons[action["icon"]], + style=btn_style) + button.config(command=lambda b=button: self._on_optional_click(b)) + button.state(state) + button.pack() + + helptext = action["helptext"] + hotkey = action["hotkey"] + helptext += "" if hotkey is None else " ({})".format(hotkey.upper()) + Tooltip(button, text=helptext) + self._optional_buttons.setdefault( + name, dict())[button] = dict(hotkey=hotkey, + group=group, + tk_var=action["tk_var"]) + self._optional_buttons[name]["frame"] = frame + self._display_optional_buttons() + + def _on_optional_click(self, button): + """ Click event for all of the optional buttons. + + Parameters + ---------- + button: str + The action name for the button that has called this event as exists in :attr:`_buttons` + """ + options = self._optional_buttons[self._selected_action.get()] + group = options[button]["group"] + for child in options["frame"].winfo_children(): + if child.winfo_class() != "TButton": + continue + child_group = options[child]["group"] + if child == button and group is not None: + child.configure(style="actions_selected.TButton") + child.state(["pressed", "focus"]) + options[child]["tk_var"].set(True) + elif child != button and group is not None and child_group == group: + child.configure(style="actions_deselected.TButton") + child.state(["!pressed", "!focus"]) + options[child]["tk_var"].set(False) + elif group is None and child_group is None: + if child.cget("style") == "actions_selected.TButton": + child.configure(style="actions_deselected.TButton") + child.state(["!pressed", "!focus"]) + options[child]["tk_var"].set(False) + else: + child.configure(style="actions_selected.TButton") + child.state(["pressed", "focus"]) + options[child]["tk_var"].set(True) + + def _display_optional_buttons(self, *args): # pylint:disable=unused-argument + """ Pack or forget the optional buttons depending on active editor """ + self._unbind_optional_hotkeys() + for editor, option in self._optional_buttons.items(): + if option is None: + continue + if editor == self._selected_action.get(): + logger.debug("Displaying optional buttons for '%s'", editor) + option["frame"].pack(side=tk.TOP, fill=tk.Y) + for child in option["frame"].winfo_children(): + if child.winfo_class() != "TButton": + continue + hotkey = option[child]["hotkey"] + if hotkey is not None: + logger.debug("Binding optional hotkey for editor '%s': %s", editor, hotkey) + self.winfo_toplevel().bind(hotkey.lower(), + lambda e, b=child: self._on_optional_click(b)) + elif option["frame"].winfo_ismapped(): + logger.debug("Hiding optional buttons for '%s'", editor) + option["frame"].pack_forget() + + def _unbind_optional_hotkeys(self): + """ Unbind all mapped optional button hotkeys """ + for editor, option in self._optional_buttons.items(): + if option is None or not option["frame"].winfo_ismapped(): + continue + for child in option["frame"].winfo_children(): + if child.winfo_class() != "TButton": + continue + hotkey = option[child]["hotkey"] + if hotkey is not None: + logger.debug("Unbinding optional hotkey for editor '%s': %s", editor, hotkey) + self.winfo_toplevel().unbind(hotkey.lower()) + + +class FrameViewer(tk.Canvas): # pylint:disable=too-many-ancestors + """ Annotation onto tkInter Canvas. + + Parameters + ---------- + parent: :class:`tkinter.ttk.Frame` + The parent frame for the canvas + tk_globals: :class:`~tools.manual.manual.TkGlobals` + The tkinter variables that apply to the whole of the GUI + detected_faces: :class:`AlignmentsData` + The alignments data for this manual session + actions: tuple + The available actions from :attr:`ActionFrame.actions` + tk_action_var: :class:`tkinter.StringVar` + The variable holding the currently selected action + """ + def __init__(self, parent, tk_globals, detected_faces, actions, tk_action_var): + logger.debug("Initializing %s: (parent: %s, tk_globals: %s, detected_faces: %s, " + "actions: %s, tk_action_var: %s)", self.__class__.__name__, + parent, tk_globals, detected_faces, actions, tk_action_var) + super().__init__(parent, bd=0, highlightthickness=0, background="black") + self.pack(side=tk.TOP, fill=tk.BOTH, expand=True, anchor=tk.E) + self._globals = tk_globals + self._det_faces = detected_faces + self._actions = actions + self._tk_action_var = tk_action_var + self._image = BackgroundImage(self) + self._editor_globals = dict(control_tk_vars=dict(), + annotation_formats=dict(), + key_bindings=dict()) + self._max_face_count = 0 + self._editors = self._get_editors() + self._add_callbacks() + self._change_active_editor() + logger.debug("Initialized %s", self.__class__.__name__) + + @property + def selected_action(self): + """str: The name of the currently selected Editor action """ + return self._tk_action_var.get() + + @property + def control_tk_vars(self): + """ dict: dictionary of tkinter variables as populated by the right hand control panel. + Tracking for all control panel variables, for access from all editors. """ + return self._editor_globals["control_tk_vars"] + + @property + def key_bindings(self): + """ dict: dictionary of key bindings for each editor for access from all editors. """ + return self._editor_globals["key_bindings"] + + @property + def annotation_formats(self): + """ dict: The selected formatting options for each annotation """ + return self._editor_globals["annotation_formats"] + + @property + def active_editor(self): + """ :class:`Editor`: The current editor in use based on :attr:`selected_action`. """ + return self._editors[self.selected_action] + + @property + def editors(self): + """ dict: All of the :class:`Editor` objects that exist """ + return self._editors + + @property + def editor_display(self): + """ dict: List of editors and any additional annotations they should display. """ + return dict(View=["BoundingBox", "ExtractBox", "Landmarks", "Mesh"], + BoundingBox=["Mesh"], + ExtractBox=["Mesh"], + Landmarks=["ExtractBox", "Mesh"], + Mask=[]) + + @property + def offset(self): + """ tuple: The (`width`, `height`) offset of the canvas based on the size of the currently + displayed image """ + frame_dims = self._globals.current_frame["display_dims"] + offset_x = (self._globals.frame_display_dims[0] - frame_dims[0]) / 2 + offset_y = (self._globals.frame_display_dims[1] - frame_dims[1]) / 2 + logger.trace("offset_x: %s, offset_y: %s", offset_x, offset_y) + return offset_x, offset_y + + def _get_editors(self): + """ Get the object editors for the canvas. + + Returns + ------ + dict + The {`action`: :class:`Editor`} dictionary of editors for :attr:`_actions` name. + """ + editors = dict() + for editor_name in self._actions + ("Mesh", ): + editor = eval(editor_name)(self, # pylint:disable=eval-used + self._det_faces) + editors[editor_name] = editor + logger.debug(editors) + return editors + + def _add_callbacks(self): + """ Add the callback trace functions to the :class:`tkinter.Variable` s + + Adds callbacks for: + :attr:`_globals.tk_update` Update the display for the current image + :attr:`__tk_action_var` Update the mouse display tracking for current action + """ + self._globals.tk_update.trace("w", self._update_display) + self._tk_action_var.trace("w", self._change_active_editor) + + def _change_active_editor(self, *args): # pylint:disable=unused-argument + """ Update the display for the active editor. + + Hide the annotations that are not relevant for the selected editor. + Set the selected editor's cursor tracking. + + Parameters + ---------- + args: tuple, unused + Required for tkinter callback but unused + """ + to_display = [self.selected_action] + self.editor_display[self.selected_action] + to_hide = [editor for editor in self._editors if editor not in to_display] + for editor in to_hide: + self._editors[editor].hide_annotation() + + self.active_editor.bind_mouse_motion() + self.active_editor.set_mouse_click_actions() + self._globals.tk_update.set(True) + + def _update_display(self, *args): # pylint:disable=unused-argument + """ Update the display on frame cache update + + Notes + ----- + A little hacky, but the editors to display or hide are processed in alphabetical + order, so that they are always processed in the same order (for tag lowering and raising) + """ + if not self._globals.tk_update.get(): + return + self._image.refresh(self.active_editor.view_mode) + to_display = sorted([self.selected_action] + self.editor_display[self.selected_action]) + self._hide_additional_faces() + for editor in to_display: + self._editors[editor].update_annotation() + self._bind_unbind_keys() + self._globals.tk_update.set(False) + self.update_idletasks() + + def _hide_additional_faces(self): + """ Hide additional faces if the number of faces on the canvas reduces on a frame + change. """ + if self._globals.is_zoomed: + current_face_count = 1 + elif self._globals.frame_index == -1: + current_face_count = 0 + else: + current_face_count = len(self._det_faces.current_faces[self._globals.frame_index]) + + if current_face_count > self._max_face_count: + # Most faces seen to date so nothing to hide. Update max count and return + logger.debug("Incrementing max face count from: %s to: %s", + self._max_face_count, current_face_count) + self._max_face_count = current_face_count + return + for idx in range(current_face_count, self._max_face_count): + tag = "face_{}".format(idx) + if any(self.itemcget(item_id, "state") != "hidden" + for item_id in self.find_withtag(tag)): + logger.debug("Hiding face tag '%s'", tag) + self.itemconfig(tag, state="hidden") + + def _bind_unbind_keys(self): + """ Bind or unbind this editor's hotkeys depending on whether it is active. """ + unbind_keys = [key for key, binding in self.key_bindings.items() + if binding["bound_to"] is not None + and binding["bound_to"] != self.selected_action] + for key in unbind_keys: + logger.debug("Unbinding key '%s'", key) + self.winfo_toplevel().unbind(key) + self.key_bindings[key]["bound_to"] = None + + bind_keys = {key: binding[self.selected_action] + for key, binding in self.key_bindings.items() + if self.selected_action in binding + and binding["bound_to"] != self.selected_action} + for key, method in bind_keys.items(): + logger.debug("Binding key '%s' to method %s", key, method) + self.winfo_toplevel().bind(key, method) + self.key_bindings[key]["bound_to"] = self.selected_action diff --git a/tools/manual/manual.py b/tools/manual/manual.py new file mode 100644 index 00000000..f55a6a74 --- /dev/null +++ b/tools/manual/manual.py @@ -0,0 +1,845 @@ +#!/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. """ +import logging +import os +import sys +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 +from lib.multithreading import MultiThread +from lib.utils import _video_extensions +from plugins.extract.pipeline import Extractor, ExtractMedia + +from .detected_faces import DetectedFaces, ThumbsCreator +from .faceviewer.frame import FacesFrame +from .frameviewer.frame import DisplayFrame + +logger = logging.getLogger(__name__) # pylint: disable=invalid-name + + +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._initialize_tkinter() + self._globals = TkGlobals(arguments.frames) + + extractor = Aligner(self._globals) + self._detected_faces = DetectedFaces(self._globals, + arguments.alignments_path, + arguments.frames, + extractor) + + video_meta_data = self._detected_faces.video_meta_data + loader = FrameLoader(self._globals, arguments.frames, video_meta_data) + + self._detected_faces.load_faces() + self._containers = self._create_containers() + self._wait_for_threads(extractor, loader, video_meta_data) + 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("", self._handle_key_press) + self._set_initial_layout() + logger.debug("Initialized %s", self.__class__.__name__) + + def _wait_for_threads(self, extractor, loader, video_meta_data): + """ 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 + video_meta_data: dict + The video meta data that exists within the alignments file + + 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 any(val is None for val in video_meta_data.values()): + 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, "") + initialize_config(self, None, 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 = tk.PanedWindow(self, + sashrelief=tk.RIDGE, + sashwidth=2, + sashpad=4, + 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 = dict(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"].sash_place(0, 1, 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 = dict() + 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=18, + 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 = dict(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 = dict() + for name in ("frame_index", "transport_index", "face_index"): + var = tk.IntVar() + var.set(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_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 + """ + def __init__(self, tk_globals): + logger.debug("Initializing: %s (tk_globals: %s)", self.__class__.__name__, tk_globals) + self._globals = tk_globals + self._aligners = {"cv2-dnn": None, "FAN": None, "mask": None} + self._aligner = "FAN" + self._detected_faces = None + self._frame_index = None + self._face_index = None + self._init_thread = self._background_init_aligner() + logger.debug("Initialized: %s", self.__class__.__name__) + + @property + def _in_queue(self): + """ :class:`queue.Queue` - The input queue to the extraction pipeline. """ + return self._aligners[self._aligner].input_queue + + @property + def _feed_face(self): + """ :class:`plugins.extract.pipeline.ExtractMedia`: The current face for feeding into the + aligner, formatted for the pipeline """ + 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: 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") + self._init_thread.check_and_raise_error() + else: + logger.trace("Aligner initialized") + self._init_thread.join() + return not thread_is_alive + + def _background_init_aligner(self): + """ Launch the aligner in a background thread so we can run other tasks whilst + waiting for initialization """ + logger.debug("Launching aligner initialization thread") + thread = MultiThread(self._init_aligner, + thread_count=1, + name="{}.init_aligner".format(self.__class__.__name__)) + thread.start() + logger.debug("Launched aligner initialization thread") + return thread + + def _init_aligner(self): + """ 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 ("mask", "cv2-dnn", "FAN"): + logger.debug("Initializing aligner: %s", model) + plugin = None if model == "mask" else model + aligner = Extractor(None, plugin, ["components", "extended"], + multiprocess=True, normalize_method="hist") + 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): + """ 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.faces_detect.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, face_index, aligner): + """ 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: ["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", + 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) + detected_face = next(self._aligners[aligner].detected_faces()).detected_faces[0] + logger.trace("landmarks: %s", detected_face.landmarks_xy) + return detected_face.landmarks_xy + + def get_masks(self, frame_index, face_index): + """ 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 + The updated masks + """ + logger.trace("frame_index: %s, face_index: %s", frame_index, face_index) + self._frame_index = frame_index + self._face_index = face_index + self._aligner = "mask" + self._in_queue.put(self._feed_face) + detected_face = next(self._aligners["mask"].detected_faces()).detected_faces[0] + logger.debug("mask: %s", detected_face.mask) + return detected_face.mask + + def set_normalization_method(self, method): + """ 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: str + The normalization method to use + """ + logger.debug("Setting normalization method to: '%s'", method) + for plugin, aligner in self._aligners.items(): + if plugin == "mask": + continue + aligner.set_aligner_normalization_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="{}.init_frames".format(self.__class__.__name__)) + 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)