mirror of
https://github.com/deepfakes/faceswap.git
synced 2026-10-29 14:09:34 +08:00
364 lines
14 KiB
Python
364 lines
14 KiB
Python
#!/usr/bin/env python3
|
|
"""Tools for manipulating the alignments using Frames as a source """
|
|
from __future__ import annotations
|
|
import logging
|
|
import os
|
|
import sys
|
|
import typing as T
|
|
|
|
from datetime import datetime
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from tqdm import tqdm
|
|
|
|
from lib.align import DetectedFace, LANDMARK_PARTS, LandmarkType
|
|
from lib.align.objects import PNGHeader, PNGSource
|
|
from lib.image import encode_image, ImagesSaver
|
|
from lib.utils import get_module_objects, deprecation_warning
|
|
from .media import ExtractedFaces, Frames
|
|
|
|
if T.TYPE_CHECKING:
|
|
from argparse import Namespace
|
|
from .media import AlignmentData
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class Draw():
|
|
"""Draws annotations onto original frames and saves into a sub-folder next to the original
|
|
frames.
|
|
|
|
Parameters
|
|
---------
|
|
alignments
|
|
The loaded alignments corresponding to the frames to be annotated
|
|
arguments
|
|
The command line arguments that have called this job
|
|
"""
|
|
def __init__(self, alignments: AlignmentData, arguments: Namespace) -> None:
|
|
logger.debug("Initializing %s: (arguments: %s)", self.__class__.__name__, arguments)
|
|
self._alignments = alignments
|
|
self._frames = Frames(arguments.frames_dir)
|
|
self._output_folder = self._set_output()
|
|
logger.debug("Initialized %s", self.__class__.__name__)
|
|
|
|
def _set_output(self) -> str:
|
|
"""Set the output folder path.
|
|
|
|
If annotating a folder of frames, output will be placed in a sub folder within the frames
|
|
folder. If annotating a video, output will be a folder next to the original video.
|
|
|
|
Returns
|
|
-------
|
|
Full path to the output folder
|
|
"""
|
|
now = datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
folder_name = f"drawn_landmarks_{now}"
|
|
if self._frames.is_video:
|
|
dest_folder = os.path.dirname(self._frames.folder)
|
|
else:
|
|
dest_folder = self._frames.folder
|
|
output_folder = os.path.join(dest_folder, folder_name)
|
|
logger.debug("Creating folder: '%s'", output_folder)
|
|
os.makedirs(output_folder)
|
|
return output_folder
|
|
|
|
def process(self) -> None:
|
|
"""Runs the process to draw face annotations onto original source frames."""
|
|
logger.info("[DRAW LANDMARKS]") # Tidy up cli output
|
|
frames_drawn = 0
|
|
for frame in tqdm(T.cast(list[dict[str, str]], self._frames.file_list_sorted),
|
|
desc="Drawing landmarks",
|
|
leave=False):
|
|
frame_name = frame["frame_fullname"]
|
|
|
|
if not self._alignments.frame_exists(frame_name):
|
|
logger.verbose("Skipping '%s' - Alignments not found", frame_name) # type:ignore
|
|
continue
|
|
|
|
self._annotate_image(frame_name)
|
|
frames_drawn += 1
|
|
logger.info("%s Frame(s) output", frames_drawn)
|
|
|
|
def _annotate_image(self, frame_name: str) -> None:
|
|
"""Annotate the frame with each face that appears in the alignments file.
|
|
|
|
Parameters
|
|
----------
|
|
frame_name
|
|
The full path to the original frame
|
|
"""
|
|
logger.trace("Annotating frame: '%s'", frame_name) # type:ignore
|
|
image = self._frames.load_image(frame_name)
|
|
|
|
for idx, alignment in enumerate(self._alignments.get_faces_in_frame(frame_name)):
|
|
face = DetectedFace()
|
|
face.from_alignment(alignment, image=image)
|
|
# Bounding Box
|
|
assert face.left is not None
|
|
assert face.top is not None
|
|
cv2.rectangle(image, (face.left, face.top), (face.right, face.bottom), (255, 0, 0), 1)
|
|
self._annotate_landmarks(image, np.rint(face.landmarks_xy).astype("int32"))
|
|
self._annotate_extract_boxes(image, face, idx)
|
|
self._annotate_pose(image, face) # Pose (head is still loaded)
|
|
|
|
self._frames.save_image(self._output_folder, frame_name, image)
|
|
|
|
def _annotate_landmarks(self, image: np.ndarray, landmarks: np.ndarray) -> None:
|
|
"""Annotate the extract boxes onto the frame.
|
|
|
|
Parameters
|
|
----------
|
|
image
|
|
The frame that extract boxes are to be annotated on to
|
|
landmarks
|
|
The facial landmarks that are to be annotated onto the frame
|
|
"""
|
|
# Mesh
|
|
for start, end, fill in LANDMARK_PARTS[LandmarkType.from_shape(landmarks.shape)].values():
|
|
cv2.polylines(image, [landmarks[start:end]], fill, (255, 255, 0), 1)
|
|
# Landmarks
|
|
for (pos_x, pos_y) in landmarks:
|
|
cv2.circle(image, (pos_x, pos_y), 1, (0, 255, 255), -1)
|
|
|
|
@classmethod
|
|
def _annotate_extract_boxes(cls, image: np.ndarray, face: DetectedFace, index: int) -> None:
|
|
"""Annotate the mesh and landmarks boxes onto the frame.
|
|
|
|
Parameters
|
|
----------
|
|
image
|
|
The frame that mesh and landmarks are to be annotated on to
|
|
face
|
|
The aligned face
|
|
index
|
|
The face index for the given face
|
|
"""
|
|
for area in T.get_args(T.Literal["face", "head"]):
|
|
face.load_aligned(image, centering=area, force=True)
|
|
color = (0, 255, 0) if area == "face" else (0, 0, 255)
|
|
top_left = face.aligned.original_roi[0]
|
|
top_left = (top_left[0], top_left[1] - 10)
|
|
cv2.putText(image, str(index), top_left, cv2.FONT_HERSHEY_DUPLEX, 1.0, color, 1)
|
|
cv2.polylines(image, [face.aligned.original_roi], True, color, 1)
|
|
|
|
@classmethod
|
|
def _annotate_pose(cls, image: np.ndarray, face: DetectedFace) -> None:
|
|
"""Annotate the pose onto the frame.
|
|
|
|
Parameters
|
|
----------
|
|
image
|
|
The frame that pose is to be annotated on to
|
|
face
|
|
The aligned face loaded for head centering
|
|
"""
|
|
center = np.array((face.aligned.size / 2,
|
|
face.aligned.size / 2)).astype("int32").reshape(1, 2)
|
|
center = np.rint(face.aligned.transform_points(center, invert=True)).astype("int32")
|
|
points = face.aligned.pose.xyz_2d * face.aligned.size
|
|
points = np.rint(face.aligned.transform_points(points, invert=True)).astype("int32")
|
|
cv2.line(image, tuple(center), tuple(points[1]), (0, 255, 0), 2)
|
|
cv2.line(image, tuple(center), tuple(points[0]), (255, 0, 0), 2)
|
|
cv2.line(image, tuple(center), tuple(points[2]), (0, 0, 255), 2)
|
|
|
|
|
|
class Extract():
|
|
"""Re-extract faces from source frames based on Alignment data
|
|
|
|
Parameters
|
|
----------
|
|
alignments
|
|
The alignments data loaded from an alignments file for this rename job
|
|
arguments
|
|
The :mod:`argparse` arguments as passed in from :mod:`tools.py`
|
|
"""
|
|
def __init__(self, alignments: AlignmentData, arguments: Namespace) -> None:
|
|
logger.debug("Initializing %s: (arguments: %s)", self.__class__.__name__, arguments)
|
|
deprecation_warning("'Extract' job", "Use 'python faceswap.py extract' instead, selecting "
|
|
"the 'file' aligner plugin.")
|
|
self._arguments = arguments
|
|
self._alignments = alignments
|
|
self._faces_dir = arguments.faces_dir
|
|
self._min_size = self._get_min_size(arguments.size, arguments.min_size)
|
|
|
|
self._frames = Frames(arguments.frames_dir, self._get_count())
|
|
self._extracted_faces = ExtractedFaces(self._frames,
|
|
self._alignments,
|
|
size=arguments.size)
|
|
self._saver: ImagesSaver | None = None
|
|
logger.debug("Initialized %s", self.__class__.__name__)
|
|
|
|
@classmethod
|
|
def _get_min_size(cls, extract_size: int, min_size: int) -> int:
|
|
"""Obtain the minimum size that a face has been resized from to be included as a valid
|
|
extract.
|
|
|
|
Parameters
|
|
----------
|
|
extract_size
|
|
The requested size of the extracted images
|
|
min_size
|
|
The percentage amount that has been supplied for valid faces (as a percentage of
|
|
extract size)
|
|
|
|
Returns
|
|
-------
|
|
The minimum size, in pixels, that a face is resized from to be considered valid
|
|
"""
|
|
retval = 0 if min_size == 0 else max(4, int(extract_size * (min_size / 100.)))
|
|
logger.debug("Extract size: %s, min percentage size: %s, min_size: %s",
|
|
extract_size, min_size, retval)
|
|
return retval
|
|
|
|
def _get_count(self) -> int | None:
|
|
"""If the alignments file has been run through the manual tool, then it will hold video
|
|
meta information, meaning that the count of frames in the alignment file can be relied
|
|
on to be accurate.
|
|
|
|
Returns
|
|
-------
|
|
For video input which contain video meta-data in the alignments file then the count of
|
|
frames is returned. In all other cases ``None`` is returned
|
|
"""
|
|
meta = self._alignments.video_meta_data
|
|
has_meta = meta is not None and all(val is not None for val in meta.values())
|
|
if has_meta:
|
|
assert meta is not None
|
|
retval: int | None = len(T.cast(dict[str, list[int]], meta["pts_time"]))
|
|
else:
|
|
retval = None
|
|
logger.debug("Frame count from alignments file: (has_meta: %s, %s", has_meta, retval)
|
|
return retval
|
|
|
|
def process(self) -> None:
|
|
"""Run the re-extraction from Alignments file process"""
|
|
logger.info("[EXTRACT FACES]") # Tidy up cli output
|
|
self._check_folder()
|
|
self._saver = ImagesSaver(self._faces_dir, as_bytes=True)
|
|
|
|
if self._min_size > 0:
|
|
logger.info("Only selecting faces that have been resized from a minimum resolution "
|
|
"of %spx", self._min_size)
|
|
|
|
self._export_faces()
|
|
|
|
def _check_folder(self) -> None:
|
|
"""Check that the faces folder doesn't pre-exist and create."""
|
|
err = None
|
|
if not self._faces_dir:
|
|
err = "ERROR: Output faces folder not provided."
|
|
elif not os.path.isdir(self._faces_dir):
|
|
logger.debug("Creating folder: '%s'", self._faces_dir)
|
|
os.makedirs(self._faces_dir)
|
|
elif os.listdir(self._faces_dir):
|
|
err = f"ERROR: Output faces folder should be empty: '{self._faces_dir}'"
|
|
if err:
|
|
logger.error(err)
|
|
sys.exit(0)
|
|
logger.verbose("Creating output folder at '%s'", self._faces_dir) # type:ignore
|
|
|
|
def _export_faces(self) -> None:
|
|
"""Export the faces to the output folder."""
|
|
extracted_faces = 0
|
|
skip_list = self._set_skip_list()
|
|
count = self._frames.count if skip_list is None else self._frames.count - len(skip_list)
|
|
|
|
for filename, image in tqdm(self._frames.stream(skip_list=skip_list),
|
|
total=count, desc="Saving extracted faces",
|
|
leave=False):
|
|
frame_name = os.path.basename(filename)
|
|
if not self._alignments.frame_exists(frame_name):
|
|
logger.verbose("Skipping '%s' - Alignments not found", frame_name) # type:ignore
|
|
continue
|
|
extracted_faces += self._output_faces(frame_name, image)
|
|
logger.info("%s face(s) extracted", extracted_faces)
|
|
|
|
def _set_skip_list(self) -> list[int] | None:
|
|
"""Set the indices for frames that should be skipped based on the `extract_every_n`
|
|
command line option.
|
|
|
|
Returns
|
|
-------
|
|
A list of indices to be skipped if extract_every_n is not `1` otherwise returns ``None``
|
|
"""
|
|
skip_num = self._arguments.extract_every_n
|
|
if skip_num == 1:
|
|
logger.debug("Not skipping any frames")
|
|
return None
|
|
skip_list = []
|
|
for idx, item in enumerate(T.cast(list[dict[str, str]], self._frames.file_list_sorted)):
|
|
if idx % skip_num != 0:
|
|
logger.trace("Adding image '%s' to skip list due to " # type:ignore
|
|
"extract_every_n = %s", item["frame_fullname"], skip_num)
|
|
skip_list.append(idx)
|
|
logger.debug("Adding skip list: %s", skip_list)
|
|
return skip_list
|
|
|
|
def _output_faces(self, filename: str, image: np.ndarray) -> int:
|
|
"""For each frame save out the faces
|
|
|
|
Parameters
|
|
----------
|
|
filename
|
|
The filename (without the full path) of the current frame
|
|
image
|
|
The full frame that faces are to be extracted from
|
|
|
|
Returns
|
|
-------
|
|
The total number of faces that have been extracted
|
|
"""
|
|
logger.trace("Outputting frame: %s", filename) # type:ignore
|
|
face_count = 0
|
|
frame_name = os.path.splitext(filename)[0]
|
|
faces = self._select_valid_faces(filename, image)
|
|
assert self._saver is not None
|
|
if not faces:
|
|
return face_count
|
|
|
|
for idx, face in enumerate(faces):
|
|
output = f"{frame_name}_{idx}.png"
|
|
meta = PNGHeader(
|
|
alignments=face.to_png_meta(),
|
|
source=PNGSource(alignments_version=self._alignments.version,
|
|
original_filename=output,
|
|
face_index=idx,
|
|
source_filename=filename,
|
|
source_is_video=self._frames.is_video,
|
|
source_frame_dims=tuple(image.shape[:2])))
|
|
assert face.aligned.face is not None
|
|
self._saver.save(output, encode_image(face.aligned.face, ".png", metadata=meta))
|
|
face_count += 1
|
|
self._saver.close()
|
|
return face_count
|
|
|
|
def _select_valid_faces(self, frame: str, image: np.ndarray) -> list[DetectedFace]:
|
|
"""Return the aligned faces from a frame that meet the selection criteria,
|
|
|
|
Parameters
|
|
----------
|
|
frame
|
|
The filename (without the full path) of the current frame
|
|
image
|
|
The full frame that faces are to be extracted from
|
|
|
|
Returns
|
|
-------
|
|
List of valid :class:`lib,align.DetectedFace` objects
|
|
"""
|
|
faces = self._extracted_faces.get_faces_in_frame(frame, image=image)
|
|
if self._min_size == 0:
|
|
valid_faces = faces
|
|
else:
|
|
sizes = self._extracted_faces.get_roi_size_for_frame(frame)
|
|
valid_faces = [faces[idx] for idx, size in enumerate(sizes)
|
|
if size >= self._min_size]
|
|
logger.trace("frame: '%s', total_faces: %s, valid_faces: %s", # type:ignore
|
|
frame, len(faces), len(valid_faces))
|
|
return valid_faces
|
|
|
|
|
|
__all__ = get_module_objects(__name__)
|