1
0
Fork 0
mirror of https://github.com/deepfakes/faceswap synced 2025-06-07 10:43:27 -04:00
faceswap/tools/manual/frameviewer/editor/extract_box.py

408 lines
17 KiB
Python

#!/usr/bin/env python3
""" Extract Box Editor for the manual adjustments tool """
import gettext
import platform
import numpy as np
from lib.align import AlignedFace
from lib.gui.custom_widgets import RightClickMenu
from lib.gui.utils import get_config
from ._base import Editor, logger
# LOCALES
_LANG = gettext.translation("tools.manual", localedir="locales", fallback=True)
_ = _LANG.gettext
class ExtractBox(Editor):
""" The Extract Box Editor.
Adjust the calculated Extract Box to shift all of the 68 point landmarks in place.
Parameters
----------
canvas: :class:`tkinter.Canvas`
The canvas that holds the image and annotations
detected_faces: :class:`~tools.manual.detected_faces.DetectedFaces`
The _detected_faces data for this manual session
"""
def __init__(self, canvas, detected_faces):
self._right_click_menu = RightClickMenu([_("Delete Face")],
[self._delete_current_face],
["Del"])
control_text = _("Extract Box Editor\nMove the extract box that has been generated by the "
"aligner. Click and drag:\n\n"
" - Inside the bounding box to relocate the landmarks.\n"
" - The corner anchors to resize the landmarks.\n"
" - Outside of the corners to rotate the landmarks.")
key_bindings = {"<Delete>": self._delete_current_face}
super().__init__(canvas, detected_faces,
control_text=control_text, key_bindings=key_bindings)
@property
def _corner_order(self):
""" dict: The position index of bounding box corners """
return {0: ("top", "left"),
3: ("top", "right"),
2: ("bottom", "right"),
1: ("bottom", "left")}
def update_annotation(self):
""" Draw the latest Extract Boxes around the faces. """
color = self._control_color
roi = self._zoomed_roi
for idx, face in enumerate(self._face_iterator):
logger.trace("Drawing Extract Box: (idx: %s)", idx)
if self._globals.is_zoomed:
box = np.array((roi[0], roi[1], roi[2], roi[1], roi[2], roi[3], roi[0], roi[3]))
else:
aligned = AlignedFace(face.landmarks_xy, centering="face")
box = self._scale_to_display(aligned.original_roi).flatten()
top_left = box[:2] - 10
kwargs = {"fill": color, "font": ('Default', 20, 'bold'), "text": str(idx)}
self._object_tracker("eb_text", "text", idx, top_left, kwargs)
kwargs = {"fill": '', "outline": color, "width": 1}
self._object_tracker("eb_box", "polygon", idx, box, kwargs)
self._update_anchor_annotation(idx, box, color)
logger.trace("Updated extract box annotations")
def _update_anchor_annotation(self, face_index, extract_box, color):
""" Update the anchor annotations for each corner of the extract box.
The anchors only display when the extract box editor is active.
Parameters
----------
face_index: int
The index of the face being annotated
extract_box: :class:`numpy.ndarray`
The scaled extract box to get the corner anchors for
color: str
The hex color of the extract box line
"""
if not self._is_active or self._globals.is_zoomed:
self.hide_annotation("eb_anc_dsp")
self.hide_annotation("eb_anc_grb")
return
fill_color = "gray"
activefill_color = "white" if self._is_active else ""
anchor_points = self._get_anchor_points((extract_box[:2],
extract_box[2:4],
extract_box[4:6],
extract_box[6:]))
for idx, (anc_dsp, anc_grb) in enumerate(zip(*anchor_points)):
dsp_kwargs = {"outline": color, "fill": fill_color, "width": 1}
grb_kwargs = {"outline": '', "fill": '', "width": 1, "activefill": activefill_color}
dsp_key = f"eb_anc_dsp_{idx}"
grb_key = f"eb_anc_grb_{idx}"
self._object_tracker(dsp_key, "oval", face_index, anc_dsp, dsp_kwargs)
self._object_tracker(grb_key, "oval", face_index, anc_grb, grb_kwargs)
logger.trace("Updated extract box anchor annotations")
# << MOUSE HANDLING >>
# Mouse cursor display
def _update_cursor(self, event):
""" Update the cursor when it is hovering over an extract box and update
:attr:`_mouse_location` with the current cursor position.
Parameters
----------
event: :class:`tkinter.Event`
The current tkinter mouse event
"""
if self._check_cursor_anchors():
return
if self._check_cursor_box():
return
if self._check_cursor_rotate(event):
return
self._canvas.config(cursor="")
self._mouse_location = None
def _check_cursor_anchors(self):
""" Check whether the cursor is over a corner anchor.
If it is, set the appropriate cursor type and set :attr:`_mouse_location` to
("anchor", `face index`, `corner_index`)
Returns
-------
bool
``True`` if cursor is over an anchor point otherwise ``False``
"""
anchors = set(self._canvas.find_withtag("eb_anc_grb"))
item_ids = set(self._canvas.find_withtag("current")).intersection(anchors)
if not item_ids:
return False
item_id = list(item_ids)[0]
tags = self._canvas.gettags(item_id)
face_idx = int(next(tag for tag in tags if tag.startswith("face_")).split("_")[-1])
corner_idx = int(next(tag for tag in tags
if tag.startswith("eb_anc_grb_")
and "face_" not in tag).split("_")[-1])
pos_x, pos_y = self._corner_order[corner_idx]
self._canvas.config(cursor=f"{pos_x}_{pos_y}_corner")
self._mouse_location = ("anchor", face_idx, corner_idx)
return True
def _check_cursor_box(self):
""" Check whether the cursor is inside an extract box.
If it is, set the appropriate cursor type and set :attr:`_mouse_location` to
("box", `face index`)
Returns
-------
bool
``True`` if cursor is over a rotate point otherwise ``False``
"""
extract_boxes = set(self._canvas.find_withtag("eb_box"))
item_ids = set(self._canvas.find_withtag("current")).intersection(extract_boxes)
if not item_ids:
return False
item_id = list(item_ids)[0]
self._canvas.config(cursor="fleur")
self._mouse_location = ("box", next(int(tag.split("_")[-1])
for tag in self._canvas.gettags(item_id)
if tag.startswith("face_")))
return True
def _check_cursor_rotate(self, event):
""" Check whether the cursor is in an area to rotate the extract box.
If it is, set the appropriate cursor type and set :attr:`_mouse_location` to
("rotate", `face index`)
Notes
-----
This code is executed after the check has been completed to see if the mouse is inside
the extract box. For this reason, we don't bother running a check to see if the mouse
is inside the box, as this code will never run if that is the case.
Parameters
----------
event: :class:`tkinter.Event`
The current tkinter mouse event
Returns
-------
bool
``True`` if cursor is over a rotate point otherwise ``False``
"""
distance = 30
boxes = np.array([np.array(self._canvas.coords(item_id)).reshape(4, 2)
for item_id in self._canvas.find_withtag("eb_box")
if self._canvas.itemcget(item_id, "state") != "hidden"])
position = np.array((event.x, event.y)).astype("float32")
for face_idx, points in enumerate(boxes):
if any(np.all(position > point - distance) and np.all(position < point + distance)
for point in points):
self._canvas.config(cursor="exchange")
self._mouse_location = ("rotate", face_idx)
return True
return False
# Mouse click actions
def set_mouse_click_actions(self):
""" Add context menu to OS specific right click action. """
super().set_mouse_click_actions()
self._canvas.bind("<Button-2>" if platform.system() == "Darwin" else "<Button-3>",
self._context_menu)
def _drag_start(self, event):
""" The action to perform when the user starts clicking and dragging the mouse.
Selects the correct extract box action based on the initial cursor position.
Parameters
----------
event: :class:`tkinter.Event`
The tkinter mouse event.
"""
if self._mouse_location is None:
self._drag_data = {}
self._drag_callback = None
return
self._drag_data["current_location"] = np.array((event.x, event.y))
callback = {"anchor": self._resize, "rotate": self._rotate, "box": self._move}
self._drag_callback = callback[self._mouse_location[0]]
def _drag_stop(self, event): # pylint:disable=unused-argument
""" Trigger a viewport thumbnail update on click + drag release
Parameters
----------
event: :class:`tkinter.Event`
The tkinter mouse event. Required but unused.
"""
if self._mouse_location is None:
return
self._det_faces.update.post_edit_trigger(self._globals.frame_index,
self._mouse_location[1])
def _move(self, event):
""" Updates the underlying detected faces landmarks based on mouse dragging delta,
which moves the Extract box on a drag event.
Parameters
----------
event: :class:`tkinter.Event`
The tkinter mouse event.
"""
if not self._drag_data:
return
shift_x = event.x - self._drag_data["current_location"][0]
shift_y = event.y - self._drag_data["current_location"][1]
scaled_shift = self.scale_from_display(np.array((shift_x, shift_y)), do_offset=False)
self._det_faces.update.landmarks(self._globals.frame_index,
self._mouse_location[1],
*scaled_shift)
self._drag_data["current_location"] = (event.x, event.y)
def _resize(self, event):
""" Resizes the landmarks contained within an extract box on a corner anchor drag event.
Parameters
----------
event: :class:`tkinter.Event`
The tkinter mouse event.
"""
face_idx = self._mouse_location[1]
face_tag = f"eb_box_face_{face_idx}"
position = np.array((event.x, event.y))
box = np.array(self._canvas.coords(face_tag))
center = np.array((sum(box[0::2]) / 4, sum(box[1::2]) / 4))
if not self._check_in_bounds(center, box, position):
logger.trace("Drag out of bounds. Not updating")
self._drag_data["current_location"] = position
return
start = self._drag_data["current_location"]
distance = ((np.linalg.norm(center - start) - np.linalg.norm(center - position))
* get_config().scaling_factor)
size = ((box[2] - box[0]) ** 2 + (box[3] - box[1]) ** 2) ** 0.5
scale = 1 - (distance / size)
logger.trace("face_index: %s, center: %s, start: %s, position: %s, distance: %s, "
"size: %s, scale: %s", face_idx, center, start, position, distance, size,
scale)
if size * scale < 20:
# Don't over shrink the box
logger.trace("Box would size to less than 20px. Not updating")
self._drag_data["current_location"] = position
return
self._det_faces.update.landmarks_scale(self._globals.frame_index,
face_idx,
scale,
self.scale_from_display(center))
self._drag_data["current_location"] = position
def _check_in_bounds(self, center, box, position):
""" Ensure that a resize drag does is not going to cross the center point from it's initial
corner location.
Parameters
----------
center: :class:`numpy.ndarray`
The (`x`, `y`) center point of the face extract box
box: :class:`numpy.ndarray`
The canvas coordinates of the extract box polygon's corners
position: : class:`numpy.ndarray`
The current (`x`, `y`) position of the mouse cursor
Returns
-------
bool
``True`` if the drag operation does not cross the center point otherwise ``False``
"""
# Generate lines that span the full frame (x and y) along the center point
center_x = np.array(((center[0], 0), (center[0], self._globals.frame_display_dims[1])))
center_y = np.array(((0, center[1]), (self._globals.frame_display_dims[0], center[1])))
# Generate a line coming from the current corner location to the current cursor position
full_line = np.array((box[self._mouse_location[2] * 2:self._mouse_location[2] * 2 + 2],
position))
logger.trace("center: %s, center_x_line: %s, center_y_line: %s, full_line: %s",
center, center_x, center_y, full_line)
# Check whether any of the generated lines intersect
for line in (center_x, center_y):
if (self._is_ccw(full_line[0], *line) != self._is_ccw(full_line[1], *line) and
self._is_ccw(*full_line, line[0]) != self._is_ccw(*full_line, line[1])):
logger.trace("line: %s crosses center: %s", full_line, center)
return False
return True
@staticmethod
def _is_ccw(point_a, point_b, point_c):
""" Check whether 3 points are counter clockwise from each other.
Parameters
----------
point_a: :class:`numpy.ndarray`
The first (`x`, `y`) point to check for counter clockwise ordering
point_b: :class:`numpy.ndarray`
The second (`x`, `y`) point to check for counter clockwise ordering
point_c: :class:`numpy.ndarray`
The third (`x`, `y`) point to check for counter clockwise ordering
Returns
-------
bool
``True`` if the 3 points are provided in counter clockwise order otherwise ``False``
"""
return ((point_c[1] - point_a[1]) * (point_b[0] - point_a[0]) >
(point_b[1] - point_a[1]) * (point_c[0] - point_a[0]))
def _rotate(self, event):
""" Rotates the landmarks contained within an extract box on a corner rotate drag event.
Parameters
----------
event: :class:`tkinter.Event`
The tkinter mouse event.
"""
face_idx = self._mouse_location[1]
face_tag = f"eb_box_face_{face_idx}"
box = np.array(self._canvas.coords(face_tag))
position = np.array((event.x, event.y))
center = np.array((sum(box[0::2]) / 4, sum(box[1::2]) / 4))
init_to_center = self._drag_data["current_location"] - center
new_to_center = position - center
angle = np.rad2deg(np.arctan2(*new_to_center) - np.arctan2(*init_to_center))
logger.trace("face_index: %s, box: %s, center: %s, init_to_center: %s, new_to_center: %s"
"center: %s, angle: %s", face_idx, box, center, init_to_center, new_to_center,
center, angle)
self._det_faces.update.landmarks_rotate(self._globals.frame_index,
face_idx,
angle,
self.scale_from_display(center))
self._drag_data["current_location"] = position
def _get_scale(self):
""" Obtain the scaling for the extract box resize """
def _context_menu(self, event):
""" Create a right click context menu to delete the alignment that is being
hovered over. """
if self._mouse_location is None or self._mouse_location[0] != "box":
return
self._right_click_menu.popup(event)
def _delete_current_face(self, *args): # pylint:disable=unused-argument
""" Called by the right click delete event. Deletes the face that the mouse is currently
over.
Parameters
----------
args: tuple (unused)
The event parameter is passed in by the hot key binding, so args is required
"""
if self._mouse_location is None or self._mouse_location[0] != "box":
return
self._det_faces.update.delete(self._globals.frame_index, self._mouse_location[1])