scripts/visualize_ckpt.py

#!/usr/bin/env python
"""Render per-decoder-stage PCA features from a pretraining checkpoint.

    python scripts/visualize_ckpt.py \
      --config configs/aff_base_pretrain_0.4ds_0.5mask_last_local.yaml \
      --checkpoint weights/pretrain/ckpt_epoch_399_affmae_fpw.pth \
      --samples docs/assets/sample{1,2,3,4}.png \
      --output output/plots/pca_decoder_stages.png

Without ``--samples`` it pulls a batch from the config's pretraining shards
instead, which is what the training loop does; naming files is usually what you
want for a figure, since the same images then appear in every run.

The model runs **unmasked**: ``forward_without_masking`` queries the decoder at
every patch position, so each stage's PCA covers the whole image with no holes.

This used to carry its own copy of the renderer, along with copies of
``compute_pca_rgb`` and ``denormalize``, and that copy ran ``_forward_internal``
instead -- so it masked the input and painted PCA only inside the masked region,
leaving the visible image showing through everywhere else. Two divergent
renderers for one figure, and the masked one is not what you want to look at:
the holes are where the model had no input, not where its features are
interesting. It now calls the same function the training loop uses.
"""

import argparse
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import torch  # noqa: E402

from affmae.config import load_config  # noqa: E402
from affmae.models.registry import get_model_spec  # noqa: E402
from affmae.utils.misc import strip_module_prefix  # noqa: E402
from affmae.viz.model_figures import run_pca_visualization  # noqa: E402



def load_samples(paths, config):
    """Preprocess image files exactly as the pretraining loader would.

    Reuses ``apply_custom_processing`` and ``create_transforms`` rather than
    ``preprocess_image``: the latter normalizes with the *finetuning* statistics
    and skips the microscope info-bar crop, and both matter here. The bar is a
    bright strip the encoder never saw during pretraining, because that crop
    removed it, so leaving it in puts out-of-distribution input at the bottom of
    every PCA panel.

    Args:
        paths: image files, one row of the figure each.
        config: Config, read for ``img_size`` and ``in_channels``.
    Returns:
        [len(paths), C, img_size, img_size] normalized tensor.
    """
    from PIL import Image

    from affmae.data.pretrain_dataset import (
        apply_custom_processing,
        create_transforms,
    )

    transform = create_transforms(config.img_size, config.in_channels)
    tensors = []
    for path in paths:
        if not os.path.isfile(path):
            raise SystemExit(f"sample not found: {path}")
        # apply_custom_processing reads the third key positionally, the way a
        # decoded WebDataset sample arrives, so the order here is load-bearing.
        sample = {"__key__": path, "__url__": "", "png": Image.open(path)}
        tensors.append(transform(apply_custom_processing(sample)["png"]))
    return torch.stack(tensors)


def parse_args():
    parser = argparse.ArgumentParser(
        description="PCA of each decoder stage's features, unmasked.")
    parser.add_argument("--config", required=True,
                        help="Pretraining config: model architecture, img_size, "
                             "and data.path when --samples is not given.")
    parser.add_argument("--samples", nargs="+", default=None,
                        help="Image files to use instead of the shards.")
    parser.add_argument("--checkpoint", required=True, help="Checkpoint (.pth).")
    parser.add_argument("--output", required=True, help="Output image (.png/.pdf).")
    parser.add_argument("--batch-skip", type=int, default=6,
                        help="Batches to discard before visualizing.")
    parser.add_argument("--num-images", type=int, default=6,
                        help="How many images to visualize.")
    parser.add_argument("--device", default=None,
                        help="Defaults to cuda when available.")
    return parser.parse_args()


def main():
    args = parse_args()

    device_name = args.device or ("cuda" if torch.cuda.is_available() else "cpu")
    if not os.path.exists(args.config):
        raise SystemExit(f"Config file not found: {args.config}")

    config = load_config(args.config)
    config.device = device_name
    if not args.samples and not getattr(config, "path", None):
        raise SystemExit(
            "without --samples this reads the pretraining shards, so it needs a "
            "pretraining config with data.path set. A finetuning config has "
            "data.base_path instead and will not work here.")

    device = torch.device(device_name)
    spec = get_model_spec(config.model_type)
    if spec.build_pretrain is None:
        raise SystemExit(f"model '{spec.name}' has no pretraining variant.")
    print(f"Initializing model type: {config.model_type}")
    model = spec.build_pretrain(config).to(device)

    if not os.path.isfile(args.checkpoint):
        raise SystemExit(f"Checkpoint file not found: {args.checkpoint}")
    checkpoint = torch.load(args.checkpoint, map_location=device, weights_only=False)
    state_dict = strip_module_prefix(
        checkpoint.get("model_state_dict", checkpoint))
    missing, unexpected = model.load_state_dict(state_dict, strict=True)
    print(f"Weights loaded. Missing: {len(missing)}, Unexpected: {len(unexpected)}")

    if args.samples:
        batch_images = load_samples(args.samples, config)
    else:
        from affmae.data.pretrain_dataset import build_pretrain_dataloader

        dataloader, _ = build_pretrain_dataloader(config)
        data_iter = iter(dataloader)
        try:
            for _ in range(args.batch_skip):
                next(data_iter)
            batch_images, _ = next(data_iter)
        except StopIteration:
            raise SystemExit(
                f"the loader ran out of data before batch {args.batch_skip}; "
                f"lower --batch-skip.")

    count = min(batch_images.shape[0], args.num_images)
    print(f"Visualizing {count} images, unmasked...")
    os.makedirs(os.path.dirname(os.path.abspath(args.output)) or ".", exist_ok=True)
    run_pca_visualization(model, batch_images[:count].to(device), args.output,
                          device)
    print(f"Wrote {args.output}")


if __name__ == "__main__":
    main()