| """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): |
| |
| |
| 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) |
| |
| |
| 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): |
| |
| |
| 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": |
| |
| 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() |
|
|
