Source code for wed.chaptors.rotating_face

# coding: utf-8
"""This submodule contains a set of functions for editing the following image:

.. image:: _images/chaptors/rotating-face.png
   :class: popup-img

"""
from typing import Tuple

import cv2
import numpy as np
import numpy.typing as npt

from ..utils._colorings import toGREEN
from ..utils._path import ROTATING_FACE_LEFT_IMAGE_PATH, ROTATING_FACE_RIGHT_IMAGE_PATH
from ..utils.image_utils import alpha_composite, arr2pil, pil2arr
from .base_editor import BaseWedOPEditor


class RotatingFaceEditor(BaseWedOPEditor):
    """Editor which in charge of editing rotating rectangle.

    .. image:: _images/chaptors/rotating-face.png
       :class: popup-img

    Args:
    """

    def __init__(
        self,
        face_left_image_path: str = ROTATING_FACE_LEFT_IMAGE_PATH,
        face_right_image_path: str = ROTATING_FACE_RIGHT_IMAGE_PATH,
        dsize: Tuple[int, int] = (200, 200),
        margin: int = 20,
    ):
        super().__init__(
            positions=(333, 374),
            image_paths=dict(
                face_left_image=face_left_image_path,
                face_right_image=face_right_image_path,
            ),
        )
        self.resize_faces(dsize=dsize, margin=margin)

    def resize_faces(self, dsize: Tuple[int, int] = (200, 200), margin: int = 20):
        self.face_left_image_pil = self.face_left_image_pil.resize(size=dsize)
        self.face_right_image_pil = self.face_right_image_pil.resize(size=dsize)
        self.logger.info(
            f"Resized both {toGREEN('face_[left/right]_image_pil')} to {toGREEN(dsize)}"
        )
        self.set_face_boxes(margin=margin)

    def set_face_boxes(self, margin: int = 20):
        W, H = self.face_left_image_pil.size
        top = (BaseWedOPEditor.FRAME_HEIGHT - H) // 2
        self.set_attribute(name="face_left_box", value=(margin, top), msg="")
        self.set_attribute(
            name="face_right_box",
            value=(BaseWedOPEditor.FRAME_WIDTH - margin - W, top),
            msg="",
        )

    def edit(self, frame: npt.NDArray[np.uint8], pos: int) -> npt.NDArray[np.uint8]:
        """Edit the image if it is an assigned chapter (``pos``)

        Args:
            frame (npt.NDArray[np.uint8]) : Current frame (BGR image) in the video.
            pos (int)                     : Current position in the video.

        Returns:
            npt.NDArray[np.uint8]: Edited frame.
        """
        frame = self.paste_faces(frame=frame, pos=pos)
        return frame

    def paste_faces(
        self, frame: npt.NDArray[np.uint8], pos: int
    ) -> npt.NDArray[np.uint8]:
        """Paste two rotating faces.

        Args:
            frame (npt.NDArray[np.uint8]) : Current frame (BGR image) in the video.
            pos (int)                     : Current position in the video.

        Returns:
            npt.NDArray[np.uint8]: Edited frame.
        """
        angle = pos * 10
        img = arr2pil(frame)
        img = alpha_composite(
            bg=img, paste=self.face_left_image_pil.rotate(angle), box=self.face_left_box
        )
        img = alpha_composite(
            bg=img,
            paste=self.face_right_image_pil.rotate(angle),
            box=self.face_right_box,
        )
        return pil2arr(img)