Source code for wed.chaptors

# coding: utf-8

from typing import Dict, List, Optional, Tuple, Union

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

from ..utils._colorings import toGREEN
from ..utils._path import OPENING_TEMPLATE_PATH
from ..utils.audio_utils import synthesize_audio
from ..utils.generic_utils import handleKeyError
from ..utils.image_utils import arr2pil, pil2arr
from ..utils.video_utils import capture2writor
from . import (
    base_editor,
    marquee,
    rotating_rectangle,
    smartphone,
    spread_tile,
    title_call,
)
from .base_editor import BaseVideoHandler, BaseWedOPEditor
from .marquee import MarqueeEditor
from .rotating_face import RotatingFaceEditor
from .rotating_rectangle import RotatingRectangleEditor
from .smartphone import SmartphoneEditor
from .spread_tile import SpreadTileEditor
from .title_call import TitleCallEditor

WED_EDITORS: Dict[str, BaseWedOPEditor] = {
    "marquee": MarqueeEditor,
    "rotating_rectangle": RotatingRectangleEditor,
    "spread_tile": SpreadTileEditor,
    "smartphone": SmartphoneEditor,
    "title_call": TitleCallEditor,
}


def get(identifier: Union[str, BaseWedOPEditor], **kwargs) -> BaseWedOPEditor:
    """Get an instance of :class:`BaseWedOPEditor <wed.chaptors.base_editor.BaseWedOPEditor>`

    Args:
        identifier (Union[str, BaseWedOPEditor]) : An identifier for :class:`BaseWedOPEditor <wed.chaptors.base_editor.BaseWedOPEditor>`.

    Raises:
        TypeError: When ``identifier`` is not a ``str`` or an instance of :class:`BaseWedOPEditor <wed.chaptors.base_editor.BaseWedOPEditor>`

    Returns:
        BaseWedOPEditor: An instance of :class:`BaseWedOPEditor <wed.chaptors.base_editor.BaseWedOPEditor>`
    """
    if isinstance(identifier, str):
        handleKeyError(lst=list(WED_EDITORS.keys()), identifier=identifier)
        instance = WED_EDITORS[identifier](**kwargs)
    elif isinstance(identifier, BaseWedOPEditor):
        instance = identifier
    else:
        raise TypeError(
            f"{toGREEN('identifier')} must be a string or an instance of {toGREEN('BaseWedOPEditor')}"
        )
    return instance


class Administorator(BaseVideoHandler):
    """Administrator in charge of Editors.

    Attributes:
        editors (List[BaseWedOPEditor]) : List containing each editor
    """

    def __init__(self, video_path: str):
        self.video_path = video_path
        self.editors: List[BaseWedOPEditor] = []

    def set_editor(self, identifier: str, **kwargs) -> None:
        """Set an editor.

        Args:
            identifier (str) : An identifier for :class:`BaseWedOPEditor <wed.chaptors.base_editor.BaseWedOPEditor>`.
        """
        self.editors.append(get(identifier=identifier, **kwargs))

    def append(self, editor: BaseWedOPEditor) -> None:
        """Append an editor.

        Args:
            editor (BaseWedOPEditor) : An instance of :class:`BaseWedOPEditor <wed.chaptors.base_editor.BaseWedOPEditor>`.
        """
        self.editors.append(editor)

    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.
        """
        for editor in self.editors:
            if editor.start_pos <= pos <= editor.end_pos:
                frame = editor.edit(frame=frame, pos=pos)
        return frame

    def export(self, out_path: Optional[str] = None, codec: str = "H264"):
        """Create a video with each editor in ``editors``.

        Args:
            out_path (Optional[str], optional) : The path to the created silence video file. Defaults to ``None``.
            codec (str, optional)              : Video codec for the created video file. Defaults to ``"H264"``.

        Returns:
            str: The path to the created video file.
        """
        cap = cv2.VideoCapture(self.video_path)
        out, out_path = capture2writor(cap=cap, out_path=out_path, codec=codec)

        for i in tqdm(
            range(int(cap.get(cv2.CAP_PROP_FRAME_COUNT))), desc="Wed DOWNTOWN OP"
        ):
            ret, frame = cap.read()
            if (not ret) or (frame is None):
                break
            frame = self.edit(frame, i)
            out.write(frame)

        out.release()
        cap.release()

        synthesized_video_path = self.synthesize_audio(out_path)
        return synthesized_video_path

    def synthesize_audio(self, out_path: str) -> str:
        """Create audio with each editor in ``editors`` and attach it to the video at ``out_path``.

        Args:
            out_path (str) : The path to the output video.

        Returns:
            str: The path to the created video file.
        """
        audio_path = self.video_path
        for editor in self.editors:
            audio_path = editor.overlayed_audio_create(video_path=audio_path)
        synthesized_video_path = synthesize_audio(
            video_path=out_path,
            audio_path=audio_path,
        )
        return synthesize_audio