Function process_with_animation

examples/visualize_features.py:301–478

Process images with radio1d_size animation. First pass: Extract PCA stats with max radio1d_size. Second pass: Generate frames for all radio1d_size values using those stats.

(loader, model, preprocessor, device, args, dirs, anim_dirs)

def process_with_animation(loader, model, preprocessor, device, args, dirs, anim_dirs):
    """
    Process images with radio1d_size animation.
    First pass: Extract PCA stats with max radio1d_size.
    Second pass: Generate frames for all radio1d_size values using those stats.
    """
    rank_print(f'Processing with radio1d animation: {args.radio1d_start} to {args.radio1d_end} step {args.radio1d_step}')

    # Collect all images first
    all_images = []
    ctr = 0
    for batches in loader:
        if ctr >= args.n:
            break
        for images, _ in batches:
            images = images.to(device=device, non_blocking=True)
            all_images.append(images)
            ctr += len(images)

    rank_print(f'Collected {len(all_images)} batches')

    # First pass: Get PCA stats from the highest radio1d_size
    rank_print(f'First pass: Computing PCA stats with radio1d_size={args.radio1d_end}')
    pca_stats_list = []

    for images in tqdm(all_images, desc='Computing PCA stats'):
        with torch.autocast(device.type, dtype=torch.bfloat16):
            p_images = preprocessor(images)

            if args.intermediates:
                outputs = model.forward_intermediates(
                    p_images,
                    indices=args.intermediates,
                    return_prefix_tokens=False,
                    norm=False,
                    stop_early=True,
                    output_fmt='NCHW',
                    intermediates_only=True,
                    aggregation=args.intermediate_aggregation,
                    norm_alpha_scheme="none",
                )
                all_feat = outputs
            else:
                kwargs = {}
                if args.neck:
                    kwargs['num_tokens'] = args.radio1d_end
                output = model(p_images, feature_fmt='NCHW', **kwargs)
                if args.adaptor_name:
                    all_feat = [
                        output['backbone'].features,
                        output[args.adaptor_name].features,
                    ]
                else:
                    all_feat = [output[1]]

        all_feat = [rearrange(f, 'b c h w -> b h w c').float() for f in all_feat]
        all_feat = list(zip(*all_feat))