| |
| """Append offset-0 RGB frames to an existing strict-causal RGB cache. |
| |
| This is faster than rebuilding the full cache with [-8,-4,-2,-1,0], |
| because the existing strict cache already stores the negative-offset history. |
| The output is intended only for current-observation ablations. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import copy |
| import json |
| import os |
| import sys |
| from collections import Counter, defaultdict |
| from typing import Any, Dict, List, Tuple |
|
|
| import cv2 |
| import torch |
|
|
| sys.path.insert(0, os.path.dirname(__file__)) |
| from build_rgb_frame_cache import parse_cell_frames, video_t_to_cell |
|
|
|
|
| def video_path(processed_root: str, boss: str, fight: int) -> str: |
| return os.path.join(processed_root, boss, f"video_fight{int(fight)}.mp4") |
|
|
|
|
| def read_decord_group( |
| path: str, |
| frame_indices: List[List[int]], |
| height: int, |
| width: int, |
| chunk_size: int, |
| ) -> torch.Tensor: |
| from decord import VideoReader, cpu |
|
|
| vr = VideoReader(path, ctx=cpu(0), width=width, height=height, num_threads=2) |
| n_frames = len(vr) |
| flat = [min(max(0, int(x)), n_frames - 1) for row in frame_indices for x in row] |
| out = torch.empty((len(flat), 3, height, width), dtype=torch.uint8) |
| for start in range(0, len(flat), chunk_size): |
| idx = flat[start:start + chunk_size] |
| batch = vr.get_batch(idx).asnumpy() |
| out[start:start + len(idx)] = torch.from_numpy(batch).permute(0, 3, 1, 2).contiguous() |
| return out.reshape(len(frame_indices), len(frame_indices[0]), 3, height, width) |
|
|
|
|
| def read_opencv_group( |
| path: str, |
| frame_indices: List[List[int]], |
| height: int, |
| width: int, |
| ) -> torch.Tensor: |
| cv2.setNumThreads(1) |
| cap = cv2.VideoCapture(path) |
| if not cap.isOpened(): |
| raise RuntimeError(f"cannot open video: {path}") |
| n_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) |
| rows = [] |
| for row in frame_indices: |
| frames = [] |
| for frame_idx in row: |
| idx = min(max(0, int(frame_idx)), max(0, n_frames - 1)) |
| cap.set(cv2.CAP_PROP_POS_FRAMES, idx) |
| ok, bgr = cap.read() |
| if not ok or bgr is None: |
| cap.release() |
| raise RuntimeError(f"cannot read frame {idx} from {path}") |
| rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) |
| if rgb.shape[0] != height or rgb.shape[1] != width: |
| rgb = cv2.resize(rgb, (width, height), interpolation=cv2.INTER_AREA) |
| frames.append(torch.from_numpy(rgb).permute(2, 0, 1).contiguous()) |
| rows.append(torch.stack(frames, dim=0)) |
| cap.release() |
| return torch.stack(rows, dim=0).to(torch.uint8) |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--base_cache", required=True) |
| ap.add_argument("--processed_root", default="data/processed") |
| ap.add_argument("--frame_in_cell", default="2") |
| ap.add_argument("--backend", choices=["decord", "opencv"], default="decord") |
| ap.add_argument("--chunk_size", type=int, default=512) |
| ap.add_argument("--out", required=True) |
| args = ap.parse_args() |
|
|
| data = torch.load(args.base_cache, map_location="cpu", weights_only=False) |
| base_rgb = data["rgb"].contiguous() |
| if base_rgb.dtype != torch.uint8: |
| raise TypeError(f"expected uint8 rgb cache, got {base_rgb.dtype}") |
| if any(int(x) >= 0 for x in data.get("history_offsets", [])): |
| raise ValueError("base_cache must be strict-causal; use build_rgb_frame_cache for custom current caches") |
| samples: List[Dict[str, Any]] = data["samples"] |
| frame_in_cell = parse_cell_frames(args.frame_in_cell) |
| n, t, c, h, w = base_rgb.shape |
| if c != 3: |
| raise ValueError(f"expected RGB channel dimension 3, got {c}") |
|
|
| groups: Dict[Tuple[str, int], List[Tuple[int, List[int]]]] = defaultdict(list) |
| all_first_indices = [] |
| for i, sample in enumerate(samples): |
| boss = sample["boss"] |
| fight = int(sample["fight"]) |
| target_cell = video_t_to_cell(sample["belief"]["time"]) |
| frame_indices = [4 * target_cell + int(in_cell) for in_cell in frame_in_cell] |
| all_first_indices.append(frame_indices[0]) |
| groups[(boss, fight)].append((i, frame_indices)) |
|
|
| current = torch.empty((n, len(frame_in_cell), 3, h, w), dtype=torch.uint8) |
| missing = 0 |
| missing_by_reason = Counter() |
| for group_id, ((boss, fight), entries) in enumerate(sorted(groups.items()), start=1): |
| path = video_path(args.processed_root, boss, fight) |
| if not os.path.exists(path): |
| missing += len(entries) |
| missing_by_reason["missing_video"] += len(entries) |
| continue |
| order = [i for i, _ in entries] |
| frame_indices = [idxs for _, idxs in entries] |
| try: |
| if args.backend == "decord": |
| frames = read_decord_group(path, frame_indices, h, w, args.chunk_size) |
| else: |
| frames = read_opencv_group(path, frame_indices, h, w) |
| except Exception as exc: |
| missing += len(entries) |
| missing_by_reason[f"read_error:{type(exc).__name__}"] += len(entries) |
| print(f"warning: failed {boss}/fight{fight}: {type(exc).__name__}: {exc}", flush=True) |
| continue |
| current[torch.tensor(order, dtype=torch.long)] = frames |
| print( |
| f"processed_groups {group_id}/{len(groups)} " |
| f"group={boss}/fight{fight} samples={len(entries)} filled={n - missing} missing={missing}", |
| flush=True, |
| ) |
|
|
| if missing: |
| raise RuntimeError(f"missing current frames: {missing} {dict(missing_by_reason)}") |
|
|
| out_rgb = torch.empty((n, t + len(frame_in_cell), 3, h, w), dtype=torch.uint8) |
| out_rgb[:, :t] = base_rgb |
| out_rgb[:, t:] = current |
|
|
| payload = dict(data) |
| payload["rgb"] = out_rgb.contiguous() |
| payload["history_offsets"] = list(data.get("history_offsets", [])) + [0] |
| payload["frame_in_cell"] = list(data.get("frame_in_cell", frame_in_cell)) |
| payload["allow_current_frame"] = True |
| payload["current_frame_source"] = "appended_from_raw_processed_fight_video" |
| payload["base_cache"] = args.base_cache |
| payload["video_backend"] = args.backend |
| payload["audit"] = copy.deepcopy(data.get("audit", {})) |
| payload["audit"]["append_current_frame"] = { |
| "base_cache": args.base_cache, |
| "requested": int(n), |
| "kept": int(n), |
| "missing": int(missing), |
| "missing_by_reason": dict(missing_by_reason), |
| "groups": int(len(groups)), |
| "frame_in_cell": list(frame_in_cell), |
| "current_frame_min": int(min(all_first_indices)) if all_first_indices else None, |
| "current_frame_max": int(max(all_first_indices)) if all_first_indices else None, |
| } |
| os.makedirs(os.path.dirname(args.out), exist_ok=True) |
| torch.save(payload, args.out) |
| print(json.dumps({ |
| "out": args.out, |
| "base_cache": args.base_cache, |
| "shape": list(out_rgb.shape), |
| "history_offsets": payload["history_offsets"], |
| "allow_current_frame": True, |
| "audit": payload["audit"]["append_current_frame"], |
| }, ensure_ascii=False, indent=2), flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
