plugins/convert/writer/patch.py - master - githubmirror/faceswap

plugins/convert/writer/patch.py
293 строки12 KB
torzdf
Extraction Overhaul (#1533)
21 мар 2026, 02:52
21 мар 2026, 02:52d217812
#!/usr/bin/env python3
""" Face patch output writer for faceswap.py converter
    Extracts the swapped Face Patch from faceswap rather than the final composited frame along with
    the transformation matrix for re-inserting the face into the original frame
"""
import json
import logging
import re
import typing as T

import os
import cv2
import numpy as np

from lib.image import encode_image, png_read_meta, tiff_read_meta
from lib.utils import get_module_objects
from ._base import Output
from . import patch_defaults as cfg

logger = logging.getLogger(__name__)


class Writer(Output):
    """ Face patch writer for outputting swapped face patches and transformation matrices

    Parameters
    ----------
    output_folder: str
        The full path to the output folder where the face patches should be saved
    patch_size: int
        The size of the face patch output from the model
    config_file: str, optional
        The full path to a custom configuration ini file. If ``None`` is passed
        then the file is loaded from the default location. Default: ``None``.
    """
    def __init__(self, output_folder: str, patch_size: int, **kwargs) -> None:
        logger.debug("patch_size: %s", patch_size)
        super().__init__(output_folder, **kwargs)
        self._extension = {"png": ".png", "tiff": ".tif"}[cfg.format()]
        self._separate_mask = cfg.separate_mask()
        self._fname_split = re.compile("[^0-9a-zA-Z]")

        if self._extension == ".png" and cfg.bit_depth() not in ("8", "16"):
            logger.warning("Patch Writer: Bit Depth '%s' is unsupported for format '%s'. "
                           "Updating to '16'", cfg.bit_depth(), cfg.format())
            cfg.bit_depth.set("16")

        self._dtype = {"8": np.uint8, "16": np.uint16, "32": np.float32}[cfg.bit_depth()]
        self._multiplier = {"8": 255., "16": 65535., "32": 1.}[cfg.bit_depth()]

        self._dummy_patch = np.zeros((1, patch_size, patch_size, 4), dtype=np.float32)

        tl_box = np.array([[0, 0], [patch_size, 0], [patch_size, patch_size], [0, patch_size]],
                          dtype=np.float32)
        self._patch_corner = {"top-left": tl_box[0],
                              "top-right": tl_box[1],
                              "bottom-right": tl_box[2],
                              "bottom-left": tl_box[3]}[cfg.origin()].copy()
        self._box = tl_box
        if cfg.origin() in ("top-right", "bottom-left"):
            self._box[[1, 3], :] = self._box[[3, 1], :]  # keep clockwise from 0,0

        self._args = self._get_save_args()
        self._matrices: dict[str, dict[str, list[list[float]]]] = {}

    def _get_save_args(self) -> tuple[int, ...]:
        """ Obtain the save parameters for the file format.

        Returns
        -------
        tuple
            The OpenCV specific arguments for the selected file format
         """
        args: tuple[int, ...] = tuple()
        if self._extension == ".png" and cfg.png_compress_level() > -1:
            args = (cv2.IMWRITE_PNG_COMPRESSION, cfg.png_compress_level())
        if self._extension == ".tif" and cfg.bit_depth() != "32":
            tiff_methods = {"none": 1, "lzw": 5, "deflate": 8}
            method = cfg.tiff_compression_method()
            method = "none" if method is None else method
            args = (cv2.IMWRITE_TIFF_COMPRESSION, tiff_methods[method])
        logger.debug(args)
        return args

    def _get_new_filename(self, filename: str, face_index: int) -> str:
        """ Obtain the filename for the output file based on the frame's filename and the user
        selected naming options

        Parameters
        ----------
        filename: str
            The original frame's filename
        face_index: int
            The index of the face within the frame

        Returns
        -------
        str
            The new filename for naming the output face patch
        """
        face_idx = str(face_index).rjust(2, "0")
        fname, ext = os.path.splitext(filename)
        fname = os.path.basename(fname)

        split_fname = self._fname_split.split(fname)
        if split_fname and split_fname[-1].isdigit():
            i_frame_no = (int(split_fname[-1]) +
                          (int(cfg.start_index()) - 1) +
                          cfg.index_offset())
            frame_no = f".{str(i_frame_no).rjust(cfg.number_padding(), '0')}"
            base_fname = fname[:-len(split_fname[-1]) - 1]
        else:
            frame_no = ""
            base_fname = fname

        retval = ""
        if cfg.include_filename():
            retval += base_fname
        if cfg.face_index_location() == "before":
            retval = f"{retval}_{face_idx}"
        retval += frame_no
        if cfg.face_index_location() == "after":
            retval = f"{retval}.{face_idx}"
        retval += ext
        logger.trace("source filename: '%s', output filename: '%s'",  # type:ignore[attr-defined]
                     filename, retval)
        return retval

    def write(self, filename: str, image: list[list[bytes]]) -> None:
        """ Write out the pre-encoded image to disk. If separate mask has been selected, write out
        the encoded mask to a sub-folder in the output directory.

        Parameters
        ----------
        filename: str
            The full path to write out the image to.
        image: list[list[bytes]]
            List of list of :class:`bytes` objects of containing all swapped faces from a frame to
            write out. The inner list will be of length 1 (mask included in the alpha channel) or
            length 2 (mask to write out separately)
        """
        logger.trace("Outputting: (filename: '%s')", filename)  # type:ignore[attr-defined]

        read_func = png_read_meta if self._extension == ".png" else tiff_read_meta
        for idx, face in enumerate(image):
            new_filename = self._get_new_filename(filename, idx)
            filenames = self.get_output_filename(new_filename, cfg.format(), self._separate_mask)
            for fname, img in zip(filenames, face):
                try:
                    with open(fname, "wb") as outfile:
                        outfile.write(img)
                except Exception as err:  # pylint:disable=broad-except
                    logger.error("Failed to save image '%s'. Original Error: %s", filename, err)
                if not cfg.json_output():
                    continue
                mat = T.cast(dict[str, list[list[float]]], read_func(img))
                self._matrices[os.path.splitext(os.path.basename(fname))[0]] = mat

    @classmethod
    def _get_inverse_matrices(cls, matrices: np.ndarray) -> np.ndarray:
        """ Obtain the inverse matrices for the given matrices. If ``None`` is supplied return a
        dummy transformation matrix that performs no action

        Parameters
        ----------
        matrices : :class:`numpy.ndarray`
            The original transform matrices that the inverse needs to be calculated for

        Returns
        -------
        :class:`numpy.ndarray`
            The inverse transformation matrices
        """
        if not np.any(matrices):
            return np.array([[[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]]], dtype=np.float32)

        identity = np.array([[[0., 0., 1.]]], dtype=np.float32)
        mat = np.concatenate([matrices, np.repeat(identity, matrices.shape[0], axis=0)], axis=1)
        retval = np.linalg.inv(mat)
        logger.trace("matrix: %s, inverse: %s", mat, retval)  # type:ignore[attr-defined]
        return retval

    def _adjust_to_origin(self, matrices: np.ndarray, canvas_size: tuple[int, int]) -> None:
        """ Adjust the transformation matrix to use the correct target coordinates system. The
        matrix adjustment is done in place, so this does not return a value

        Parameters
        ----------
        matrices: :class:`numpy.ndarray`
            The transformation matrices to be adjusted
        canvas_size: tuple[int, int]
            The size of the canvas width, height) that the transformation matrix applies to.
        """
        if cfg.origin() == "top-left":
            return

        for mat in matrices:
            og_cnr = cv2.transform(self._patch_corner[None, None], mat[:2, ...]).squeeze()
            x_shift, y_shift = og_cnr
            if cfg.origin().split("-")[-1] == "right":
                x_shift = canvas_size[0] - x_shift
            if cfg.origin().split("-")[0] == "bottom":
                y_shift = canvas_size[1] - y_shift
            mat[:2, 2] = [x_shift, y_shift]

        if cfg.origin() in ("top-right", "bottom-left"):
            matrices[..., :2, :2] *= [[[1, -1], [-1, 1]]]  # switch shear

    def _get_roi(self, matrices: np.ndarray) -> np.ndarray:
        """ Obtain the (x, y) ROI points of the patch in the original frame. Points are returned
        in clockwise order from the origin location

        Parameters
        ----------
        matrices: :class:`numpy.ndarray`
            The transformation matrices for the current frame

        Returns
        -------
        np.ndarray
            The ROI of the patches in original frame co-ordinates in clockwise order from the
            origin point
        """
        retval = [cv2.transform(np.expand_dims(self._box, axis=1), mat[:2, ...]).squeeze()
                  for mat in matrices]
        return np.array(retval, dtype=np.float32)

    def pre_encode(self, image: np.ndarray, **kwargs) -> list[list[bytes]]:
        """ Pre_encode the image in lib/convert.py threads as it is a LOT quicker.

        Parameters
        ----------
        image: :class:`numpy.ndarray`
            A 3 or 4 channel BGR swapped face batch as float32
        canvas_size: tuple[int, int]
            The size of the canvas (x, y) that the transformation matrix applies to.
        matrices: :class:`numpy.ndarray`, optional
            The transformation matrices for extracting the face patches from the original frame.
            Must be provided if an image is provided, otherwise ``None`` to insert a dummy matrix

        Returns
        -------
        list
            List of :class:`bytes` objects ready for writing. The list will be of length 1 with
            image bytes object as the only member unless separate mask has been requested, in which
            case it will be length 2 with the image in position 0 and mask in position 1
         """
        logger.trace("Pre-encoding image")  # type:ignore[attr-defined]
        retval = []
        canvas_size: tuple[int, int] = kwargs.get("canvas_size", (1, 1))
        matrices: np.ndarray = kwargs.get("matrices", np.array([]))

        if not np.any(image) and cfg.empty_frames() == "blank":
            image = self._dummy_patch

        matrices = self._get_inverse_matrices(matrices)
        self._adjust_to_origin(matrices, canvas_size)
        rois = self._get_roi(matrices)
        patches = (image * self._multiplier).astype(self._dtype)

        for patch, matrix, roi in zip(patches, matrices, rois):
            this_face = []
            mat = json.dumps({"transform_matrix": matrix.tolist(), "roi": roi.tolist()},
                             ensure_ascii=True).encode("ascii")
            if self._separate_mask:
                mask = patch[..., -1]
                face = patch[..., :3]

                this_face.append(encode_image(mask,
                                              self._extension,
                                              encoding_args=self._args,
                                              metadata=mat))
            else:
                face = patch

            this_face.insert(0, encode_image(face,
                                             self._extension,
                                             encoding_args=self._args,
                                             metadata=mat))
            retval.append(this_face)
        return retval

    def close(self) -> None:
        """ Outputs json file if requested """
        if not cfg.json_output():
            return
        fname = os.path.join(self.output_folder, "matrices.json")
        with open(fname, "w", encoding="utf-8") as ofile:
            json.dump(self._matrices, ofile, indent=2, sort_keys=True)
        logger.info("Patch matrices written to: '%s'", fname)


__all__ = get_module_objects(__name__)