| |
| """Visualize the global/high INS encoder branches on a RenderMe view. |
| |
| The images are activation diagnostics. They show spatial response, PCA |
| projections, and representative channels; they do not assign semantic labels |
| to the features. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import sys |
| from pathlib import Path |
|
|
| 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, |
| read_checkpoint, |
| ) |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser( |
| description="Visualize global and high-resolution INS branch activations." |
| ) |
| 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("--resolution", type=int, default=512) |
| parser.add_argument("--device", default="cuda:0") |
| parser.add_argument("--amp", choices=("none", "bf16"), default="none") |
| parser.add_argument( |
| "--output-dir", |
| type=Path, |
| default=Path("/workspace/cpfs-data/tmp/ins_branch_viz"), |
| ) |
| parser.add_argument("--top-channels", type=int, default=12) |
| parser.add_argument("--save-features", action="store_true") |
| args = parser.parse_args() |
| if not args.checkpoint.is_file(): |
| raise FileNotFoundError(args.checkpoint) |
| if args.resolution <= 0 or args.top_channels <= 0: |
| raise ValueError("resolution and top-channels must be positive.") |
| return args |
|
|
|
|
| def robust_normalize(value: Tensor) -> Tensor: |
| value = value.float() |
| low = torch.quantile(value, 0.01) |
| high = torch.quantile(value, 0.99) |
| return ((value - low) / (high - low).clamp_min(1e-8)).clamp(0.0, 1.0) |
|
|
|
|
| def feature_mean_map(features: Tensor) -> Tensor: |
| value = robust_normalize(features.abs().mean(dim=1, keepdim=True)) |
| return value.expand(-1, 3, -1, -1) |
|
|
|
|
| def feature_pca_rgb(features: Tensor) -> Tensor: |
| features = features.float() |
| batch, channels, height, width = features.shape |
| if batch != 1: |
| raise ValueError("Branch visualization expects batch size one.") |
| samples = features[0].permute(1, 2, 0).reshape(height * width, channels) |
| samples = samples - samples.mean(dim=0, keepdim=True) |
| _, _, vectors = torch.pca_lowrank(samples, q=3, center=False, niter=2) |
| projected = samples @ vectors[:, :3] |
| projected = projected.reshape(height, width, 3).permute(2, 0, 1).unsqueeze(0) |
| return robust_normalize(projected) |
|
|
|
|
| def top_channel_grid(features: Tensor, count: int) -> tuple[Tensor, list[int]]: |
| features = features.float() |
| count = min(count, features.shape[1]) |
| scores = features.abs().flatten(2).mean(dim=2)[0] |
| indices = torch.topk(scores, k=count).indices.tolist() |
| maps = [ |
| robust_normalize(features[:, index : index + 1]).expand(-1, 3, -1, -1) |
| for index in indices |
| ] |
| return make_grid(torch.cat(maps, dim=0), nrow=min(4, count), padding=2), indices |
|
|
|
|
| def gradient_energy(features: Tensor) -> Tensor: |
| features = features.float() |
| horizontal = (features[:, :, :, 1:] - features[:, :, :, :-1]).abs().mean() |
| vertical = (features[:, :, 1:, :] - features[:, :, :-1, :]).abs().mean() |
| return 0.5 * (horizontal + vertical) |
|
|
|
|
| def capture_branches(model) -> tuple[dict[str, Tensor], list[object]]: |
| captured: dict[str, Tensor] = {} |
|
|
| def make_hook(name: str): |
| def hook(_module, _inputs, output): |
| captured[name] = output.detach().float() |
|
|
| return hook |
|
|
| handles = [ |
| model.encoder.f_encoder.register_forward_hook(make_hook("global")), |
| model.encoder.high_encoder.register_forward_hook(make_hook("high")), |
| model.encoder.final_encoder.register_forward_hook(make_hook("fused")), |
| ] |
| return captured, handles |
|
|
|
|
| def resolve_data_config(args: argparse.Namespace, checkpoint: dict) -> dict: |
| saved = checkpoint.get("args", {}) |
| if not isinstance(saved, dict): |
| saved = {} |
| 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 |
| source = args.source_camera_id |
| if source is None: |
| source = saved.get("source_camera_id") |
| canonical_radius = float(saved.get("canonical_radius", 2.7)) |
| canonical_focus = tuple(float(v) for v in saved.get("canonical_focus", (0.0, 0.0, 0.2))) |
| background = tuple(float(v) for v in saved.get("background", (1.0, 1.0, 1.0))) |
| return { |
| "root": root, |
| "frame_index": frame_index, |
| "source": None if source is None else str(source), |
| "canonical_radius": canonical_radius, |
| "canonical_focus": canonical_focus, |
| "background": background, |
| } |
|
|
|
|
| @torch.inference_mode() |
| def visualize(args: argparse.Namespace) -> dict[str, object]: |
| device = torch.device(args.device) |
| if device.type != "cuda" or not torch.cuda.is_available(): |
| raise RuntimeError("Branch visualization requires CUDA.") |
| torch.cuda.set_device(device) |
| checkpoint = read_checkpoint(args.checkpoint) |
| model = build_model_from_checkpoint(checkpoint, args.resolution, device) |
| config = resolve_data_config(args, checkpoint) |
|
|
| |
| camera_id = args.camera_id or config["source"] |
| dataset = RenderMeINSViewDataset( |
| root=config["root"], |
| frame_index=config["frame_index"], |
| image_size=512, |
| camera_ids=None if camera_id is None else [camera_id], |
| source_camera_id=config["source"], |
| canonical_radius=config["canonical_radius"], |
| canonical_focus=config["canonical_focus"], |
| background=config["background"], |
| ) |
| camera_id = str(camera_id or dataset.source_camera_id) |
| sample = dataset.get_view(camera_id) |
| image = sample["image"].unsqueeze(0).to(device) |
| captured, handles = capture_branches(model) |
| with torch.autocast( |
| device_type="cuda", |
| dtype=torch.bfloat16, |
| enabled=args.amp == "bf16", |
| ): |
| triplanes = model.encoder(image) |
| for handle in handles: |
| handle.remove() |
| if set(captured) != {"global", "high", "fused"}: |
| raise RuntimeError(f"Missing branch activations: {sorted(captured)}") |
|
|
| args.output_dir.mkdir(parents=True, exist_ok=True) |
| source = ((image.float() + 1.0) * 0.5).clamp(0.0, 1.0) |
| global_features = captured["global"] |
| high_features = captured["high"] |
| fused_features = captured["fused"] |
| global_mean = feature_mean_map(global_features) |
| high_mean = feature_mean_map(high_features) |
| global_pca = feature_pca_rgb(global_features) |
| high_pca = feature_pca_rgb(high_features) |
| fused_mean = feature_mean_map(fused_features) |
| global_channels, global_indices = top_channel_grid(global_features, args.top_channels) |
| high_channels, high_indices = top_channel_grid(high_features, args.top_channels) |
|
|
| source_small = F.interpolate( |
| source, size=global_mean.shape[-2:], mode="bilinear", align_corners=False |
| ) |
| overview = make_grid( |
| torch.cat( |
| [source_small, global_mean, global_pca, high_mean, high_pca, fused_mean], |
| dim=0, |
| ), |
| nrow=3, |
| padding=3, |
| ) |
| overview_path = args.output_dir / "branch_overview.png" |
| global_path = args.output_dir / "global_top_channels.png" |
| high_path = args.output_dir / "high_top_channels.png" |
| save_image(overview, overview_path) |
| save_image(global_channels, global_path) |
| save_image(high_channels, high_path) |
|
|
| stats: dict[str, object] = { |
| "checkpoint": str(args.checkpoint.resolve()), |
| "camera_id": camera_id, |
| "global_shape": list(global_features.shape), |
| "high_shape": list(high_features.shape), |
| "fused_shape": list(fused_features.shape), |
| "triplane_shape": list(triplanes.shape), |
| "global_top_channels": global_indices, |
| "high_top_channels": high_indices, |
| "global_mean_abs": float(global_features.abs().mean()), |
| "high_mean_abs": float(high_features.abs().mean()), |
| "global_gradient_energy": float(gradient_energy(global_features)), |
| "high_gradient_energy": float(gradient_energy(high_features)), |
| "global_spatial_std": float(global_features.flatten(2).std(dim=2).mean()), |
| "high_spatial_std": float(high_features.flatten(2).std(dim=2).mean()), |
| "outputs": { |
| "overview": str(overview_path), |
| "global_channels": str(global_path), |
| "high_channels": str(high_path), |
| }, |
| } |
| if args.save_features: |
| feature_path = args.output_dir / "branch_features.pt" |
| torch.save( |
| { |
| "global": global_features.cpu(), |
| "high": high_features.cpu(), |
| "fused": fused_features.cpu(), |
| }, |
| feature_path, |
| ) |
| stats["outputs"]["features"] = str(feature_path) |
| with (args.output_dir / "stats.json").open("w", encoding="utf-8") as file: |
| json.dump(stats, file, indent=2) |
| file.write("\n") |
| print(json.dumps(stats, indent=2), flush=True) |
| return stats |
|
|
|
|
| def main() -> int: |
| visualize(parse_args()) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
