faceswap/plugins/extract/detect/_base.py

#!/usr/bin/env python3
""" Base class for Face Detector plugins
    Plugins should inherit from this class

    See the override methods for which methods are
    required.

    For each source frame, the plugin must pass a dict to finalize containing:
    {"filename": <filename of source frame>,
     "image": <source image>,
     "detected_faces": <list of dlib.rectangles>}
    """

import logging
import os
import traceback
from io import StringIO

import cv2
import dlib
from math import sqrt

from lib.gpu_stats import GPUStats
from lib.utils import rotate_landmarks
from plugins.extract._config import Config

logger = logging.getLogger(__name__)  # pylint: disable=invalid-name


def get_config(plugin_name):
    """ Return the config for the requested model """
    return Config(plugin_name).config_dict


class Detector():
    """ Detector object """
    def __init__(self, loglevel, rotation=None):
        logger.debug("Initializing %s: (rotation: %s)", self.__class__.__name__, rotation)
        self.config = get_config(".".join(self.__module__.split(".")[-2:]))
        self.loglevel = loglevel
        self.cachepath = os.path.join(os.path.dirname(__file__), ".cache")
        self.rotation = self.get_rotation_angles(rotation)
        self.parent_is_pool = False
        self.init = None

        # The input and output queues for the plugin.
        # See lib.queue_manager.QueueManager for getting queues
        self.queues = {"in": None, "out": None}

        #  Path to model if required
        self.model_path = self.set_model_path()

        # Target image size for passing images through the detector
        # Set to tuple of dimensions (x, y) or int of pixel count
        self.target = None

        # Approximate VRAM used for the set target. Used to calculate
        # how many parallel processes / batches can be run.
        # Be conservative to avoid OOM.
        self.vram = None

        # For detectors that support batching, this should be set to
        # the calculated batch size that the amount of available VRAM
        # will support. It is also used for holding the number of threads/
        # processes for parallel processing plugins
        self.batch_size = 1
        logger.debug("Initialized _base %s", self.__class__.__name__)

    # <<< OVERRIDE METHODS >>> #
    # These methods must be overriden when creating a plugin
    @staticmethod
    def set_model_path():
        """ path to data file/models
            override for specific detector """
        raise NotImplementedError()

    def initialize(self, *args, **kwargs):
        """ Inititalize the detector
            Tasks to be run before any detection is performed.
            Override for specific detector """
        logger_init = kwargs["log_init"]
        log_queue = kwargs["log_queue"]
        logger_init(self.loglevel, log_queue)
        logger.debug("initialize %s (PID: %s, args: %s, kwargs: %s)",
                     self.__class__.__name__, os.getpid(), args, kwargs)
        self.init = kwargs.get("event", False)
        self.queues["in"] = kwargs["in_queue"]
        self.queues["out"] = kwargs["out_queue"]

    def detect_faces(self, *args, **kwargs):
        """ Detect faces in rgb image
            Override for specific detector
            Must return a list of dlib rects"""
        try:
            if not self.init:
                self.initialize(*args, **kwargs)
        except ValueError as err:
            logger.error(err)
            exit(1)
        logger.debug("Detecting Faces (args: %s, kwargs: %s)", args, kwargs)

    # <<< DETECTION WRAPPER >>> #
    def run(self, *args, **kwargs):
        """ Parent detect process.
            This should always be called as the entry point so exceptions
            are passed back to parent.
            Do not override """
        try:
            self.detect_faces(*args, **kwargs)
        except Exception:  # pylint: disable=broad-except
            logger.error("Caught exception in child process: %s", os.getpid())
            # Display traceback if in initialization stage 
            if not self.init.is_set():
                logger.exception("Traceback:")
            tb_buffer = StringIO()
            traceback.print_exc(file=tb_buffer)
            logger.trace(tb_buffer.getvalue())
            exception = {"exception": (os.getpid(), tb_buffer)}
            self.queues["out"].put(exception)
            exit(1)

    # <<< FINALIZE METHODS>>> #
    def finalize(self, output):
        """ This should be called as the final task of each plugin
            Performs fianl processing and puts to the out queue """
        if isinstance(output, dict):
            logger.trace("Item out: %s", {key: val
                                          for key, val in output.items()
                                          if key != "image"})
        else:
            logger.trace("Item out: %s", output)
        self.queues["out"].put(output)

    # <<< DETECTION IMAGE COMPILATION METHODS >>> #
    def compile_detection_image(self, image, is_square, scale_up):
        """ Compile the detection image """
        scale = self.set_scale(image, is_square=is_square, scale_up=scale_up)
        return [self.set_detect_image(image, scale), scale]

    def set_scale(self, image, is_square=False, scale_up=False):
        """ Set the scale factor for incoming image """
        height, width = image.shape[:2]
        if is_square:
            if isinstance(self.target, int):
                dims = (self.target ** 0.5, self.target ** 0.5)
                self.target = dims
            source = max(height, width)
            target = max(self.target)
        else:
            if isinstance(self.target, tuple):
                self.target = self.target[0] * self.target[1]
            source = width * height
            target = self.target

        if scale_up or target < source:
            scale = sqrt(target / source)
        else:
            scale = 1.0
        logger.trace("Detector scale: %s", scale)
        
        return scale

    def set_detect_image(self, input_image, scale):
        """ Convert the image to RGB and scale """
        # pylint: disable=no-member
        image = input_image[:, :, ::-1].copy()
        if scale == 1.0:
            return image

        height, width = image.shape[:2]
        interpln = cv2.INTER_LINEAR if scale > 1.0 else cv2.INTER_AREA
        dims = (int(width * scale), int(height * scale))

        if scale < 1.0:
            logger.verbose("Resizing image from %sx%s to %s.",
                           width, height, "x".join(str(i) for i in dims))

        image = cv2.resize(image, dims, interpolation=interpln)
        return image

    # <<< IMAGE ROTATION METHODS >>> #
    @staticmethod
    def get_rotation_angles(rotation):
        """ Set the rotation angles. Includes backwards compatibility for the
            'on' and 'off' options:
                - 'on' - increment 90 degrees
                - 'off' - disable
                - 0 is prepended to the list, as whatever happens, we want to
                  scan the image in it's upright state """
        rotation_angles = [0]

        if not rotation or rotation.lower() == "off":
            logger.debug("Not setting rotation angles")
            return rotation_angles

        if rotation.lower() == "on":
            rotation_angles.extend(range(90, 360, 90))
        else:
            passed_angles = [int(angle)
                             for angle in rotation.split(",")]
            if len(passed_angles) == 1:
                rotation_step_size = passed_angles[0]
                rotation_angles.extend(range(rotation_step_size,
                                             360,
                                             rotation_step_size))
            elif len(passed_angles) > 1:
                rotation_angles.extend(passed_angles)

        logger.debug("Rotation Angles: %s", rotation_angles)
        return rotation_angles

    def rotate_image(self, image, angle):
        """ Rotate the image by given angle and return
            Image with rotation matrix """
        if angle == 0:
            return image, None
        return self.rotate_image_by_angle(image, angle)

    @staticmethod
    def rotate_rect(d_rect, rotation_matrix):
        """ Rotate a dlib rect based on the rotation_matrix"""
        logger.trace("Rotating d_rectangle")
        d_rect = rotate_landmarks(d_rect, rotation_matrix)
        return d_rect

    @staticmethod
    def rotate_image_by_angle(image, angle,
                              rotated_width=None, rotated_height=None):
        """ Rotate an image by a given angle.
            From: https://stackoverflow.com/questions/22041699 """

        logger.trace("Rotating image: (angle: %s, rotated_width: %s, rotated_height: %s)",
                     angle, rotated_width, rotated_height)
        height, width = image.shape[:2]
        image_center = (width/2, height/2)
        rotation_matrix = cv2.getRotationMatrix2D(  # pylint: disable=no-member
            image_center, -1.*angle, 1.)
        if rotated_width is None or rotated_height is None:
            abs_cos = abs(rotation_matrix[0, 0])
            abs_sin = abs(rotation_matrix[0, 1])
            if rotated_width is None:
                rotated_width = int(height*abs_sin + width*abs_cos)
            if rotated_height is None:
                rotated_height = int(height*abs_cos + width*abs_sin)
        rotation_matrix[0, 2] += rotated_width/2 - image_center[0]
        rotation_matrix[1, 2] += rotated_height/2 - image_center[1]
        logger.trace("Rotated image: (rotation_matrix: %s", rotation_matrix)
        return (cv2.warpAffine(image,  # pylint: disable=no-member
                               rotation_matrix,
                               (rotated_width, rotated_height)),
                rotation_matrix)

    # << QUEUE METHODS >> #
    def get_item(self):
        """ Yield one item from the queue """
        item = self.queues["in"].get()
        if isinstance(item, dict):
            logger.trace("Item in: %s", item["filename"])
        else:
            logger.trace("Item in: %s", item)
        if item == "EOF":
            logger.debug("In Queue Exhausted")
            # Re-put EOF into queue for other threads
            self.queues["in"].put(item)
        return item

    def get_batch(self):
        """ Get items from the queue in batches of
            self.batch_size

            First item in output tuple indicates whether the
            queue is exhausted.
            Second item is the batch

            Remember to put "EOF" to the out queue after processing
            the final batch """
        exhausted = False
        batch = list()
        for _ in range(self.batch_size):
            item = self.get_item()
            if item == "EOF":
                exhausted = True
                break
            batch.append(item)
        logger.trace("Returning batch size: %s", len(batch))
        return (exhausted, batch)

    # <<< DLIB RECTANGLE METHODS >>> #
    @staticmethod
    def is_mmod_rectangle(d_rectangle):
        """ Return whether the passed in object is
            a dlib.mmod_rectangle """
        return isinstance(
            d_rectangle,
            dlib.mmod_rectangle)  # pylint: disable=c-extension-no-member

    def convert_to_dlib_rectangle(self, d_rect):
        """ Convert detected mmod_rects to dlib_rectangle """
        if self.is_mmod_rectangle(d_rect):
            return d_rect.rect
        return d_rect

    # <<< MISC METHODS >>> #
    @staticmethod
    def get_vram_free():
        """ Return total free VRAM on largest card """
        stats = GPUStats()
        vram = stats.get_card_most_free()
        logger.verbose("Using device %s with %sMB free of %sMB",
                       vram["device"],
                       int(vram["free"]),
                       int(vram["total"]))
        return int(vram["free"])

    @staticmethod
    def set_predetected(width, height):
        """ Set a dlib rectangle for predetected faces """
        # Predetected_face is used for sort tool.
        # Landmarks should not be extracted again from predetected faces,
        # because face data is lost, resulting in a large variance
        # against extract from original image
        logger.debug("Setting predetected face")
        return [dlib.rectangle(0, 0, width, height)]  # pylint: disable=c-extension-no-member