# coding: utf-8
"""This submodule contains a set of functions for editing the following image:
.. image:: _images/chaptors/freedom1.png
:class: popup-img
"""
import math
import os
from typing import Any, Dict, List, Optional, Tuple, Union
import cv2
import numpy as np
import numpy.typing as npt
from PIL import Image
from ..utils._colorings import toBLUE, toGREEN
from ..utils._path import (
SMARTPHONE_IMAGE_PATH,
SMARTPHONE_TOUCH_IMAGE_PATH,
SMARTPHONE_VIDEO_PATH,
)
from ..utils.image_utils import alpha_composite, arr2pil, pil2arr
from .base_editor import BaseWedOPEditor
[docs]class SmartphoneEditor(BaseWedOPEditor):
def __init__(
self,
video_path: str = SMARTPHONE_VIDEO_PATH,
image_path: str = SMARTPHONE_IMAGE_PATH,
touch_image_path: str = SMARTPHONE_TOUCH_IMAGE_PATH,
touch_timings: Dict[int, Tuple[int]] = {
158: (140, 220),
167: (140, 300),
179: (125, 290),
205: (225, 280),
211: (280, 100),
},
touch_start_pos: Tuple[int, int] = (430, 200),
touch_end_pos: Tuple[int, int] = (430, 200),
positions: Tuple[int, int] = (152, 211),
smartphone_tblr: Tuple[int, int, int, int] = [72, 536, 17, 283],
transition: int = 3,
smartphone_pos_margin_top: int = 10,
):
super().__init__(
positions=positions,
image_paths=dict(
smartphone_image=image_path,
smartphone_touch_image=touch_image_path,
),
)
smartphone_video = cv2.VideoCapture(video_path)
t, b, l, r = smartphone_tblr
self.smartphone_tblr = smartphone_tblr
self.set_attribute(name="smartphone_tblr", value=smartphone_tblr, msg="")
self.set_attribute(name="smartphone_size", value=(r - l, b - t), msg="")
smartphone_video_length = smartphone_video.get(cv2.CAP_PROP_FRAME_COUNT)
if smartphone_video_length < self.duration:
self.logger.warning(
f"Video length should be greater than or equal to the {toGREEN('duration')} of the chapter in charge of this editor, but {smartphone_video_length } > {self.duration}"
)
self.smartphone_video = smartphone_video
self.set_attribute(
name="smartphone_video_length", value=smartphone_video_length, msg=""
)
timings = {self.start_pos: touch_start_pos, self.end_pos: touch_end_pos}
timings.update(touch_timings)
self.touch_timings = dict(sorted(timings.items()))
self.transition = transition
self.set_smartphone_box(margin_top=smartphone_pos_margin_top)
[docs] def edit(
self, frame: npt.NDArray[np.uint8], pos: int, span: int = 20
) -> 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.
"""
if self.start_pos <= pos <= self.end_pos:
frame = self.paste_smartphone_movie(
frame=frame, pos=pos, box=self.smartphone_box
)
else:
diff = min(abs(self.start_pos - pos), abs(self.end_pos - pos))
if diff <= self.transition:
l, t = self.smartphone_box
t_max = BaseWedOPEditor.FRAME_HEIGHT
t = int(
(t * (self.transition - diff) + t_max * (diff)) / self.transition
)
frame = self.paste_smartphone_movie(frame=frame, pos=pos, box=(l, t))
return frame
[docs] def set_smartphone_box(self, margin_top: int = 10) -> None:
w, h = self.smartphone_image_pil.size
W = BaseWedOPEditor.FRAME_WIDTH
self.smartphone_box: Tuple[int, int] = ((W - w) // 2, margin_top)
[docs] def get_touch_pos(self, pos: int) -> Tuple[int, int]:
prev = (0, (0, 0))
for p, (w, h) in self.touch_timings.items():
if pos < p:
pp, (pw, ph) = prev
w = (pw * (p - pos) + w * (pos - pp)) / (p - pp)
h = (ph * (p - pos) + h * (pos - pp)) / (p - pp)
break
prev = (p, (w, h))
return (int(w), int(h))
[docs] def paste_smartphone_movie(
self,
frame: npt.NDArray[np.uint8],
pos: int,
box: Tuple[int, int] = (160, 10),
) -> npt.NDArray[np.uint8]:
"""Paste a smartphone movie.
Args:
frame (npt.NDArray[np.uint8]) : Input frame (BGR ndarray image). NOT USED.
pos (int) : The position of the frame to get.
Returns:
npt.NDArray[np.uint8]: ``frame`` with ``text`` drawn.
"""
self.smartphone_video.set(
cv2.CAP_PROP_POS_FRAMES,
min(self.smartphone_video_length - 1, max(0, (pos - self.start_pos))),
)
is_ok, paste_frame = self.smartphone_video.read()
if is_ok:
img = arr2pil(frame)
t, b, l, r = self.smartphone_tblr
smartphone_image = alpha_composite(
bg=self.smartphone_image_pil,
paste=arr2pil(paste_frame).resize(self.smartphone_size),
box=(l, t),
)
img = alpha_composite(bg=img, paste=smartphone_image, box=box)
img = alpha_composite(
bg=img,
paste=self.smartphone_touch_image_pil,
box=self.get_touch_pos(pos),
)
frame = pil2arr(img)
return frame
