neural_codec/canvas_assembler.py

"""Patch selection + canvas assembly for the DCVC-RT release backend.

Turns per-frame DCVC-RT bit-cost bitmaps into the exact on-disk contract the
release codec path (``codec_video_processing_magevl.py``) consumes:

    <out_dir>/canvas_000.jpg ...      RGB canvases, all identical size
    <out_dir>/src_patch_position.npy  int32 [total_patches, 3] = (t, h, w)
    <out_dir>/meta.json               {"fps": ..., "canvas_files": [...]}

Layout invariants required by the downstream release code
---------------------------------------------------------
* All canvases share one size, so ``total_patches`` is divisible by the canvas
  count (``drop_padding_canvases``).
* Canvas patch grid ``(Gh, Gw)`` is even in both dims (``spatial_merge_size=2``
  block reorder in ``codec_positions_for_processor``).
* We select at **2x2-patch (28px) block granularity** and pack blocks in
  time-sorted order, so after the block reorder the ``t`` column forms
  consecutive runs whose lengths are multiples of 4 -- required for
  ``rewrite_text_with_codec_positions`` (``count // merge**2`` token counts).
* ``t`` is the *original video frame index* so ``t / fps`` is the real
  timestamp (same convention the HEVC/cv-preinfer contract uses).
* We never emit fully-padding canvases; a short final canvas is repeat-filled
  with real blocks (so ``drop_padding_canvases`` is a no-op).

The importance signal (which blocks to keep) is DCVC-RT bits; everything else
mirrors the dev codec's ``pack_topk`` selection (mandatory full first frame,
then global top-k over the rest).
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple

import cv2
import numpy as np
from PIL import Image


@dataclass
class AssembleConfig:
    # Must match the release image processor (preprocessor_config.json):
    # patch_size=16, merge_size=2 -> selectable unit is 32px. (The codec.patch=14
    # field in that config is a cv-preinfer internal; the Qwen2VLImageProcessor
    # grid that codec_positions_for_processor reads uses 16.)
    patch: int = 16
    spatial_merge_size: int = 2
    target_canvas: int = 32
    seq_len_frames: int = 64
    max_pixels: int = 150000
    # Canvas is square with ``canvas_token_side`` tokens per side (a token = one
    # 2x2 patch block = ``patch*merge`` px). If None, derived from max_pixels.
    canvas_token_side: Optional[int] = None
    mandatory_first_frame: bool = True

    @property
    def unit(self) -> int:
        return int(self.patch) * int(self.spatial_merge_size)

    def token_side(self) -> int:
        if self.canvas_token_side is not None:
            return int(self.canvas_token_side)
        side_px = math.sqrt(float(self.max_pixels))
        return max(2, int(round(side_px / self.unit)))


def _resize_longer_pad_square(img: np.ndarray, side_px: int, is_map: bool = False):
    """Resize longer side to ``side_px`` (keep aspect) and center-pad to a
    square. Returns (square_array, pad_info). ``is_map`` uses float + linear.
    """
    H, W = img.shape[:2]
    scale = float(side_px) / float(max(H, W))
    Hn = max(1, int(round(H * scale)))
    Wn = max(1, int(round(W * scale)))
    interp = cv2.INTER_LINEAR
    resized = cv2.resize(img.astype(np.float32) if is_map else img, (Wn, Hn), interpolation=interp)
    pad_top = (side_px - Hn) // 2
    pad_left = (side_px - Wn) // 2
    if is_map:
        out = np.zeros((side_px, side_px), dtype=np.float32)
        out[pad_top:pad_top + Hn, pad_left:pad_left + Wn] = resized
    else:
        out = np.zeros((side_px, side_px, 3), dtype=np.uint8)
        out[pad_top:pad_top + Hn, pad_left:pad_left + Wn] = resized
    info = dict(scale=scale, Hn=Hn, Wn=Wn, pad_top=pad_top, pad_left=pad_left)
    return out, info


def _token_valid_mask(side_px: int, unit: int, info: dict) -> np.ndarray:
    """Boolean (Ts, Ts): a token is valid if fully inside the non-padded region."""
    ts = side_px // unit
    top, left = info["pad_top"], info["pad_left"]
    bottom, right = top + info["Hn"], left + info["Wn"]
    mask = np.zeros((ts, ts), dtype=bool)
    for r in range(ts):
        y0, y1 = r * unit, (r + 1) * unit
        if y0 < top or y1 > bottom:
            continue
        for c in range(ts):
            x0, x1 = c * unit, (c + 1) * unit
            if x0 < left or x1 > right:
                continue
            mask[r, c] = True
    return mask


def _pool_map_to_tokens(square_map: np.ndarray, unit: int) -> np.ndarray:
    """Sum a (side, side) map into (Ts, Ts) token scores."""
    side = square_map.shape[0]
    ts = side // unit
    m = square_map[: ts * unit, : ts * unit]
    return m.reshape(ts, unit, ts, unit).sum(axis=(1, 3))


def assemble_canvases(
    sampled_frames_rgb: List[np.ndarray],
    frame_ids: List[int],
    bitmaps: Dict[int, np.ndarray],
    cfg: AssembleConfig,
) -> Tuple[List[Image.Image], np.ndarray]:
    """Select high-bit blocks and pack them into release-contract canvases.

    ``sampled_frames_rgb[i]`` is the decoded RGB frame for ``frame_ids[i]``;
    ``bitmaps[frame_ids[i]]`` is its DCVC-RT bit map (H/16, W/16). Returns
    (list of PIL RGB canvases, src_positions int32 [total_patches, 3]).
    """
    unit = cfg.unit
    ts = cfg.token_side()
    side_px = ts * unit
    seq_len = len(sampled_frames_rgb)

    # Per-frame square frames + token score maps + validity.
    frames_sq: List[np.ndarray] = []
    scores = np.full((seq_len, ts, ts), -np.inf, dtype=np.float32)
    for i in range(seq_len):
        fsq, info = _resize_longer_pad_square(sampled_frames_rgb[i], side_px, is_map=False)
        frames_sq.append(fsq)
        bm = bitmaps.get(int(frame_ids[i]))
        if bm is None:
            bm = np.zeros((max(1, side_px // 16), max(1, side_px // 16)), dtype=np.float32)
        bm_sq, _ = _resize_longer_pad_square(bm, side_px, is_map=True)
        vmask = _token_valid_mask(side_px, unit, info)
        tok = _pool_map_to_tokens(bm_sq, unit)
        scores[i][vmask] = tok[vmask]

    tokens_per_canvas = ts * ts
    target_tokens = int(cfg.target_canvas) * tokens_per_canvas

    # --- selection: mandatory full first frame, then global top-k over rest ---
    selected: List[Tuple[int, int, int]] = []  # (t_orig, th, tw)
    if cfg.mandatory_first_frame and seq_len > 0:
        for r in range(ts):
            for c in range(ts):
                selected.append((int(frame_ids[0]), r, c))

    remaining = target_tokens - len(selected)
    if remaining > 0 and seq_len > 1:
        rest = scores[1:].copy()  # (seq_len-1, ts, ts)
        flat = rest.reshape(-1)
        finite = np.isfinite(flat)
        n_avail = int(finite.sum())
        k = min(remaining, n_avail)
        if k > 0:
            order = np.argsort(-flat, kind="stable")[:k]
            for idx in order:
                fi = idx // (ts * ts) + 1
                rem = idx % (ts * ts)
                selected.append((int(frame_ids[fi]), rem // ts, rem % ts))

    if not selected:
        selected.append((int(frame_ids[0]) if frame_ids else 0, 0, 0))

    # Time-sorted packing so per-canvas ``t`` runs are contiguous.
    selected.sort(key=lambda x: (x[0], x[1], x[2]))

    # Repeat-fill to a whole number of canvases (>=1), never half-pad.
    if len(selected) < tokens_per_canvas:
        n_canvas = 1
    else:
        n_canvas = min(int(cfg.target_canvas), len(selected) // tokens_per_canvas)
        n_canvas = max(1, n_canvas)
    target = n_canvas * tokens_per_canvas
    if len(selected) < target:
        selected = selected + [selected[-1]] * (target - len(selected))
    else:
        selected = selected[:target]

    # index sampled frame_id -> its square frame (dup ids collapse fine)
    fid_to_sq = {int(frame_ids[i]): frames_sq[i] for i in range(seq_len)}

    Gh = Gw = ts * cfg.spatial_merge_size  # canvas patch grid
    sms = int(cfg.spatial_merge_size)
    images: List[Image.Image] = []
    all_positions: List[np.ndarray] = []

    for ci in range(n_canvas):
        canvas = np.zeros((side_px, side_px, 3), dtype=np.uint8)
        positions = np.zeros((Gh * Gw, 3), dtype=np.int32)
        chunk = selected[ci * tokens_per_canvas:(ci + 1) * tokens_per_canvas]
        for p, (t_orig, th, tw) in enumerate(chunk):
            rc, cc = p // ts, p % ts  # canvas token cell
            src = fid_to_sq.get(int(t_orig))
            if src is not None:
                canvas[rc * unit:(rc + 1) * unit, cc * unit:(cc + 1) * unit] = \
                    src[th * unit:(th + 1) * unit, tw * unit:(tw + 1) * unit]
            pr, pc = sms * rc, sms * cc          # canvas patch coords (top-left)
            sph, spw = sms * th, sms * tw        # source patch coords (top-left)
            for dy in range(sms):
                for dx in range(sms):
                    row = (pr + dy) * Gw + (pc + dx)
                    positions[row] = (int(t_orig), sph + dy, spw + dx)
        images.append(Image.fromarray(canvas))
        all_positions.append(positions)

    src_positions = np.concatenate(all_positions, axis=0).astype(np.int32)
    return images, src_positions