AI Flow: Reverse Engineering

"""The renderer: scenes made of shots, each shot a slow camera move over one or more real pictures, with flat layers on top.

A frame is a pure function of time, so the frames can be made in any order and by several workers. Frames go to ffmpeg as
raw RGB through a pipe; no browser is involved.
"""
import math
import os
import subprocess
import numpy as np
import cv2
from style import W, H, FPS, BG

cv2.setNumThreads(1)
XF = 0.5        # crossfade between scenes, seconds
SHOT_XF = 0.45  # crossfade between shots of one scene


def ease(u):
    u = min(1.0, max(0.0, u))
    return u * u * (3 - 2 * u)


class Plate:
    """A picture the camera moves over, with smaller copies so a wide view is not sampled coarsely."""
    MAX_W = 3200

    def __init__(self, rgb):
        if rgb.shape[1] > self.MAX_W:
            s = self.MAX_W / rgb.shape[1]
            rgb = cv2.resize(rgb, (self.MAX_W, int(round(rgb.shape[0] * s))), interpolation=cv2.INTER_AREA)
        self.levels = [np.ascontiguousarray(rgb)]
        self.w, self.h = rgb.shape[1], rgb.shape[0]

    def level(self, k):
        while len(self.levels) <= k:
            self.levels.append(cv2.pyrDown(self.levels[-1]))
        return self.levels[k]

    def view(self, cam, vw, vh):
        """The part of the plate the camera (centre x, centre y, width in plate pixels) sees, drawn vw x vh."""
        cx, cy, cw = cam
        s = vw / cw
        k = 0
        while s < 0.5 and min(self.w, self.h) / (2 ** (k + 1)) > 64:
            s *= 2
            k += 1
        f = 2 ** k
        m = np.array([[s, 0, vw / 2 - s * cx / f], [0, s, vh / 2 - s * cy / f]], np.float64)
        return cv2.warpAffine(self.level(k), m, (vw, vh), flags=cv2.INTER_LINEAR if s <= 1.0 else cv2.INTER_CUBIC,
                              borderMode=cv2.BORDER_CONSTANT, borderValue=BG)


def cam_fit(plate, vw, vh, rect=None, margin=0.04):
    """The camera that shows rect (x0, y0, x1, y1 in plate pixels; the whole plate when None) whole, with a margin."""
    x0, y0, x1, y1 = rect if rect else (0, 0, plate.w, plate.h)
    cw = max((x1 - x0), (y1 - y0) * vw / vh) * (1 + 2 * margin)
    return ((x0 + x1) / 2, (y0 + y1) / 2, cw)


def cam_fill(plate, vw, vh, zoom=1.0, at=(0.5, 0.5)):
    """The camera that fills the viewport with the plate (cropping what does not fit), zoomed in by `zoom` about `at`."""
    cw = min(plate.w, plate.h * vw / vh) / zoom
    ch = cw * vh / vw
    cx = min(max(at[0] * plate.w, cw / 2), plate.w - cw / 2)
    cy = min(max(at[1] * plate.h, ch / 2), plate.h - ch / 2)
    return (cx, cy, cw)


class View:
    """One viewport on the frame: a stack of plates of the same size (each fading in at its time) under one camera move."""

    def __init__(self, plates, viewport, cam0, cam1=None, fades=None):
        self.plates = plates if isinstance(plates, list) else [plates]
        self.viewport = viewport
        self.cam0, self.cam1 = cam0, cam1 or cam0
        self.fades = fades or [0.0] * len(self.plates)   # seconds into the shot at which each plate starts to fade in
        self.line = (255, 184, 48)

    def draw(self, frame, t, dur):
        x, y, vw, vh = self.viewport
        u = ease(t / dur) if dur > 0 else 0
        cam = tuple(a + (b - a) * u for a, b in zip(self.cam0, self.cam1))
        out = None
        for p, t_in in zip(self.plates, self.fades):
            a = 1.0 if t_in <= 0 else min(1.0, max(0.0, (t - t_in) / 0.6))
            if a <= 0:
                continue
            img = p.view(cam, vw, vh)
            out = img if out is None or a >= 1 else cv2.addWeighted(img, a, out, 1 - a, 0)
        if out is not None:
            frame[y:y + vh, x:x + vw] = out


class WipeView(View):
    """Two plates of the same size under one camera: the first fills the viewport and the second is uncovered from the
    right edge by a line that crosses the picture and comes back, so the same place is seen in both."""

    def draw(self, frame, t, dur):
        x, y, vw, vh = self.viewport
        u = ease(t / dur) if dur > 0 else 0
        cam = tuple(a + (b - a) * u for a, b in zip(self.cam0, self.cam1))
        a = self.plates[0].view(cam, vw, vh)
        b = self.plates[1].view(cam, vw, vh)
        # the line starts at the right, sweeps to 12% and returns to the middle
        k = t / dur if dur > 0 else 0
        pos = 1.0 - 0.88 * ease(k / 0.55) if k < 0.55 else 0.12 + 0.38 * ease((k - 0.55) / 0.45)
        cut = int(round(vw * pos))
        a[:, cut:] = b[:, cut:]
        a[:, max(0, cut - 2):cut + 2] = self.line
        frame[y:y + vh, x:x + vw] = a


class Layer:
    """A flat RGBA picture on the frame, with its own fade in (and out) in seconds from the start of its shot or scene."""

    def __init__(self, rgba, x, y, t_in=0.0, t_out=None, fade=0.4):
        self.x, self.y, self.t_in, self.t_out, self.fade = int(x), int(y), t_in, t_out, fade
        self.rgb = np.ascontiguousarray(rgba[:, :, :3])
        a = rgba[:, :, 3]
        self.opaque = bool((a == 255).all())
        self.alpha = (a.astype(np.float32) / 255.0)[:, :, None]

    def draw(self, frame, t):
        k = 1.0
        if self.t_in > 0 or self.fade > 0:
            k = min(1.0, max(0.0, (t - self.t_in) / self.fade)) if self.fade > 0 else float(t >= self.t_in)
        if self.t_out is not None:
            k = min(k, max(0.0, (self.t_out - t) / self.fade)) if self.fade > 0 else k * float(t < self.t_out)
        if k <= 0:
            return
        h, w = self.rgb.shape[:2]
        x0, y0 = max(0, self.x), max(0, self.y)
        x1, y1 = min(W, self.x + w), min(H, self.y + h)
        if x1 <= x0 or y1 <= y0:
            return
        src = self.rgb[y0 - self.y:y1 - self.y, x0 - self.x:x1 - self.x]
        dst = frame[y0:y1, x0:x1]
        if self.opaque and k >= 1:
            dst[:] = src
            return
        a = self.alpha[y0 - self.y:y1 - self.y, x0 - self.x:x1 - self.x] * k
        dst[:] = (dst * (1 - a) + src * a + 0.5).astype(np.uint8)


class Shot:
    def __init__(self, key, views, layers=None, ground=None, min_len=0.0, files=None):
        self.key, self.views, self.layers = key, views, layers or []
        self.ground = ground       # a full-frame RGB picture under the views, or None for the plain ground
        self.min_len = min_len
        self.files = files or []
        self.start = self.end = 0.0

    def draw(self, t):
        frame = self.ground.copy() if self.ground is not None else np.empty((H, W, 3), np.uint8)
        if self.ground is None:
            frame[:] = BG
        d = self.end - self.start
        for v in self.views:
            v.draw(frame, t, d + SHOT_XF)
        for l in self.layers:
            l.draw(frame, t)
        return frame


class Scene:
    """One chapter: its shots in order, layers that stay for the whole scene, and the narration beats that time it."""

    def __init__(self, key, title, shots, beats, layers=None, numbers=None, weight=1.0):
        self.key, self.title, self.shots, self.beats = key, title, shots, beats
        self.layers = layers or []
        self.numbers = numbers or []
        self.weight = weight
        self.start = self.dur = 0.0

    def frame(self, t):
        cur = self.shots[0]
        for s in self.shots:
            if t >= s.start:
                cur = s
        f = cur.draw(t - cur.start)
        k = self.shots.index(cur)
        if k > 0 and t - cur.start < SHOT_XF:
            prev = self.shots[k - 1]
            a = ease((t - cur.start) / SHOT_XF)
            f = cv2.addWeighted(f, a, prev.draw(t - prev.start), 1 - a, 0)
        for l in self.layers:
            l.draw(f, t)
        return f


class Timeline:
    def __init__(self, scenes):
        self.scenes = scenes
        t = 0.0
        for s in scenes:
            s.start = t
            t += s.dur - XF
        self.total = t + XF
        self.frames = int(round(self.total * FPS))

    def frame(self, n):
        t = n / FPS
        live = [s for s in self.scenes if s.start <= t < s.start + s.dur]
        if not live:
            live = [self.scenes[-1]]
        f = live[0].frame(min(t - live[0].start, live[0].dur))
        if len(live) > 1:
            b = live[1]
            a = ease((t - b.start) / XF)
            f = cv2.addWeighted(b.frame(t - b.start), a, f, 1 - a, 0)
        # the video opens from the ground and closes to it
        edge = min(t / 0.4, (self.total - t) / 0.6)
        if edge < 1:
            g = np.empty_like(f)
            g[:] = BG
            f = cv2.addWeighted(f, max(0.0, edge), g, 1 - max(0.0, edge), 0)
        return f


def ffmpeg_exe():
    return os.environ.get("ADOM_AIFLOW_FFMPEG") or "ffmpeg"


def encode_range(tl, a, b, path, threads=1, crf=22):
    """Frames a..b of the timeline into one H.264 segment. The rate is capped so two minutes stay a file that is easy to
    send (about 50 MB at most), and a keyframe comes every two seconds so a player can seek."""
    cmd = [ffmpeg_exe(), "-y", "-v", "error", "-f", "rawvideo", "-pix_fmt", "rgb24", "-s", f"{W}x{H}", "-r", str(FPS), "-i", "-",
           "-c:v", "libx264", "-preset", "fast", "-crf", str(crf), "-maxrate", "3500k", "-bufsize", "7000k", "-pix_fmt", "yuv420p",
           "-g", str(FPS * 2), "-threads", str(threads), path]
    p = subprocess.Popen(cmd, stdin=subprocess.PIPE, stderr=subprocess.PIPE)
    try:
        for n in range(a, b):
            p.stdin.write(tl.frame(n).tobytes())
        p.stdin.close()
    except BrokenPipeError:
        pass
    err = p.stderr.read().decode(errors="replace")
    if p.wait() != 0:
        raise RuntimeError(f"ffmpeg failed encoding frames {a}..{b}: {err.strip()[-400:]}")
    return path


def split_ranges(n, k):
    step = math.ceil(n / k)
    return [(i, min(n, i + step)) for i in range(0, n, step)]