analysis/config.py

"""
Конфигурация SAE-визуализаций.
"""

from __future__ import annotations

import json
import os
from dataclasses import dataclass, field
from typing import Dict, List, Mapping, Sequence, Tuple

DEFAULT_SAE_VIS_CONFIG_PATH = os.path.join(
    os.path.dirname(__file__), '..', 'default_configs', 'sae_vis_config.json'
)

SUPPORTED_DATASETS: Tuple[str, ...] = ('kadid10k', 'local_kadid', 'QGround', 'SRGround')

_DATASET_IMAGE_SUBDIRS = {dataset: dataset for dataset in SUPPORTED_DATASETS}


def dataset_images_root(datasets_root: str, dataset: str) -> str:
    """Filesystem root for a dataset's image files (under DATASETS_ROOT)."""
    subdir = _DATASET_IMAGE_SUBDIRS.get(str(dataset), str(dataset))
    return os.path.join(datasets_root, subdir)


@dataclass(frozen=True)
class DatasetCachePaths:
    dataset: str
    cache_dir: str
    dataset_cache_dir: str
    acts_cache_path: str
    corr_group_cache_path: str
    corr_type_cache_path: str
    corr_group_patch_cache_path: str
    corr_type_patch_cache_path: str
    mi_group_cache_path: str
    mi_type_cache_path: str
    mi_group_patch_cache_path: str
    mi_type_patch_cache_path: str
    auc_group_cache_path: str
    auc_type_cache_path: str
    auc_group_patch_cache_path: str
    auc_type_patch_cache_path: str
    precision_group_cache_path: str
    precision_type_cache_path: str
    precision_group_patch_cache_path: str
    precision_type_patch_cache_path: str
    recall_group_cache_path: str
    recall_type_cache_path: str
    recall_group_patch_cache_path: str
    recall_type_patch_cache_path: str
    iou_group_cache_path: str
    iou_type_cache_path: str


@dataclass
class SaeVisConfig:
    """Runtime configuration loaded from SAE vis JSON (+ optional sae_config.json)."""

    SAE_CHECKPOINT_PATH: str
    SAE_CONFIG_PATH: str
    DATASETS_ROOT: str
    IN_DIM: int
    INNER_DIM: int
    LAYER_NUM: int
    IQA_METRIC: str
    SWIN_NUM: int
    CROP_SIZE: int
    DOWNSCALE_FACTOR: int
    KADID_MIN_DISTORTION_LEVEL: int
    KADID_MAX_DISTORTION_LEVEL: int
    SRGROUND_INCLUDE_SR_ARTIFACT: bool
    SCALING_FACTOR: float
    BATCH_SIZE: int
    NUM_WORKERS: int
    DEVICE: str
    ACTIVATION_STEPS_TO_KEEP: List[int]
    CORR_TOP_K: int
    HEATMAP_CRITERION: str
    N_TOP_FEATURES_HEATMAPS: int
    N_IMAGES_PER_FEATURE: int
    HEATMAP_AGGREGATIONS: List[str]
    FEATURE_FILTERS: List[dict]
    SELECTOR_CONFIGS: List[dict]
    SCATTER_TOP_K_PATCHES: int
    SCATTER_GROUP_NAME: str
    SCATTER_BACKEND: str
    SCATTER_METRICS: List[str]
    SCATTER_ACTIVATION_THRESHOLD: float
    SCATTER_METRIC_LEVEL: str
    BUILD_CACHE_IF_MISSING: bool
    SAVE_TABULAR_ARTIFACTS: bool
    SHOW_PROGRESS_BARS: bool
    CACHE_DIR: str
    DATASET: str
    SUPPORTED_DATASETS: Tuple[str, ...]
    DATASET_CACHE_CONFIGS: Dict[str, DatasetCachePaths]
    CACHE_PATHS: DatasetCachePaths
    KADID_IMAGES_PATH: str
    config_path: str = field(repr=False, compare=False, default='')

    @property
    def DATASET_ROOT(self) -> str:
        """Filesystem root for the active ``DATASET`` (``DATASETS_ROOT/<dataset>``)."""
        return dataset_images_root(self.DATASETS_ROOT, self.DATASET)

    def as_cache_params(self) -> Dict[str, object]:
        """Dict for ensure_activation_cache / build_activation_cache call sites."""
        return {
            'SAE_CHECKPOINT_PATH': self.SAE_CHECKPOINT_PATH,
            'LAYER_NUM': self.LAYER_NUM,
            'IQA_METRIC': self.IQA_METRIC,
            'SWIN_NUM': self.SWIN_NUM,
            'DEVICE': self.DEVICE,
            'BATCH_SIZE': self.BATCH_SIZE,
            'NUM_WORKERS': self.NUM_WORKERS,
            'CROP_SIZE': self.CROP_SIZE,
            'SCALING_FACTOR': self.SCALING_FACTOR,
            'KADID_MAX_DISTORTION_LEVEL': self.KADID_MAX_DISTORTION_LEVEL,
            'SRGROUND_INCLUDE_SR_ARTIFACT': self.SRGROUND_INCLUDE_SR_ARTIFACT,
        }


def _read_json_object(path: str) -> Dict[str, object]:
    if not os.path.exists(path):
        raise FileNotFoundError(
            f'SAE visualization config not found: {path}. '
            f'Please create JSON config or set SAE_VIS_CONFIG_PATH.'
        )
    with open(path, 'r', encoding='utf-8') as f:
        payload = json.load(f)
    if not isinstance(payload, dict):
        raise ValueError(f'SAE visualization config must be a JSON object: {path}')
    return payload


def _cfg_get(cfg: Mapping[str, object], key: str, default: object = None) -> object:
    return cfg.get(key, default)


def resolve_sae_config_path(checkpoint_path: str) -> str:
    if os.path.isdir(checkpoint_path):
        return os.path.join(os.path.dirname(checkpoint_path), 'sae_config.json')
    return os.path.join(os.path.dirname(os.path.dirname(checkpoint_path)), 'sae_config.json')


def build_dataset_cache_paths(
    cache_dir: str,
    datasets: Sequence[str],
) -> Dict[str, DatasetCachePaths]:
    dataset_cache_configs: Dict[str, DatasetCachePaths] = {}
    acts_filenames = {
        'kadid10k': 'kadid_acts.feather',
        'local_kadid': 'local_kadid_acts.feather',
        'QGround': 'QGround_acts.feather',
        'SRGround': 'SRGround_acts.feather',
    }

    common_cache_filenames = {
        'corr_group': 'corr_group.parquet',
        'corr_type': 'corr_type.parquet',
        'corr_group_patch': 'corr_group_patch.parquet',
        'corr_type_patch': 'corr_type_patch.parquet',
        'mi_group': 'mi_group.parquet',
        'mi_type': 'mi_type.parquet',
        'mi_group_patch': 'mi_group_patch.parquet',
        'mi_type_patch': 'mi_type_patch.parquet',
        'auc_group': 'auc_group.parquet',
        'auc_type': 'auc_type.parquet',
        'auc_group_patch': 'auc_group_patch.parquet',
        'auc_type_patch': 'auc_type_patch.parquet',
        'precision_group': 'precision_group.parquet',
        'precision_type': 'precision_type.parquet',
        'precision_group_patch': 'precision_group_patch.parquet',
        'precision_type_patch': 'precision_type_patch.parquet',
        'recall_group': 'recall_group.parquet',
        'recall_type': 'recall_type.parquet',
        'recall_group_patch': 'recall_group_patch.parquet',
        'recall_type_patch': 'recall_type_patch.parquet',
        'iou_group': 'iou_group.parquet',
        'iou_type': 'iou_type.parquet',
    }

    for dataset in datasets:
        dataset_cache_dir = os.path.join(cache_dir, dataset)
        acts_filename = acts_filenames.get(dataset, f'{dataset}_acts.feather')

        dataset_paths: Dict[str, str] = {
            'acts': os.path.join(dataset_cache_dir, acts_filename),
        }
        for key, filename in common_cache_filenames.items():
            dataset_paths[key] = os.path.join(dataset_cache_dir, filename)

        dataset_cache_configs[dataset] = DatasetCachePaths(
            dataset=dataset,
            cache_dir=cache_dir,
            dataset_cache_dir=dataset_cache_dir,
            acts_cache_path=dataset_paths['acts'],
            corr_group_cache_path=dataset_paths['corr_group'],
            corr_type_cache_path=dataset_paths['corr_type'],
            corr_group_patch_cache_path=dataset_paths['corr_group_patch'],
            corr_type_patch_cache_path=dataset_paths['corr_type_patch'],
            mi_group_cache_path=dataset_paths['mi_group'],
            mi_type_cache_path=dataset_paths['mi_type'],
            mi_group_patch_cache_path=dataset_paths['mi_group_patch'],
            mi_type_patch_cache_path=dataset_paths['mi_type_patch'],
            auc_group_cache_path=dataset_paths['auc_group'],
            auc_type_cache_path=dataset_paths['auc_type'],
            auc_group_patch_cache_path=dataset_paths['auc_group_patch'],
            auc_type_patch_cache_path=dataset_paths['auc_type_patch'],
            precision_group_cache_path=dataset_paths['precision_group'],
            precision_type_cache_path=dataset_paths['precision_type'],
            precision_group_patch_cache_path=dataset_paths['precision_group_patch'],
            precision_type_patch_cache_path=dataset_paths['precision_type_patch'],
            recall_group_cache_path=dataset_paths['recall_group'],
            recall_type_cache_path=dataset_paths['recall_type'],
            recall_group_patch_cache_path=dataset_paths['recall_group_patch'],
            recall_type_patch_cache_path=dataset_paths['recall_type_patch'],
            iou_group_cache_path=dataset_paths['iou_group'],
            iou_type_cache_path=dataset_paths['iou_type'],
        )

    return dataset_cache_configs


def load_sae_vis_config(path: str | None = None) -> SaeVisConfig:
    """Load and validate SAE visualization config from JSON."""
    config_path = path or os.environ.get('SAE_VIS_CONFIG_PATH', DEFAULT_SAE_VIS_CONFIG_PATH)
    # config_path = "/home/28m_gov@lab.graphicon.ru/SAE/xIQA/logs/arniqa_logs/ARNIQA_layer5_lambda5e-4_scale1_exp37/vis/vis_srground/config.json"
    vis_cfg = _read_json_object(config_path)

    sae_checkpoint_path = str(_cfg_get(vis_cfg, 'SAE_CHECKPOINT_PATH', '')).strip()
    if not sae_checkpoint_path:
        raise ValueError('SAE_CHECKPOINT_PATH must be set in SAE vis config JSON')

    sae_config_path = resolve_sae_config_path(sae_checkpoint_path)
    sae_cfg: Dict[str, object] = _read_json_object(sae_config_path) if os.path.exists(sae_config_path) else {}

    datasets_root = str(_cfg_get(vis_cfg, 'DATASETS_ROOT', '')).strip()
    if not datasets_root:
        legacy_kadid_images_path = str(_cfg_get(vis_cfg, 'KADID_IMAGES_PATH', '')).strip()
        if legacy_kadid_images_path:
            datasets_root = os.path.dirname(legacy_kadid_images_path)
    if not datasets_root:
        raise ValueError('DATASETS_ROOT must be set in SAE vis config JSON')

    kadid_min_distortion_level = int(_cfg_get(vis_cfg, 'KADID_MIN_DISTORTION_LEVEL', 1))
    kadid_max_distortion_level = int(_cfg_get(vis_cfg, 'KADID_MAX_DISTORTION_LEVEL', 5))
    if not (1 <= kadid_min_distortion_level <= 5):
        raise ValueError('KADID_MIN_DISTORTION_LEVEL must be in [1, 5]')
    if not (1 <= kadid_max_distortion_level <= 5):
        raise ValueError('KADID_MAX_DISTORTION_LEVEL must be in [1, 5]')
    if kadid_min_distortion_level > kadid_max_distortion_level:
        raise ValueError('KADID_MIN_DISTORTION_LEVEL must be <= KADID_MAX_DISTORTION_LEVEL')

    activation_steps_to_keep = [int(step) for step in _cfg_get(vis_cfg, 'ACTIVATION_STEPS_TO_KEEP', [])]
    if any(step <= 0 for step in activation_steps_to_keep):
        raise ValueError('ACTIVATION_STEPS_TO_KEEP must contain only positive integers')

    dataset = str(_cfg_get(vis_cfg, 'DATASET', 'kadid10k')).strip()
    if dataset not in SUPPORTED_DATASETS:
        raise ValueError(f'Unsupported DATASET={dataset!r}; expected one of {SUPPORTED_DATASETS}')

    cache_dir = str(
        _cfg_get(
            vis_cfg,
            'CACHE_DIR',
            os.path.join(os.path.dirname(os.path.dirname(sae_checkpoint_path)), 'cache/'),
        )
    )
    dataset_cache_configs = build_dataset_cache_paths(cache_dir, SUPPORTED_DATASETS)

    raw_selector_configs = _cfg_get(vis_cfg, 'SELECTOR_CONFIGS', [])
    if raw_selector_configs is None:
        raw_selector_configs = []

    feature_filters = _cfg_get(
        vis_cfg,
        'FEATURE_FILTERS',
        [
            {'name': 'nonzero_max', 'params': {}},
            {
                'name': 'kruskal_wallis',
                'params': {'alpha': 0.05, 'group_col': 'dist_type', 'min_group_size': 3},
            },
        ],
    )

    dataset_images_subdir = _DATASET_IMAGE_SUBDIRS[dataset]

    return SaeVisConfig(
        SAE_CHECKPOINT_PATH=sae_checkpoint_path,
        SAE_CONFIG_PATH=sae_config_path,
        DATASETS_ROOT=datasets_root,
        IN_DIM=int(_cfg_get(sae_cfg, 'sae_input_dim', 64)),
        INNER_DIM=int(_cfg_get(sae_cfg, 'inner_dim', 6400)),
        LAYER_NUM=int(_cfg_get(sae_cfg, 'layer_num', 3)),
        IQA_METRIC=str(_cfg_get(sae_cfg, 'iqa_metric', 'arniqa-kadid')),
        SWIN_NUM=int(_cfg_get(sae_cfg, 'swin_num', 2)),
        CROP_SIZE=int(_cfg_get(vis_cfg, 'CROP_SIZE', 224)),
        DOWNSCALE_FACTOR=int(_cfg_get(vis_cfg, 'DOWNSCALE_FACTOR', 2)),
        KADID_MIN_DISTORTION_LEVEL=kadid_min_distortion_level,
        KADID_MAX_DISTORTION_LEVEL=kadid_max_distortion_level,
        SRGROUND_INCLUDE_SR_ARTIFACT=bool(_cfg_get(vis_cfg, 'SRGROUND_INCLUDE_SR_ARTIFACT', False)),
        SCALING_FACTOR=float(_cfg_get(sae_cfg, 'scaling_factor', 1.0)),
        BATCH_SIZE=int(_cfg_get(vis_cfg, 'BATCH_SIZE', 32)),
        NUM_WORKERS=int(_cfg_get(vis_cfg, 'NUM_WORKERS', 4)),
        DEVICE=str(_cfg_get(vis_cfg, 'DEVICE', 'cuda')),
        ACTIVATION_STEPS_TO_KEEP=activation_steps_to_keep,
        CORR_TOP_K=int(_cfg_get(vis_cfg, 'CORR_TOP_K', 30)),
        HEATMAP_CRITERION=str(_cfg_get(vis_cfg, 'HEATMAP_CRITERION', 'max')),
        N_TOP_FEATURES_HEATMAPS=int(_cfg_get(vis_cfg, 'N_TOP_FEATURES_HEATMAPS', 3)),
        N_IMAGES_PER_FEATURE=int(_cfg_get(vis_cfg, 'N_IMAGES_PER_FEATURE', 5)),
        HEATMAP_AGGREGATIONS=[str(value) for value in _cfg_get(vis_cfg, 'HEATMAP_AGGREGATIONS', ['max', 'mean_acts', 'sum'])],
        FEATURE_FILTERS=list(feature_filters),
        SELECTOR_CONFIGS=list(raw_selector_configs),
        SCATTER_TOP_K_PATCHES=int(_cfg_get(vis_cfg, 'SCATTER_TOP_K_PATCHES', 1000)),
        SCATTER_GROUP_NAME=str(_cfg_get(vis_cfg, 'SCATTER_GROUP_NAME', 'blur')),
        SCATTER_BACKEND=str(_cfg_get(vis_cfg, 'SCATTER_BACKEND', 'matplotlib')),
        SCATTER_METRICS=[str(value) for value in _cfg_get(vis_cfg, 'SCATTER_METRICS', ['entropy', 'iou', 'roc_auc', 'precision', 'recall'])],
        SCATTER_ACTIVATION_THRESHOLD=float(_cfg_get(vis_cfg, 'SCATTER_ACTIVATION_THRESHOLD', 0.0)),
        SCATTER_METRIC_LEVEL=str(_cfg_get(vis_cfg, 'SCATTER_METRIC_LEVEL', 'patch')),
        BUILD_CACHE_IF_MISSING=bool(_cfg_get(vis_cfg, 'BUILD_CACHE_IF_MISSING', True)),
        SAVE_TABULAR_ARTIFACTS=bool(_cfg_get(vis_cfg, 'SAVE_TABULAR_ARTIFACTS', False)),
        SHOW_PROGRESS_BARS=bool(_cfg_get(vis_cfg, 'SHOW_PROGRESS_BARS', True)),
        CACHE_DIR=cache_dir,
        DATASET=dataset,
        SUPPORTED_DATASETS=SUPPORTED_DATASETS,
        DATASET_CACHE_CONFIGS=dataset_cache_configs,
        CACHE_PATHS=dataset_cache_configs[dataset],
        KADID_IMAGES_PATH=os.path.join(datasets_root, dataset_images_subdir),
        config_path=config_path,
    )