shared/score_case.py

"""Score staged baseline predictions with the established WorldMem metric contracts."""

import argparse
import importlib
import importlib.util
import json
from pathlib import Path
import sys
import types

sys.path.insert(0, str(Path(__file__).resolve().parent / "metric_vendor"))

import numpy as np
import torch
import torch.nn.functional as F
import torchmetrics
from torchmetrics.functional.image.ssim import structural_similarity_index_measure
from torchmetrics.image.fid import FrechetInceptionDistance
from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity

from protocol import case_key, load_re10k_case


def load_source(path, name):
    spec = importlib.util.spec_from_file_location(name, path)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def load_oasis_vae(worldmem_root, checkpoint, device):
    # A namespace import keeps the authoritative VAE's relative imports intact
    # without importing WorldMem's training/experiment package.
    package = types.ModuleType("_benchmark_worldmem_models")
    package.__path__ = [str(Path(worldmem_root) / "algorithms/worldmem/models")]
    sys.modules[package.__name__] = package
    module = importlib.import_module(package.__name__ + ".vae")
    vae = module.VAE_models["vit-l-20-shallow-encoder"]()
    payload = torch.load(checkpoint, map_location="cpu", weights_only=False)
    state = payload.get("state_dict", payload)
    state = {key.removeprefix("vae."): value for key, value in state.items()}
    vae.load_state_dict(state, strict=True)
    return vae.requires_grad_(False).eval().to(device)


@torch.inference_mode()
def minecraft_reference(row, feature_root, vae, frames, device):
    relative = Path(row["relative_path"])
    feature = Path(feature_root) / relative.parent / (relative.stem + "_vae_feature.npy")
    start = int(row["target_start"])
    latents = np.load(feature, mmap_mode="r")[start:start + frames]
    decoded = []
    for offset in range(0, frames, 16):
        z = torch.from_numpy(np.array(latents[offset:offset + 16], copy=True)).float().to(device)
        # The existing cache stores scaled C x18x32 latents; decode matches
        # DeMemWM.decode, including float32 scaling before fp16 autocast.
        z = z.flatten(2).transpose(1, 2) / 0.07843137255
        with torch.autocast(device_type=device.type, dtype=torch.float16, enabled=device.type == "cuda"):
            rgb = (vae.decode(z) + 1) / 2
        decoded.append(rgb.float().clamp(0, 1).cpu())
    return torch.cat(decoded), str(feature)


@torch.inference_mode()
def frame_metrics(prediction, reference, dataset, lpips, device, batch_size=16):
    values = {name: [] for name in (["mse", "psnr", "lpips", "ssim"] if dataset == "re10k" else ["mse", "psnr", "lpips"])}
    for start in range(0, len(prediction), batch_size):
        pred, real = (video[start:start + batch_size].float().to(device).clamp(0, 1)
                      for video in (prediction, reference))
        mse = (pred - real).square().flatten(1).mean(1)
        epsilon = 1e-12 if dataset == "re10k" else torch.finfo(mse.dtype).eps
        values["mse"].append(mse.cpu())
        values["psnr"].append((-10 * mse.clamp_min(epsilon).log10()).cpu())
        values["lpips"].append(lpips.net(pred, real, normalize=True).flatten().cpu())
        if dataset == "re10k":
            values["ssim"].append(structural_similarity_index_measure(
                pred, real, data_range=1.0, reduction="none").cpu())
    return {name: torch.cat(parts) for name, parts in values.items()}


@torch.inference_mode()
def inception_features(frames, inception, device):
    # FrechetInceptionDistance(normalize=True).update performs this exact
    # truncating uint8 conversion before running torch-fidelity Inception.
    return torch.cat([inception((batch.to(device) * 255).byte()).reshape(len(batch), -1).float().cpu()
                      for batch in frames.split(16)])


@torch.inference_mode()
def video_inception_features(path, legacy_fid, inception, count, device):
    return torch.cat([inception(batch.to(device)).reshape(len(batch), -1).float().cpu()
                      for batch in legacy_fid.frame_batches(path, 0, count, 16)])


def metric_ranges(dataset, frames):
    width = 25 if dataset == "re10k" else 100
    ranges = [(f"frames_{start:03d}_{min(start + width, frames):03d}", start, min(start + width, frames))
              for start in range(0, frames, width)]
    ranges.append(("overall", 0, frames))
    if dataset == "re10k":
        ranges.append(("forward", 0, min(100, frames)))
        if frames > 100:
            ranges.append(("return", 100, frames))
    return ranges


@torch.inference_mode()
def score_case(case_dir, dataset, data_root, worldmem_root, vae_checkpoint=None,
               feature_root=None, keep_tensor=False, device="cuda"):
    case_dir, data_root, worldmem_root = map(Path, (case_dir, data_root, worldmem_root))
    device = torch.device(device)
    metadata = json.loads((case_dir / "case.json").read_text())
    row = metadata["case"]
    prediction_path = case_dir / "prediction.pt"
    prediction = torch.load(prediction_path, map_location="cpu", weights_only=True).float()
    expected_size = (256, 256) if dataset == "re10k" else (360, 640)
    max_frames = 300 if dataset == "re10k" else 500
    errors = []
    if torchmetrics.__version__ != "0.11.4":
        errors.append(f"Legacy contract requires torchmetrics0.11.4, found {torchmetrics.__version__}")
    if prediction.ndim != 4 or prediction.shape[1] != 3 or not 0 < len(prediction) <= max_frames:
        errors.append(f"Invalid generated RGB suffix shape {tuple(prediction.shape)}")
    if not torch.isfinite(prediction).all():
        errors.append("Non-finite prediction")
    if errors:
        raise ValueError("; ".join(errors))
    native_size = list(prediction.shape[-2:])
    if tuple(native_size) != expected_size:
        prediction = torch.cat([F.interpolate(batch.to(device), size=expected_size, mode="bicubic",
                                             align_corners=False, antialias=True).cpu()
                                for batch in prediction.split(16)])
    prediction.clamp_(0, 1)
    if dataset == "re10k":
        reference = load_re10k_case(row, data_root)["rgb"][100:100 + len(prediction)]
        reference_source = "raw DFoT test_256 RGB at matching source indices"
    else:
        vae = load_oasis_vae(worldmem_root, vae_checkpoint, device)
        reference, reference_source = minecraft_reference(
            row, feature_root or data_root / "vae_features", vae, len(prediction), device)
        del vae
        torch.cuda.empty_cache()
    if reference.shape != prediction.shape or not torch.isfinite(reference).all():
        raise ValueError("Reference frames do not match the complete finite prediction population")
    lpips = LearnedPerceptualImagePatchSimilarity(net_type="alex", normalize=True).eval().to(device)
    values = frame_metrics(prediction, reference, dataset, lpips, device)
    del lpips
    legacy = load_source(worldmem_root / "evaluation/long_horizon.py", "_benchmark_long_horizon")
    for name, frames in (("prediction", prediction), ("reference", reference)):
        legacy._write_video(case_dir / f"{name}.mp4", legacy._uint8_video(frames), fps=15)
    fid = FrechetInceptionDistance(feature=2048, normalize=dataset == "re10k").eval().to(device)
    if dataset == "minecraft":
        # calculate_fid imports only the authoritative long_horizon helpers.
        sys.path.insert(0, str(worldmem_root))
        legacy_fid = load_source(worldmem_root / "calculate_fid.py", "_benchmark_legacy_fid")
        fake_features = video_inception_features(case_dir / "prediction.mp4", legacy_fid, fid.inception, len(prediction), device)
        real_features = video_inception_features(case_dir / "reference.mp4", legacy_fid, fid.inception, len(reference), device)
    else:
        fake_features = inception_features(prediction, fid.inception, device)
        real_features = inception_features(reference, fid.inception, device)
    if any(not torch.isfinite(value).all() for value in (*values.values(), fake_features, real_features)):
        raise ValueError("Non-finite metric or FID feature; preserving prediction.pt for diagnosis")
    np.savez(case_dir / "frame_metrics.npz", **{key: value.numpy() for key, value in values.items()})
    np.save(case_dir / "fake_features.npy", fake_features.numpy())
    np.save(case_dir / "real_features.npy", real_features.numpy())
    summary = {
        "case": row, "sample_seed": int(metadata["seed"]), "dataset": dataset,
        "frames": len(prediction), "native_resolution": native_size, "metric_resolution": list(expected_size),
        "resize": "none" if tuple(native_size) == expected_size else "bicubic align_corners=False antialias=True",
        "reference_source": reference_source, "vae_checkpoint": str(vae_checkpoint) if dataset == "minecraft" else None,
        "torch": torch.__version__, "torchmetrics": torchmetrics.__version__,
        "lpips": "alex; normalize=True; mean of individual frame values",
        "psnr_epsilon": 1e-12 if dataset == "re10k" else torch.finfo(torch.float32).eps,
        "fid": "direct float01 -> truncating uint8" if dataset == "re10k" else "OpenCV decoding of libx264 quality8 yuv420p uint8 videos",
        "video": {"codec": "libx264", "quality": 8, "pixelformat": "yuv420p", "macro_block_size": 1, "fps": 15},
        "ranges": {name: {metric: float(value[start:end].mean()) for metric, value in values.items()}
                   for name, start, end in metric_ranges(dataset, len(prediction))},
    }
    (case_dir / "metrics.json").write_text(json.dumps(summary, indent=2) + "\n")
    if not keep_tensor:
        prediction_path.unlink()
    return summary


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--case-dir", type=Path, required=True)
    parser.add_argument("--dataset", choices=["minecraft", "re10k"], required=True)
    parser.add_argument("--data-root", type=Path, required=True)
    parser.add_argument("--worldmem-root", type=Path, required=True)
    parser.add_argument("--vae-checkpoint", type=Path)
    parser.add_argument("--feature-root", type=Path)
    parser.add_argument("--keep-tensor", action="store_true")
    parser.add_argument("--device", default="cuda")
    args = parser.parse_args()
    result = score_case(args.case_dir, args.dataset, args.data_root, args.worldmem_root,
                        args.vae_checkpoint, args.feature_root, args.keep_tensor, args.device)
    print(json.dumps({"case": case_key(result["case"]), "frames": result["frames"], "metrics": result["ranges"]["overall"]}), flush=True)


if __name__ == "__main__":
    main()