"""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: /canvas_000.jpg ... RGB canvases, all identical size /src_patch_position.npy int32 [total_patches, 3] = (t, h, w) /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