analysis/tools/cache_activations.py

#!/usr/bin/env python
"""Cache SAE activations for analysis notebooks and dashboard.

Usage:
    python analysis/tools/cache_activations.py
    python analysis/tools/cache_activations.py --checkpoint /path/to/checkpoint
    python analysis/tools/cache_activations.py --max-batches 10
    python analysis/tools/cache_activations.py --max-memory-gb 8

Paths and defaults are read from ``default_configs/sae_vis_config.json``
(or ``SAE_VIS_CONFIG_PATH`` / ``--sae-vis-config-path``).
"""
from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from analysis.cache_utils import build_activation_cache
from analysis.config import load_sae_vis_config
from log_config import get_logger, setup_logging

logger = get_logger(__name__)


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description='Cache KADID-10k SAE activations')
    p.add_argument(
        '--sae-vis-config-path',
        type=str,
        default=None,
        help='Path to sae_vis_config.json (default: template or SAE_VIS_CONFIG_PATH)',
    )
    p.add_argument(
        '--dataset',
        type=str,
        default=None,
        choices=['kadid', 'local_kadid'],
        help='Explicit dataset choice for caching',
    )
    p.add_argument('--checkpoint', type=str, default=None, help='Override SAE_CHECKPOINT_PATH')
    p.add_argument('--kadid-path', type=str, default=None, help='Override KADID_IMAGES_PATH')
    p.add_argument('--layer', type=int, default=None, help='Override LAYER_NUM')
    p.add_argument(
        '--iqa-metric',
        type=str,
        default=None,
        choices=['arniqa-kadid', 'maniqa', 'qalign', 'liqe', 'liqe_mix', 'topiq_nr'],
        help='Override IQA_METRIC',
    )
    p.add_argument(
        '--swin-num',
        type=int,
        default=None,
        choices=[1, 2],
        help='Override SWIN_NUM (MANIQA only)',
    )
    p.add_argument('--device', type=str, default=None, help='Override DEVICE (cuda / cpu)')
    p.add_argument('--batch-size', type=int, default=None, help='Override BATCH_SIZE')
    p.add_argument('--num-workers', type=int, default=None, help='Override NUM_WORKERS')
    p.add_argument('--max-batches', type=int, default=None, help='Limit batches (debug)')
    p.add_argument(
        '--max-memory-gb',
        type=float,
        default=30.0,
        help='Memory limit for accumulated sparse matrices, GB',
    )
    p.add_argument(
        '--min-distortion-level',
        type=int,
        default=None,
        help='Minimum KADID distortion level (1..5)',
    )
    p.add_argument(
        '--srground-include-sr-artifact',
        action='store_true',
        help='Include SR artifacts in SRGround mask (default: real distortions only)',
    )
    return p.parse_args()


def main() -> None:
    setup_logging()
    args = parse_args()
    cfg = load_sae_vis_config(args.sae_vis_config_path)

    dataset_arg = args.dataset
    if dataset_arg is None:
        dataset_arg = 'local_kadid' if cfg.DATASET == 'local_kadid' else 'kadid'
    dataset_key = 'local_kadid' if dataset_arg == 'local_kadid' else 'kadid10k'

    checkpoint_path = args.checkpoint or cfg.SAE_CHECKPOINT_PATH
    dataset_root = args.kadid_path or cfg.DATASET_ROOT
    layer_num = args.layer if args.layer is not None else cfg.LAYER_NUM
    iqa_metric = args.iqa_metric or cfg.IQA_METRIC
    swin_num = args.swin_num if args.swin_num is not None else cfg.SWIN_NUM
    device = args.device or cfg.DEVICE
    batch_size = args.batch_size if args.batch_size is not None else cfg.BATCH_SIZE
    num_workers = args.num_workers if args.num_workers is not None else cfg.NUM_WORKERS
    min_distortion_level = (
        args.min_distortion_level
        if args.min_distortion_level is not None
        else cfg.KADID_MIN_DISTORTION_LEVEL
    )

    if not (1 <= min_distortion_level <= 5):
        raise ValueError('--min-distortion-level must be in [1, 5]')

    if args.checkpoint is not None:
        cache_dir = os.path.join(os.path.dirname(checkpoint_path), 'activation_cache')
        acts_filename = os.path.basename(cfg.DATASET_CACHE_CONFIGS[dataset_key].acts_cache_path)
        cache_path = os.path.join(cache_dir, dataset_key, acts_filename)
        os.makedirs(os.path.dirname(cache_path), exist_ok=True)
    else:
        cache_path = cfg.DATASET_CACHE_CONFIGS[dataset_key].acts_cache_path

    logger.info('=' * 60)
    logger.info('  SAE checkpoint : %s', checkpoint_path)
    logger.info('  Dataset        : %s', dataset_arg)
    logger.info('  Dataset root   : %s', dataset_root)
    logger.info('  Layer          : %s', layer_num)
    logger.info('  IQA metric     : %s', iqa_metric)
    if iqa_metric == 'maniqa':
        logger.info('  Swin num       : %s', swin_num)
    logger.info('  Device         : %s', device)
    logger.info('  Batch size     : %s', batch_size)
    logger.info('  Min dist level : %s', min_distortion_level)
    logger.info('  Cache output   : %s', cache_path)
    if args.max_batches:
        logger.debug('  Max batches    : %s', args.max_batches)
    logger.info('=' * 60)

    logger.info('Starting caching...')
    build_info = build_activation_cache(
        dataset=dataset_key,
        cache_path=cache_path,
        checkpoint_path=checkpoint_path,
        dataset_root=dataset_root,
        layer_num=layer_num,
        iqa_metric=iqa_metric,
        swin_num=swin_num,
        device=device,
        batch_size=batch_size,
        num_workers=num_workers,
        crop_size=cfg.CROP_SIZE,
        scaling_factor=cfg.SCALING_FACTOR,
        min_distortion_level=min_distortion_level,
        max_batches=args.max_batches,
        max_memory_gb=args.max_memory_gb,
        add_patch_mask_stats=True,
        include_pristine=True,
        show_progress_bars=cfg.SHOW_PROGRESS_BARS,
        srground_include_sr_artifact=(
            args.srground_include_sr_artifact or cfg.SRGROUND_INCLUDE_SR_ARTIFACT
        ),
    )

    layer_name = str(build_info['layer_name'])
    patch_grid_shape = tuple(build_info['patch_grid_shape'])
    patches_per_image = int(build_info['patches_per_image'])
    sae_config = build_info.get('sae_config')

    logger.info('  Hooked layer: %r', layer_name)
    logger.info('  Grid=%s, patches/image=%s', patch_grid_shape, patches_per_image)
    if sae_config:
        logger.info('  SAE type: %s', sae_config.get('sae_type', 'unknown'))
        logger.info('  Lambda param: %s', sae_config.get('lambda_param', 'unknown'))
        logger.info('  Inner dim: %s', sae_config.get('inner_dim', 'unknown'))

    run_meta = {
        'dataset_type': dataset_arg,
        'iqa_metric': iqa_metric,
        'swin_num': swin_num,
        'layer': layer_num,
        'batch_size': batch_size,
        'crop_size': cfg.CROP_SIZE,
        'min_distortion_level': min_distortion_level,
        'checkpoint_path': checkpoint_path,
        'cache_path': cache_path,
    }
    run_meta_path = os.path.splitext(cache_path)[0] + '_run_meta.json'
    with open(run_meta_path, 'w', encoding='utf-8') as f:
        json.dump(run_meta, f, indent=2, ensure_ascii=True)
    logger.info('Run metadata saved -> %s', run_meta_path)

    logger.info('Done.')


if __name__ == '__main__':
    main()