mirror of
https://github.com/deepfakes/faceswap
synced 2025-06-07 19:05:02 -04:00
27a685383e
* Separate predict and implement pool * Add check and raise error to multithreading Box functions to config. Add crop box option. * All masks to mask module. Refactor convert masks Predicted mask passed from model. Cli update * Intesect box with mask and fixes * Use raw NN output for convert Use raw mask for face adjustments. Split adjustments to pre and post warp * Separate out adjustments. Add unmask sharpen * Set sensible defaults. Pre PR Testing * Fix queue sizes. Move masked.py to lib * Fix Skip Frames. Fix GUI Config popup * Sensible queue limits. Add a less resource intensive single processing mode * Fix predicted mask. Amend smooth box defaults * Deterministic ordering for video output * Video to Video convert implemented * Fixups - Remove defaults from folders across all stages - Move match-hist and aca into color adjustments selectable - Handle crashes properly for pooled processes - Fix output directory does not exist error when creating a new output folder - Force output to frames if input is not a video * Add Color Transfer adjustment method Wrap info text in GUI plugin configure popup * Refactor image adjustments. Easier to create plugins Start implementing config options for video encoding * Add video encoding config options Allow video encoding for frames input (must pass in a reference video) Move video and image output writers to plugins * Image writers config options Move scaling to cli Move draw_transparent to images config Add config options for cv2 writer Add Pillow image writer * Add gif filetype to Pillow. Fix draw transparent for Pillow * Add Animated GIF writer standardize opencv/pillow defaults * [speedup] Pre-encode supported writers in the convert pool (opencv, pillow) Move scaling to convert pool Remove dfaker mask * Fix default writer * Bugfixes * Better custom argparse formatting
217 lines
7.8 KiB
Python
217 lines
7.8 KiB
Python
#!/usr/bin python3
|
|
""" Utilities available across all scripts """
|
|
|
|
import logging
|
|
import os
|
|
import warnings
|
|
|
|
from hashlib import sha1
|
|
from pathlib import Path
|
|
from re import finditer
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
import dlib
|
|
|
|
from lib.faces_detect import DetectedFace
|
|
from lib.logger import get_loglevel
|
|
|
|
|
|
logger = logging.getLogger(__name__) # pylint: disable=invalid-name
|
|
|
|
# Global variables
|
|
_image_extensions = [ # pylint: disable=invalid-name
|
|
".bmp", ".jpeg", ".jpg", ".png", ".tif", ".tiff"]
|
|
_video_extensions = [ # pylint: disable=invalid-name
|
|
".avi", ".flv", ".mkv", ".mov", ".mp4", ".mpeg", ".webm"]
|
|
|
|
|
|
def get_folder(path, make_folder=True):
|
|
""" Return a path to a folder, creating it if it doesn't exist """
|
|
logger.debug("Requested path: '%s'", path)
|
|
output_dir = Path(path)
|
|
if not make_folder and not output_dir.exists():
|
|
logger.debug("%s does not exist", path)
|
|
return None
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
logger.debug("Returning: '%s'", output_dir)
|
|
return output_dir
|
|
|
|
|
|
def get_image_paths(directory):
|
|
""" Return a list of images that reside in a folder """
|
|
image_extensions = _image_extensions
|
|
dir_contents = list()
|
|
|
|
if not os.path.exists(directory):
|
|
logger.debug("Creating folder: '%s'", directory)
|
|
directory = get_folder(directory)
|
|
|
|
dir_scanned = sorted(os.scandir(directory), key=lambda x: x.name)
|
|
logger.debug("Scanned Folder contains %s files", len(dir_scanned))
|
|
logger.trace("Scanned Folder Contents: %s", dir_scanned)
|
|
|
|
for chkfile in dir_scanned:
|
|
if any([chkfile.name.lower().endswith(ext)
|
|
for ext in image_extensions]):
|
|
logger.trace("Adding '%s' to image list", chkfile.path)
|
|
dir_contents.append(chkfile.path)
|
|
|
|
logger.debug("Returning %s images", len(dir_contents))
|
|
return dir_contents
|
|
|
|
|
|
def hash_image_file(filename):
|
|
""" Return an image file's sha1 hash """
|
|
img = cv2.imread(filename) # pylint: disable=no-member
|
|
img_hash = sha1(img).hexdigest()
|
|
logger.trace("filename: '%s', hash: %s", filename, img_hash)
|
|
return img_hash
|
|
|
|
|
|
def hash_encode_image(image, extension):
|
|
""" Encode the image, get the hash and return the hash with
|
|
encoded image """
|
|
img = cv2.imencode(extension, image)[1] # pylint: disable=no-member
|
|
f_hash = sha1(
|
|
cv2.imdecode(img, cv2.IMREAD_UNCHANGED)).hexdigest() # pylint: disable=no-member
|
|
return f_hash, img
|
|
|
|
|
|
def backup_file(directory, filename):
|
|
""" Backup a given file by appending .bk to the end """
|
|
logger.trace("Backing up: '%s'", filename)
|
|
origfile = os.path.join(directory, filename)
|
|
backupfile = origfile + '.bk'
|
|
if os.path.exists(backupfile):
|
|
logger.trace("Removing existing file: '%s'", backup_file)
|
|
os.remove(backupfile)
|
|
if os.path.exists(origfile):
|
|
logger.trace("Renaming: '%s' to '%s'", origfile, backup_file)
|
|
os.rename(origfile, backupfile)
|
|
|
|
|
|
def set_system_verbosity(loglevel):
|
|
""" Set the verbosity level of tensorflow and suppresses
|
|
future and deprecation warnings from any modules
|
|
From:
|
|
https://stackoverflow.com/questions/35911252/disable-tensorflow-debugging-information
|
|
Can be set to:
|
|
0 - all logs shown
|
|
1 - filter out INFO logs
|
|
2 - filter out WARNING logs
|
|
3 - filter out ERROR logs """
|
|
|
|
numeric_level = get_loglevel(loglevel)
|
|
loglevel = "2" if numeric_level > 15 else "0"
|
|
logger.debug("System Verbosity level: %s", loglevel)
|
|
os.environ['TF_CPP_MIN_LOG_LEVEL'] = loglevel
|
|
if loglevel != '0':
|
|
for warncat in (FutureWarning, DeprecationWarning, UserWarning):
|
|
warnings.simplefilter(action='ignore', category=warncat)
|
|
|
|
|
|
def rotate_landmarks(face, rotation_matrix):
|
|
# pylint: disable=c-extension-no-member
|
|
""" Rotate the landmarks and bounding box for faces
|
|
found in rotated images.
|
|
Pass in a DetectedFace object, Alignments dict or DLib rectangle"""
|
|
logger.trace("Rotating landmarks: (rotation_matrix: %s, type(face): %s",
|
|
rotation_matrix, type(face))
|
|
if isinstance(face, DetectedFace):
|
|
bounding_box = [[face.x, face.y],
|
|
[face.x + face.w, face.y],
|
|
[face.x + face.w, face.y + face.h],
|
|
[face.x, face.y + face.h]]
|
|
landmarks = face.landmarksXY
|
|
|
|
elif isinstance(face, dict):
|
|
bounding_box = [[face.get("x", 0), face.get("y", 0)],
|
|
[face.get("x", 0) + face.get("w", 0),
|
|
face.get("y", 0)],
|
|
[face.get("x", 0) + face.get("w", 0),
|
|
face.get("y", 0) + face.get("h", 0)],
|
|
[face.get("x", 0),
|
|
face.get("y", 0) + face.get("h", 0)]]
|
|
landmarks = face.get("landmarksXY", list())
|
|
|
|
elif isinstance(face,
|
|
dlib.rectangle): # pylint: disable=c-extension-no-member
|
|
bounding_box = [[face.left(), face.top()],
|
|
[face.right(), face.top()],
|
|
[face.right(), face.bottom()],
|
|
[face.left(), face.bottom()]]
|
|
landmarks = list()
|
|
else:
|
|
raise ValueError("Unsupported face type")
|
|
|
|
logger.trace("Original landmarks: %s", landmarks)
|
|
|
|
rotation_matrix = cv2.invertAffineTransform( # pylint: disable=no-member
|
|
rotation_matrix)
|
|
rotated = list()
|
|
for item in (bounding_box, landmarks):
|
|
if not item:
|
|
continue
|
|
points = np.array(item, np.int32)
|
|
points = np.expand_dims(points, axis=0)
|
|
transformed = cv2.transform(points, # pylint: disable=no-member
|
|
rotation_matrix).astype(np.int32)
|
|
rotated.append(transformed.squeeze())
|
|
|
|
# Bounding box should follow x, y planes, so get min/max
|
|
# for non-90 degree rotations
|
|
pt_x = min([pnt[0] for pnt in rotated[0]])
|
|
pt_y = min([pnt[1] for pnt in rotated[0]])
|
|
pt_x1 = max([pnt[0] for pnt in rotated[0]])
|
|
pt_y1 = max([pnt[1] for pnt in rotated[0]])
|
|
|
|
if isinstance(face, DetectedFace):
|
|
face.x = int(pt_x)
|
|
face.y = int(pt_y)
|
|
face.w = int(pt_x1 - pt_x)
|
|
face.h = int(pt_y1 - pt_y)
|
|
face.r = 0
|
|
if len(rotated) > 1:
|
|
rotated_landmarks = [tuple(point) for point in rotated[1].tolist()]
|
|
face.landmarksXY = rotated_landmarks
|
|
elif isinstance(face, dict):
|
|
face["x"] = int(pt_x)
|
|
face["y"] = int(pt_y)
|
|
face["w"] = int(pt_x1 - pt_x)
|
|
face["h"] = int(pt_y1 - pt_y)
|
|
face["r"] = 0
|
|
if len(rotated) > 1:
|
|
rotated_landmarks = [tuple(point) for point in rotated[1].tolist()]
|
|
face["landmarksXY"] = rotated_landmarks
|
|
else:
|
|
rotated_landmarks = dlib.rectangle( # pylint: disable=c-extension-no-member
|
|
int(pt_x), int(pt_y), int(pt_x1), int(pt_y1))
|
|
face = rotated_landmarks
|
|
|
|
logger.trace("Rotated landmarks: %s", rotated_landmarks)
|
|
return face
|
|
|
|
|
|
def camel_case_split(identifier):
|
|
""" Split a camel case name
|
|
from: https://stackoverflow.com/questions/29916065 """
|
|
matches = finditer(
|
|
".+?(?:(?<=[a-z])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])|$)",
|
|
identifier)
|
|
return [m.group(0) for m in matches]
|
|
|
|
|
|
def safe_shutdown():
|
|
""" Close queues, threads and processes in event of crash """
|
|
logger.debug("Safely shutting down")
|
|
from lib.queue_manager import queue_manager
|
|
from lib.multithreading import terminate_processes
|
|
queue_manager.terminate_queues()
|
|
terminate_processes()
|
|
logger.debug("Cleanup complete. Shutting down queue manager and exiting")
|
|
queue_manager._log_queue.put(None) # pylint: disable=protected-access
|
|
while not queue_manager._log_queue.empty(): # pylint: disable=protected-access
|
|
continue
|
|
queue_manager.manager.shutdown()
|