Source code for wed.chaptors.rotating_rectangle

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

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

"""
from typing import Tuple

import cv2
import numpy as np
import numpy.typing as npt
import scipy as sp
import scipy.ndimage

from ..utils._path import ROTATING_SQUARE_IMAGE_PATH
from .base_editor import BaseWedOPEditor


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

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

    Args:
        square_image_path (str, optional)  : The path to the image for embedding. Defaults to ``ROTATING_SQUARE_IMAGE_PATH``.
        atol (int, optional)               : Absolute tolerance to recognize where to embed a square image. Defaults to ``50``.
        rotating_start_pos (int, optional) : The number of the frame where the rotation starts. Defaults to ``220``.

    Attributes:
        MASKED_COLOR (Tuple[int,int,int]) : BGR Color of the part to be masked.
    """

    MASKED_COLOR: Tuple[int, int, int] = (248, 232, 217)

    def __init__(
        self,
        square_image_path: str = ROTATING_SQUARE_IMAGE_PATH,
        atol: int = 50,
        rotating_start_pos: int = 220,
    ):
        super().__init__(
            positions=(212, 246), image_paths=dict(square_image=square_image_path)
        )
        self.atol = atol
        self.rotating_start_pos = rotating_start_pos

    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.embed_image(frame=frame, pos=pos)
        return frame

    def embed_image(
        self, frame: npt.NDArray[np.uint8], pos: int
    ) -> npt.NDArray[np.uint8]:
        """Embed a square image in the following gray square:

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

        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.
        """
        H, W = frame.shape[:2]
        h, w = cH, cW = (H // 2, W // 2)
        flag = (
            np.sum(
                np.abs(
                    frame - np.asarray([[RotatingRectangleEditor.MASKED_COLOR]]),
                    dtype=int,
                ),
                axis=2,
            )
            < self.atol
        )
        while True:
            if not flag[h, w]:
                if np.any(flag[h - 1, w : w + 3]):
                    w += 1
                else:
                    break
            else:
                h += 1
        l = int(np.sqrt((h - cH) ** 2 + (w - cW) ** 2))
        mask = np.zeros_like(flag)
        mask[cH - l : cH + l, cW - l : cW + l] = flag[cH - l : cH + l, cW - l : cW + l]
        paste = np.zeros_like(frame)
        paste[cH - l : cH + l, cW - l : cW + l] = cv2.resize(
            src=sp.ndimage.rotate(
                self.square_image_arr,
                angle=int(45 + max(0, 20 / 26 * (pos - self.rotating_start_pos))),
            ),
            dsize=(int(2 * l), int(2 * l)),
        )
        return np.where(mask[:, :, None], paste, frame)