OSSForks/faceswap
1
0
Fork 0
mirror of https://github.com/zebrajr/faceswap.git synced 2025-12-06 00:20:09 +01:00
Code Issues Actions Packages Projects Releases Wiki Activity

bugfix:

Browse Source 
- 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
This commit is contained in:
torzdf 2023-01-17 15:03:29 +00:00
parent 80f63280ca
commit 34b558426e
14 changed files with 1789 additions and 1178 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
docs/full/tests/tests.rst
View File

@ -14,3 +14,4 @@ Subpackages
:maxdepth: 1
lib
tools

15
docs/full/tests/tools.preview.rst Normal file
View File

@ -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:

14
docs/full/tests/tools.rst Normal file
View File

@ -0,0 +1,14 @@
*************
tools package
*************
.. contents:: Contents
:local:
Subpackages
===========
.. toctree::
:maxdepth: 1
tools.preview

44
docs/full/tools/preview.rst Normal file
View File

@ -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:

26
docs/full/tools/tools.rst
View File

@ -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:

35
lib/gui/utils/config.py
View File

@ -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

60
locales/tools.preview.pot
View File

@ -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 <EMAIL@ADDRESS>, 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 <EMAIL@ADDRESS>\n"
"Language-Team: LANGUAGE <LL@li.org>\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 ""

0
tests/tools/__init__.py Normal file
View File

0
tests/tools/preview/__init__.py Normal file
View File

480
tests/tools/preview/viewer_test.py Normal file
View File

@ -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()

22
tools/preview/cli.py
View File

@ -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,

667
tools/preview/control_panels.py Normal file
View File

@ -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")

1308
tools/preview/preview.py
View File

File diff suppressed because it is too large Load Diff

295
tools/preview/viewer.py Normal file
View File

@ -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("<Configure>", 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")