Masked source frame exporter

export_source_masked.py

'''
Created on Oct 23, 2019

@author: kilroythethird
'''

import cv2
import numpy as np

from os import path as ospath
import os
from lib.logger import log_setup
log_setup("INFO", None, "test", False)

from lib.alignments import Alignments
from lib.faces_detect import DetectedFace
from lib.image import ImagesLoader
from lib.convert import Converter
import argparse

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Saves the frames with masked faces as pngs.")
    parser.add_argument("-al", "--alignment", required=False, help="The .fsa file containing your alignments. Must be extracted with a mask.")
    parser.add_argument("-i", "--input", required=False, help="The frame folder or video the alignments are extracted from.")
    parser.add_argument("-o", "--output", required=False, help="The folder you want to save the masked frames in.")
    parser.add_argument("-b", "--blur", required=False, help="The ammount of blur you want.")
    parser.add_argument("-c", "--crop-face", action="store_true", help="Crop the faces.")
    parser.add_argument("-s", "--output-size", required=False, help="Resizes the faces to a specific size. Use 0 to not resize.")
    parser.add_argument("-co", "--coverage", required=False, help="Coverage value to use. 0.5 - 1.0")
    args = parser.parse_args()

    frames = args.input
    out_dir = args.output
    alignment = args.alignment
    blur = args.blur
    osize = args.output_size
    coverage = args.coverage
    if frames is None:
        frames = input("Source frame folder or video: ").strip('\'" ')
    if alignment is None:
        alignment = input("Alignemnt (.fsa) file containing masks: ").strip('\'" ')
    if out_dir is None:
        out_dir = input("Output folder for masked frames: ").strip('\'" ')
    if blur is None:
        blur = int(input("Blur to use in pixel [Default: 0]: ") or 0)
    else:
        blur = int(blur)
    if args.crop_face:
        if osize is None:
            osize = input("Face output size. Use 0 to NOT resize. [Default: 256]: ")
            osize = osize if len(osize.strip()) else 256
        osize = int(osize)
        if coverage is None:
            coverage = input("Coverage from 0.1 - 1.0. [Default: 0.9]: ") or "0.9"
        coverage = float(coverage)

    print("Source: '%s'" % frames)
    print("Alignment: '%s'" % alignment)
    print("Output directory: '%s'" % out_dir)
    print("blur: '%s'" % blur)
    print("Output size: '%s'" % osize)
    print("Coverage: '%s'" % coverage)

    os.makedirs(out_dir, exist_ok=True)
    a = Alignments(frames, filename=alignment)
    alignments = a.data
    loader = ImagesLoader(frames)
    for fname, img in loader.load():
        _, fname = os.path.split(fname)
        out_name = ospath.join(out_dir, fname)
        background = np.zeros(img.shape[:2] + (4,), dtype="float32")
        background[..., :3] = img
        faces = list()
        for i, face in enumerate(alignments.get(fname, [])):
            detface = DetectedFace()
            detface.from_alignment(face, img)
            mask = list(detface.mask.values())[0]
            mask.set_blur_kernel_and_threshold(blur, 0)
            if args.crop_face:
                out_name = "%s_%i.png" % (ospath.splitext(out_name)[0], i)
                if osize:
                    detface.load_feed_face(img, osize, coverage)
                    face = detface.feed_face
                    mask = mask.mask
                    mask = cv2.resize(mask, face.shape[:2], interpolation=cv2.INTER_AREA)
                    mask = np.expand_dims(mask, axis=-1)
                    face = np.concatenate([face, mask], axis=-1)
                else:
                    mask = mask.get_full_frame_mask(img.shape[1], img.shape[0])
                    face = background.copy()
                    face[..., 3] = mask
                    face = face[max(0, detface.top):max(0, detface.bottom), max(0, detface.left):max(0, detface.right), ...]
                cv2.imwrite(out_name, face)
            else:
                mask = mask.get_full_frame_mask(img.shape[1], img.shape[0])
                background[..., 3] = mask
        if not args.crop_face:
            out_name = "%s.png" % ospath.splitext(out_name)[0]
            cv2.imwrite(out_name, background)

    input("All done. Press enter to exit")