| |
| """Causal global/high branch ablation for a RenderMe INS checkpoint.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import sys |
| from contextlib import contextmanager |
| from pathlib import Path |
| from types import SimpleNamespace |
| from typing import Iterator, Optional |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch import Tensor |
| from torchvision.utils import make_grid, save_image |
|
|
| REPOSITORY_ROOT = Path(__file__).resolve().parents[1] |
| if str(REPOSITORY_ROOT) not in sys.path: |
| sys.path.insert(0, str(REPOSITORY_ROOT)) |
|
|
| from datasets.renderme_ins_dataset import RenderMeINSViewDataset |
| from scripts.evaluate_ins_stage1_renderme import ( |
| build_model_from_checkpoint, |
| foreground_metrics, |
| read_checkpoint, |
| ) |
| from scripts.train_ins_stage1 import ( |
| DistributedContext, |
| build_auxiliary_losses, |
| resize_target, |
| ) |
|
|
|
|
| CONDITIONS = ("baseline", "global_zero", "high_zero") |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser( |
| description="Render baseline/global=0/high=0 from one INS checkpoint." |
| ) |
| parser.add_argument("--checkpoint", type=Path, required=True) |
| parser.add_argument("--data-root", type=Path, default=None) |
| parser.add_argument("--frame-index", type=int, default=None) |
| parser.add_argument("--camera-id", default=None) |
| parser.add_argument("--source-camera-id", default=None) |
| parser.add_argument( |
| "--mode", |
| choices=("auto", "same-view", "fixed-source-reuse"), |
| default="auto", |
| ) |
| parser.add_argument("--resolution", type=int, default=512) |
| parser.add_argument("--amp", choices=("auto", "none", "bf16"), default="auto") |
| parser.add_argument("--lpips", action=argparse.BooleanOptionalAction, default=True) |
| parser.add_argument( |
| "--lpips-net", choices=("alex", "squeeze", "vgg"), default="alex" |
| ) |
| parser.add_argument("--device", default="cuda:0") |
| parser.add_argument( |
| "--output-dir", |
| type=Path, |
| default=Path("/workspace/cpfs-data/tmp/ins_branch_ablation"), |
| ) |
| args = parser.parse_args() |
| if not args.checkpoint.is_file(): |
| raise FileNotFoundError(args.checkpoint) |
| if args.data_root is not None and not args.data_root.is_dir(): |
| raise FileNotFoundError(args.data_root) |
| if args.frame_index is not None and args.frame_index < 0: |
| raise ValueError("--frame-index cannot be negative.") |
| if args.resolution <= 0: |
| raise ValueError("--resolution must be positive.") |
| return args |
|
|
|
|
| def saved_args(checkpoint: dict[str, object]) -> dict[str, object]: |
| value = checkpoint.get("args", {}) |
| return value if isinstance(value, dict) else {} |
|
|
|
|
| def resolve_config( |
| args: argparse.Namespace, |
| checkpoint: dict[str, object], |
| ) -> dict[str, object]: |
| saved = saved_args(checkpoint) |
| root = args.data_root or Path( |
| saved.get( |
| "data_root", |
| "/workspace/cpfs-data/data/renderme_all/0026/0026_e0", |
| ) |
| ) |
| frame_index = ( |
| int(saved.get("frame_index", 0)) |
| if args.frame_index is None |
| else args.frame_index |
| ) |
| mode = str(saved.get("mode", "same-view")) if args.mode == "auto" else args.mode |
| if mode not in {"same-view", "fixed-source-reuse"}: |
| raise ValueError(f"Unsupported checkpoint mode: {mode}") |
| source_camera_id = args.source_camera_id |
| if source_camera_id is None: |
| source_camera_id = saved.get("source_camera_id") |
| checkpoint_amp = str(saved.get("amp", "none")) |
| amp = checkpoint_amp if args.amp == "auto" else args.amp |
| if amp not in {"none", "bf16"}: |
| raise ValueError( |
| f"Ablation supports none/bf16, but checkpoint requested {amp!r}." |
| ) |
| return { |
| "root": Path(root), |
| "frame_index": frame_index, |
| "mode": mode, |
| "source_camera_id": ( |
| None if source_camera_id is None else str(source_camera_id) |
| ), |
| "canonical_radius": float(saved.get("canonical_radius", 2.7)), |
| "canonical_focus": tuple( |
| float(value) |
| for value in saved.get("canonical_focus", (0.0, 0.0, 0.2)) |
| ), |
| "background": tuple( |
| float(value) for value in saved.get("background", (1.0, 1.0, 1.0)) |
| ), |
| "amp": amp, |
| } |
|
|
|
|
| @contextmanager |
| def ablate_branch(model, condition: str) -> Iterator[None]: |
| """Temporarily replace one branch output with an all-zero tensor.""" |
| if condition not in CONDITIONS: |
| raise ValueError(f"Unknown branch condition: {condition}") |
| if condition == "baseline": |
| yield |
| return |
| module = ( |
| model.encoder.f_encoder |
| if condition == "global_zero" |
| else model.encoder.high_encoder |
| ) |
|
|
| def zero_output(_module, _inputs, output): |
| if not torch.is_tensor(output): |
| raise TypeError("INS encoder branch output must be a tensor.") |
| return torch.zeros_like(output) |
|
|
| handle = module.register_forward_hook(zero_output) |
| try: |
| yield |
| finally: |
| handle.remove() |
|
|
|
|
| def masked_mean(value: Tensor, mask: Tensor) -> Tensor: |
| if mask.shape[1] == 1 and value.shape[1] != 1: |
| mask = mask.expand(-1, value.shape[1], -1, -1) |
| return (value * mask).sum() / mask.sum().clamp_min(1.0) |
|
|
|
|
| @torch.inference_mode() |
| def evaluate(args: argparse.Namespace) -> dict[str, object]: |
| device = torch.device(args.device) |
| if device.type != "cuda" or not torch.cuda.is_available(): |
| raise RuntimeError("INS branch ablation requires CUDA.") |
| torch.cuda.set_device(device) |
| checkpoint = read_checkpoint(args.checkpoint) |
| config = resolve_config(args, checkpoint) |
| model = build_model_from_checkpoint(checkpoint, args.resolution, device) |
|
|
| dataset = RenderMeINSViewDataset( |
| root=config["root"], |
| frame_index=config["frame_index"], |
| image_size=512, |
| camera_ids=None if args.camera_id is None else [args.camera_id], |
| source_camera_id=config["source_camera_id"], |
| canonical_radius=config["canonical_radius"], |
| canonical_focus=config["canonical_focus"], |
| background=config["background"], |
| ) |
| target_camera_id = str(args.camera_id or dataset.source_camera_id) |
| target_sample = dataset.get_view(target_camera_id) |
| source_camera_id = str(dataset.source_camera_id) |
| source_sample = dataset.get_view(source_camera_id) |
| if config["mode"] == "same-view": |
| encoder_sample = target_sample |
| else: |
| encoder_sample = source_sample |
|
|
| encoder_image = encoder_sample["image"].unsqueeze(0).to(device) |
| encoder_camera = encoder_sample["camera"].unsqueeze(0).to(device) |
| render_camera = target_sample["camera"].unsqueeze(0).to(device) |
| target_source = target_sample["image"].unsqueeze(0).to(device) |
| target = resize_target(target_source, args.resolution) |
| foreground_mask = target_sample["foreground_mask"].unsqueeze(0).to(device) |
| background = torch.tensor(config["background"], device=device) |
|
|
| context = DistributedContext( |
| rank=0, |
| local_rank=device.index or 0, |
| world_size=1, |
| distributed=False, |
| device=device, |
| ) |
| metric_args = SimpleNamespace( |
| lpips_weight=1.0 if args.lpips else 0.0, |
| id_weight=0.0, |
| lpips_net=args.lpips_net, |
| ) |
| lpips_metric, _ = build_auxiliary_losses(metric_args, context) |
|
|
| predictions: dict[str, Tensor] = {} |
| masks: dict[str, Tensor] = {} |
| metrics: dict[str, dict[str, float]] = {} |
| for condition in CONDITIONS: |
| with ablate_branch(model, condition): |
| with torch.autocast( |
| device_type="cuda", |
| dtype=torch.bfloat16, |
| enabled=config["amp"] == "bf16", |
| ): |
| output = model( |
| encoder_image, |
| render_camera, |
| sampling_camera=encoder_camera, |
| encoder_camera=encoder_camera, |
| render_resolution=args.resolution, |
| deterministic_sampling=True, |
| ) |
| raw_prediction = output["image"].float().clamp(0.0, 1.0) |
| condition_metrics, composited, resized_mask = foreground_metrics( |
| raw_prediction, |
| target, |
| foreground_mask, |
| background, |
| lpips_metric, |
| identity_metric=None, |
| ) |
| gaussians = output["gaussians"] |
| condition_metrics.update( |
| { |
| "opacity_mean": float(gaussians["opacities"].float().mean()), |
| "scale_mean": float(gaussians["scales"].float().mean()), |
| "scale_max": float(gaussians["scales"].float().max()), |
| } |
| ) |
| predictions[condition] = composited |
| masks[condition] = resized_mask |
| metrics[condition] = condition_metrics |
|
|
| baseline = predictions["baseline"] |
| mask = masks["baseline"] |
| for condition in ("global_zero", "high_zero"): |
| metrics[condition]["foreground_render_change_mae"] = float( |
| masked_mean((predictions[condition] - baseline).abs(), mask) |
| ) |
| metrics[condition]["delta_foreground_psnr"] = ( |
| metrics[condition]["foreground_psnr"] |
| - metrics["baseline"]["foreground_psnr"] |
| ) |
| if "lpips" in metrics[condition]: |
| metrics[condition]["delta_lpips"] = ( |
| metrics[condition]["lpips"] - metrics["baseline"]["lpips"] |
| ) |
|
|
| args.output_dir.mkdir(parents=True, exist_ok=True) |
| encoder_01 = resize_target(encoder_image, args.resolution).cpu() |
| target_cpu = target.cpu() |
| mask_rgb = mask.expand(-1, 3, -1, -1).cpu() |
| prediction_row = [predictions[name].cpu() for name in CONDITIONS] |
| error_row = [(prediction - target_cpu).abs() * mask_rgb for prediction in prediction_row] |
| overview = make_grid( |
| torch.cat( |
| [ |
| encoder_01, |
| target_cpu, |
| prediction_row[0], |
| prediction_row[1], |
| mask_rgb, |
| error_row[0], |
| error_row[1], |
| error_row[2], |
| prediction_row[2], |
| ], |
| dim=0, |
| ), |
| nrow=4, |
| padding=3, |
| ) |
| |
| |
| |
| overview_path = args.output_dir / "ablation_overview.png" |
| save_image(overview, overview_path) |
| for name, prediction in predictions.items(): |
| save_image(prediction.cpu(), args.output_dir / f"{name}.png") |
| save_image( |
| (prediction.cpu() - target_cpu).abs() * mask_rgb, |
| args.output_dir / f"{name}_error.png", |
| ) |
|
|
| result: dict[str, object] = { |
| "checkpoint": str(args.checkpoint.resolve()), |
| "checkpoint_step": int(checkpoint.get("global_step", 0)), |
| "protocol": config["mode"], |
| "amp": config["amp"], |
| "encoder_camera_id": str(encoder_sample["camera_id"]), |
| "target_camera_id": target_camera_id, |
| "resolution": args.resolution, |
| "metrics": metrics, |
| "interpretation": { |
| "foreground_render_change_mae": ( |
| "Mean absolute foreground render change relative to baseline." |
| ), |
| "delta_foreground_psnr": ( |
| "A large negative value means the ablated branch was important." |
| ), |
| "delta_lpips": ( |
| "A large positive value means the ablated branch was important." |
| ), |
| }, |
| "outputs": { |
| "overview": str(overview_path), |
| **{ |
| name: str(args.output_dir / f"{name}.png") |
| for name in CONDITIONS |
| }, |
| }, |
| } |
| metrics_path = args.output_dir / "metrics.json" |
| with metrics_path.open("w", encoding="utf-8") as file: |
| json.dump(result, file, indent=2) |
| file.write("\n") |
| print(json.dumps(result, indent=2), flush=True) |
| print(f"overview: {overview_path}", flush=True) |
| print(f"metrics: {metrics_path}", flush=True) |
| return result |
|
|
|
|
| def main() -> int: |
| evaluate(parse_args()) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
