From 34b558426e516cd08004a69fd4261a0f7086a32c Mon Sep 17 00:00:00 2001 From: torzdf <36920800+torzdf@users.noreply.github.com> Date: Tue, 17 Jan 2023 15:03:29 +0000 Subject: [PATCH] bugfix: - preview tool - Prevent tkinter variables from exiting main thread refactor: - preview tool - Split module to smaller sub-modules + docs, locales Typing: - tools.preview.cli Unit test: - tools.preview.viewer --- docs/full/tests/tests.rst | 1 + docs/full/tests/tools.preview.rst | 15 + docs/full/tests/tools.rst | 14 + docs/full/tools/preview.rst | 44 + docs/full/tools/tools.rst | 26 +- lib/gui/utils/config.py | 35 +- locales/tools.preview.pot | 60 +- tests/tools/__init__.py | 0 tests/tools/preview/__init__.py | 0 tests/tools/preview/viewer_test.py | 480 ++++++++++ tools/preview/cli.py | 22 +- tools/preview/control_panels.py | 667 ++++++++++++++ tools/preview/preview.py | 1308 +++++----------------------- tools/preview/viewer.py | 295 +++++++ 14 files changed, 1789 insertions(+), 1178 deletions(-) create mode 100644 docs/full/tests/tools.preview.rst create mode 100644 docs/full/tests/tools.rst create mode 100644 docs/full/tools/preview.rst create mode 100644 tests/tools/__init__.py create mode 100644 tests/tools/preview/__init__.py create mode 100644 tests/tools/preview/viewer_test.py create mode 100644 tools/preview/control_panels.py create mode 100644 tools/preview/viewer.py diff --git a/docs/full/tests/tests.rst b/docs/full/tests/tests.rst index 0785dd7..a24f36a 100644 --- a/docs/full/tests/tests.rst +++ b/docs/full/tests/tests.rst @@ -14,3 +14,4 @@ Subpackages :maxdepth: 1 lib + tools diff --git a/docs/full/tests/tools.preview.rst b/docs/full/tests/tools.preview.rst new file mode 100644 index 0000000..7c744b1 --- /dev/null +++ b/docs/full/tests/tools.preview.rst @@ -0,0 +1,15 @@ +*************** +preview package +*************** + +.. contents:: Contents + :local: + +viewer_test module +****************** +Unittests for the :class:`~tools.preview.viewer` module + +.. automodule:: tests.tools.preview.viewer_test + :members: + :undoc-members: + :show-inheritance: \ No newline at end of file diff --git a/docs/full/tests/tools.rst b/docs/full/tests/tools.rst new file mode 100644 index 0000000..881e7c7 --- /dev/null +++ b/docs/full/tests/tools.rst @@ -0,0 +1,14 @@ +************* +tools package +************* + +.. contents:: Contents + :local: + +Subpackages +=========== + +.. toctree:: + :maxdepth: 1 + + tools.preview diff --git a/docs/full/tools/preview.rst b/docs/full/tools/preview.rst new file mode 100644 index 0000000..5c0c76d --- /dev/null +++ b/docs/full/tools/preview.rst @@ -0,0 +1,44 @@ +*************** +preview package +*************** + +.. contents:: Contents + :local: + + +preview module +============== +The Preview Module is the main entry point into the Preview Tool. + +.. automodule:: tools.preview.preview + :members: + :undoc-members: + :show-inheritance: + + +cli module +========== + +.. automodule:: tools.preview.cli + :members: + :undoc-members: + :show-inheritance: + + +control_panels module +===================== + +.. automodule:: tools.preview.control_panels + :members: + :undoc-members: + :show-inheritance: + + +viewer module +============= + +.. automodule:: tools.preview.viewer + :members: + :undoc-members: + :show-inheritance: + diff --git a/docs/full/tools/tools.rst b/docs/full/tools/tools.rst index d14a72a..0381c4b 100644 --- a/docs/full/tools/tools.rst +++ b/docs/full/tools/tools.rst @@ -15,6 +15,7 @@ Subpackages alignments manual + preview sort mask module @@ -32,28 +33,3 @@ model module :members: :undoc-members: :show-inheritance: - -preview module -=============== - -.. rubric:: Module Summary - -.. autosummary:: - :nosignatures: - - ~tools.preview.preview.ActionFrame - ~tools.preview.preview.ConfigFrame - ~tools.preview.preview.ConfigTools - ~tools.preview.preview.FacesDisplay - ~tools.preview.preview.ImagesCanvas - ~tools.preview.preview.OptionsBook - ~tools.preview.preview.Patch - ~tools.preview.preview.Preview - ~tools.preview.preview.Samples - -.. rubric:: Module - -.. automodule:: tools.preview.preview - :members: - :undoc-members: - :show-inheritance: diff --git a/lib/gui/utils/config.py b/lib/gui/utils/config.py index 1c8dd19..3d4096a 100644 --- a/lib/gui/utils/config.py +++ b/lib/gui/utils/config.py @@ -26,8 +26,8 @@ _CONFIG: Optional["Config"] = None def initialize_config(root: tk.Tk, - cli_opts: "CliOptions", - statusbar: "StatusBar") -> Optional["Config"]: + cli_opts: Optional["CliOptions"], + statusbar: Optional["StatusBar"]) -> Optional["Config"]: """ Initialize the GUI Master :class:`Config` and add to global constant. This should only be called once on first GUI startup. Future access to :class:`Config` @@ -37,10 +37,10 @@ def initialize_config(root: tk.Tk, ---------- root: :class:`tkinter.Tk` The root Tkinter object - cli_opts: :class:`lib.gui.options.CliOptions` - The command line options object - statusbar: :class:`lib.gui.custom_widgets.StatusBar` - The GUI Status bar + cli_opts: :class:`lib.gui.options.CliOptions` or ``None`` + The command line options object. Must be provided for main GUI. Must be ``None`` for tools + statusbar: :class:`lib.gui.custom_widgets.StatusBar` or ``None`` + The GUI Status bar. Must be provided for main GUI. Must be ``None`` for tools Returns ------- @@ -145,11 +145,11 @@ class GlobalVariables(): @dataclass class _GuiObjects: """ Data class for commonly accessed GUI Objects """ - cli_opts: "CliOptions" + cli_opts: Optional["CliOptions"] tk_vars: GlobalVariables project: Project tasks: Tasks - status_bar: "StatusBar" + status_bar: Optional["StatusBar"] default_options: Dict[str, Dict[str, Any]] = field(default_factory=dict) command_notebook: Optional["CommandNotebook"] = None @@ -165,12 +165,15 @@ class Config(): ---------- root: :class:`tkinter.Tk` The root Tkinter object - cli_opts: :class:`lib.gui.options.CliOpts` - The command line options object - statusbar: :class:`lib.gui.custom_widgets.StatusBar` - The GUI Status bar + cli_opts: :class:`lib.gui.options.CliOptions` or ``None`` + The command line options object. Must be provided for main GUI. Must be ``None`` for tools + statusbar: :class:`lib.gui.custom_widgets.StatusBar` or ``None`` + The GUI Status bar. Must be provided for main GUI. Must be ``None`` for tools """ - def __init__(self, root: tk.Tk, cli_opts: "CliOptions", statusbar: "StatusBar") -> None: + def __init__(self, + root: tk.Tk, + cli_opts: Optional["CliOptions"], + statusbar: Optional["StatusBar"]) -> None: logger.debug("Initializing %s: (root %s, cli_opts: %s, statusbar: %s)", self.__class__.__name__, root, cli_opts, statusbar) self._default_font = cast(dict, tk.font.nametofont("TkDefaultFont").configure())["family"] @@ -210,6 +213,9 @@ class Config(): @property def cli_opts(self) -> "CliOptions": """ :class:`lib.gui.options.CliOptions`: The command line options for this GUI Session. """ + # This should only be None when a separate tool (not main GUI) is used, at which point + # cli_opts do not exist + assert self._gui_objects.cli_opts is not None return self._gui_objects.cli_opts @property @@ -236,6 +242,9 @@ class Config(): def statusbar(self) -> "StatusBar": """ :class:`lib.gui.custom_widgets.StatusBar`: The GUI StatusBar :class:`tkinter.ttk.Frame`. """ + # This should only be None when a separate tool (not main GUI) is used, at which point + # this statusbar does not exist + assert self._gui_objects.status_bar is not None return self._gui_objects.status_bar @property diff --git a/locales/tools.preview.pot b/locales/tools.preview.pot index 3d98ad6..aa2e650 100644 --- a/locales/tools.preview.pot +++ b/locales/tools.preview.pot @@ -1,72 +1,80 @@ # SOME DESCRIPTIVE TITLE. -# Copyright (C) YEAR ORGANIZATION +# Copyright (C) YEAR THE PACKAGE'S COPYRIGHT HOLDER +# This file is distributed under the same license as the PACKAGE package. # FIRST AUTHOR , YEAR. # +#, fuzzy msgid "" msgstr "" "Project-Id-Version: PACKAGE VERSION\n" -"POT-Creation-Date: 2021-03-10 16:51-0000\n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2023-01-16 12:27+0000\n" "PO-Revision-Date: YEAR-MO-DA HO:MI+ZONE\n" "Last-Translator: FULL NAME \n" "Language-Team: LANGUAGE \n" +"Language: \n" "MIME-Version: 1.0\n" -"Content-Type: text/plain; charset=cp1252\n" +"Content-Type: text/plain; charset=CHARSET\n" "Content-Transfer-Encoding: 8bit\n" -"Generated-By: pygettext.py 1.5\n" - -#: ./tools/preview\cli.py:13 +#: tools/preview/cli.py:14 msgid "This command allows you to preview swaps to tweak convert settings." msgstr "" -#: ./tools/preview\cli.py:22 +#: tools/preview/cli.py:29 msgid "" "Preview tool\n" "Allows you to configure your convert settings with a live preview" msgstr "" -#: ./tools/preview\cli.py:32 ./tools/preview\cli.py:41 -#: ./tools/preview\cli.py:48 +#: tools/preview/cli.py:46 tools/preview/cli.py:55 tools/preview/cli.py:62 msgid "data" msgstr "" -#: ./tools/preview\cli.py:34 -msgid "Input directory or video. Either a directory containing the image files you wish to process or path to a video file." +#: tools/preview/cli.py:48 +msgid "" +"Input directory or video. Either a directory containing the image files you " +"wish to process or path to a video file." msgstr "" -#: ./tools/preview\cli.py:43 -msgid "Path to the alignments file for the input, if not at the default location" +#: tools/preview/cli.py:57 +msgid "" +"Path to the alignments file for the input, if not at the default location" msgstr "" -#: ./tools/preview\cli.py:50 -msgid "Model directory. A directory containing the trained model you wish to process." +#: tools/preview/cli.py:64 +msgid "" +"Model directory. A directory containing the trained model you wish to " +"process." msgstr "" -#: ./tools/preview\cli.py:57 +#: tools/preview/cli.py:71 msgid "Swap the model. Instead of A -> B, swap B -> A" msgstr "" -#: ./tools/preview\preview.py:1303 +#: tools/preview/control_panels.py:496 msgid "Save full config" msgstr "" -#: ./tools/preview\preview.py:1306 +#: tools/preview/control_panels.py:499 msgid "Reset full config to default values" msgstr "" -#: ./tools/preview\preview.py:1309 +#: tools/preview/control_panels.py:502 msgid "Reset full config to saved values" msgstr "" -#: ./tools/preview\preview.py:1453 -msgid "Save {} config" +#: tools/preview/control_panels.py:653 +#, python-brace-format +msgid "Save {title} config" msgstr "" -#: ./tools/preview\preview.py:1456 -msgid "Reset {} config to default values" +#: tools/preview/control_panels.py:656 +#, python-brace-format +msgid "Reset {title} config to default values" msgstr "" -#: ./tools/preview\preview.py:1459 -msgid "Reset {} config to saved values" +#: tools/preview/control_panels.py:659 +#, python-brace-format +msgid "Reset {title} config to saved values" msgstr "" - diff --git a/tests/tools/__init__.py b/tests/tools/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/tools/preview/__init__.py b/tests/tools/preview/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/tools/preview/viewer_test.py b/tests/tools/preview/viewer_test.py new file mode 100644 index 0000000..96a9bb7 --- /dev/null +++ b/tests/tools/preview/viewer_test.py @@ -0,0 +1,480 @@ +#!/usr/bin python3 +""" Pytest unit tests for :mod:`tools.preview.viewer` """ +import tkinter as tk +from tkinter import ttk + +from typing import cast, TYPE_CHECKING +from unittest.mock import MagicMock + +import pytest +import pytest_mock +import numpy as np +from PIL import ImageTk + +from lib.logger import log_setup +log_setup("DEBUG", "", "PyTest, False") # Need to setup logging to avoid trace/verbose errors + +from lib.utils import get_backend # pylint:disable=wrong-import-position # noqa +from tools.preview.viewer import _Faces, FacesDisplay, ImagesCanvas # pylint:disable=wrong-import-position # noqa + +if TYPE_CHECKING: + from lib.align.aligned_face import CenteringType + + +# pylint:disable=protected-access + + +def test__faces(): + """ Test the :class:`~tools.preview.viewer._Faces dataclass initializes correctly """ + faces = _Faces() + assert faces.filenames == [] + assert faces.matrix == [] + assert faces.src == [] + assert faces.dst == [] + + +_PARAMS = [(3, 448), (4, 333), (5, 254), (6, 128)] # columns/face_size +_IDS = [f"cols:{c},size:{s}[{get_backend().upper()}]" for c, s in _PARAMS] + + +class TestFacesDisplay(): + """ Test :class:`~tools.preview.viewer.FacesDisplay """ + _padding = 64 + + def get_faces_display_instance(self, columns: int = 5, face_size: int = 256) -> FacesDisplay: + """ Obtain an instance of :class:`~tools.preview.viewer.FacesDisplay` with the given column + and face size layout. + + Parameters + ---------- + columns: int, optional + The number of columns to display in the viewer, default: 5 + face_size: int, optional + The size of each face image to be displayed in the viewer, default: 256 + + Returns + ------- + :class:`~tools.preview.viewer.FacesDisplay` + An instance of the FacesDisplay class at the given settings + """ + app = MagicMock() + retval = FacesDisplay(app, face_size, self._padding) + retval._faces = _Faces( + matrix=[np.random.rand(2, 3) for _ in range(columns)], + src=[np.random.rand(face_size, face_size, 3) for _ in range(columns)], + dst=[np.random.rand(face_size, face_size, 3) for _ in range(columns)]) + return retval + + def test_init(self) -> None: + """ Test :class:`~tools.preview.viewer.FacesDisplay` __init__ method """ + f_display = self.get_faces_display_instance(face_size=256) + assert f_display._size == 256 + assert f_display._padding == self._padding + assert isinstance(f_display._app, MagicMock) + + assert f_display._display_dims == (1, 1) + assert isinstance(f_display._faces, _Faces) + + assert f_display._centering is None + assert f_display._faces_source.size == 0 + assert f_display._faces_dest.size == 0 + assert f_display._tk_image is None + assert f_display.update_source is False + assert not f_display.source and isinstance(f_display.source, list) + assert not f_display.destination and isinstance(f_display.destination, list) + + @pytest.mark.parametrize("columns, face_size", _PARAMS, ids=_IDS) + def test__total_columns(self, columns: int, face_size: int) -> None: + """ Test :class:`~tools.preview.viewer.FacesDisplay` _total_columns property is correctly + calculated + + Parameters + ---------- + columns: int + The number of columns to display in the viewer + face_size: int + The size of each face image to be displayed in the viewer + """ + f_display = self.get_faces_display_instance(columns, face_size) + f_display.source = [None for _ in range(columns)] # type:ignore + assert f_display._total_columns == columns + + def test_set_centering(self) -> None: + """ Test :class:`~tools.preview.viewer.FacesDisplay` set_centering method """ + f_display = self.get_faces_display_instance() + assert f_display._centering is None + centering: "CenteringType" = "legacy" + f_display.set_centering(centering) + assert f_display._centering == centering + + def test_set_display_dimensions(self) -> None: + """ Test :class:`~tools.preview.viewer.FacesDisplay` set_display_dimensions method """ + f_display = self.get_faces_display_instance() + assert f_display._display_dims == (1, 1) + dimensions = (800, 600) + f_display.set_display_dimensions(dimensions) + assert f_display._display_dims == dimensions + + @pytest.mark.parametrize("columns, face_size", _PARAMS, ids=_IDS) + def test_update_tk_image(self, + columns: int, + face_size: int, + mocker: pytest_mock.MockerFixture) -> None: + """ Test :class:`~tools.preview.viewer.FacesDisplay` update_tk_image method + + Parameters + ---------- + columns: int + The number of columns to display in the viewer + face_size: int + The size of each face image to be displayed in the viewer + mocker: :class:`pytest_mock.MockerFixture` + Mocker for checking _build_faces_image method called + """ + f_display = self.get_faces_display_instance(columns, face_size) + f_display._build_faces_image = cast(MagicMock, mocker.MagicMock()) # type:ignore + f_display._get_scale_size = cast(MagicMock, # type:ignore + mocker.MagicMock(return_value=(128, 128))) + f_display._faces_source = np.zeros((face_size, face_size, 3), dtype=np.uint8) + f_display._faces_dest = np.zeros((face_size, face_size, 3), dtype=np.uint8) + + tk.Tk() # tkinter instance needed for image creation + f_display.update_tk_image() + + f_display._build_faces_image.assert_called_once() + f_display._get_scale_size.assert_called_once() + assert isinstance(f_display._tk_image, ImageTk.PhotoImage) + assert f_display._tk_image.width() == 128 + assert f_display._tk_image.height() == 128 + assert f_display.tk_image == f_display._tk_image # public property test + + @pytest.mark.parametrize("columns, face_size", _PARAMS, ids=_IDS) + def test_get_scale_size(self, columns: int, face_size: int) -> None: + """ Test :class:`~tools.preview.viewer.FacesDisplay` get_scale_size method + + Parameters + ---------- + columns: int + The number of columns to display in the viewer + face_size: int + The size of each face image to be displayed in the viewer + """ + f_display = self.get_faces_display_instance(columns, face_size) + f_display.set_display_dimensions((800, 600)) + + img = np.zeros((face_size, face_size, 3), dtype=np.uint8) + size = f_display._get_scale_size(img) + assert size == (600, 600) + + @pytest.mark.parametrize("columns, face_size", _PARAMS, ids=_IDS) + def test__build_faces_image(self, + columns: int, + face_size: int, + mocker: pytest_mock.MockerFixture) -> None: + """ Test :class:`~tools.preview.viewer.FacesDisplay` _build_faces_image method + + Parameters + ---------- + columns: int + The number of columns to display in the viewer + face_size: int + The size of each face image to be displayed in the viewer + mocker: :class:`pytest_mock.MockerFixture` + Mocker for checking internal methods called + """ + header_size = 32 + + f_display = self.get_faces_display_instance(columns, face_size) + f_display._faces_from_frames = cast(MagicMock, mocker.MagicMock()) # type:ignore + f_display._header_text = cast( # type:ignore + MagicMock, + mocker.MagicMock(return_value=np.random.rand(header_size, face_size * columns, 3))) + f_display._draw_rect = cast(MagicMock, # type:ignore + mocker.MagicMock(side_effect=lambda x: x)) + + # Test full update + f_display.update_source = True + f_display._build_faces_image() + + f_display._faces_from_frames.assert_called_once() + f_display._header_text.assert_called_once() + assert f_display._draw_rect.call_count == columns * 2 # src + dst + assert f_display._faces_source.shape == (face_size + header_size, face_size * columns, 3) + assert f_display._faces_dest.shape == (face_size, face_size * columns, 3) + + f_display._faces_from_frames.reset_mock() + f_display._header_text.reset_mock() + f_display._draw_rect.reset_mock() + + # Test dst update only + f_display.update_source = False + f_display._build_faces_image() + + f_display._faces_from_frames.assert_called_once() + assert not f_display._header_text.called + assert f_display._draw_rect.call_count == columns # dst only + assert f_display._faces_dest.shape == (face_size, face_size * columns, 3) + + @pytest.mark.parametrize("columns, face_size", _PARAMS, ids=_IDS) + def test_faces__from_frames(self, + columns, + face_size, + mocker: pytest_mock.MockerFixture) -> None: + """ Test :class:`~tools.preview.viewer.FacesDisplay` _from_frames method + + Parameters + ---------- + columns: int + The number of columns to display in the viewer + face_size: int + The size of each face image to be displayed in the viewer + mocker: :class:`pytest_mock.MockerFixture` + Mocker for checking _build_faces_image method called + """ + f_display = self.get_faces_display_instance(columns, face_size) + f_display.source = [mocker.MagicMock() for _ in range(3)] + f_display.destination = [np.random.rand(face_size, face_size, 3) for _ in range(3)] + f_display._crop_source_faces = cast(MagicMock, mocker.MagicMock()) # type:ignore + f_display._crop_destination_faces = cast(MagicMock, mocker.MagicMock()) # type:ignore + + # Both src + dst + f_display.update_source = True + f_display._faces_from_frames() + f_display._crop_source_faces.assert_called_once() + f_display._crop_destination_faces.assert_called_once() + + f_display._crop_source_faces.reset_mock() + f_display._crop_destination_faces.reset_mock() + + # Just dst + f_display.update_source = False + f_display._faces_from_frames() + assert not f_display._crop_source_faces.called + f_display._crop_destination_faces.assert_called_once() + + @pytest.mark.parametrize("columns, face_size", _PARAMS, ids=_IDS) + def test__crop_source_faces(self, + columns: int, + face_size: int, + monkeypatch: pytest.MonkeyPatch, + mocker: pytest_mock.MockerFixture) -> None: + """ Test :class:`~tools.preview.viewer.FacesDisplay` _crop_source_faces method + + Parameters + ---------- + columns: int + The number of columns to display in the viewer + face_size: int + The size of each face image to be displayed in the viewer + monkeypatch: :class:`pytest.MonkeyPatch` + For patching the transform_image function + mocker: :class:`pytest_mock.MockerFixture` + Mocker for mocking various internal methods + """ + f_display = self.get_faces_display_instance(columns, face_size) + f_display._centering = "face" + f_display.update_source = True + f_display._faces.src = [] + + transform_image_mock = mocker.MagicMock() + monkeypatch.setattr("tools.preview.viewer.transform_image", transform_image_mock) + + f_display.source = [mocker.MagicMock() for _ in range(columns)] + for idx, mock in enumerate(f_display.source): + assert isinstance(mock, MagicMock) + mock.inbound.detected_faces.__getitem__ = lambda self, x, y=mock: y + mock.aligned.matrix = f"test_matrix_{idx}" + mock.inbound.filename = f"test_filename_{idx}.txt" + + f_display._crop_source_faces() + + assert len(f_display._faces.filenames) == columns + assert len(f_display._faces.matrix) == columns + assert len(f_display._faces.src) == columns + assert not f_display.update_source + assert transform_image_mock.call_count == columns + + for idx in range(columns): + assert f_display._faces.filenames[idx] == f"test_filename_{idx}" + assert f_display._faces.matrix[idx] == f"test_matrix_{idx}" + + @pytest.mark.parametrize("columns, face_size", _PARAMS, ids=_IDS) + def test__crop_destination_faces(self, + columns: int, + face_size: int, + mocker: pytest_mock.MockerFixture) -> None: + """ Test :class:`~tools.preview.viewer.FacesDisplay` _crop_destination_faces method + + Parameters + ---------- + columns: int + The number of columns to display in the viewer + face_size: int + The size of each face image to be displayed in the viewer + mocker: :class:`pytest_mock.MockerFixture` + Mocker for dummying in full frames + """ + f_display = self.get_faces_display_instance(columns, face_size) + f_display._centering = "face" + f_display._faces.dst = [] # empty object and test populated correctly + + f_display.source = [mocker.MagicMock() for _ in range(columns)] + for item in f_display.source: # type ignore + item.inbound.image = np.random.rand(1280, 720, 3) # type:ignore + + f_display._crop_destination_faces() + assert len(f_display._faces.dst) == columns + assert all(f.shape == (face_size, face_size, 3) for f in f_display._faces.dst) + + @pytest.mark.parametrize("columns, face_size", _PARAMS, ids=_IDS) + def test__header_text(self, + columns: int, + face_size: int, + mocker: pytest_mock.MockerFixture) -> None: + """ Test :class:`~tools.preview.viewer.FacesDisplay` _header_text method + + Parameters + ---------- + columns: int + The number of columns to display in the viewer + face_size: int + The size of each face image to be displayed in the viewer + mocker: :class:`pytest_mock.MockerFixture` + Mocker for dummying in cv2 calls + """ + f_display = self.get_faces_display_instance(columns, face_size) + f_display.source = [None for _ in range(columns)] # type:ignore + f_display._faces.filenames = [f"filename_{idx}.png" for idx in range(columns)] + + cv2_mock = mocker.patch("tools.preview.viewer.cv2") + text_width, text_height = (100, 32) + cv2_mock.getTextSize.return_value = [(text_width, text_height), ] + + header_box = f_display._header_text() + assert cv2_mock.getTextSize.call_count == columns + assert cv2_mock.putText.call_count == columns + assert header_box.shape == (face_size // 8, face_size * columns, 3) + + @pytest.mark.parametrize("columns, face_size", _PARAMS, ids=_IDS) + def test__draw_rect_text(self, + columns: int, + face_size: int, + mocker: pytest_mock.MockerFixture) -> None: + """ Test :class:`~tools.preview.viewer.FacesDisplay` _draw_rect method + + Parameters + ---------- + columns: int + The number of columns to display in the viewer + face_size: int + The size of each face image to be displayed in the viewer + mocker: :class:`pytest_mock.MockerFixture` + Mocker for dummying in cv2 calls + """ + f_display = self.get_faces_display_instance(columns, face_size) + cv2_mock = mocker.patch("tools.preview.viewer.cv2") + + image = (np.random.rand(face_size, face_size, 3) * 255.0) + 50 + assert image.max() > 255.0 + output = f_display._draw_rect(image) + cv2_mock.rectangle.assert_called_once() + assert output.max() == 255.0 # np.clip + + +class TestImagesCanvas: + """ Test :class:`~tools.preview.viewer.ImagesCanvas` """ + + @pytest.fixture + def parent(self) -> MagicMock: + """ Mock object to act as the parent widget to the ImagesCanvas + + Returns + -------- + :class:`unittest.mock.MagicMock` + The mocked ttk.PanedWindow widget + """ + retval = MagicMock(spec=ttk.PanedWindow) + retval.tk = retval + retval._w = "mock_ttkPanedWindow" + retval.children = {} + retval.call = retval + retval.createcommand = retval + retval.preview_display = MagicMock(spec=FacesDisplay) + return retval + + @pytest.fixture(name="images_canvas_instance") + def images_canvas_fixture(self, parent) -> ImagesCanvas: + """ Fixture for creating a testing :class:`~tools.preview.viewer.ImagesCanvas` instance + + Parameters + ---------- + parent: :class:`unittest.mock.MagicMock` + The mocked ttk.PanedWindow parent + + Returns + ------- + :class:`~tools.preview.viewer.ImagesCanvas` + The class instance for testing + """ + app = MagicMock() + return ImagesCanvas(app, parent) + + def test_init(self, images_canvas_instance: ImagesCanvas, parent: MagicMock) -> None: + """ Test :class:`~tools.preview.viewer.ImagesCanvas` __init__ method + + Parameters + ---------- + images_canvas_instance: :class:`~tools.preview.viewer.ImagesCanvas` + The class instance to test + parent: :class:`unittest.mock.MagicMock` + The mocked parent ttk.PanedWindow + """ + assert images_canvas_instance._display == parent.preview_display + assert isinstance(images_canvas_instance._canvas, tk.Canvas) + assert images_canvas_instance._canvas.master == images_canvas_instance + assert images_canvas_instance._canvas.winfo_ismapped() + + def test_resize(self, + images_canvas_instance: ImagesCanvas, + parent: MagicMock, + mocker: pytest_mock.MockerFixture) -> None: + """ Test :class:`~tools.preview.viewer.ImagesCanvas` resize method + + Parameters + ---------- + images_canvas_instance: :class:`~tools.preview.viewer.ImagesCanvas` + The class instance to test + parent: :class:`unittest.mock.MagicMock` + The mocked parent ttk.PanedWindow + mocker: :class:`pytest_mock.MockerFixture` + Mocker for dummying in tk calls + """ + event_mock = mocker.MagicMock(spec=tk.Event, width=100, height=200) + images_canvas_instance.reload = cast(MagicMock, mocker.MagicMock()) # type:ignore + + images_canvas_instance._resize(event_mock) + + parent.preview_display.set_display_dimensions.assert_called_once_with((100, 200)) + images_canvas_instance.reload.assert_called_once() + + def test_reload(self, + images_canvas_instance: ImagesCanvas, + parent: MagicMock, + mocker: pytest_mock.MockerFixture) -> None: + """ Test :class:`~tools.preview.viewer.ImagesCanvas` reload method + + Parameters + ---------- + images_canvas_instance: :class:`~tools.preview.viewer.ImagesCanvas` + The class instance to test + parent: :class:`unittest.mock.MagicMock` + The mocked parent ttk.PanedWindow + mocker: :class:`pytest_mock.MockerFixture` + Mocker for dummying in tk calls + """ + itemconfig_mock = mocker.patch.object(tk.Canvas, "itemconfig") + + images_canvas_instance.reload() + + parent.preview_display.update_tk_image.assert_called_once() + itemconfig_mock.assert_called_once() diff --git a/tools/preview/cli.py b/tools/preview/cli.py index 3ee985d..7c32475 100644 --- a/tools/preview/cli.py +++ b/tools/preview/cli.py @@ -1,6 +1,7 @@ #!/usr/bin/env python3 """ Command Line Arguments for tools """ import gettext +from typing import Any, List, Dict from lib.cli.args import FaceSwapArgs from lib.cli.actions import DirOrFileFullPaths, DirFullPaths, FileFullPaths @@ -17,13 +18,26 @@ class PreviewArgs(FaceSwapArgs): """ Class to parse the command line arguments for Preview (Convert Settings) tool """ @staticmethod - def get_info(): - """ Return command information """ + def get_info() -> str: + """ Return command information + + Returns + ------- + str + Top line information about the Preview tool + """ return _("Preview tool\nAllows you to configure your convert settings with a live preview") - def get_argument_list(self): + @staticmethod + def get_argument_list() -> List[Dict[str, Any]]: + """ Put the arguments in a list so that they are accessible from both argparse and gui - argument_list = list() + Returns + ------- + list[dict[str, Any]] + Top command line options for the preview tool + """ + argument_list = [] argument_list.append(dict( opts=("-i", "--input-dir"), action=DirOrFileFullPaths, diff --git a/tools/preview/control_panels.py b/tools/preview/control_panels.py new file mode 100644 index 0000000..9b811d2 --- /dev/null +++ b/tools/preview/control_panels.py @@ -0,0 +1,667 @@ +#!/usr/bin/env python3 +""" Manages the widgets that hold the bottom 'control' area of the preview tool """ +import gettext +import logging +import tkinter as tk + +from tkinter import ttk +from configparser import ConfigParser +from typing import Any, Callable, cast, Dict, List, Optional, TYPE_CHECKING, Union + +from lib.gui.custom_widgets import Tooltip +from lib.gui.control_helper import ControlPanel, ControlPanelOption +from lib.gui.utils import get_images +from plugins.plugin_loader import PluginLoader +from plugins.convert._config import Config + +if TYPE_CHECKING: + from .preview import Preview + +logger = logging.getLogger(__name__) + +# LOCALES +_LANG = gettext.translation("tools.preview", localedir="locales", fallback=True) +_ = _LANG.gettext + + +class ConfigTools(): + """ Tools for loading, saving, setting and retrieving configuration file values. + + Attributes + ---------- + tk_vars: dict + Global tkinter variables. `Refresh` and `Busy` :class:`tkinter.BooleanVar` + """ + def __init__(self) -> None: + self._config = Config(None) + self.tk_vars: Dict[str, Dict[str, Union[tk.BooleanVar, + tk.StringVar, + tk.IntVar, + tk.DoubleVar]]] = {} + self._config_dicts = self._get_config_dicts() # Holds currently saved config + + @property + def config(self) -> Config: + """ :class:`plugins.convert._config.Config` The convert configuration """ + return self._config + + @property + def config_dicts(self) -> Dict[str, Any]: + """ dict: The convert configuration options in dictionary form.""" + return self._config_dicts + + @property + def sections(self) -> List[str]: + """ list: The sorted section names that exist within the convert Configuration options. """ + return sorted(set(plugin.split(".")[0] for plugin in self._config.config.sections() + if plugin.split(".")[0] != "writer")) + + @property + def plugins_dict(self) -> Dict[str, List[str]]: + """ dict: Dictionary of configuration option sections as key with a list of containing + plugins as the value """ + return {section: sorted([plugin.split(".")[1] for plugin in self._config.config.sections() + if plugin.split(".")[0] == section]) + for section in self.sections} + + def update_config(self) -> None: + """ Update :attr:`config` with the currently selected values from the GUI. """ + for section, items in self.tk_vars.items(): + for item, value in items.items(): + try: + new_value = str(value.get()) + except tk.TclError as err: + # When manually filling in text fields, blank values will + # raise an error on numeric data types so return 0 + logger.debug("Error getting value. Defaulting to 0. Error: %s", str(err)) + new_value = str(0) + old_value = self._config.config[section][item] + if new_value != old_value: + logger.trace("Updating config: %s, %s from %s to %s", # type: ignore + section, item, old_value, new_value) + self._config.config[section][item] = new_value + + def _get_config_dicts(self) -> Dict[str, Dict[str, Any]]: + """ Obtain a custom configuration dictionary for convert configuration items in use + by the preview tool formatted for control helper. + + Returns + ------- + dict + Each configuration section as keys, with the values as a dict of option: + :class:`lib.gui.control_helper.ControlOption` pairs. """ + logger.debug("Formatting Config for GUI") + config_dicts: Dict[str, Dict[str, Any]] = {} + for section in self._config.config.sections(): + if section.startswith("writer."): + continue + for key, val in self._config.defaults[section].items(): + if key == "helptext": + config_dicts.setdefault(section, {})[key] = val + continue + cp_option = ControlPanelOption(title=key, + dtype=val["type"], + group=val["group"], + default=val["default"], + initial_value=self._config.get(section, key), + choices=val["choices"], + is_radio=val["gui_radio"], + rounding=val["rounding"], + min_max=val["min_max"], + helptext=val["helptext"]) + self.tk_vars.setdefault(section, {})[key] = cp_option.tk_var + config_dicts.setdefault(section, {})[key] = cp_option + logger.debug("Formatted Config for GUI: %s", config_dicts) + return config_dicts + + def reset_config_to_saved(self, section: Optional[str] = None) -> None: + """ Reset the GUI parameters to their saved values within the configuration file. + + Parameters + ---------- + section: str, optional + The configuration section to reset the values for, If ``None`` provided then all + sections are reset. Default: ``None`` + """ + logger.debug("Resetting to saved config: %s", section) + sections = [section] if section is not None else list(self.tk_vars.keys()) + for config_section in sections: + for item, options in self._config_dicts[config_section].items(): + if item == "helptext": + continue + val = options.value + if val != self.tk_vars[config_section][item].get(): + self.tk_vars[config_section][item].set(val) + logger.debug("Setting %s - %s to saved value %s", config_section, item, val) + logger.debug("Reset to saved config: %s", section) + + def reset_config_to_default(self, section: Optional[str] = None) -> None: + """ Reset the GUI parameters to their default configuration values. + + Parameters + ---------- + section: str, optional + The configuration section to reset the values for, If ``None`` provided then all + sections are reset. Default: ``None`` + """ + logger.debug("Resetting to default: %s", section) + sections = [section] if section is not None else list(self.tk_vars.keys()) + for config_section in sections: + for item, options in self._config_dicts[config_section].items(): + if item == "helptext": + continue + default = options.default + if default != self.tk_vars[config_section][item].get(): + self.tk_vars[config_section][item].set(default) + logger.debug("Setting %s - %s to default value %s", + config_section, item, default) + logger.debug("Reset to default: %s", section) + + def save_config(self, section: Optional[str] = None) -> None: + """ Save the configuration ``.ini`` file with the currently stored values. + + Notes + ----- + We cannot edit the existing saved config as comments tend to get removed, so we create + a new config and populate that. + + Parameters + ---------- + section: str, optional + The configuration section to save, If ``None`` provided then all sections are saved. + Default: ``None`` + """ + logger.debug("Saving %s config", section) + + new_config = ConfigParser(allow_no_value=True) + + for config_section, items in self._config.defaults.items(): + logger.debug("Adding section: '%s')", config_section) + self._config.insert_config_section(config_section, + items["helptext"], + config=new_config) + for item, options in items.items(): + if item == "helptext": + continue # helptext already written at top + if ((section is not None and config_section != section) + or config_section not in self.tk_vars): + # retain saved values that have not been updated + new_opt = self._config.get(config_section, item) + logger.debug("Retaining option: (item: '%s', value: '%s')", item, new_opt) + else: + new_opt = self.tk_vars[config_section][item].get() + logger.debug("Setting option: (item: '%s', value: '%s')", item, new_opt) + + # Set config_dicts value to new saved value + self._config_dicts[config_section][item].set_initial_value(new_opt) + + helptext = self._config.format_help(options["helptext"], is_section=False) + new_config.set(config_section, helptext) + new_config.set(config_section, item, str(new_opt)) + + self._config.config = new_config + self._config.save_config() + logger.info("Saved config: '%s'", self._config.configfile) + + +class BusyProgressBar(): + """ An infinite progress bar for when a thread is running to swap/patch a group of samples """ + def __init__(self, parent: ttk.Frame) -> None: + self._progress_bar = self._add_busy_indicator(parent) + + def _add_busy_indicator(self, parent: ttk.Frame) -> ttk.Progressbar: + """ Place progress bar into bottom bar to indicate when processing. + + Parameters + ---------- + parent: tkinter object + The tkinter object that holds the busy indicator + + Returns + ------- + ttk.Progressbar + A Progress bar to indicate that the Preview tool is busy + """ + logger.debug("Placing busy indicator") + pbar = ttk.Progressbar(parent, mode="indeterminate") + pbar.pack(side=tk.LEFT) + pbar.pack_forget() + return pbar + + def stop(self) -> None: + """ Stop and hide progress bar """ + logger.debug("Stopping busy indicator") + if not self._progress_bar.winfo_ismapped(): + logger.debug("busy indicator already hidden") + return + self._progress_bar.stop() + self._progress_bar.pack_forget() + + def start(self) -> None: + """ Start and display progress bar """ + logger.debug("Starting busy indicator") + if self._progress_bar.winfo_ismapped(): + logger.debug("busy indicator already started") + return + + self._progress_bar.pack(side=tk.LEFT, padx=5, pady=(5, 10), fill=tk.X, expand=True) + self._progress_bar.start(25) + + +class ActionFrame(ttk.Frame): # pylint: disable=too-many-ancestors + """ Frame that holds the left hand side options panel containing the command line options. + + Parameters + ---------- + app: :class:`Preview` + The main tkinter Preview app + parent: tkinter object + The parent tkinter object that holds the Action Frame + """ + def __init__(self, app: 'Preview', parent: ttk.Frame) -> None: + logger.debug("Initializing %s: (app: %s, parent: %s)", + self.__class__.__name__, app, parent) + self._app = app + + super().__init__(parent) + self.pack(side=tk.LEFT, anchor=tk.N, fill=tk.Y) + self._tk_vars: Dict[str, tk.StringVar] = {} + + self._options = dict( + color=app._patch.converter.cli_arguments.color_adjustment.replace("-", "_"), + mask_type=app._patch.converter.cli_arguments.mask_type.replace("-", "_")) + defaults = {opt: self._format_to_display(val) + for opt, val in self._options.items()} + self._busy_bar = self._build_frame(defaults, + app._samples.generate, + app._refresh, + app._samples.available_masks, + app._samples.predictor.has_predicted_mask) + + @property + def convert_args(self) -> Dict[str, Any]: + """ dict: Currently selected Command line arguments from the :class:`ActionFrame`. """ + return {opt if opt != "color" else "color_adjustment": + self._format_from_display(self._tk_vars[opt].get()) + for opt in self._options} + + @property + def busy_progress_bar(self) -> BusyProgressBar: + """ :class:`BusyProgressBar`: The progress bar that appears on the left hand side whilst a + swap/patch is being applied """ + return self._busy_bar + + @staticmethod + def _format_from_display(var: str) -> str: + """ Format a variable from the display version to the command line action version. + + Parameters + ---------- + var: str + The variable name to format + + Returns + ------- + str + The formatted variable name + """ + return var.replace(" ", "_").lower() + + @staticmethod + def _format_to_display(var: str) -> str: + """ Format a variable from the command line action version to the display version. + Parameters + ---------- + var: str + The variable name to format + + Returns + ------- + str + The formatted variable name + """ + return var.replace("_", " ").replace("-", " ").title() + + def _build_frame(self, + defaults: Dict[str, Any], + refresh_callback: Callable[[], None], + patch_callback: Callable[[], None], + available_masks: List[str], + has_predicted_mask: bool) -> BusyProgressBar: + """ Build the :class:`ActionFrame`. + + Parameters + ---------- + defaults: dict + The default command line options + patch_callback: python function + The function to execute when a patch callback is received + refresh_callback: python function + The function to execute when a refresh callback is received + available_masks: list + The available masks that exist within the alignments file + has_predicted_mask: bool + Whether the model was trained with a mask + + Returns + ------- + ttk.Progressbar + A Progress bar to indicate that the Preview tool is busy + """ + logger.debug("Building Action frame") + + bottom_frame = ttk.Frame(self) + bottom_frame.pack(side=tk.BOTTOM, fill=tk.X, anchor=tk.S) + top_frame = ttk.Frame(self) + top_frame.pack(side=tk.TOP, fill=tk.BOTH, anchor=tk.N, expand=True) + + self._add_cli_choices(top_frame, defaults, available_masks, has_predicted_mask) + + busy_indicator = BusyProgressBar(bottom_frame) + self._add_refresh_button(bottom_frame, refresh_callback) + self._add_patch_callback(patch_callback) + self._add_actions(bottom_frame) + logger.debug("Built Action frame") + return busy_indicator + + def _add_cli_choices(self, + parent: ttk.Frame, + defaults: Dict[str, Any], + available_masks: List[str], + has_predicted_mask: bool) -> None: + """ Create :class:`lib.gui.control_helper.ControlPanel` object for the command + line options. + + parent: :class:`ttk.Frame` + The frame to hold the command line choices + defaults: dict + The default command line options + available_masks: list + The available masks that exist within the alignments file + has_predicted_mask: bool + Whether the model was trained with a mask + """ + cp_options = self._get_control_panel_options(defaults, available_masks, has_predicted_mask) + panel_kwargs = dict(blank_nones=False, label_width=10, style="CPanel") + ControlPanel(parent, cp_options, header_text=None, **panel_kwargs) + + def _get_control_panel_options(self, + defaults: Dict[str, Any], + available_masks: List[str], + has_predicted_mask: bool) -> List[ControlPanelOption]: + """ Create :class:`lib.gui.control_helper.ControlPanelOption` objects for the command + line options. + + defaults: dict + The default command line options + available_masks: list + The available masks that exist within the alignments file + has_predicted_mask: bool + Whether the model was trained with a mask + + Returns + ------- + list + The list of `lib.gui.control_helper.ControlPanelOption` objects for the Action Frame + """ + cp_options: List[ControlPanelOption] = [] + for opt in self._options: + if opt == "mask_type": + choices = self._create_mask_choices(defaults, available_masks, has_predicted_mask) + else: + choices = PluginLoader.get_available_convert_plugins(opt, True) + cp_option = ControlPanelOption(title=opt, + dtype=str, + default=defaults[opt], + initial_value=defaults[opt], + choices=choices, + group="Command Line Choices", + is_radio=False) + self._tk_vars[opt] = cp_option.tk_var + cp_options.append(cp_option) + return cp_options + + def _create_mask_choices(self, + defaults: Dict[str, Any], + available_masks: List[str], + has_predicted_mask: bool) -> List[str]: + """ Set the mask choices and default mask based on available masks. + + Parameters + ---------- + defaults: dict + The default command line options + available_masks: list + The available masks that exist within the alignments file + has_predicted_mask: bool + Whether the model was trained with a mask + + Returns + ------- + list + The masks that are available to use from the alignments file + """ + logger.debug("Initial mask choices: %s", available_masks) + if has_predicted_mask: + available_masks += ["predicted"] + if "none" not in available_masks: + available_masks += ["none"] + if self._format_from_display(defaults["mask_type"]) not in available_masks: + logger.debug("Setting default mask to first available: %s", available_masks[0]) + defaults["mask_type"] = available_masks[0] + logger.debug("Final mask choices: %s", available_masks) + return available_masks + + @classmethod + def _add_refresh_button(cls, + parent: ttk.Frame, + refresh_callback: Callable[[], None]) -> None: + """ Add a button to refresh the images. + + Parameters + ---------- + refresh_callback: python function + The function to execute when the refresh button is pressed + """ + btn = ttk.Button(parent, text="Update Samples", command=refresh_callback) + btn.pack(padx=5, pady=5, side=tk.TOP, fill=tk.X, anchor=tk.N) + + def _add_patch_callback(self, patch_callback: Callable[[], None]) -> None: + """ Add callback to re-patch images on action option change. + + Parameters + ---------- + patch_callback: python function + The function to execute when the images require patching + """ + for tk_var in self._tk_vars.values(): + tk_var.trace("w", patch_callback) + + def _add_actions(self, parent: ttk.Frame) -> None: + """ Add Action Buttons to the :class:`ActionFrame` + + Parameters + ---------- + parent: tkinter object + The tkinter object that holds the action buttons + """ + logger.debug("Adding util buttons") + frame = ttk.Frame(parent) + frame.pack(padx=5, pady=(5, 10), side=tk.RIGHT, fill=tk.X, anchor=tk.E) + + for utl in ("save", "clear", "reload"): + logger.debug("Adding button: '%s'", utl) + img = get_images().icons[utl] + if utl == "save": + text = _("Save full config") + action = self._app.config_tools.save_config + elif utl == "clear": + text = _("Reset full config to default values") + action = self._app.config_tools.reset_config_to_default + elif utl == "reload": + text = _("Reset full config to saved values") + action = self._app.config_tools.reset_config_to_saved + + btnutl = ttk.Button(frame, + image=img, + command=action) + btnutl.pack(padx=2, side=tk.RIGHT) + Tooltip(btnutl, text=text, wrap_length=200) + logger.debug("Added util buttons") + + +class OptionsBook(ttk.Notebook): # pylint:disable=too-many-ancestors + """ The notebook that holds the Convert configuration options. + + Parameters + ---------- + parent: tkinter object + The parent tkinter object that holds the Options book + config_tools: :class:`ConfigTools` + Tools for loading and saving configuration files + patch_callback: python function + The function to execute when a patch callback is received + + Attributes + ---------- + config_tools: :class:`ConfigTools` + Tools for loading and saving configuration files + """ + def __init__(self, + parent: ttk.Frame, + config_tools: ConfigTools, + patch_callback: Callable[[], None]) -> None: + logger.debug("Initializing %s: (parent: %s, config: %s)", + self.__class__.__name__, parent, config_tools) + super().__init__(parent) + self.pack(side=tk.RIGHT, anchor=tk.N, fill=tk.BOTH, expand=True) + self.config_tools = config_tools + + self._tabs: Dict[str, Dict[str, Union[ttk.Notebook, ConfigFrame]]] = {} + self._build_tabs() + self._build_sub_tabs() + self._add_patch_callback(patch_callback) + logger.debug("Initialized %s", self.__class__.__name__) + + def _build_tabs(self) -> None: + """ Build the notebook tabs for the each configuration section. """ + logger.debug("Build Tabs") + for section in self.config_tools.sections: + tab = ttk.Notebook(self) + self._tabs[section] = {"tab": tab} + self.add(tab, text=section.replace("_", " ").title()) + + def _build_sub_tabs(self) -> None: + """ Build the notebook sub tabs for each convert section's plugin. """ + for section, plugins in self.config_tools.plugins_dict.items(): + for plugin in plugins: + config_key = ".".join((section, plugin)) + config_dict = self.config_tools.config_dicts[config_key] + tab = ConfigFrame(self, config_key, config_dict) + self._tabs[section][plugin] = tab + text = plugin.replace("_", " ").title() + cast(ttk.Notebook, self._tabs[section]["tab"]).add(tab, text=text) + + def _add_patch_callback(self, patch_callback: Callable[[], None]) -> None: + """ Add callback to re-patch images on configuration option change. + + Parameters + ---------- + patch_callback: python function + The function to execute when the images require patching + """ + for plugins in self.config_tools.tk_vars.values(): + for tk_var in plugins.values(): + tk_var.trace("w", patch_callback) + + +class ConfigFrame(ttk.Frame): # pylint: disable=too-many-ancestors + """ Holds the configuration options for a convert plugin inside the :class:`OptionsBook`. + + Parameters + ---------- + parent: tkinter object + The tkinter object that will hold this configuration frame + config_key: str + The section/plugin key for these configuration options + options: dict + The options for this section/plugin + """ + + def __init__(self, + parent: OptionsBook, + config_key: str, + options: Dict[str, Any]): + logger.debug("Initializing %s", self.__class__.__name__) + super().__init__(parent) + self.pack(side=tk.TOP, fill=tk.BOTH, expand=True) + + self._options = options + + self._action_frame = ttk.Frame(self) + self._action_frame.pack(padx=0, pady=(0, 5), side=tk.BOTTOM, fill=tk.X, anchor=tk.E) + self._add_frame_separator() + + self._build_frame(parent, config_key) + logger.debug("Initialized %s", self.__class__.__name__) + + def _build_frame(self, parent: OptionsBook, config_key: str) -> None: + """ Build the options frame for this command + + Parameters + ---------- + parent: tkinter object + The tkinter object that will hold this configuration frame + config_key: str + The section/plugin key for these configuration options + """ + logger.debug("Add Config Frame") + panel_kwargs = dict(columns=2, option_columns=2, blank_nones=False, style="CPanel") + frame = ttk.Frame(self) + frame.pack(side=tk.TOP, fill=tk.BOTH, expand=True) + cp_options = [opt for key, opt in self._options.items() if key != "helptext"] + ControlPanel(frame, cp_options, header_text=None, **panel_kwargs) + self._add_actions(parent, config_key) + logger.debug("Added Config Frame") + + def _add_frame_separator(self) -> None: + """ Add a separator between top and bottom frames. """ + logger.debug("Add frame seperator") + sep = ttk.Frame(self._action_frame, height=2, relief=tk.RIDGE) + sep.pack(fill=tk.X, pady=5, side=tk.TOP) + logger.debug("Added frame seperator") + + def _add_actions(self, parent: OptionsBook, config_key: str) -> None: + """ Add Action Buttons. + + Parameters + ---------- + parent: tkinter object + The tkinter object that will hold this configuration frame + config_key: str + The section/plugin key for these configuration options + """ + logger.debug("Adding util buttons") + + title = config_key.split(".")[1].replace("_", " ").title() + btn_frame = ttk.Frame(self._action_frame) + btn_frame.pack(padx=5, side=tk.BOTTOM, fill=tk.X) + for utl in ("save", "clear", "reload"): + logger.debug("Adding button: '%s'", utl) + img = get_images().icons[utl] + if utl == "save": + text = _(f"Save {title} config") + action = parent.config_tools.save_config + elif utl == "clear": + text = _(f"Reset {title} config to default values") + action = parent.config_tools.reset_config_to_default + elif utl == "reload": + text = _(f"Reset {title} config to saved values") + action = parent.config_tools.reset_config_to_saved + + btnutl = ttk.Button(btn_frame, + image=img, + command=lambda cmd=action: cmd(config_key)) # type: ignore + btnutl.pack(padx=2, side=tk.RIGHT) + Tooltip(btnutl, text=text, wrap_length=200) + logger.debug("Added util buttons") diff --git a/tools/preview/preview.py b/tools/preview/preview.py index 712b08a..a5e83df 100644 --- a/tools/preview/preview.py +++ b/tools/preview/preview.py @@ -1,38 +1,35 @@ #!/usr/bin/env python3 """ Tool to preview swaps and tweak configuration prior to running a convert """ -from dataclasses import dataclass, field import gettext import logging import random import tkinter as tk from tkinter import ttk -from typing import Any, Callable, cast, Dict, List, Optional, Tuple, TYPE_CHECKING, Union +from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING, Union import os import sys -from configparser import ConfigParser + from threading import Event, Lock, Thread -import cv2 import numpy as np -from PIL import Image, ImageTk -from lib.align import DetectedFace, transform_image +from lib.align import DetectedFace from lib.cli.args import ConvertArgs from lib.gui.utils import get_images, get_config, initialize_config, initialize_images -from lib.gui.custom_widgets import Tooltip -from lib.gui.control_helper import ControlPanel, ControlPanelOption from lib.convert import Converter from lib.utils import FaceswapError from lib.queue_manager import queue_manager from scripts.fsmedia import Alignments, Images from scripts.convert import Predict, ConvertItem -from plugins.plugin_loader import PluginLoader -from plugins.convert._config import Config from plugins.extract.pipeline import ExtractMedia +from .control_panels import ActionFrame, ConfigTools, OptionsBook +from .viewer import FacesDisplay, ImagesCanvas + + if sys.version_info < (3, 8): from typing_extensions import Literal else: @@ -40,8 +37,8 @@ else: if TYPE_CHECKING: from argparse import Namespace - from lib.align.aligned_face import CenteringType from lib.queue_manager import EventQueue + from .control_panels import BusyProgressBar logger = logging.getLogger(__name__) # pylint: disable=invalid-name @@ -70,23 +67,10 @@ class Preview(tk.Tk): # pylint:disable=too-few-public-methods super().__init__() self._config_tools = ConfigTools() self._lock = Lock() - - self._tk_vars: Dict[Literal["refresh", "busy"], - tk.BooleanVar] = dict(refresh=tk.BooleanVar(), busy=tk.BooleanVar()) - for val in self._tk_vars.values(): - val.set(False) - self._display = FacesDisplay(256, 64, self._tk_vars) - - trigger_patch = Event() - self._samples = Samples(arguments, 5, self._display, self._lock, trigger_patch) - self._patch = Patch(arguments, - self._available_masks, - self._samples, - self._display, - self._lock, - trigger_patch, - self._config_tools, - self._tk_vars) + self._dispatcher = Dispatcher(self) + self._display = FacesDisplay(self, 256, 64) + self._samples = Samples(self, arguments, 5) + self._patch = Patch(self, arguments) self._initialize_tkinter() self._image_canvas: Optional[ImagesCanvas] = None @@ -95,12 +79,41 @@ class Preview(tk.Tk): # pylint:disable=too-few-public-methods logger.debug("Initialized %s", self.__class__.__name__) @property - def _available_masks(self) -> List[str]: - """ list: The mask names that are available for every face in the alignments file """ - retval = [key - for key, val in self._samples.alignments.mask_summary.items() - if val == self._samples.alignments.faces_count] - return retval + def config_tools(self) -> "ConfigTools": + """ :class:`ConfigTools`: The object responsible for parsing configuration options and + updating to/from the GUI """ + return self._config_tools + + @property + def dispatcher(self) -> "Dispatcher": + """ :class:`Dispatcher`: The object responsible for triggering events and variables and + handling global GUI state """ + return self._dispatcher + + @property + def display(self) -> FacesDisplay: + """ :class:`~tools.preview.viewer.FacesDisplay`: The object that holds the sample, + converted and patched faces """ + return self._display + + @property + def lock(self) -> Lock: + """ :class:`threading.Lock`: The threading lock object for the Preview GUI """ + return self._lock + + @property + def progress_bar(self) -> "BusyProgressBar": + """ :class:`~tools.preview.control_panels.BusyProgressBar`: The progress bar that indicates + a swap/patch thread is running """ + assert self._cli_frame is not None + return self._cli_frame.busy_progress_bar + + def update_display(self): + """ Update the images in the canvas and redraw """ + if not hasattr(self, "_image_canvas"): # On first call object not yet created + return + assert self._image_canvas is not None + self._image_canvas.reload() def _initialize_tkinter(self) -> None: """ Initialize a standalone tkinter instance. """ @@ -125,22 +138,22 @@ class Preview(tk.Tk): # pylint:disable=too-few-public-methods self.mainloop() def _refresh(self, *args) -> None: - """ Load new faces to display in preview. + """ Patch faces with current convert settings. Parameters ---------- *args: tuple Unused, but required for tkinter callback. """ - logger.trace("Refreshing swapped faces. args: %s", args) # type: ignore - self._tk_vars["busy"].set(True) + logger.debug("Patching swapped faces. args: %s", args) + self._dispatcher.set_busy() self._config_tools.update_config() with self._lock: assert self._cli_frame is not None self._patch.converter_arguments = self._cli_frame.convert_args - self._patch.current_config = self._config_tools.config - self._patch.trigger.set() - logger.trace("Refreshed swapped faces") # type: ignore + + self._dispatcher.set_needs_patch() + logger.debug("Patched swapped faces") def _build_ui(self) -> None: """ Build the elements for displaying preview images and options panels. """ @@ -148,20 +161,11 @@ class Preview(tk.Tk): # pylint:disable=too-few-public-methods orient=tk.VERTICAL) container.pack(fill=tk.BOTH, expand=True) setattr(container, "preview_display", self._display) # TODO subclass not setattr - self._image_canvas = ImagesCanvas(container, self._tk_vars) + self._image_canvas = ImagesCanvas(self, container) container.add(self._image_canvas, weight=3) options_frame = ttk.Frame(container) - self._cli_frame = ActionFrame( - options_frame, - self._available_masks, - self._samples.predictor.has_predicted_mask, - self._patch.converter.cli_arguments.color_adjustment.replace("-", "_"), - self._patch.converter.cli_arguments.mask_type.replace("-", "_"), - self._config_tools, - self._refresh, - self._samples.generate, - self._tk_vars) + self._cli_frame = ActionFrame(self, options_frame) self._opts_book = OptionsBook(options_frame, self._config_tools, self._refresh) @@ -170,6 +174,90 @@ class Preview(tk.Tk): # pylint:disable=too-few-public-methods container.sashpos(0, int(400 * get_config().scaling_factor)) +class Dispatcher(): + """ Handles the app level tk.Variables and the threading events. Dispatches events to the + correct location and handles GUI state whilst events are handled + + Parameters + ---------- + app: :class:`Preview` + The main tkinter Preview app + """ + def __init__(self, app: Preview): + logger.debug("Initializing %s: (app: %s)", self.__class__.__name__, app) + self._app = app + self._tk_busy = tk.BooleanVar(value=False) + self._evnt_needs_patch = Event() + self._is_updating = False + self._stacked_event = False + logger.debug("Initialized %s", self.__class__.__name__) + + @property + def needs_patch(self) -> Event: + """:class:`threading.Event`. Set by the parent and cleared by the child. Informs the child + patching thread that a run needs to be processed """ + return self._evnt_needs_patch + + # TKInter Variables + def set_busy(self) -> None: + """ Set the tkinter busy variable to ``True`` and display the busy progress bar """ + if self._tk_busy.get(): + logger.debug("Busy event is already set. Doing nothing") + return + if not hasattr(self._app, "progress_bar"): + logger.debug("Not setting busy during initial startup") + return + + logger.debug("Setting busy event to True") + self._tk_busy.set(True) + self._app.progress_bar.start() + self._app.update_idletasks() + + def _unset_busy(self) -> None: + """ Set the tkinter busy variable to ``False`` and hide the busy progress bar """ + self._is_updating = False + if not self._tk_busy.get(): + logger.debug("busy unset when already unset. Doing nothing") + return + logger.debug("Setting busy event to False") + self._tk_busy.set(False) + self._app.progress_bar.stop() + self._app.update_idletasks() + + # Threading Events + def _wait_for_patch(self) -> None: + """ Wait for a patch thread to complete before triggering a display refresh and unsetting + the busy indicators """ + logger.debug("Checking for patch completion...") + if self._evnt_needs_patch.is_set(): + logger.debug("Samples not patched. Waiting...") + self._app.after(1000, self._wait_for_patch) + return + + logger.debug("Patch completion detected") + self._app.update_display() + self._unset_busy() + + if self._stacked_event: + logger.debug("Processing last stacked event") + self.set_busy() + self._stacked_event = False + self.set_needs_patch() + return + + def set_needs_patch(self) -> None: + """ Sends a trigger to the patching thread that it needs to be run. Waits for the patching + to complete prior to triggering a display refresh and unsetting the busy indicators """ + if self._is_updating: + logger.debug("Request to run patch when it is already running. Adding stacked event.") + self._stacked_event = True + return + self._is_updating = True + logger.debug("Triggering patch") + self._evnt_needs_patch.set() + self._wait_for_patch() + + class Samples(): """ The display samples. @@ -182,31 +270,19 @@ class Samples(): Parameters ---------- + app: :class:`Preview` + The main tkinter Preview app arguments: :class:`argparse.Namespace` The :mod:`argparse` arguments as passed in from :mod:`tools.py` sample_size: int The number of samples to take from the input video/images - display: :class:`FacesDisplay` - The display section of the Preview GUI. - lock: :class:`threading.Lock` - A threading lock to prevent multiple GUI updates at the same time. - trigger_patch: :class:`threading.Event` - An event to indicate that a converter patch should be run """ - def __init__(self, - arguments: "Namespace", - sample_size: int, - display: "FacesDisplay", - lock: Lock, - trigger_patch: Event) -> None: - logger.debug("Initializing %s: (arguments: '%s', sample_size: %s, display: %s, lock: %s, " - "trigger_patch: %s)", self.__class__.__name__, arguments, sample_size, - display, lock, trigger_patch) + def __init__(self, app: Preview, arguments: "Namespace", sample_size: int) -> None: + logger.debug("Initializing %s: (app: %s, arguments: '%s', sample_size: %s)", + self.__class__.__name__, app, arguments, sample_size) self._sample_size = sample_size - self._display = display - self._lock = lock - self._trigger_patch = trigger_patch + self._app = app self._input_images: List[ConvertItem] = [] self._predicted_images: List[Tuple[ConvertItem, np.ndarray]] = [] @@ -228,11 +304,19 @@ class Samples(): self._predictor = Predict(queue_manager.get_queue("preview_predict_in"), sample_size, arguments) - self._display.set_centering(self._predictor.centering) + self._app._display.set_centering(self._predictor.centering) self.generate() logger.debug("Initialized %s", self.__class__.__name__) + @property + def available_masks(self) -> List[str]: + """ list: The mask names that are available for every face in the alignments file """ + retval = [key + for key, val in self.alignments.mask_summary.items() + if val == self.alignments.faces_count] + return retval + @property def sample_size(self) -> int: """ int: The number of samples to take from the input video/images """ @@ -319,9 +403,12 @@ class Samples(): Selects :attr:`sample_size` random faces. Runs them through prediction to obtain the swap, then trigger the patch event to run the faces through patching. """ + logger.debug("Generating new random samples") + self._app.dispatcher.set_busy() self._load_frames() self._predict() - self._trigger_patch.set() + self._app.dispatcher.set_needs_patch() + logger.debug("Generated new random samples") def _load_frames(self) -> None: """ Load a sample of random frames. @@ -344,8 +431,8 @@ class Samples(): detected_face.from_alignment(face, image=image) inbound = ExtractMedia(filename=filename, image=image, detected_faces=[detected_face]) self._input_images.append(ConvertItem(inbound=inbound)) - self._display.source = self._input_images - self._display.update_source = True + self._app.display.source = self._input_images + self._app.display.update_source = True logger.debug("Selected frames: %s", [frame.inbound.filename for frame in self._input_images]) @@ -355,7 +442,7 @@ class Samples(): With a threading lock (to prevent stacking), run the selected faces through the Faceswap model predict function and add the output to :attr:`predicted` """ - with self._lock: + with self._app.lock: self._predicted_images = [] for frame in self._input_images: self._predictor.in_queue.put(frame) @@ -375,7 +462,7 @@ class Samples(): logger.debug("Predicted faces") -class Patch(): +class Patch(): # pylint:disable=too-few-public-methods """ The Patch pipeline Runs in it's own thread. Takes the output from the Faceswap model predictor and runs the faces @@ -383,79 +470,42 @@ class Patch(): Parameters ---------- + app: :class:`Preview` + The main tkinter Preview app arguments: :class:`argparse.Namespace` The :mod:`argparse` arguments as passed in from :mod:`tools.py` - available_masks: list - The masks that are available for convert - samples: :class:`Samples` - The Samples for display. - display: :class:`FacesDisplay` - The display section of the Preview GUI. - lock: :class:`threading.Lock` - A threading lock to prevent multiple GUI updates at the same time. - trigger: :class:`threading.Event` - An event to indicate that a converter patch should be run - config_tools: :class:`ConfigTools` - Tools for loading and saving configuration files - tk_vars: dict - Global tkinter variables. `Refresh` and `Busy` :class:`tkinter.BooleanVar` Attributes ---------- converter_arguments: dict The currently selected converter command line arguments for the patch queue - current_config::class:`lib.config.FaceswapConfig` - The currently set configuration for the patch queue """ - def __init__(self, - arguments: "Namespace", - available_masks: List[str], - samples: Samples, - display: "FacesDisplay", - lock: Lock, - trigger: Event, - config_tools: "ConfigTools", - tk_vars: Dict[Literal["refresh", "busy"], tk.BooleanVar]) -> None: - logger.debug("Initializing %s: (arguments: '%s', available_masks: %s, samples: %s, " - "display: %s, lock: %s, trigger: %s, config_tools: %s, tk_vars %s)", - self.__class__.__name__, arguments, available_masks, samples, display, lock, - trigger, config_tools, tk_vars) - self._samples = samples + def __init__(self, app: Preview, arguments: "Namespace") -> None: + logger.debug("Initializing %s: (app: %s, arguments: '%s')", + self.__class__.__name__, app, arguments) + self._app = app self._queue_patch_in = queue_manager.get_queue("preview_patch_in") - self._display = display - self._lock = lock - self._trigger = trigger - self.current_config = config_tools.config self.converter_arguments: Optional[Dict[str, Any]] = None # Updated converter args dict configfile = arguments.configfile if hasattr(arguments, "configfile") else None - self._converter = Converter(output_size=self._samples.predictor.output_size, - coverage_ratio=self._samples.predictor.coverage_ratio, - centering=self._samples.predictor.centering, + self._converter = Converter(output_size=app._samples.predictor.output_size, + coverage_ratio=app._samples.predictor.coverage_ratio, + centering=app._samples.predictor.centering, draw_transparent=False, pre_encode=None, - arguments=self._generate_converter_arguments(arguments, - available_masks), + arguments=self._generate_converter_arguments( + arguments, + app._samples.available_masks), configfile=configfile) - self._shutdown = Event() - self._thread = Thread(target=self._process, name="patch_thread", - args=(self._trigger, - self._shutdown, - self._queue_patch_in, - self._samples, - tk_vars), + args=(self._queue_patch_in, + self._app.dispatcher.needs_patch, + app._samples), daemon=True) self._thread.start() logger.debug("Initializing %s", self.__class__.__name__) - @property - def trigger(self) -> Event: - """ :class:`threading.Event`: The trigger to indicate that a patching run should - commence. """ - return self._trigger - @property def converter(self) -> Converter: """ :class:`lib.convert.Converter`: The converter to use for patching the images. """ @@ -499,11 +549,9 @@ class Patch(): return arguments def _process(self, - trigger_event: Event, - shutdown_event: Event, patch_queue_in: "EventQueue", - samples: Samples, - tk_vars: Dict[Literal["refresh", "busy"], tk.BooleanVar]) -> None: + trigger_event: Event, + samples: Samples) -> None: """ The face patching process. Runs in a thread, and waits for an event to be set. Once triggered, runs a patching @@ -511,40 +559,32 @@ class Patch(): Parameters ---------- - trigger_event: :class:`threading.Event` - Set by parent process when a patching run should be executed - shutdown_event :class:`threading.Event` - Set by parent process if a shutdown has been requested patch_queue_in: :class:`~lib.queue_manager.EventQueue` The input queue for the patching process + trigger_event: :class:`threading.Event` + The event that indicates a patching run needs to be processed samples: :class:`Samples` The Samples for display. - tk_vars: dict - Global tkinter variables. `Refresh` and `Busy` :class:`tkinter.BooleanVar` """ - logger.debug("Launching patch process thread: (trigger_event: %s, shutdown_event: %s, " - "patch_queue_in: %s, samples: %s, tk_vars: %s)", trigger_event, - shutdown_event, patch_queue_in, samples, tk_vars) + logger.debug("Launching patch process thread: (patch_queue_in: %s, trigger_event: %s, " + "samples: %s)", patch_queue_in, trigger_event, samples) patch_queue_out = queue_manager.get_queue("preview_patch_out") while True: trigger = trigger_event.wait(1) - if shutdown_event.is_set(): - logger.debug("Shutdown received") - break if not trigger: continue - # Clear trigger so calling process can set it during this run - trigger_event.clear() + logger.debug("Patch Triggered") queue_manager.flush_queue("preview_patch_in") self._feed_swapped_faces(patch_queue_in, samples) - with self._lock: + with self._app.lock: self._update_converter_arguments() - self._converter.reinitialize(config=self.current_config) + self._converter.reinitialize(config=self._app.config_tools.config) swapped = self._patch_faces(patch_queue_in, patch_queue_out, samples.sample_size) - with self._lock: - self._display.destination = swapped - tk_vars["refresh"].set(True) - tk_vars["busy"].set(False) + with self._app.lock: + self._app.display.destination = swapped + + logger.debug("Patch complete") + trigger_event.clear() logger.debug("Closed patch process thread") @@ -570,12 +610,12 @@ class Patch(): samples: :class:`Samples` The Samples for display. """ - logger.trace("feeding swapped faces to converter") # type: ignore + logger.debug("feeding swapped faces to converter") for item in samples.predicted_images: patch_queue_in.put(item) - logger.trace("fed %s swapped faces to converter", # type: ignore + logger.debug("fed %s swapped faces to converter", len(samples.predicted_images)) - logger.trace("Putting EOF to converter") # type: ignore + logger.debug("Putting EOF to converter") patch_queue_in.put("EOF") def _patch_faces(self, @@ -598,967 +638,15 @@ class Patch(): list The swapped faces patched with the selected convert settings """ - logger.trace("Patching faces") # type: ignore + logger.debug("Patching faces") self._converter.process(queue_in, queue_out) swapped = [] idx = 0 while idx < sample_size: - logger.trace("Patching image %s of %s", idx + 1, sample_size) # type: ignore + logger.debug("Patching image %s of %s", idx + 1, sample_size) item = queue_out.get() swapped.append(item[1]) - logger.trace("Patched image %s of %s", idx + 1, sample_size) # type: ignore + logger.debug("Patched image %s of %s", idx + 1, sample_size) idx += 1 - logger.trace("Patched faces") # type: ignore + logger.debug("Patched faces") return swapped - - -@dataclass -class _Faces: - """ Dataclass for holding faces """ - filenames: List[str] = field(default_factory=list) - matrix: List[np.ndarray] = field(default_factory=list) - src: List[np.ndarray] = field(default_factory=list) - dst: List[np.ndarray] = field(default_factory=list) - - -class FacesDisplay(): - """ Compiles the 2 rows of sample faces (original and swapped) into a single image - - Parameters - ---------- - size: int - The size of each individual face sample in pixels - padding: int - The amount of extra padding to apply to the outside of the face - tk_vars: dict - Global tkinter variables. `Refresh` and `Busy` :class:`tkinter.BooleanVar` - - Attributes - ---------- - update_source: bool - Flag to indicate that the source images for the preview have been updated, so the preview - should be recompiled. - source: list - The list of :class:`numpy.ndarray` source preview images for top row of display - destination: list - The list of :class:`numpy.ndarray` swapped and patched preview images for bottom row of - display - """ - def __init__(self, - size: int, - padding: int, - tk_vars: Dict[Literal["refresh", "busy"], tk.BooleanVar]) -> None: - logger.trace("Initializing %s: (size: %s, padding: %s, tk_vars: %s)", # type: ignore - self.__class__.__name__, size, padding, tk_vars) - self._size = size - self._display_dims = (1, 1) - self._tk_vars = tk_vars - self._padding = padding - - self._faces = _Faces() - self._centering: Optional["CenteringType"] = None - self._faces_source: np.ndarray = np.array([]) - self._faces_dest: np.ndarray = np.array([]) - self._tk_image: Optional[ImageTk.PhotoImage] = None - - # Set from Samples - self.update_source = False - self.source: List[ConvertItem] = [] # Source images, filenames + detected faces - # Set from Patch - self.destination: List[np.ndarray] = [] # Swapped + patched images - - logger.trace("Initialized %s", self.__class__.__name__) # type: ignore - - @property - def tk_image(self) -> Optional[ImageTk.PhotoImage]: - """ :class:`PIL.ImageTk.PhotoImage`: The compiled preview display in tkinter display - format """ - return self._tk_image - - @property - def _total_columns(self) -> int: - """ int: The total number of images that are being displayed """ - return len(self.source) - - def set_centering(self, centering: "CenteringType") -> None: - """ The centering that the model uses is not known at initialization time. - Set :attr:`_centering` when the model has been loaded. - - Parameters - ---------- - centering: str - The centering that the model was trained on - """ - self._centering = centering - - def set_display_dimensions(self, dimensions: Tuple[int, int]) -> None: - """ Adjust the size of the frame that will hold the preview samples. - - Parameters - ---------- - dimensions: tuple - The (`width`, `height`) of the frame that holds the preview - """ - self._display_dims = dimensions - - def update_tk_image(self) -> None: - """ Build the full preview images and compile :attr:`tk_image` for display. """ - logger.trace("Updating tk image") # type: ignore - self._build_faces_image() - img = np.vstack((self._faces_source, self._faces_dest)) - size = self._get_scale_size(img) - img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) - pilimg = Image.fromarray(img) - pilimg = pilimg.resize(size, Image.ANTIALIAS) - self._tk_image = ImageTk.PhotoImage(pilimg) - self._tk_vars["refresh"].set(False) - logger.trace("Updated tk image") # type: ignore - - def _get_scale_size(self, image: np.ndarray) -> Tuple[int, int]: - """ Get the size that the full preview image should be resized to fit in the - display window. - - Parameters - ---------- - image: :class:`numpy.ndarray` - The full sized compiled preview image - - Returns - ------- - tuple - The (`width`, `height`) that the display image should be sized to fit in the display - window - """ - frameratio = float(self._display_dims[0]) / float(self._display_dims[1]) - imgratio = float(image.shape[1]) / float(image.shape[0]) - - if frameratio <= imgratio: - scale = self._display_dims[0] / float(image.shape[1]) - size = (self._display_dims[0], max(1, int(image.shape[0] * scale))) - else: - scale = self._display_dims[1] / float(image.shape[0]) - size = (max(1, int(image.shape[1] * scale)), self._display_dims[1]) - logger.trace("scale: %s, size: %s", scale, size) # type: ignore - return size - - def _build_faces_image(self) -> None: - """ Compile the source and destination rows of the preview image. """ - logger.trace("Building Faces Image") # type: ignore - update_all = self.update_source - self._faces_from_frames() - if update_all: - header = self._header_text() - source = np.hstack([self._draw_rect(face) for face in self._faces.src]) - self._faces_source = np.vstack((header, source)) - self._faces_dest = np.hstack([self._draw_rect(face) for face in self._faces.dst]) - logger.debug("source row shape: %s, swapped row shape: %s", - self._faces_dest.shape, self._faces_source.shape) - - def _faces_from_frames(self) -> None: - """ Extract the preview faces from the source frames and apply the requisite padding. """ - logger.debug("Extracting faces from frames: Number images: %s", len(self.source)) - if self.update_source: - self._crop_source_faces() - self._crop_destination_faces() - logger.debug("Extracted faces from frames: %s", - {k: len(v) for k, v in self._faces.__dict__.items()}) - - def _crop_source_faces(self) -> None: - """ Extract the source faces from the source frames, along with their filenames and the - transformation matrix used to extract the faces. """ - logger.debug("Updating source faces") - self._faces = _Faces() # Init new class - for item in self.source: - detected_face = item.inbound.detected_faces[0] - src_img = item.inbound.image - detected_face.load_aligned(src_img, - size=self._size, - centering=cast("CenteringType", self._centering)) - matrix = detected_face.aligned.matrix - self._faces.filenames.append(os.path.splitext(item.inbound.filename)[0]) - self._faces.matrix.append(matrix) - self._faces.src.append(transform_image(src_img, matrix, self._size, self._padding)) - self.update_source = False - logger.debug("Updated source faces") - - def _crop_destination_faces(self) -> None: - """ Extract the swapped faces from the swapped frames using the source face destination - matrices. """ - logger.debug("Updating destination faces") - self._faces.dst = [] - destination = self.destination if self.destination else [np.ones_like(src.inbound.image) - for src in self.source] - for idx, image in enumerate(destination): - self._faces.dst.append(transform_image(image, - self._faces.matrix[idx], - self._size, - self._padding)) - logger.debug("Updated destination faces") - - def _header_text(self) -> np.ndarray: - """ Create the header text displaying the frame name for each preview column. - - Returns - ------- - :class:`numpy.ndarray` - The header row of the preview image containing the frame names for each column - """ - font_scale = self._size / 640 - height = self._size // 8 - font = cv2.FONT_HERSHEY_SIMPLEX - # Get size of placed text for positioning - text_sizes = [cv2.getTextSize(self._faces.filenames[idx], - font, - font_scale, - 1)[0] - for idx in range(self._total_columns)] - # Get X and Y co-ordinates for each text item - text_y = int((height + text_sizes[0][1]) / 2) - text_x = [int((self._size - text_sizes[idx][0]) / 2) + self._size * idx - for idx in range(self._total_columns)] - logger.debug("filenames: %s, text_sizes: %s, text_x: %s, text_y: %s", - self._faces.filenames, text_sizes, text_x, text_y) - header_box = np.ones((height, self._size * self._total_columns, 3), np.uint8) * 255 - for idx, text in enumerate(self._faces.filenames): - cv2.putText(header_box, - text, - (text_x[idx], text_y), - font, - font_scale, - (0, 0, 0), - 1, - lineType=cv2.LINE_AA) - logger.debug("header_box.shape: %s", header_box.shape) - return header_box - - def _draw_rect(self, image: np.ndarray) -> np.ndarray: - """ Place a white border around a given image. - - Parameters - ---------- - image: :class:`numpy.ndarray` - The image to place a border on to - Returns - ------- - :class:`numpy.ndarray` - The given image with a border drawn around the outside - """ - cv2.rectangle(image, (0, 0), (self._size - 1, self._size - 1), (255, 255, 255), 1) - image = np.clip(image, 0.0, 255.0) - return image.astype("uint8") - - -class ConfigTools(): - """ Tools for loading, saving, setting and retrieving configuration file values. - - Attributes - ---------- - tk_vars: dict - Global tkinter variables. `Refresh` and `Busy` :class:`tkinter.BooleanVar` - """ - def __init__(self) -> None: - self._config = Config(None) - self.tk_vars: Dict[str, Dict[str, Union[tk.BooleanVar, - tk.StringVar, - tk.IntVar, - tk.DoubleVar]]] = {} - self._config_dicts = self._get_config_dicts() # Holds currently saved config - - @property - def config(self) -> Config: - """ :class:`plugins.convert._config.Config` The convert configuration """ - return self._config - - @property - def config_dicts(self) -> Dict[str, Any]: - """ dict: The convert configuration options in dictionary form.""" - return self._config_dicts - - @property - def sections(self) -> List[str]: - """ list: The sorted section names that exist within the convert Configuration options. """ - return sorted(set(plugin.split(".")[0] for plugin in self._config.config.sections() - if plugin.split(".")[0] != "writer")) - - @property - def plugins_dict(self) -> Dict[str, List[str]]: - """ dict: Dictionary of configuration option sections as key with a list of containing - plugins as the value """ - return {section: sorted([plugin.split(".")[1] for plugin in self._config.config.sections() - if plugin.split(".")[0] == section]) - for section in self.sections} - - def update_config(self) -> None: - """ Update :attr:`config` with the currently selected values from the GUI. """ - for section, items in self.tk_vars.items(): - for item, value in items.items(): - try: - new_value = str(value.get()) - except tk.TclError as err: - # When manually filling in text fields, blank values will - # raise an error on numeric data types so return 0 - logger.debug("Error getting value. Defaulting to 0. Error: %s", str(err)) - new_value = str(0) - old_value = self._config.config[section][item] - if new_value != old_value: - logger.trace("Updating config: %s, %s from %s to %s", # type: ignore - section, item, old_value, new_value) - self._config.config[section][item] = new_value - - def _get_config_dicts(self) -> Dict[str, Dict[str, Any]]: - """ Obtain a custom configuration dictionary for convert configuration items in use - by the preview tool formatted for control helper. - - Returns - ------- - dict - Each configuration section as keys, with the values as a dict of option: - :class:`lib.gui.control_helper.ControlOption` pairs. """ - logger.debug("Formatting Config for GUI") - config_dicts: Dict[str, Dict[str, Any]] = {} - for section in self._config.config.sections(): - if section.startswith("writer."): - continue - for key, val in self._config.defaults[section].items(): - if key == "helptext": - config_dicts.setdefault(section, {})[key] = val - continue - cp_option = ControlPanelOption(title=key, - dtype=val["type"], - group=val["group"], - default=val["default"], - initial_value=self._config.get(section, key), - choices=val["choices"], - is_radio=val["gui_radio"], - rounding=val["rounding"], - min_max=val["min_max"], - helptext=val["helptext"]) - self.tk_vars.setdefault(section, {})[key] = cp_option.tk_var - config_dicts.setdefault(section, {})[key] = cp_option - logger.debug("Formatted Config for GUI: %s", config_dicts) - return config_dicts - - def reset_config_to_saved(self, section: Optional[str] = None) -> None: - """ Reset the GUI parameters to their saved values within the configuration file. - - Parameters - ---------- - section: str, optional - The configuration section to reset the values for, If ``None`` provided then all - sections are reset. Default: ``None`` - """ - logger.debug("Resetting to saved config: %s", section) - sections = [section] if section is not None else list(self.tk_vars.keys()) - for config_section in sections: - for item, options in self._config_dicts[config_section].items(): - if item == "helptext": - continue - val = options.value - if val != self.tk_vars[config_section][item].get(): - self.tk_vars[config_section][item].set(val) - logger.debug("Setting %s - %s to saved value %s", config_section, item, val) - logger.debug("Reset to saved config: %s", section) - - def reset_config_to_default(self, section: Optional[str] = None) -> None: - """ Reset the GUI parameters to their default configuration values. - - Parameters - ---------- - section: str, optional - The configuration section to reset the values for, If ``None`` provided then all - sections are reset. Default: ``None`` - """ - logger.debug("Resetting to default: %s", section) - sections = [section] if section is not None else list(self.tk_vars.keys()) - for config_section in sections: - for item, options in self._config_dicts[config_section].items(): - if item == "helptext": - continue - default = options.default - if default != self.tk_vars[config_section][item].get(): - self.tk_vars[config_section][item].set(default) - logger.debug("Setting %s - %s to default value %s", - config_section, item, default) - logger.debug("Reset to default: %s", section) - - def save_config(self, section: Optional[str] = None) -> None: - """ Save the configuration ``.ini`` file with the currently stored values. - - Notes - ----- - We cannot edit the existing saved config as comments tend to get removed, so we create - a new config and populate that. - - Parameters - ---------- - section: str, optional - The configuration section to save, If ``None`` provided then all sections are saved. - Default: ``None`` - """ - logger.debug("Saving %s config", section) - - new_config = ConfigParser(allow_no_value=True) - - for config_section, items in self._config.defaults.items(): - logger.debug("Adding section: '%s')", config_section) - self._config.insert_config_section(config_section, - items["helptext"], - config=new_config) - for item, options in items.items(): - if item == "helptext": - continue # helptext already written at top - if ((section is not None and config_section != section) - or config_section not in self.tk_vars): - # retain saved values that have not been updated - new_opt = self._config.get(config_section, item) - logger.debug("Retaining option: (item: '%s', value: '%s')", item, new_opt) - else: - new_opt = self.tk_vars[config_section][item].get() - logger.debug("Setting option: (item: '%s', value: '%s')", item, new_opt) - - # Set config_dicts value to new saved value - self._config_dicts[config_section][item].set_initial_value(new_opt) - - helptext = self._config.format_help(options["helptext"], is_section=False) - new_config.set(config_section, helptext) - new_config.set(config_section, item, str(new_opt)) - - self._config.config = new_config - self._config.save_config() - logger.info("Saved config: '%s'", self._config.configfile) - - -class ImagesCanvas(ttk.Frame): # pylint:disable=too-many-ancestors - """ tkinter Canvas that holds the preview images. - - Parameters - ---------- - parent: tkinter object - The parent tkinter object that holds the canvas - tk_vars: dict - Global tkinter variables. `Refresh` and `Busy` :class:`tkinter.BooleanVar` - """ - def __init__(self, - parent: ttk.PanedWindow, - tk_vars: Dict[Literal["refresh", "busy"], tk.BooleanVar]) -> None: - logger.debug("Initializing %s: (parent: %s, tk_vars: %s)", - self.__class__.__name__, parent, tk_vars) - super().__init__(parent) - self.pack(expand=True, fill=tk.BOTH, padx=2, pady=2) - - self._refresh_display_trigger = tk_vars["refresh"] - self._refresh_display_trigger.trace("w", self._refresh_display_callback) - self._display: FacesDisplay = parent.preview_display # type: ignore - self._canvas = tk.Canvas(self, bd=0, highlightthickness=0) - self._canvas.pack(side=tk.TOP, fill=tk.BOTH, expand=True) - self._displaycanvas = self._canvas.create_image(0, 0, - image=self._display.tk_image, - anchor=tk.NW) - self.bind("", self._resize) - logger.debug("Initialized %s", self.__class__.__name__) - - def _refresh_display_callback(self, *args) -> None: - """ Add a trace to refresh display on callback """ - if not self._refresh_display_trigger.get(): - return - logger.trace("Refresh display trigger received: %s", args) # type: ignore - self._reload() - - def _resize(self, event: tk.Event) -> None: - """ Resize the image to fit the frame, maintaining aspect ratio """ - logger.trace("Resizing preview image") # type: ignore - framesize = (event.width, event.height) - self._display.set_display_dimensions(framesize) - self._reload() - - def _reload(self) -> None: - """ Reload the preview image """ - logger.trace("Reloading preview image") # type: ignore - self._display.update_tk_image() - self._canvas.itemconfig(self._displaycanvas, image=self._display.tk_image) - - -class ActionFrame(ttk.Frame): # pylint: disable=too-many-ancestors - """ Frame that holds the left hand side options panel containing the command line options. - - Parameters - ---------- - parent: tkinter object - The parent tkinter object that holds the Action Frame - available_masks: list - The available masks that exist within the alignments file - has_predicted_mask: bool - Whether the model was trained with a mask - selected_color: str - The selected color adjustment type - selected_mask_type: str - The selected mask type - config_tools: :class:`ConfigTools` - Tools for loading and saving configuration files - patch_callback: python function - The function to execute when a patch callback is received - refresh_callback: python function - The function to execute when a refresh callback is received - tk_vars: dict - Global tkinter variables. `Refresh` and `Busy` :class:`tkinter.BooleanVar` - """ - def __init__(self, - parent: ttk.Frame, - available_masks: List[str], - has_predicted_mask: bool, - selected_color: str, - selected_mask_type: str, - config_tools: ConfigTools, - patch_callback: Callable[[], None], - refresh_callback: Callable[[], None], - tk_vars: Dict[Literal["refresh", "busy"], tk.BooleanVar]) -> None: - logger.debug("Initializing %s: (available_masks: %s, has_predicted_mask: %s, " - "selected_color: %s, selected_mask_type: %s, patch_callback: %s, " - "refresh_callback: %s, tk_vars: %s)", - self.__class__.__name__, available_masks, has_predicted_mask, selected_color, - selected_mask_type, patch_callback, refresh_callback, tk_vars) - self._config_tools = config_tools - - super().__init__(parent) - self.pack(side=tk.LEFT, anchor=tk.N, fill=tk.Y) - self._options = ["color", "mask_type"] - self._busy_tkvar = tk_vars["busy"] - self._tk_vars: Dict[str, tk.StringVar] = {} - - d_locals = locals() - defaults = {opt: self._format_to_display(d_locals[f"selected_{opt}"]) - for opt in self._options} - self._busy_indicator = self._build_frame(defaults, - refresh_callback, - patch_callback, - available_masks, - has_predicted_mask) - - @property - def convert_args(self) -> Dict[str, Any]: - """ dict: Currently selected Command line arguments from the :class:`ActionFrame`. """ - return {opt if opt != "color" else "color_adjustment": - self._format_from_display(self._tk_vars[opt].get()) - for opt in self._options} - - @staticmethod - def _format_from_display(var: str) -> str: - """ Format a variable from the display version to the command line action version. - - Parameters - ---------- - var: str - The variable name to format - - Returns - ------- - str - The formatted variable name - """ - return var.replace(" ", "_").lower() - - @staticmethod - def _format_to_display(var: str) -> str: - """ Format a variable from the command line action version to the display version. - Parameters - ---------- - var: str - The variable name to format - - Returns - ------- - str - The formatted variable name - """ - return var.replace("_", " ").replace("-", " ").title() - - def _build_frame(self, - defaults: Dict[str, Any], - refresh_callback: Callable[[], None], - patch_callback: Callable[[], None], - available_masks: List[str], - has_predicted_mask: bool) -> ttk.Progressbar: - """ Build the :class:`ActionFrame`. - - Parameters - ---------- - defaults: dict - The default command line options - patch_callback: python function - The function to execute when a patch callback is received - refresh_callback: python function - The function to execute when a refresh callback is received - available_masks: list - The available masks that exist within the alignments file - has_predicted_mask: bool - Whether the model was trained with a mask - - Returns - ------- - ttk.Progressbar - A Progress bar to indicate that the Preview tool is busy - """ - logger.debug("Building Action frame") - - bottom_frame = ttk.Frame(self) - bottom_frame.pack(side=tk.BOTTOM, fill=tk.X, anchor=tk.S) - top_frame = ttk.Frame(self) - top_frame.pack(side=tk.TOP, fill=tk.BOTH, anchor=tk.N, expand=True) - - self._add_cli_choices(top_frame, defaults, available_masks, has_predicted_mask) - - busy_indicator = self._add_busy_indicator(bottom_frame) - self._add_refresh_button(bottom_frame, refresh_callback) - self._add_patch_callback(patch_callback) - self._add_actions(bottom_frame) - logger.debug("Built Action frame") - return busy_indicator - - def _add_cli_choices(self, - parent: ttk.Frame, - defaults: Dict[str, Any], - available_masks: List[str], - has_predicted_mask: bool) -> None: - """ Create :class:`lib.gui.control_helper.ControlPanel` object for the command - line options. - - parent: :class:`ttk.Frame` - The frame to hold the command line choices - defaults: dict - The default command line options - available_masks: list - The available masks that exist within the alignments file - has_predicted_mask: bool - Whether the model was trained with a mask - """ - cp_options = self._get_control_panel_options(defaults, available_masks, has_predicted_mask) - panel_kwargs = dict(blank_nones=False, label_width=10, style="CPanel") - ControlPanel(parent, cp_options, header_text=None, **panel_kwargs) - - def _get_control_panel_options(self, - defaults: Dict[str, Any], - available_masks: List[str], - has_predicted_mask: bool) -> List[ControlPanelOption]: - """ Create :class:`lib.gui.control_helper.ControlPanelOption` objects for the command - line options. - - defaults: dict - The default command line options - available_masks: list - The available masks that exist within the alignments file - has_predicted_mask: bool - Whether the model was trained with a mask - - Returns - ------- - list - The list of `lib.gui.control_helper.ControlPanelOption` objects for the Action Frame - """ - cp_options: List[ControlPanelOption] = [] - for opt in self._options: - if opt == "mask_type": - choices = self._create_mask_choices(defaults, available_masks, has_predicted_mask) - else: - choices = PluginLoader.get_available_convert_plugins(opt, True) - cp_option = ControlPanelOption(title=opt, - dtype=str, - default=defaults[opt], - initial_value=defaults[opt], - choices=choices, - group="Command Line Choices", - is_radio=False) - self._tk_vars[opt] = cp_option.tk_var - cp_options.append(cp_option) - return cp_options - - @classmethod - def _create_mask_choices(cls, - defaults: Dict[str, Any], - available_masks: List[str], - has_predicted_mask: bool) -> List[str]: - """ Set the mask choices and default mask based on available masks. - - Parameters - ---------- - defaults: dict - The default command line options - available_masks: list - The available masks that exist within the alignments file - has_predicted_mask: bool - Whether the model was trained with a mask - - Returns - ------- - list - The masks that are available to use from the alignments file - """ - logger.debug("Initial mask choices: %s", available_masks) - if has_predicted_mask: - available_masks += ["predicted"] - if "none" not in available_masks: - available_masks += ["none"] - if defaults["mask_type"] not in available_masks: - logger.debug("Setting default mask to first available: %s", available_masks[0]) - defaults["mask_type"] = available_masks[0] - logger.debug("Final mask choices: %s", available_masks) - return available_masks - - @classmethod - def _add_refresh_button(cls, - parent: ttk.Frame, - refresh_callback: Callable[[], None]) -> None: - """ Add a button to refresh the images. - - Parameters - ---------- - refresh_callback: python function - The function to execute when the refresh button is pressed - """ - btn = ttk.Button(parent, text="Update Samples", command=refresh_callback) - btn.pack(padx=5, pady=5, side=tk.TOP, fill=tk.X, anchor=tk.N) - - def _add_patch_callback(self, patch_callback: Callable[[], None]) -> None: - """ Add callback to re-patch images on action option change. - - Parameters - ---------- - patch_callback: python function - The function to execute when the images require patching - """ - for tk_var in self._tk_vars.values(): - tk_var.trace("w", patch_callback) - - def _add_busy_indicator(self, parent: ttk.Frame) -> ttk.Progressbar: - """ Place progress bar into bottom bar to indicate when processing. - - Parameters - ---------- - parent: tkinter object - The tkinter object that holds the busy indicator - - Returns - ------- - ttk.Progressbar - A Progress bar to indicate that the Preview tool is busy - """ - logger.debug("Placing busy indicator") - pbar = ttk.Progressbar(parent, mode="indeterminate") - pbar.pack(side=tk.LEFT) - pbar.pack_forget() - self._busy_tkvar.trace("w", self._busy_indicator_trace) - return pbar - - def _busy_indicator_trace(self, *args) -> None: - """ Show or hide busy indicator based on whether the preview is updating. - - Parameters - ---------- - args: unused - Required for tkinter event, but unused - """ - logger.trace("Busy indicator trace: %s", args) # type: ignore - if self._busy_tkvar.get(): - self._start_busy_indicator() - else: - self._stop_busy_indicator() - - def _stop_busy_indicator(self) -> None: - """ Stop and hide progress bar """ - logger.debug("Stopping busy indicator") - self._busy_indicator.stop() - self._busy_indicator.pack_forget() - - def _start_busy_indicator(self) -> None: - """ Start and display progress bar """ - logger.debug("Starting busy indicator") - self._busy_indicator.pack(side=tk.LEFT, padx=5, pady=(5, 10), fill=tk.X, expand=True) - self._busy_indicator.start() - - def _add_actions(self, parent: ttk.Frame) -> None: - """ Add Action Buttons to the :class:`ActionFrame` - - Parameters - ---------- - parent: tkinter object - The tkinter object that holds the action buttons - """ - logger.debug("Adding util buttons") - frame = ttk.Frame(parent) - frame.pack(padx=5, pady=(5, 10), side=tk.RIGHT, fill=tk.X, anchor=tk.E) - - for utl in ("save", "clear", "reload"): - logger.debug("Adding button: '%s'", utl) - img = get_images().icons[utl] - if utl == "save": - text = _("Save full config") - action = self._config_tools.save_config - elif utl == "clear": - text = _("Reset full config to default values") - action = self._config_tools.reset_config_to_default - elif utl == "reload": - text = _("Reset full config to saved values") - action = self._config_tools.reset_config_to_saved - - btnutl = ttk.Button(frame, - image=img, - command=action) - btnutl.pack(padx=2, side=tk.RIGHT) - Tooltip(btnutl, text=text, wrap_length=200) - logger.debug("Added util buttons") - - -class OptionsBook(ttk.Notebook): # pylint:disable=too-many-ancestors - """ The notebook that holds the Convert configuration options. - - Parameters - ---------- - parent: tkinter object - The parent tkinter object that holds the Options book - config_tools: :class:`ConfigTools` - Tools for loading and saving configuration files - patch_callback: python function - The function to execute when a patch callback is received - - Attributes - ---------- - config_tools: :class:`ConfigTools` - Tools for loading and saving configuration files - """ - def __init__(self, - parent: ttk.Frame, - config_tools: ConfigTools, - patch_callback: Callable[[], None]) -> None: - logger.debug("Initializing %s: (parent: %s, config: %s)", - self.__class__.__name__, parent, config_tools) - super().__init__(parent) - self.pack(side=tk.RIGHT, anchor=tk.N, fill=tk.BOTH, expand=True) - self.config_tools = config_tools - - self._tabs: Dict[str, Dict[str, Union[ttk.Notebook, ConfigFrame]]] = {} - self._build_tabs() - self._build_sub_tabs() - self._add_patch_callback(patch_callback) - logger.debug("Initialized %s", self.__class__.__name__) - - def _build_tabs(self) -> None: - """ Build the notebook tabs for the each configuration section. """ - logger.debug("Build Tabs") - for section in self.config_tools.sections: - tab = ttk.Notebook(self) - self._tabs[section] = {"tab": tab} - self.add(tab, text=section.replace("_", " ").title()) - - def _build_sub_tabs(self) -> None: - """ Build the notebook sub tabs for each convert section's plugin. """ - for section, plugins in self.config_tools.plugins_dict.items(): - for plugin in plugins: - config_key = ".".join((section, plugin)) - config_dict = self.config_tools.config_dicts[config_key] - tab = ConfigFrame(self, config_key, config_dict) - self._tabs[section][plugin] = tab - text = plugin.replace("_", " ").title() - cast(ttk.Notebook, self._tabs[section]["tab"]).add(tab, text=text) - - def _add_patch_callback(self, patch_callback: Callable[[], None]) -> None: - """ Add callback to re-patch images on configuration option change. - - Parameters - ---------- - patch_callback: python function - The function to execute when the images require patching - """ - for plugins in self.config_tools.tk_vars.values(): - for tk_var in plugins.values(): - tk_var.trace("w", patch_callback) - - -class ConfigFrame(ttk.Frame): # pylint: disable=too-many-ancestors - """ Holds the configuration options for a convert plugin inside the :class:`OptionsBook`. - - Parameters - ---------- - parent: tkinter object - The tkinter object that will hold this configuration frame - config_key: str - The section/plugin key for these configuration options - options: dict - The options for this section/plugin - """ - - def __init__(self, - parent: OptionsBook, - config_key: str, - options: Dict[str, Any]): - logger.debug("Initializing %s", self.__class__.__name__) - super().__init__(parent) - self.pack(side=tk.TOP, fill=tk.BOTH, expand=True) - - self._options = options - - self._action_frame = ttk.Frame(self) - self._action_frame.pack(padx=0, pady=(0, 5), side=tk.BOTTOM, fill=tk.X, anchor=tk.E) - self._add_frame_separator() - - self._build_frame(parent, config_key) - logger.debug("Initialized %s", self.__class__.__name__) - - def _build_frame(self, parent: OptionsBook, config_key: str) -> None: - """ Build the options frame for this command - - Parameters - ---------- - parent: tkinter object - The tkinter object that will hold this configuration frame - config_key: str - The section/plugin key for these configuration options - """ - logger.debug("Add Config Frame") - panel_kwargs = dict(columns=2, option_columns=2, blank_nones=False, style="CPanel") - frame = ttk.Frame(self) - frame.pack(side=tk.TOP, fill=tk.BOTH, expand=True) - cp_options = [opt for key, opt in self._options.items() if key != "helptext"] - ControlPanel(frame, cp_options, header_text=None, **panel_kwargs) - self._add_actions(parent, config_key) - logger.debug("Added Config Frame") - - def _add_frame_separator(self) -> None: - """ Add a separator between top and bottom frames. """ - logger.debug("Add frame seperator") - sep = ttk.Frame(self._action_frame, height=2, relief=tk.RIDGE) - sep.pack(fill=tk.X, pady=5, side=tk.TOP) - logger.debug("Added frame seperator") - - def _add_actions(self, parent: OptionsBook, config_key: str) -> None: - """ Add Action Buttons. - - Parameters - ---------- - parent: tkinter object - The tkinter object that will hold this configuration frame - config_key: str - The section/plugin key for these configuration options - """ - logger.debug("Adding util buttons") - - title = config_key.split(".")[1].replace("_", " ").title() - btn_frame = ttk.Frame(self._action_frame) - btn_frame.pack(padx=5, side=tk.BOTTOM, fill=tk.X) - for utl in ("save", "clear", "reload"): - logger.debug("Adding button: '%s'", utl) - img = get_images().icons[utl] - if utl == "save": - text = _(f"Save {title} config") - action = parent.config_tools.save_config - elif utl == "clear": - text = _(f"Reset {title} config to default values") - action = parent.config_tools.reset_config_to_default - elif utl == "reload": - text = _(f"Reset {title} config to saved values") - action = parent.config_tools.reset_config_to_saved - - btnutl = ttk.Button(btn_frame, - image=img, - command=lambda cmd=action: cmd(config_key)) # type: ignore - btnutl.pack(padx=2, side=tk.RIGHT) - Tooltip(btnutl, text=text, wrap_length=200) - logger.debug("Added util buttons") diff --git a/tools/preview/viewer.py b/tools/preview/viewer.py new file mode 100644 index 0000000..b187603 --- /dev/null +++ b/tools/preview/viewer.py @@ -0,0 +1,295 @@ +#!/usr/bin/env python3 +""" Manages the widgets that hold the top 'viewer' area of the preview tool """ +import logging +import os +import tkinter as tk +from tkinter import ttk + +from dataclasses import dataclass, field +from typing import cast, List, Optional, Tuple, TYPE_CHECKING + +import cv2 +import numpy as np +from PIL import Image, ImageTk + +from lib.align import transform_image +from lib.align.aligned_face import CenteringType +from scripts.convert import ConvertItem + + +if TYPE_CHECKING: + from .preview import Preview + +logger = logging.getLogger(__name__) + + +@dataclass +class _Faces: + """ Dataclass for holding faces """ + filenames: List[str] = field(default_factory=list) + matrix: List[np.ndarray] = field(default_factory=list) + src: List[np.ndarray] = field(default_factory=list) + dst: List[np.ndarray] = field(default_factory=list) + + +class FacesDisplay(): + """ Compiles the 2 rows of sample faces (original and swapped) into a single image + + Parameters + ---------- + app: :class:`Preview` + The main tkinter Preview app + size: int + The size of each individual face sample in pixels + padding: int + The amount of extra padding to apply to the outside of the face + + Attributes + ---------- + update_source: bool + Flag to indicate that the source images for the preview have been updated, so the preview + should be recompiled. + source: list + The list of :class:`numpy.ndarray` source preview images for top row of display + destination: list + The list of :class:`numpy.ndarray` swapped and patched preview images for bottom row of + display + """ + def __init__(self, app: 'Preview', size: int, padding: int) -> None: + logger.trace("Initializing %s: (app: %s, size: %s, padding: %s)", # type: ignore + self.__class__.__name__, app, size, padding) + self._size = size + self._display_dims = (1, 1) + self._app = app + self._padding = padding + + self._faces = _Faces() + self._centering: Optional[CenteringType] = None + self._faces_source: np.ndarray = np.array([]) + self._faces_dest: np.ndarray = np.array([]) + self._tk_image: Optional[ImageTk.PhotoImage] = None + + # Set from Samples + self.update_source = False + self.source: List[ConvertItem] = [] # Source images, filenames + detected faces + # Set from Patch + self.destination: List[np.ndarray] = [] # Swapped + patched images + + logger.trace("Initialized %s", self.__class__.__name__) # type: ignore + + @property + def tk_image(self) -> Optional[ImageTk.PhotoImage]: + """ :class:`PIL.ImageTk.PhotoImage`: The compiled preview display in tkinter display + format """ + return self._tk_image + + @property + def _total_columns(self) -> int: + """ int: The total number of images that are being displayed """ + return len(self.source) + + def set_centering(self, centering: CenteringType) -> None: + """ The centering that the model uses is not known at initialization time. + Set :attr:`_centering` when the model has been loaded. + + Parameters + ---------- + centering: str + The centering that the model was trained on + """ + self._centering = centering + + def set_display_dimensions(self, dimensions: Tuple[int, int]) -> None: + """ Adjust the size of the frame that will hold the preview samples. + + Parameters + ---------- + dimensions: tuple + The (`width`, `height`) of the frame that holds the preview + """ + self._display_dims = dimensions + + def update_tk_image(self) -> None: + """ Build the full preview images and compile :attr:`tk_image` for display. """ + logger.trace("Updating tk image") # type: ignore + self._build_faces_image() + img = np.vstack((self._faces_source, self._faces_dest)) + size = self._get_scale_size(img) + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + pilimg = Image.fromarray(img) + pilimg = pilimg.resize(size, Image.ANTIALIAS) + self._tk_image = ImageTk.PhotoImage(pilimg) + logger.trace("Updated tk image") # type: ignore + + def _get_scale_size(self, image: np.ndarray) -> Tuple[int, int]: + """ Get the size that the full preview image should be resized to fit in the + display window. + + Parameters + ---------- + image: :class:`numpy.ndarray` + The full sized compiled preview image + + Returns + ------- + tuple + The (`width`, `height`) that the display image should be sized to fit in the display + window + """ + frameratio = float(self._display_dims[0]) / float(self._display_dims[1]) + imgratio = float(image.shape[1]) / float(image.shape[0]) + + if frameratio <= imgratio: + scale = self._display_dims[0] / float(image.shape[1]) + size = (self._display_dims[0], max(1, int(image.shape[0] * scale))) + else: + scale = self._display_dims[1] / float(image.shape[0]) + size = (max(1, int(image.shape[1] * scale)), self._display_dims[1]) + logger.trace("scale: %s, size: %s", scale, size) # type: ignore + return size + + def _build_faces_image(self) -> None: + """ Compile the source and destination rows of the preview image. """ + logger.trace("Building Faces Image") # type: ignore + update_all = self.update_source + self._faces_from_frames() + if update_all: + header = self._header_text() + source = np.hstack([self._draw_rect(face) for face in self._faces.src]) + self._faces_source = np.vstack((header, source)) + self._faces_dest = np.hstack([self._draw_rect(face) for face in self._faces.dst]) + logger.debug("source row shape: %s, swapped row shape: %s", + self._faces_dest.shape, self._faces_source.shape) + + def _faces_from_frames(self) -> None: + """ Extract the preview faces from the source frames and apply the requisite padding. """ + logger.debug("Extracting faces from frames: Number images: %s", len(self.source)) + if self.update_source: + self._crop_source_faces() + self._crop_destination_faces() + logger.debug("Extracted faces from frames: %s", + {k: len(v) for k, v in self._faces.__dict__.items()}) + + def _crop_source_faces(self) -> None: + """ Extract the source faces from the source frames, along with their filenames and the + transformation matrix used to extract the faces. """ + logger.debug("Updating source faces") + self._faces = _Faces() # Init new class + for item in self.source: + detected_face = item.inbound.detected_faces[0] + src_img = item.inbound.image + detected_face.load_aligned(src_img, + size=self._size, + centering=cast(CenteringType, self._centering)) + matrix = detected_face.aligned.matrix + self._faces.filenames.append(os.path.splitext(item.inbound.filename)[0]) + self._faces.matrix.append(matrix) + self._faces.src.append(transform_image(src_img, matrix, self._size, self._padding)) + self.update_source = False + logger.debug("Updated source faces") + + def _crop_destination_faces(self) -> None: + """ Extract the swapped faces from the swapped frames using the source face destination + matrices. """ + logger.debug("Updating destination faces") + self._faces.dst = [] + destination = self.destination if self.destination else [np.ones_like(src.inbound.image) + for src in self.source] + for idx, image in enumerate(destination): + self._faces.dst.append(transform_image(image, + self._faces.matrix[idx], + self._size, + self._padding)) + logger.debug("Updated destination faces") + + def _header_text(self) -> np.ndarray: + """ Create the header text displaying the frame name for each preview column. + + Returns + ------- + :class:`numpy.ndarray` + The header row of the preview image containing the frame names for each column + """ + font_scale = self._size / 640 + height = self._size // 8 + font = cv2.FONT_HERSHEY_SIMPLEX + # Get size of placed text for positioning + text_sizes = [cv2.getTextSize(self._faces.filenames[idx], + font, + font_scale, + 1)[0] + for idx in range(self._total_columns)] + # Get X and Y co-ordinates for each text item + text_y = int((height + text_sizes[0][1]) / 2) + text_x = [int((self._size - text_sizes[idx][0]) / 2) + self._size * idx + for idx in range(self._total_columns)] + logger.debug("filenames: %s, text_sizes: %s, text_x: %s, text_y: %s", + self._faces.filenames, text_sizes, text_x, text_y) + header_box = np.ones((height, self._size * self._total_columns, 3), np.uint8) * 255 + for idx, text in enumerate(self._faces.filenames): + cv2.putText(header_box, + text, + (text_x[idx], text_y), + font, + font_scale, + (0, 0, 0), + 1, + lineType=cv2.LINE_AA) + logger.debug("header_box.shape: %s", header_box.shape) + return header_box + + def _draw_rect(self, image: np.ndarray) -> np.ndarray: + """ Place a white border around a given image. + + Parameters + ---------- + image: :class:`numpy.ndarray` + The image to place a border on to + Returns + ------- + :class:`numpy.ndarray` + The given image with a border drawn around the outside + """ + cv2.rectangle(image, (0, 0), (self._size - 1, self._size - 1), (255, 255, 255), 1) + image = np.clip(image, 0.0, 255.0) + return image.astype("uint8") + + +class ImagesCanvas(ttk.Frame): # pylint:disable=too-many-ancestors + """ tkinter Canvas that holds the preview images. + + Parameters + ---------- + app: :class:`Preview` + The main tkinter Preview app + parent: tkinter object + The parent tkinter object that holds the canvas + """ + def __init__(self, app: 'Preview', parent: ttk.PanedWindow) -> None: + logger.debug("Initializing %s: (app: %s, parent: %s)", + self.__class__.__name__, app, parent) + super().__init__(parent) + self.pack(expand=True, fill=tk.BOTH, padx=2, pady=2) + + self._display: FacesDisplay = parent.preview_display # type: ignore + self._canvas = tk.Canvas(self, bd=0, highlightthickness=0) + self._canvas.pack(side=tk.TOP, fill=tk.BOTH, expand=True) + self._displaycanvas = self._canvas.create_image(0, 0, + image=self._display.tk_image, + anchor=tk.NW) + self.bind("", self._resize) + logger.debug("Initialized %s", self.__class__.__name__) + + def _resize(self, event: tk.Event) -> None: + """ Resize the image to fit the frame, maintaining aspect ratio """ + logger.debug("Resizing preview image") + framesize = (event.width, event.height) + self._display.set_display_dimensions(framesize) + self.reload() + + def reload(self) -> None: + """ Update the images in the canvas and redraw """ + logger.debug("Reloading preview image") + self._display.update_tk_image() + self._canvas.itemconfig(self._displaycanvas, image=self._display.tk_image) + logger.debug("Reloaded preview image")