Blame for code/vbench_prompts_compile/scripts/cache_matrix_transforms.py

#!/usr/bin/env python3
"""Cache physical video mirrors and controlled occlusions on the full frozen plan.

Every source/variant retains all 16 frame slots. Per-frame missing localization,
impossible controls, failed pixel parity and missing backend results remain in the
denominator. Geometry coverage never substitutes for independent visibility labels.
"""
from __future__ import annotations
import argparse
from collections import Counter, defaultdict
import hashlib
import json
from pathlib import Path
import shutil
import subprocess
import sys
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / 'src'))
from vbench_prompts_compile import annotation_jobs as J
from vbench_prompts_compile.official_replay import UPSTREAM_SHA
from vbench_prompts_compile.visual_transforms import mirror_detections, patch_frame
from cache_backend_outputs import sha


class BackendUnavailable(RuntimeError):
    """Stop a batch on model initialization failure instead of repeating it."""


def evidence_id(row):
    fields = [row['task'], row['relative_path'], row['evidence_kind'], row.get('axis'),
              row.get('occlusion_control'), row.get('level')]
    return hashlib.sha256(json.dumps(fields).encode()).hexdigest()[:24]


def unique_plan(rows):
    result = {}
    for row in rows:
        if row['evidence_kind'] == 'original':
            continue
        key = evidence_id(row)
        result.setdefault(key, {**row, 'evidence_id': key, 'matrix_ids': []})['matrix_ids'].append(row['sample_id'])
    return list(result.values())


def frame_hash(frame):
    return hashlib.sha256(frame.tobytes()).hexdigest()


def frames_for_backend(path, load_video):
    import torchvision.transforms as tv
    video = load_video(str(path), num_frames=16)
    _, _, height, width = video.size()
    if min(height, width) > 768:
        scale = 720. / min(height, width)
        video = tv.Resize(size=(int(scale * height), int(scale * width)))(video)
    return video.permute(0, 2, 3, 1).numpy()


def physical_mirror(source, target, axis, load_video):
    """Lossless RGB encode; GIFs use the official PIL decoder and frame indices."""
    import numpy as np
    target.parent.mkdir(parents=True, exist_ok=True)
    input_bytes = None
    timing = 'source timestamps passed through'
    input_args = ['-i', str(source)]
    if source.suffix.lower() == '.gif':
        # FFmpeg's GIF decoder can discard zero-duration frames that the official
        # PIL loader includes. Decode with that same loader before lossless encode.
        from fractions import Fraction
        from PIL import Image, ImageSequence
        with Image.open(source) as image:
            durations = [f.info.get('duration') for f in ImageSequence.Iterator(image)]
        known = set(durations) - {None, 0}
        if len(known) > 1 or (known and any(d is None or d == 0 for d in durations)):
            raise ValueError('variable GIF timing requires an explicit timestamp-preserving encoder')
        rate = str(Fraction(1000, next(iter(known)))) if known else '10'
        timing = 'uniform GIF duration preserved' if known else 'GIF has no duration metadata; 10 fps representation, all PIL frame indices preserved'
        full = load_video(str(source), return_tensor=False)
        _, height, width, _ = full.shape
        input_bytes = full.tobytes()
        input_args = ['-f', 'rawvideo', '-pixel_format', 'rgb24', '-video_size', f'{width}x{height}',
                      '-framerate', rate, '-i', 'pipe:0']
    command = ['ffmpeg', '-v', 'error', '-nostdin', '-y', *input_args,
        '-map', '0:v:0', '-vf', axis, '-c:v', 'libx264rgb', '-crf', '0', '-preset', 'veryfast',
        '-pix_fmt', 'rgb24', '-fps_mode', 'passthrough', '-an', '-f', 'mp4']
    if not target.exists():
        temporary = target.with_suffix('.partial.mp4')
        run = subprocess.run(command + [str(temporary)], input=input_bytes, capture_output=True)
        if run.returncode:
            raise RuntimeError('ffmpeg: ' + run.stderr[-800:].decode(errors='replace'))
        temporary.replace(target)
    before = load_video(str(source), num_frames=16, return_tensor=False)
    after = load_video(str(target), num_frames=16, return_tensor=False)
    expected = np.flip(before, axis=2 if axis == 'hflip' else 1)
    parity = before.shape == after.shape and np.array_equal(expected, after)
    if not parity:
        difference = float(np.max(np.abs(expected.astype('float32') - after.astype('float32')))) if expected.shape == after.shape else None
        raise ValueError(f'lossless mirror sample mismatch: {before.shape} vs {after.shape}, max error {difference}')
    return {'transformed_video': str(target), 'transformed_video_sha256': sha(target),
            'timing_policy': timing,
            'sampled_pixels_equal_exact_mirror': True, 'source_sample_sha256': [frame_hash(f) for f in before],
            'transformed_sample_sha256': [frame_hash(f) for f in after], 'ffmpeg_command': command + [str(target)]}


def detect(model, frame):
    import torch
    with torch.no_grad():
        result = model.run_caption_tensor(frame)
    detections = [{'label': r[0], 'box': [float(v) for v in r[1][:4]]} for r in result[0]]
    labels = sorted(set(str(v) for v in result[0][0][2])) if result[0] else []
    return detections, labels


def parse_args(argv=None):
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument('--task', choices=['spatial','objects'], required=True)
    p.add_argument('--plan', required=True)
    p.add_argument('--cache', required=True, help='complete original backend export for this task')
    p.add_argument('--video-root', required=True)
    p.add_argument('--out', required=True)
    p.add_argument('--artifact-root', required=True)
    p.add_argument('--vbench-root', default='/root/wenbiao_zhao/VBench')
    p.add_argument('--grit-weights', default='/root/.cache/vbench/grit_model/grit_b_densecap_objectdet.pth')
    p.add_argument('--device', default='cuda:0')
    p.add_argument('--max-videos', type=int, help='pilot; keeps complete plan')
    p.add_argument('--reserve-disk-gb', type=float, default=20)
    p.add_argument('--dry-run', action='store_true')
    return p.parse_args(argv)


def run(args, out):
    plan = unique_plan([r for r in J.read_jsonl(Path(args.plan)) if r['task'] == args.task])
    if not plan:
        raise ValueError('empty transform plan')
    cache = {r['relative_path']: r for r in J.read_jsonl(Path(args.cache))}
    upstream = subprocess.check_output(['git','-C',args.vbench_root,'rev-parse','HEAD'], text=True).strip()
    if upstream != UPSTREAM_SHA:
        raise ValueError('upstream commit mismatch')
    identity = {'schema': 1, 'task': args.task, 'plan_sha256': sha(args.plan), 'cache_sha256': sha(args.cache),
        'code_sha256': {n: sha(ROOT / n) for n in ['scripts/cache_matrix_transforms.py', 'src/vbench_prompts_compile/visual_transforms.py']},
        'upstream_sha': upstream, 'weights_sha256': sha(args.grit_weights), 'frames': 16,
        'variants': len(plan), 'artifact_root': args.artifact_root,
        'note': 'occlusion geometry verified, semantic visibility unverified'}
    J.bind_job(out.with_suffix('.manifest.json'), identity)
    if args.dry_run:
        print(json.dumps(identity, indent=2)); return 0
    journal = out.with_suffix('.attempts.jsonl')
    latest = {r['evidence_id']: r for r in J.read_jsonl(journal)}
    grouped = defaultdict(list)
    for row in plan:
        if row['evidence_id'] not in latest:
            grouped[row['relative_path']].append(row)
    # GPU imports are only needed for pending physical variants.
    sys.path.insert(0, str(Path(args.vbench_root).resolve()))
    import numpy as np
    import torch
    from PIL import Image
    from vbench.utils import load_video
    if args.device.startswith('cuda') and not torch.cuda.is_available():
        raise RuntimeError('requested CUDA backend unavailable')
    model = None
    artifact_root = Path(args.artifact_root)
    artifact_root.mkdir(parents=True, exist_ok=True)
    videos_attempted = 0
    stopped = None
    for relative, variants in grouped.items():
        if args.max_videos is not None and videos_attempted >= args.max_videos:
            break
        if shutil.disk_usage(artifact_root).free < args.reserve_disk_gb * 1024 ** 3:
            stopped = 'disk_reserve_reached'; break
        base = cache.get(relative, {})
        source = Path(args.video_root) / relative
        video_uid = hashlib.sha256(relative.encode()).hexdigest()[:24]
        frames = None
        frames_error = None
        memo = {}
        if base.get('status') == 'ok' and any(v['eligible'] and v['evidence_kind'] != 'evidence_mirror' for v in variants):
            try:
                if sha(source) != base.get('video_sha256'):
                    raise ValueError('source video changed since original cache')
                frames = frames_for_backend(source, load_video)
                if len(frames) != 16 or [frame_hash(f) for f in frames] != base['frame_sha256']:
                    raise ValueError('original frame sampling/preprocessing changed')
                for frame, detections, labels in zip(frames, base['frame_detections'], base['frame_labels']):
                    memo[frame_hash(frame)] = (detections, labels)
            except Exception as error:
                frames_error = f'{type(error).__name__}: {error}'
        for row in variants:
            record = {**row, 'source_video_sha256': base.get('video_sha256'), 'frame_size': base.get('frame_size')}
            try:
                if not row['eligible']:
                    record['status'] = 'ineligible'
                elif base.get('status') != 'ok':
                    record.update(status='source_cache_missing', source_status=base.get('status'))
                elif row['evidence_kind'] == 'evidence_mirror':
                    record.update(status='ok', frame_detections=mirror_detections(base['frame_detections'], base['frame_size'], row['axis']),
                                  frame_labels=base['frame_labels'], frame_status=['ok'] * 16,
                                  source_frame_sha256=base['frame_sha256'])
                else:
                    if frames_error or frames is None:
                        raise ValueError(frames_error or 'no source frames')
                    if row['evidence_kind'] == 'video_mirror':
                        target = artifact_root / 'spatial' / (video_uid + '-' + row['axis'] + '.mp4')
                        record.update(physical_mirror(source, target, row['axis'], load_video))
                        transformed = frames_for_backend(target, load_video)
                        patch_details = [{'status': 'ok', 'kind': row['axis']} for _ in transformed]
                    else:
                        required = row['base_official_target'].split(' and ')
                        patched = [patch_frame(frame, detections, required, row['occlusion_control'], row['level'])
                                   for frame, detections in zip(frames, base['frame_detections'])]
                        transformed = [pair[0] for pair in patched]
                        patch_details = [pair[1] for pair in patched]
                    detections, labels, hashes, errors = [], [], [], []
                    for i, (frame, details) in enumerate(zip(transformed, patch_details)):
                        if frame is None:
                            detections.append(None); labels.append(None); hashes.append(None); errors.append(details['status'])
                            continue
                        digest = frame_hash(frame)
                        hashes.append(digest)
                        try:
                            if digest not in memo:
                                if model is None:
                                    try:
                                        from vbench.third_party.grit_model import DenseCaptioning
                                        candidate = DenseCaptioning(torch.device(args.device))
                                        candidate.initialize_model_det(model_weight=args.grit_weights)
                                        model = candidate
                                    except Exception as error:
                                        raise BackendUnavailable(str(error)) from error
                                memo[digest] = detect(model, frame)
                            d, lab = memo[digest]
                            detections.append(d); labels.append(lab); errors.append('ok')
                        except BackendUnavailable:
                            raise
                        except Exception as error:
                            detections.append(None); labels.append(None); errors.append('backend:' + type(error).__name__)
                            details['backend_error'] = str(error)[:300]
                        if args.task == 'objects' and row.get('occlusion_control') == 'target' and row.get('level') == 1:
                            png = artifact_root / 'objects' / video_uid / f'target-full-frame-{i:02d}.png'
                            png.parent.mkdir(parents=True, exist_ok=True)
                            Image.fromarray(frame.astype('uint8')).save(png)
                            details['endpoint_image'] = str(png)
                            details['endpoint_image_sha256'] = sha(png)
                    record.update(status='ok' if all(e == 'ok' for e in errors) else 'partial',
                        frame_detections=detections, frame_labels=labels, frame_sha256=hashes, frame_status=errors,
                        frame_transform=patch_details, source_frame_sha256=base['frame_sha256'])
            except BackendUnavailable:
                raise
            except Exception as error:
                record.update(status='error:' + type(error).__name__, error=str(error)[:800])
            J.append_jsonl(journal, record)
            latest[row['evidence_id']] = record
            print(json.dumps({'cached_variants':len(latest), 'planned_variants':len(plan), 'last_status':record['status'],
                              'task':args.task, 'relative_path':relative, 'transform':row['transform']}), flush=True)
        videos_attempted += 1
    records = [latest.get(row['evidence_id'], {**row, 'status':'not_run'}) for row in plan]
    temporary = out.with_suffix('.jsonl.tmp')
    with temporary.open('w') as f:
        for row in records:
            f.write(json.dumps(row, ensure_ascii=False) + '\n')
    temporary.replace(out)
    statuses = dict(Counter(r['status'] for r in records))
    report = {'planned':len(plan), 'attempted':len(latest), 'status':statuses,
        'all_attempted':len(latest) == len(plan), 'all_eligible_complete':all(r['status']=='ok' for r in records if r['eligible']),
        'frame_status':dict(Counter(s for r in records for s in r.get('frame_status', []))),
        'stopped':stopped, 'visibility_verified':0, 'note':'full box coverage is not a visibility annotation'}
    J.atomic_json(out.with_suffix('.report.json'), report)
    print(json.dumps(report), flush=True)
    return 0 if report['all_attempted'] else 2


def main(argv=None):
    args = parse_args(argv)
    out = Path(args.out)
    with J.job_lock(out.with_suffix('.lock')):
        return run(args, out)


if __name__ == '__main__':
    raise SystemExit(main())