| |
| """Render segmentation and token-layout figures for every image in a test split. |
| |
| python scripts/render_test_set.py \ |
| --config configs/aff_base_finetune_512_fpw.yaml \ |
| --checkpoint weights/segmentation/fpw_aff_base_ft_512_slits_pgbmi.pth --tag ft512 --img-size 512 |
| |
| Writes ``<out>/<tag>/segmentation/NNNN_<name>.png`` and |
| ``<out>/<tag>/tokens/NNNN_<name>.png``. |
| """ |
|
|
| import argparse |
| import os |
| import sys |
|
|
| import torch |
|
|
| sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) |
|
|
| from affmae.config import load_config |
| from affmae.data.finetune_dataset import build_test_dataloader |
| from affmae.eval.loader import load_for_eval |
| from affmae.utils.dist import resolve_device, unwrap_model |
|
|
|
|
| def stage_positions(model, images): |
| """Per-stage token positions for a preprocessed batch. |
| |
| Mirrors ``AFFMAE.token_layout`` but takes an already-normalized batch, so |
| the dataloader's own preprocessing is not applied twice. |
| |
| Args: |
| model: an AFF segmentation model. |
| images: [B, C, H, W] normalized batch on the model's device. |
| Returns: |
| List of [B, N_s, 2] position tensors, or None if this encoder keeps a |
| fixed grid. |
| """ |
| bare = unwrap_model(model) |
| encoder = getattr(bare, "encoder", None) |
| if encoder is None or not hasattr(encoder, "forward_with_pos"): |
| return None |
| |
| pos, feat, height, width = encoder.patch_embed(images, ids_masked=None) |
| return encoder.forward_with_pos(feat, pos, height, width) |
|
|
|
|
| def build_parser(): |
| parser = argparse.ArgumentParser( |
| description=__doc__.split("\n")[0], |
| formatter_class=argparse.RawDescriptionHelpFormatter, epilog=__doc__) |
| parser.add_argument("--config", required=True) |
| parser.add_argument("--checkpoint", required=True) |
| parser.add_argument("--tag", required=True, help="Subdirectory name.") |
| parser.add_argument("--img-size", type=int, default=None, |
| help="Override the config's resolution to match the " |
| "checkpoint it was trained at.") |
| parser.add_argument("--num-classes", type=int, default=None, |
| help="Override the config's class count. The dataset has " |
| "3 (background, PGBMI, slit), which is what every " |
| "config declares, so this is only for a checkpoint " |
| "trained against a different labelling.") |
| parser.add_argument("--out", default="output/test_set_renders") |
| parser.add_argument("--limit", type=int, default=None, |
| help="Stop after this many images (for a smoke test).") |
| parser.add_argument("--device", default=None) |
| return parser |
|
|
|
|
| def main() -> None: |
| args = build_parser().parse_args() |
| from affmae.viz import PAPER, render_segmentation, render_token_layout |
|
|
| cfg = load_config(args.config) |
| if args.img_size: |
| cfg.img_size = args.img_size |
| if args.num_classes: |
| cfg.num_classes = args.num_classes |
| cfg.device = resolve_device(args.device or getattr(cfg, "device", None)) |
|
|
| seg_dir = os.path.join(args.out, args.tag, "segmentation") |
| tok_dir = os.path.join(args.out, args.tag, "tokens") |
| os.makedirs(seg_dir, exist_ok=True) |
| os.makedirs(tok_dir, exist_ok=True) |
|
|
| model = load_for_eval(cfg, args.checkpoint) |
| loader = build_test_dataloader(cfg) |
| print(f"[{args.tag}] {cfg.img_size}px, {cfg.num_classes} classes, " |
| f"{len(loader)} images -> {args.out}/{args.tag}/") |
|
|
| written, skipped_tokens = 0, 0 |
| with torch.no_grad(): |
| for index, (images, targets, paths) in enumerate(loader): |
| if args.limit and index >= args.limit: |
| break |
| images = images.to(cfg.device) |
| targets = targets.to(cfg.device).long() |
| outputs = model(images) |
| logits = outputs[-1] if isinstance(outputs, (list, tuple)) else outputs |
|
|
| raw = paths[0] if isinstance(paths, (list, tuple)) else paths |
| name = os.path.splitext(os.path.basename(str( |
| raw[0] if isinstance(raw, (list, tuple)) else raw)))[0] |
| |
| safe = "".join(c if c.isalnum() or c in "-_." else "_" |
| for c in name)[:80] |
|
|
| render_segmentation( |
| images.cpu(), logits.float().cpu(), cfg.num_classes, |
| os.path.join(seg_dir, f"{index:04d}_{safe}.png"), |
| targets=targets.cpu(), config=PAPER) |
|
|
| positions = stage_positions(model, images) |
| if positions is None: |
| skipped_tokens += 1 |
| else: |
| render_token_layout( |
| images.cpu(), [p.float().cpu() for p in positions], |
| cfg.patch_size, |
| os.path.join(tok_dir, f"{index:04d}_{safe}.png"), |
| config=PAPER) |
| written += 1 |
| if written % 25 == 0: |
| print(f" {written}/{len(loader)}") |
|
|
| print(f"[{args.tag}] wrote {written} segmentation figure(s)" |
| + (f", token layout skipped for {skipped_tokens} (fixed-grid encoder)" |
| if skipped_tokens else f" and {written} token figure(s)")) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
