mirror of
https://github.com/deepfakes/faceswap
synced 2025-06-08 20:13:52 -04:00
4153a7ea0d
* Restructure tools to allow dynamic (plugable) loading * Update gui cli parsing to allow new tools structure
939 lines
38 KiB
Python
939 lines
38 KiB
Python
#!/usr/bin/env python3
|
|
""" Manual processing of alignments """
|
|
|
|
import logging
|
|
import platform
|
|
import sys
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from lib.faces_detect import DetectedFace
|
|
from lib.queue_manager import queue_manager
|
|
from plugins.extract.pipeline import Extractor, ExtractMedia
|
|
from .annotate import Annotate
|
|
from .media import ExtractedFaces, Frames
|
|
|
|
logger = logging.getLogger(__name__) # pylint: disable=invalid-name
|
|
|
|
|
|
class Interface():
|
|
""" Key controls and interfacing options for OpenCV """
|
|
def __init__(self, alignments, frames):
|
|
logger.debug("Initializing %s: (alignments: %s, frames: %s)",
|
|
self.__class__.__name__, alignments, frames)
|
|
self.alignments = alignments
|
|
self.frames = frames
|
|
self.controls = self.set_controls()
|
|
self.state = self.set_state()
|
|
self.skip_mode = {1: "Standard",
|
|
2: "No Faces",
|
|
3: "Multi-Faces",
|
|
4: "Has Faces"}
|
|
logger.debug("Initialized %s", self.__class__.__name__)
|
|
|
|
def set_controls(self):
|
|
""" Set keyboard controls, destination and help text """
|
|
controls = {"z": {"action": self.iterate_frame,
|
|
"args": ("navigation", - 1),
|
|
"help": "Previous Frame"},
|
|
"x": {"action": self.iterate_frame,
|
|
"args": ("navigation", 1),
|
|
"help": "Next Frame"},
|
|
"[": {"action": self.iterate_frame,
|
|
"args": ("navigation", - 100),
|
|
"help": "100 Frames Back"},
|
|
"]": {"action": self.iterate_frame,
|
|
"args": ("navigation", 100),
|
|
"help": "100 Frames Forward"},
|
|
"{": {"action": self.iterate_frame,
|
|
"args": ("navigation", "first"),
|
|
"help": "Go to First Frame"},
|
|
"}": {"action": self.iterate_frame,
|
|
"args": ("navigation", "last"),
|
|
"help": "Go to Last Frame"},
|
|
27: {"action": "quit",
|
|
"key_text": "ESC",
|
|
"args": ("navigation", None),
|
|
"help": "Exit",
|
|
"key_type": ord},
|
|
"/": {"action": self.iterate_state,
|
|
"args": ("navigation", "frame-size"),
|
|
"help": "Cycle Frame Zoom"},
|
|
"s": {"action": self.iterate_state,
|
|
"args": ("navigation", "skip-mode"),
|
|
"help": ("Skip Mode (All, No Faces, Multi Faces, Has Faces)")},
|
|
" ": {"action": self.save_alignments,
|
|
"key_text": "SPACE",
|
|
"args": ("edit", None),
|
|
"help": "Save Alignments"},
|
|
"r": {"action": self.reload_alignments,
|
|
"args": ("edit", None),
|
|
"help": "Reload Alignments (Discard all changes)"},
|
|
"d": {"action": self.delete_alignment,
|
|
"args": ("edit", None),
|
|
"help": "Delete Selected Alignment"},
|
|
"m": {"action": self.toggle_state,
|
|
"args": ("edit", "active"),
|
|
"help": "Change Mode (View, Edit)"},
|
|
range(10): {"action": self.set_state_value,
|
|
"key_text": "0 to 9",
|
|
"args": ["edit", "selected"],
|
|
"help": "Select/Deselect Face at this Index",
|
|
"key_type": range},
|
|
"c": {"action": self.copy_alignments,
|
|
"args": ("edit", -1),
|
|
"help": "Copy Alignments from Previous Frame with Alignments"},
|
|
"v": {"action": self.copy_alignments,
|
|
"args": ("edit", 1),
|
|
"help": "Copy Alignments from Next Frame with Alignments"},
|
|
"y": {"action": self.toggle_state,
|
|
"args": ("image", "display"),
|
|
"help": "Toggle Image"},
|
|
"u": {"action": self.iterate_state,
|
|
"args": ("bounding_box", "color"),
|
|
"help": "Cycle Bounding Box Color"},
|
|
"i": {"action": self.iterate_state,
|
|
"args": ("extract_box", "color"),
|
|
"help": "Cycle Extract Box Color"},
|
|
"o": {"action": self.iterate_state,
|
|
"args": ("landmarks", "color"),
|
|
"help": "Cycle Landmarks Color"},
|
|
"p": {"action": self.iterate_state,
|
|
"args": ("landmarks_mesh", "color"),
|
|
"help": "Cycle Landmarks Mesh Color"},
|
|
"h": {"action": self.iterate_state,
|
|
"args": ("bounding_box", "size"),
|
|
"help": "Cycle Bounding Box thickness"},
|
|
"j": {"action": self.iterate_state,
|
|
"args": ("extract_box", "size"),
|
|
"help": "Cycle Extract Box thickness"},
|
|
"k": {"action": self.iterate_state,
|
|
"args": ("landmarks", "size"),
|
|
"help": "Cycle Landmarks - point size"},
|
|
"l": {"action": self.iterate_state,
|
|
"args": ("landmarks_mesh", "size"),
|
|
"help": "Cycle Landmarks Mesh - thickness"}}
|
|
|
|
logger.debug("Controls: %s", controls)
|
|
return controls
|
|
|
|
@staticmethod
|
|
def set_state():
|
|
""" Set the initial display state """
|
|
state = {"bounding_box": dict(),
|
|
"extract_box": dict(),
|
|
"landmarks": dict(),
|
|
"landmarks_mesh": dict(),
|
|
"image": dict(),
|
|
"navigation": {"skip-mode": 1,
|
|
"frame-size": 1,
|
|
"frame_idx": 0,
|
|
"max_frame": 0,
|
|
"last_request": 0,
|
|
"frame_name": None},
|
|
"edit": {"updated": False,
|
|
"update_faces": False,
|
|
"selected": None,
|
|
"active": 0,
|
|
"redraw": False}}
|
|
|
|
# See lib_alignments/annotate.py for color mapping
|
|
color = 0
|
|
for key in sorted(state.keys()):
|
|
if key not in ("bounding_box", "extract_box", "landmarks", "landmarks_mesh", "image"):
|
|
continue
|
|
state[key]["display"] = True
|
|
if key == "image":
|
|
continue
|
|
color += 1
|
|
state[key]["size"] = 1
|
|
state[key]["color"] = color
|
|
logger.debug("State: %s", state)
|
|
return state
|
|
|
|
def save_alignments(self, *args): # pylint: disable=unused-argument
|
|
""" Save alignments """
|
|
logger.debug("Saving Alignments")
|
|
if not self.state["edit"]["updated"]:
|
|
logger.debug("Save received, but state not updated. Not saving")
|
|
return
|
|
self.alignments.save()
|
|
self.state["edit"]["updated"] = False
|
|
self.set_redraw(True)
|
|
|
|
def reload_alignments(self, *args): # pylint: disable=unused-argument
|
|
""" Reload alignments """
|
|
logger.debug("Reloading Alignments")
|
|
if not self.state["edit"]["updated"]:
|
|
logger.debug("Reload received, but state not updated. Not reloading")
|
|
return
|
|
self.alignments.reload()
|
|
self.state["edit"]["updated"] = False
|
|
self.state["edit"]["update_faces"] = True
|
|
self.set_redraw(True)
|
|
|
|
def delete_alignment(self, *args): # pylint: disable=unused-argument
|
|
""" Save alignments """
|
|
logger.debug("Deleting Alignments")
|
|
selected_face = self.get_selected_face_id()
|
|
if self.get_edit_mode() == "View" or selected_face is None:
|
|
logger.debug("Delete received, but edit mode is 'View'. Not deleting")
|
|
return
|
|
frame = self.get_frame_name()
|
|
if self.alignments.delete_face_at_index(frame, selected_face):
|
|
self.state["edit"]["selected"] = None
|
|
self.state["edit"]["updated"] = True
|
|
self.state["edit"]["update_faces"] = True
|
|
self.set_redraw(True)
|
|
|
|
def copy_alignments(self, *args):
|
|
""" Copy the alignments from the previous or next frame
|
|
to the current frame """
|
|
logger.debug("Copying Alignments")
|
|
if self.get_edit_mode() != "Edit":
|
|
logger.debug("Copy received, but edit mode is not 'Edit'. Not copying")
|
|
return
|
|
frame_id = self.get_next_face_idx(args[1])
|
|
if not 0 <= frame_id <= self.state["navigation"]["max_frame"]:
|
|
return
|
|
current_frame = self.get_frame_name()
|
|
get_frame = self.frames.file_list_sorted[frame_id]["frame_fullname"]
|
|
alignments = self.alignments.get_faces_in_frame(get_frame)
|
|
for alignment in alignments:
|
|
self.alignments. add_face(current_frame, alignment)
|
|
self.state["edit"]["updated"] = True
|
|
self.state["edit"]["update_faces"] = True
|
|
self.set_redraw(True)
|
|
|
|
def toggle_state(self, item, category):
|
|
""" Toggle state of requested item """
|
|
logger.debug("Toggling state: (item: %s, category: %s)", item, category)
|
|
self.state[item][category] = not self.state[item][category]
|
|
logger.debug("State toggled: (item: %s, category: %s, value: %s)",
|
|
item, category, self.state[item][category])
|
|
self.set_redraw(True)
|
|
|
|
def iterate_state(self, item, category):
|
|
""" Cycle through options (6 possible or 3 currently supported) """
|
|
logger.debug("Cycling state: (item: %s, category: %s)", item, category)
|
|
if category == "color":
|
|
max_val = 7
|
|
elif category == "frame-size":
|
|
max_val = 6
|
|
elif category == "skip-mode":
|
|
max_val = 4
|
|
else:
|
|
max_val = 3
|
|
val = self.state[item][category]
|
|
val = val + 1 if val != max_val else 1
|
|
self.state[item][category] = val
|
|
logger.debug("Cycled state: (item: %s, category: %s, value: %s)",
|
|
item, category, self.state[item][category])
|
|
self.set_redraw(True)
|
|
|
|
def set_state_value(self, item, category, value):
|
|
""" Set state of requested item or toggle off """
|
|
logger.debug("Setting state value: (item: %s, category: %s, value: %s)",
|
|
item, category, value)
|
|
state = self.state[item][category]
|
|
value = str(value) if value is not None else value
|
|
if state == value:
|
|
self.state[item][category] = None
|
|
else:
|
|
self.state[item][category] = value
|
|
logger.debug("Setting state value: (item: %s, category: %s, value: %s)",
|
|
item, category, self.state[item][category])
|
|
self.set_redraw(True)
|
|
|
|
def iterate_frame(self, *args):
|
|
""" Iterate frame up or down, stopping at either end """
|
|
logger.debug("Iterating frame: (args: %s)", args)
|
|
iteration = args[1]
|
|
max_frame = self.state["navigation"]["max_frame"]
|
|
if iteration in ("first", "last"):
|
|
next_frame = 0 if iteration == "first" else max_frame
|
|
self.state["navigation"]["frame_idx"] = next_frame
|
|
self.state["navigation"]["last_request"] = 0
|
|
self.set_redraw(True)
|
|
return
|
|
|
|
current_frame = self.state["navigation"]["frame_idx"]
|
|
next_frame = current_frame + iteration
|
|
end = 0 if iteration < 0 else max_frame
|
|
if (max_frame == 0 or
|
|
(end > 0 and next_frame >= end) or
|
|
(end == 0 and next_frame <= end)):
|
|
next_frame = end
|
|
self.state["navigation"]["frame_idx"] = next_frame
|
|
self.state["navigation"]["last_request"] = iteration
|
|
self.set_state_value("edit", "selected", None)
|
|
|
|
def get_color(self, item):
|
|
""" Return color for selected item """
|
|
return self.state[item]["color"]
|
|
|
|
def get_size(self, item):
|
|
""" Return size for selected item """
|
|
return self.state[item]["size"]
|
|
|
|
def get_frame_scaling(self):
|
|
""" Return frame scaling factor for requested item """
|
|
factors = (1, 1.25, 1.5, 2, 0.5, 0.75)
|
|
idx = self.state["navigation"]["frame-size"] - 1
|
|
return factors[idx]
|
|
|
|
def get_edit_mode(self):
|
|
""" Return text version and border color for edit mode """
|
|
if self.state["edit"]["active"]:
|
|
return "Edit"
|
|
return "View"
|
|
|
|
def get_skip_mode(self):
|
|
""" Return text version of skip mode """
|
|
return self.skip_mode[self.state["navigation"]["skip-mode"]]
|
|
|
|
def get_state_color(self):
|
|
""" Return a color based on current state
|
|
white - View Mode
|
|
yellow - Edit Mode
|
|
red - Unsaved alignments """
|
|
color = (255, 255, 255)
|
|
if self.state["edit"]["updated"]:
|
|
color = (0, 0, 255)
|
|
elif self.state["edit"]["active"]:
|
|
color = (0, 255, 255)
|
|
return color
|
|
|
|
def get_frame_name(self):
|
|
""" Return the current frame number """
|
|
return self.state["navigation"]["frame_name"]
|
|
|
|
def get_selected_face_id(self):
|
|
""" Return the index of the currently selected face """
|
|
try:
|
|
return int(self.state["edit"]["selected"])
|
|
except TypeError:
|
|
return None
|
|
|
|
def redraw(self):
|
|
""" Return whether a redraw is required """
|
|
return self.state["edit"]["redraw"]
|
|
|
|
def set_redraw(self, request):
|
|
""" Turn redraw requirement on or off """
|
|
self.state["edit"]["redraw"] = request
|
|
|
|
def get_next_face_idx(self, increment):
|
|
"""Get the index of the previous or next frame which has a face"""
|
|
navigation = self.state["navigation"]
|
|
frame_list = self.frames.file_list_sorted
|
|
frame_idx = navigation["frame_idx"] + increment
|
|
while True:
|
|
if not 0 <= frame_idx <= navigation["max_frame"]:
|
|
break
|
|
frame = frame_list[frame_idx]["frame_fullname"]
|
|
if not self.alignments.frame_has_faces(frame):
|
|
frame_idx += increment
|
|
else:
|
|
break
|
|
return frame_idx
|
|
|
|
|
|
class Help():
|
|
""" Generate and display help in cli and in window """
|
|
def __init__(self, interface):
|
|
logger.debug("Initializing %s: (interface: %s)", self.__class__.__name__, interface)
|
|
self.interface = interface
|
|
self.helptext = self.generate()
|
|
logger.debug("Initialized %s", self.__class__.__name__)
|
|
|
|
def generate(self):
|
|
""" Generate help output """
|
|
logger.debug("Generating help")
|
|
sections = ("navigation", "display", "color", "size", "edit")
|
|
helpout = {section: list() for section in sections}
|
|
helptext = ""
|
|
for key, val in self.interface.controls.items():
|
|
logger.trace("Generating help for:(key: '%s', val: '%s'", key, val)
|
|
help_section = val["args"][0]
|
|
if help_section not in ("navigation", "edit"):
|
|
help_section = val["args"][1]
|
|
key_text = val.get("key_text", None)
|
|
key_text = key_text if key_text else key
|
|
logger.trace("Adding help for:(section: '%s', val: '%s', text: '%s'",
|
|
help_section, val["help"], key_text)
|
|
helpout[help_section].append((val["help"], key_text))
|
|
|
|
helpout["edit"].append(("Bounding Box - Move", "Left Click"))
|
|
helpout["edit"].append(("Bounding Box - Resize", "Middle Click"))
|
|
|
|
for section in sections:
|
|
spacer = "=" * int((40 - len(section)) / 2)
|
|
display = "\n{0} {1} {0}\n".format(spacer, section.upper())
|
|
helpsection = sorted(helpout[section])
|
|
if section == "navigation":
|
|
helpsection = sorted(helpout[section], reverse=True)
|
|
display += "\n".join(" - '{}': {}".format(item[1], item[0])
|
|
for item in helpsection)
|
|
|
|
helptext += display
|
|
logger.debug("Added helptext: '%s'", helptext)
|
|
return helptext
|
|
|
|
def render(self):
|
|
""" Render help text to image window """
|
|
# pylint: disable=no-member
|
|
logger.trace("Rendering help text")
|
|
image = self.background()
|
|
display_text = self.helptext + self.compile_status()
|
|
self.text_to_image(image, display_text)
|
|
cv2.namedWindow("Help")
|
|
cv2.imshow("Help", image)
|
|
logger.trace("Rendered help text")
|
|
|
|
def background(self):
|
|
""" Create an image to hold help text """
|
|
# pylint: disable=no-member
|
|
logger.trace("Creating help text canvas")
|
|
height = 880
|
|
width = 480
|
|
image = np.zeros((height, width, 3), np.uint8)
|
|
color = self.interface.get_state_color()
|
|
cv2.rectangle(image, (0, 0), (width - 1, height - 1), color, 2)
|
|
logger.trace("Created help text canvas")
|
|
return image
|
|
|
|
def compile_status(self):
|
|
""" Render the status text """
|
|
logger.trace("Compiling Status text")
|
|
status = "\n=== STATUS\n"
|
|
navigation = self.interface.state["navigation"]
|
|
frame_scale = int(self.interface.get_frame_scaling() * 100)
|
|
status += " File: {}\n".format(self.interface.get_frame_name())
|
|
status += " Frame: {} / {}\n".format(
|
|
navigation["frame_idx"] + 1, navigation["max_frame"] + 1)
|
|
status += " Frame Size: {}%\n".format(frame_scale)
|
|
status += " Skip Mode: {}\n".format(self.interface.get_skip_mode())
|
|
status += " View Mode: {}\n".format(self.interface.get_edit_mode())
|
|
if self.interface.get_selected_face_id() is not None:
|
|
status += " Selected Face Index: {}\n".format(self.interface.get_selected_face_id())
|
|
if self.interface.state["edit"]["updated"]:
|
|
status += " Warning: There are unsaved changes\n"
|
|
|
|
logger.trace("Compiled Status text")
|
|
return status
|
|
|
|
@staticmethod
|
|
def text_to_image(image, display_text):
|
|
""" Write out and format help text to image """
|
|
# pylint: disable=no-member
|
|
logger.trace("Converting help text to image")
|
|
pos_y = 0
|
|
for line in display_text.split("\n"):
|
|
if line.startswith("==="):
|
|
pos_y += 10
|
|
line = line.replace("=", "").strip()
|
|
line = line.replace("- '", "[ ").replace("':", " ]")
|
|
cv2.putText(image, line, (20, pos_y),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.43, (255, 255, 255), 1)
|
|
pos_y += 20
|
|
logger.trace("Converted help text to image")
|
|
|
|
|
|
class Manual():
|
|
""" Manually adjust or create landmarks data """
|
|
def __init__(self, alignments, arguments):
|
|
logger.debug("Initializing %s: (alignments: %s, arguments: %s)",
|
|
self.__class__.__name__, alignments, arguments)
|
|
self.arguments = arguments
|
|
self.alignments = alignments
|
|
self.frames = Frames(arguments.frames_dir)
|
|
self.extracted_faces = None
|
|
self.interface = None
|
|
self.help = None
|
|
self.mouse_handler = None
|
|
logger.debug("Initialized %s", self.__class__.__name__)
|
|
|
|
def process(self):
|
|
""" Process manual extraction """
|
|
logger.info("[MANUAL PROCESSING]") # Tidy up cli output
|
|
self.extracted_faces = ExtractedFaces(self.frames, self.alignments, size=256)
|
|
self.interface = Interface(self.alignments, self.frames)
|
|
self.help = Help(self.interface)
|
|
self.mouse_handler = MouseHandler(self.interface, self.arguments.loglevel)
|
|
|
|
print(self.help.helptext)
|
|
max_idx = self.frames.count - 1
|
|
self.interface.state["navigation"]["max_frame"] = max_idx
|
|
self.display_frames()
|
|
|
|
def display_frames(self):
|
|
""" Iterate through frames """
|
|
# pylint: disable=no-member
|
|
logger.debug("Display frames")
|
|
is_windows = platform.system() == "Windows"
|
|
is_conda = "conda" in sys.version.lower()
|
|
logger.debug("is_windows: %s, is_conda: %s", is_windows, is_conda)
|
|
cv2.namedWindow("Frame")
|
|
cv2.namedWindow("Faces")
|
|
cv2.setMouseCallback('Frame', self.mouse_handler.on_event)
|
|
|
|
frame, faces = self.get_frame()
|
|
press = self.get_keys()
|
|
self.interface.set_redraw(True)
|
|
|
|
while True:
|
|
if self.interface.redraw():
|
|
self.help.render()
|
|
cv2.imshow("Frame", frame)
|
|
cv2.imshow("Faces", faces)
|
|
self.interface.set_redraw(False)
|
|
key = cv2.waitKey(1000)
|
|
|
|
if self.window_closed(is_windows, is_conda, key):
|
|
queue_manager.terminate_queues()
|
|
break
|
|
|
|
if key and key != -1:
|
|
logger.trace("Keypress received: '%s'", key)
|
|
if key in press.keys():
|
|
action = press[key]["action"]
|
|
logger.debug("Keypress action: key: ('%s', action: '%s')", key, action)
|
|
if action == "quit":
|
|
break
|
|
|
|
if press[key].get("key_type") == range:
|
|
args = press[key]["args"] + [chr(key)]
|
|
else:
|
|
args = press[key]["args"]
|
|
action(*args)
|
|
|
|
if not self.interface.redraw():
|
|
continue
|
|
|
|
logger.trace("Redraw requested")
|
|
frame, faces = self.get_frame()
|
|
|
|
cv2.destroyAllWindows()
|
|
|
|
def window_closed(self, is_windows, is_conda, key):
|
|
""" Check whether the window has been closed
|
|
|
|
MS Windows doesn't appear to read the window state property
|
|
properly, so we check for a negative key press.
|
|
|
|
Conda (tested on Windows) doesn't appear to read the window
|
|
state property or negative key press properly, so we arbitrarily
|
|
use another property """
|
|
# pylint: disable=no-member
|
|
logger.trace("Commencing closed window check")
|
|
closed = False
|
|
prop_autosize = cv2.getWindowProperty('Frame', cv2.WND_PROP_AUTOSIZE)
|
|
prop_visible = cv2.getWindowProperty('Frame', cv2.WND_PROP_VISIBLE)
|
|
if self.arguments.disable_monitor:
|
|
closed = False
|
|
elif is_conda and prop_autosize < 1:
|
|
closed = True
|
|
elif is_windows and not is_conda and key == -1:
|
|
closed = True
|
|
elif not is_windows and not is_conda and prop_visible < 1:
|
|
closed = True
|
|
logger.trace("Completed closed window check. Closed is %s", closed)
|
|
if closed:
|
|
logger.debug("Window closed detected")
|
|
return closed
|
|
|
|
def get_keys(self):
|
|
""" Convert keys dict into something useful
|
|
for OpenCV """
|
|
keys = dict()
|
|
for key, val in self.interface.controls.items():
|
|
if val.get("key_type", str) == range:
|
|
for range_key in key:
|
|
keys[ord(str(range_key))] = val
|
|
elif val.get("key_type", str) == ord:
|
|
keys[key] = val
|
|
else:
|
|
keys[ord(key)] = val
|
|
|
|
return keys
|
|
|
|
def get_frame(self):
|
|
""" Compile the frame and get faces """
|
|
image = self.frame_selector()
|
|
frame_name = self.interface.get_frame_name()
|
|
logger.debug("Frame Name: '%s'", frame_name)
|
|
alignments = self.alignments.get_faces_in_frame(frame_name)
|
|
faces_updated = self.interface.state["edit"]["update_faces"]
|
|
logger.debug("Faces Updated: %s", faces_updated)
|
|
self.extracted_faces.get_faces(frame_name)
|
|
roi = [face.original_roi for face in self.extracted_faces.faces]
|
|
|
|
if faces_updated:
|
|
self.interface.state["edit"]["update_faces"] = False
|
|
|
|
frame = FrameDisplay(image, alignments, roi, self.interface).image
|
|
faces = self.set_faces(frame_name).image
|
|
return frame, faces
|
|
|
|
def frame_selector(self):
|
|
""" Return frame at given index """
|
|
navigation = self.interface.state["navigation"]
|
|
frame_list = self.frames.file_list_sorted
|
|
frame = frame_list[navigation["frame_idx"]]["frame_fullname"]
|
|
skip_mode = self.interface.get_skip_mode().lower()
|
|
logger.debug("navigation: %s, frame: '%s', skip_mode: '%s'", navigation, frame, skip_mode)
|
|
|
|
while True:
|
|
if navigation["last_request"] == 0:
|
|
break
|
|
if navigation["frame_idx"] in (0, navigation["max_frame"]):
|
|
break
|
|
if skip_mode == "standard":
|
|
break
|
|
if skip_mode == "no faces" and not self.alignments.frame_has_faces(frame):
|
|
break
|
|
if skip_mode == "multi-faces" and self.alignments.frame_has_multiple_faces(frame):
|
|
break
|
|
if skip_mode == "has faces" and self.alignments.frame_has_faces(frame):
|
|
break
|
|
self.interface.iterate_frame("navigation", navigation["last_request"])
|
|
frame = frame_list[navigation["frame_idx"]]["frame_fullname"]
|
|
|
|
image = self.frames.load_image(frame)
|
|
navigation["last_request"] = 0
|
|
navigation["frame_name"] = frame
|
|
return image
|
|
|
|
def set_faces(self, frame):
|
|
""" Pass the current frame faces to faces window """
|
|
faces = self.extracted_faces.get_faces_in_frame(frame)
|
|
landmarks = [{"landmarks_xy": face.aligned_landmarks}
|
|
for face in self.extracted_faces.faces]
|
|
return FacesDisplay(faces, landmarks, self.extracted_faces.size, self.interface)
|
|
|
|
|
|
class FrameDisplay():
|
|
"""" Window that holds the frame """
|
|
def __init__(self, image, alignments, roi, interface):
|
|
logger.trace("Initializing %s: (alignments: %s, roi: %s, interface: %s)",
|
|
self.__class__.__name__, alignments, roi, interface)
|
|
self.image = image
|
|
self.roi = roi
|
|
self.alignments = alignments
|
|
self.interface = interface
|
|
self.annotate_frame()
|
|
logger.trace("Initialized %s", self.__class__.__name__)
|
|
|
|
def annotate_frame(self):
|
|
""" Annotate the frame """
|
|
state = self.interface.state
|
|
logger.trace("State: %s", state)
|
|
annotate = Annotate(self.image, self.alignments, self.roi)
|
|
if not state["image"]["display"]:
|
|
annotate.draw_black_image()
|
|
|
|
for item in ("bounding_box", "extract_box", "landmarks", "landmarks_mesh"):
|
|
color = self.interface.get_color(item)
|
|
size = self.interface.get_size(item)
|
|
state[item]["display"] = color != 7
|
|
if not state[item]["display"]:
|
|
continue
|
|
logger.trace("Annotating: '%s'", item)
|
|
annotation = getattr(annotate, "draw_{}".format(item))
|
|
annotation(color, size)
|
|
|
|
selected_face = self.interface.get_selected_face_id()
|
|
if (selected_face is not None and
|
|
int(selected_face) < len(self.alignments)):
|
|
annotate.draw_grey_out_faces(selected_face)
|
|
|
|
self.image = self.resize_frame(annotate.image)
|
|
|
|
def resize_frame(self, image):
|
|
""" Set the displayed frame size and add state border"""
|
|
# pylint: disable=no-member
|
|
logger.trace("Resizing frame")
|
|
height, width = image.shape[:2]
|
|
color = self.interface.get_state_color()
|
|
cv2.rectangle(image, (0, 0), (width - 1, height - 1), color, 1)
|
|
scaling = self.interface.get_frame_scaling()
|
|
image = cv2.resize(image, (0, 0), fx=scaling, fy=scaling)
|
|
logger.trace("Resized frame")
|
|
return image
|
|
|
|
|
|
class FacesDisplay():
|
|
""" Window that holds faces thumbnail """
|
|
def __init__(self, extracted_faces, landmarks, size, interface):
|
|
logger.trace("Initializing %s: (extracted_faces: %s, landmarks: %s, size: %s, "
|
|
"interface: %s)", self.__class__.__name__, extracted_faces,
|
|
landmarks, size, interface)
|
|
self.row_length = 4
|
|
self.faces = self.copy_faces(extracted_faces)
|
|
self.roi = self.set_full_roi(size)
|
|
self.landmarks = landmarks
|
|
self.interface = interface
|
|
|
|
self.annotate_faces()
|
|
|
|
self.image = self.build_faces_image(size)
|
|
logger.trace("Initialized %s", self.__class__.__name__)
|
|
|
|
@staticmethod
|
|
def copy_faces(faces):
|
|
""" Copy the extracted faces so as not to save the annotations back """
|
|
return [face.aligned_face.copy() for face in faces]
|
|
|
|
@staticmethod
|
|
def set_full_roi(size):
|
|
""" ROI is the full frame for faces, so set based on size """
|
|
return [np.array([[(0, 0), (0, size - 1), (size - 1, size - 1), (size - 1, 0)]], np.int32)]
|
|
|
|
def annotate_faces(self):
|
|
""" Annotate each of the faces """
|
|
state = self.interface.state
|
|
selected_face = self.interface.get_selected_face_id()
|
|
logger.trace("State: %s, Selected Face ID: %s", state, selected_face)
|
|
for idx, face in enumerate(self.faces):
|
|
annotate = Annotate(face, [self.landmarks[idx]], self.roi)
|
|
if not state["image"]["display"]:
|
|
annotate.draw_black_image()
|
|
|
|
for item in ("landmarks", "landmarks_mesh"):
|
|
if not state[item]["display"]:
|
|
continue
|
|
logger.trace("Annotating: '%s'", item)
|
|
color = self.interface.get_color(item)
|
|
size = self.interface.get_size(item)
|
|
annotation = getattr(annotate, "draw_{}".format(item))
|
|
annotation(color, size)
|
|
|
|
if (selected_face is not None
|
|
and int(selected_face) < len(self.faces)
|
|
and int(selected_face) != idx):
|
|
annotate.draw_grey_out_faces(1)
|
|
|
|
self.faces[idx] = annotate.image
|
|
|
|
def build_faces_image(self, size):
|
|
""" Display associated faces """
|
|
total_faces = len(self.faces)
|
|
logger.trace("Building faces panel. (total_faces: %s)", total_faces)
|
|
if not total_faces:
|
|
logger.trace("Returning empty row")
|
|
image = self.build_faces_row(list(), size)
|
|
return image
|
|
total_rows = int(total_faces / self.row_length) + 1
|
|
for idx in range(total_rows):
|
|
logger.trace("Building row %s", idx)
|
|
face_idx = idx * self.row_length
|
|
row_faces = self.faces[face_idx:face_idx + self.row_length]
|
|
if not row_faces:
|
|
break
|
|
row = self.build_faces_row(row_faces, size)
|
|
image = row if idx == 0 else np.concatenate((image, row), axis=0)
|
|
return image
|
|
|
|
def build_faces_row(self, faces, size):
|
|
""" Build a row of 4 faces """
|
|
# pylint: disable=no-member
|
|
logger.trace("Building row for %s faces", len(faces))
|
|
if len(faces) != 4:
|
|
remainder = 4 - (len(faces) % self.row_length)
|
|
for _ in range(remainder):
|
|
faces.append(np.zeros((size, size, 3), np.uint8))
|
|
for idx, face in enumerate(faces):
|
|
color = self.interface.get_state_color()
|
|
cv2.rectangle(face, (0, 0), (size - 1, size - 1),
|
|
color, 1)
|
|
if idx == 0:
|
|
row = face
|
|
else:
|
|
row = np.concatenate((row, face), axis=1)
|
|
return row
|
|
|
|
|
|
class MouseHandler():
|
|
""" Manual Extraction """
|
|
def __init__(self, interface, loglevel):
|
|
logger.debug("Initializing %s: (interface: %s, loglevel: %s)",
|
|
self.__class__.__name__, interface, loglevel)
|
|
self.interface = interface
|
|
self.alignments = interface.alignments
|
|
self.frames = interface.frames
|
|
|
|
self.queues = dict()
|
|
self.extractor = self.init_extractor()
|
|
|
|
self.mouse_state = None
|
|
self.last_move = None
|
|
self.center = None
|
|
self.dims = None
|
|
self.media = {"frame_id": None,
|
|
"image": None,
|
|
"bounding_box": list(),
|
|
"bounding_last": list(),
|
|
"bounding_box_orig": list()}
|
|
logger.debug("Initialized %s", self.__class__.__name__)
|
|
|
|
def init_extractor(self):
|
|
""" Initialize Aligner """
|
|
logger.debug("Initialize Extractor")
|
|
extractor = Extractor(None, "fan", None, multiprocess=True, normalize_method="hist")
|
|
self.queues["in"] = extractor.input_queue
|
|
# Set the batchsize to 1
|
|
extractor.set_batchsize("align", 1)
|
|
extractor.launch()
|
|
logger.debug("Initialized Extractor")
|
|
return extractor
|
|
|
|
def on_event(self, event, x, y, flags, param): # pylint: disable=unused-argument,invalid-name
|
|
""" Handle the mouse events """
|
|
# pylint: disable=no-member
|
|
if self.interface.get_edit_mode() != "Edit":
|
|
return
|
|
logger.trace("Mouse event: (event: %s, x: %s, y: %s, flags: %s, param: %s",
|
|
event, x, y, flags, param)
|
|
if not self.mouse_state and event not in (cv2.EVENT_LBUTTONDOWN, cv2.EVENT_MBUTTONDOWN):
|
|
return
|
|
|
|
self.initialize()
|
|
|
|
if event in (cv2.EVENT_LBUTTONUP, cv2.EVENT_MBUTTONUP):
|
|
self.mouse_state = None
|
|
self.last_move = None
|
|
elif event == cv2.EVENT_LBUTTONDOWN:
|
|
self.mouse_state = "left"
|
|
self.set_bounding_box(x, y)
|
|
elif event == cv2.EVENT_MBUTTONDOWN:
|
|
self.mouse_state = "middle"
|
|
self.set_bounding_box(x, y)
|
|
elif event == cv2.EVENT_MOUSEMOVE:
|
|
if self.mouse_state == "left":
|
|
self.move_bounding_box(x, y)
|
|
elif self.mouse_state == "middle":
|
|
self.resize_bounding_box(x, y)
|
|
|
|
def initialize(self):
|
|
""" Update changed parameters """
|
|
frame = self.interface.get_frame_name()
|
|
if frame == self.media["frame_id"]:
|
|
return
|
|
logger.debug("Initialize frame: '%s'", frame)
|
|
self.media["frame_id"] = frame
|
|
self.media["image"] = self.frames.load_image(frame)
|
|
self.dims = None
|
|
self.center = None
|
|
self.last_move = None
|
|
self.mouse_state = None
|
|
self.media["bounding_box"] = DetectedFace()
|
|
self.media["bounding_box_orig"] = None
|
|
|
|
def set_bounding_box(self, pt_x, pt_y):
|
|
""" Select or create bounding box """
|
|
if self.interface.get_selected_face_id() is None:
|
|
self.check_click_location(pt_x, pt_y)
|
|
|
|
if self.interface.get_selected_face_id() is not None:
|
|
self.dims_from_alignment()
|
|
else:
|
|
self.dims_from_image()
|
|
|
|
self.move_bounding_box(pt_x, pt_y)
|
|
|
|
def check_click_location(self, pt_x, pt_y):
|
|
""" Check whether the point clicked is within an existing
|
|
bounding box and set face_id """
|
|
frame = self.media["frame_id"]
|
|
alignments = self.alignments.get_faces_in_frame(frame)
|
|
scale = self.interface.get_frame_scaling()
|
|
pt_x = int(pt_x / scale)
|
|
pt_y = int(pt_y / scale)
|
|
|
|
for idx, alignment in enumerate(alignments):
|
|
left = alignment["x"]
|
|
right = alignment["x"] + alignment["w"]
|
|
top = alignment["y"]
|
|
bottom = alignment["y"] + alignment["h"]
|
|
|
|
if left <= pt_x <= right and top <= pt_y <= bottom:
|
|
self.interface.set_state_value("edit", "selected", idx)
|
|
break
|
|
|
|
def dims_from_alignment(self):
|
|
""" Set the height and width of bounding box from alignment """
|
|
frame = self.media["frame_id"]
|
|
face_id = self.interface.get_selected_face_id()
|
|
alignment = self.alignments.get_faces_in_frame(frame)[face_id]
|
|
self.dims = (alignment["w"], alignment["h"])
|
|
|
|
def dims_from_image(self):
|
|
""" Set the height and width of bounding
|
|
box at 10% of longest axis """
|
|
size = max(self.media["image"].shape[:2])
|
|
dim = int(size / 10.00)
|
|
self.dims = (dim, dim)
|
|
|
|
def bounding_from_center(self):
|
|
""" Get bounding X Y from center """
|
|
pt_x, pt_y = self.center
|
|
width, height = self.dims
|
|
scale = self.interface.get_frame_scaling()
|
|
self.media["bounding_box"].x = int((pt_x / scale) - width / 2)
|
|
self.media["bounding_box"].y = int((pt_y / scale) - height / 2)
|
|
self.media["bounding_box"].w = width
|
|
self.media["bounding_box"].h = height
|
|
|
|
def move_bounding_box(self, pt_x, pt_y):
|
|
""" Move the bounding box """
|
|
self.center = (pt_x, pt_y)
|
|
self.bounding_from_center()
|
|
self.update_landmarks()
|
|
|
|
def resize_bounding_box(self, pt_x, pt_y):
|
|
""" Resize the bounding box """
|
|
scale = self.interface.get_frame_scaling()
|
|
if not self.last_move:
|
|
self.last_move = (pt_x, pt_y)
|
|
self.media["bounding_box_orig"] = self.media["bounding_box"]
|
|
|
|
move_x = int(pt_x - self.last_move[0])
|
|
move_y = int(self.last_move[1] - pt_y)
|
|
|
|
original = self.media["bounding_box_orig"]
|
|
updated = self.media["bounding_box"]
|
|
|
|
minsize = int(20 / scale)
|
|
center = (int(self.center[0] / scale), int(self.center[1] / scale))
|
|
updated.x = min(center[0] - (minsize // 2), original.x - move_x)
|
|
updated.y = min(center[1] - (minsize // 2), original.y - move_y)
|
|
updated.w = max(minsize, original.w + move_x)
|
|
updated.h = max(minsize, original.h + move_y)
|
|
self.update_landmarks()
|
|
self.last_move = (pt_x, pt_y)
|
|
|
|
def update_landmarks(self):
|
|
""" Update the landmarks """
|
|
feed = ExtractMedia(self.media["frame_id"],
|
|
self.media["image"],
|
|
detected_faces=[self.media["bounding_box"]])
|
|
self.queues["in"].put(feed)
|
|
detected_face = next(self.extractor.detected_faces()).detected_faces[0]
|
|
alignment = detected_face.to_alignment()
|
|
# Mask will now be incorrect for updated landmarks so delete
|
|
alignment["mask"] = dict()
|
|
|
|
frame = self.media["frame_id"]
|
|
|
|
if self.interface.get_selected_face_id() is None:
|
|
idx = self.alignments.add_face(frame, alignment)
|
|
self.interface.set_state_value("edit", "selected", idx)
|
|
else:
|
|
self.alignments.update_face(frame,
|
|
self.interface.get_selected_face_id(),
|
|
alignment)
|
|
self.interface.set_redraw(True)
|
|
|
|
self.interface.state["edit"]["updated"] = True
|
|
self.interface.state["edit"]["update_faces"] = True
|