style(dashboard): translate commentaries

.gitignore

@@ -1,2 +1,4 @@
 data/
-tests/
+tests/
+__pycache__/
+.venv/

analysis/cache_utils.py

@@ -1,5 +1,5 @@
 """
-Утилиты для кэширования и загрузки активаций SAE.
+Utilities for caching and loading SAE activations.
 """
 
 from pathlib import Path
@@ -17,20 +17,20 @@ from analysis.models import _iqa_activations
 
 
 def _df_mb(df: pd.DataFrame) -> float:
-    """Объём памяти DataFrame в МБ."""
+    """DataFrame memory usage in MB."""
     return df.memory_usage(deep=True).sum() / 1024 ** 2
 
 
 def _sparse_mb(mat: sp.csr_matrix) -> float:
-    """Объём памяти CSR-матрицы (данные + индексы) в МБ."""
+    """CSR matrix memory usage (data + indices) in MB."""
     return (mat.data.nbytes + mat.indices.nbytes + mat.indptr.nbytes) / 1024 ** 2
 
 
 def _cache_paths(base_path: str) -> Tuple[str, str, str]:
     """
-    Возвращает пути к трём файлам кэша из базового пути.
+    Return paths to the three cache files derived from a base path.
 
-    Пример:
+    Example:
         'cache/kadid_acts.feather'
           → ('cache/kadid_acts_meta.feather', 'cache/kadid_acts_codes.npz',
               'cache/kadid_acts_steps.npz')
@@ -273,8 +273,8 @@ def collect_and_cache(
     sparse_mb = _sparse_mb(codes_csr)
     steps_mb = _sparse_mb(steps_csr)
 
-    print(f'Activations: shape={codes_csr.shape}, {sparse_mb:.1f} МБ (sparse)')
-    print(f'Activation steps: shape={steps_csr.shape}, {steps_mb:.1f} МБ (sparse)')
+    print(f'Activations: shape={codes_csr.shape}, {sparse_mb:.1f} MB (sparse)')
+    print(f'Activation steps: shape={steps_csr.shape}, {steps_mb:.1f} MB (sparse)')
 
     meta_path, codes_path, steps_path = _cache_paths(output_path)
 
@@ -285,9 +285,9 @@ def collect_and_cache(
     print(f'Saved metadata ({len(meta_df)} rows) -> {meta_path}')
     print(f'Saved activations -> {codes_path}')
     print(f'Saved activation steps -> {steps_path}')
-    print(f'  Metadata:    {_df_mb(meta_df):.1f} МБ')
-    print(f'  Activations: {sparse_mb:.1f} МБ')
-    print(f'  Steps:       {steps_mb:.1f} МБ')
+    print(f'  Metadata:    {_df_mb(meta_df):.1f} MB')
+    print(f'  Activations: {sparse_mb:.1f} MB')
+    print(f'  Steps:       {steps_mb:.1f} MB')
 
     if pristine_dataloader is not None:
         print('\nProcessing pristine dataset...')
@@ -318,8 +318,8 @@ def collect_and_cache(
         sp.save_npz(pristine_codes_path, pristine_codes)
         sp.save_npz(pristine_steps_path, pristine_steps)
 
-        print(f'Pristine activations: shape={pristine_codes.shape}, {pristine_sparse_mb:.1f} МБ')
-        print(f'Pristine steps: shape={pristine_steps.shape}, {pristine_steps_mb:.1f} МБ')
+        print(f'Pristine activations: shape={pristine_codes.shape}, {pristine_sparse_mb:.1f} MB')
+        print(f'Pristine steps: shape={pristine_steps.shape}, {pristine_steps_mb:.1f} MB')
 
         print(f'Saved pristine metadata ({len(pristine_meta)} rows) -> {pristine_meta_path}')
         print(f'Saved pristine activations -> {pristine_codes_path}')
@@ -333,7 +333,7 @@ def build_activation_cache(
     dataset: str,
     cache_path: str,
     checkpoint_path: str,
-    kadid_path: str,
+    dataset_root: str,
     layer_num: int,
     iqa_metric: str,
     swin_num: int,
@@ -348,6 +348,7 @@ def build_activation_cache(
     add_patch_mask_stats: bool = True,
     include_pristine: bool = True,
     show_progress_bars: bool = True,
+    srground_include_sr_artifact: bool = False,
 ) -> Dict[str, Any]:
     """Build activation cache end-to-end for KADID/local-KADID datasets."""
     from .datasets import (
@@ -377,7 +378,7 @@ def build_activation_cache(
     label_to_dist_group = None
 
     if dataset == 'local_kadid':
-        data = LocalKadidPresavedDataset(root=kadid_path, crop_size=crop_size)
+        data = LocalKadidPresavedDataset(root=dataset_root, crop_size=crop_size)
         collate_fn = local_kadid_collate_fn
         meta_keys = [
             'dist_type',
@@ -392,11 +393,11 @@ def build_activation_cache(
             'distorted_img_path',
             'original_img_path',
         ]
-        pristine_data = LocalKadidPristineDataset(root=kadid_path, crop_size=crop_size) if include_pristine else None
+        pristine_data = LocalKadidPristineDataset(root=dataset_root, crop_size=crop_size) if include_pristine else None
         pristine_collate = local_kadid_pristine_collate_fn
     elif dataset in {'kadid10k', 'kadid'}:
         data = Kadid10kDataset(
-            root=kadid_path,
+            root=dataset_root,
             crop_size=crop_size,
             min_distortion_level=min_distortion_level or 1,
         )
@@ -409,11 +410,11 @@ def build_activation_cache(
             'distorted_img_path',
             'original_img_path',
         ]
-        pristine_data = KadidPristineDataset(root=kadid_path, crop_size=crop_size) if include_pristine else None
+        pristine_data = KadidPristineDataset(root=dataset_root, crop_size=crop_size) if include_pristine else None
         pristine_collate = kadid_pristine_collate_fn
     elif dataset == 'QGround':
         data = QGroundDataset(
-            root=kadid_path,
+            root=dataset_root,
             split='test',
             crop_size=crop_size,
         )
@@ -438,9 +439,10 @@ def build_activation_cache(
         pristine_collate = None
     elif dataset == 'SRGround':
         data = SRGroundSmallDataset(
-            root=kadid_path,
+            root=dataset_root,
             json_path='srground_train.json',
             crop_size=crop_size,
+            include_sr_artifact=srground_include_sr_artifact,
         )
         collate_fn = srground_collate_fn
         meta_keys = [
@@ -538,11 +540,10 @@ def load_cache(
     Tuple[pd.DataFrame, sp.csr_matrix],
     Tuple[pd.DataFrame, sp.csr_matrix, sp.csr_matrix],
 ]:
-    """Загружает кэш активаций SAE из раздельных файлов.
+    """Load SAE activation cache from separate files.
 
-    Если ``return_activation_steps=True``, дополнительно возвращает CSR-матрицу
-    порядка активаций, где значение ``0`` означает отсутствие активации,
-    а ``k>0`` соответствует шагу ``k`` в pursuit.
+    If ``return_activation_steps=True``, also returns a CSR matrix of activation
+    order, where ``0`` means no activation and ``k>0`` corresponds to pursuit step ``k``.
     """
     meta_path, codes_path, steps_path = _cache_paths(path)
     meta = pd.read_feather(meta_path)
@@ -550,8 +551,8 @@ def load_cache(
 
     print(f'Loaded from {meta_path}: {meta.shape[0]} rows × {meta.shape[1]} cols')
     print(f'Loaded from {codes_path}: shape={codes.shape}, dtype={codes.dtype}')
-    print(f'  Metadata:    {_df_mb(meta):.1f} МБ')
-    print(f'  Activations: {_sparse_mb(codes):.1f} МБ (sparse)')
+    print(f'  Metadata:    {_df_mb(meta):.1f} MB')
+    print(f'  Activations: {_sparse_mb(codes):.1f} MB (sparse)')
 
     keep_idx: Optional[np.ndarray] = None
     if min_distortion_level is not None or max_distortion_level is not None:
@@ -569,7 +570,7 @@ def load_cache(
             keep_idx = np.flatnonzero(keep_mask.to_numpy())
         else: 
             keep_mask = (meta['dist_level'] >= -1000) 
-            keep_idx = np.flatnonzero(keep_mask.to_numpy()) # Костыль -- поправить
+            keep_idx = np.flatnonzero(keep_mask.to_numpy())  # Temporary workaround -- fix later
  
     if return_activation_steps:
         if Path(steps_path).exists():
@@ -592,7 +593,7 @@ def load_cache(
                 f'{meta.shape[0]} rows kept'
             )
 
-        print(f'  Steps:       {_sparse_mb(steps):.1f} МБ (sparse)')
+        print(f'  Steps:       {_sparse_mb(steps):.1f} MB (sparse)')
         return meta, codes, steps
 
     if keep_idx is not None:
@@ -620,8 +621,8 @@ def load_pristine_cache(
 
     print(f'Loaded pristine from {meta_path}: {meta.shape[0]} rows × {meta.shape[1]} cols')
     print(f'Loaded pristine from {codes_path}: shape={codes.shape}, dtype={codes.dtype}')
-    print(f'  Metadata:    {_df_mb(meta):.1f} МБ')
-    print(f'  Activations: {_sparse_mb(codes):.1f} МБ (sparse)')
+    print(f'  Metadata:    {_df_mb(meta):.1f} MB')
+    print(f'  Activations: {_sparse_mb(codes):.1f} MB (sparse)')
 
     if return_activation_steps:
         if Path(steps_path).exists():
@@ -635,7 +636,7 @@ def load_pristine_cache(
             print('No pristine steps cache found. Using all-zero activation steps.')
             steps = sp.csr_matrix(codes.shape, dtype=np.int32)
 
-        print(f'  Steps:       {_sparse_mb(steps):.1f} МБ (sparse)')
+        print(f'  Steps:       {_sparse_mb(steps):.1f} MB (sparse)')
         return meta, codes, steps
 
     return meta, codes
@@ -649,12 +650,12 @@ def ensure_cache_ready(
     load_cache_kwargs: Optional[Dict[str, Any]] = None,
     build_cache_fn: Optional[Callable[[], None]] = None,
 ) -> None:
-    """Проверяет доступность кэша и при необходимости собирает его.
+    """Check cache availability and rebuild it if needed.
 
-    Поведение:
-    - пытается загрузить кэш через ``load_cache``;
-    - если кэш отсутствует или выставлен ``force_recache=True``, запускает сборку;
-    - если сборка отключена, пробрасывает ``FileNotFoundError``.
+    Behavior:
+    - try loading the cache via ``load_cache``;
+    - if the cache is missing or ``force_recache=True``, start a rebuild;
+    - if building is disabled, raise ``FileNotFoundError``.
     """
     needs_rebuild = force_recache
     if not needs_rebuild:
@@ -688,18 +689,18 @@ def zero_codes_outside_activation_steps(
     activation_steps_csr: sp.csr_matrix,
     activation_steps_to_keep: List[int],
 ) -> sp.csr_matrix:
-    """Обнуляет активации, шаг появления которых не входит в allow-list.
+    """Zero out activations whose appearance step is not in the allow-list.
 
-    Параметры
+    Parameters
     ----------
-    codes_csr : CSR-матрица активаций SAE.
-    activation_steps_csr : CSR-матрица шагов активаций (0 = не активирован).
-    activation_steps_to_keep : список шагов, которые нужно сохранить.
-
-    Возвращает
-    ----------
-    CSR-матрицу той же формы, где вне указанных шагов значения занулены.
-    Если список шагов пуст, возвращается исходная матрица без изменений.
+    codes_csr : CSR matrix of SAE activations.
+    activation_steps_csr : CSR matrix of activation steps (0 = not activated).
+    activation_steps_to_keep : list of steps to keep.
+
+    Returns
+    -------
+    CSR matrix of the same shape with values outside the specified steps zeroed.
+    If the step list is empty, returns the original matrix unchanged.
     """
     if not activation_steps_to_keep:
         return codes_csr
@@ -739,7 +740,7 @@ def zero_codes_outside_activation_steps(
 def ensure_activation_cache(
     dataset: str,
     acts_cache_path: str,
-    kadid_images_path: str,
+    dataset_root: str,
     min_distortion_level: int,
     params: dict,
     include_pristine_cache: Optional[bool] = None,
@@ -769,7 +770,7 @@ def ensure_activation_cache(
         dataset=dataset,
         cache_path=acts_cache_path,
         checkpoint_path=params.get('SAE_CHECKPOINT_PATH'),
-        kadid_path=kadid_images_path,
+        dataset_root=dataset_root,
         layer_num=params.get('LAYER_NUM'),
         iqa_metric=params.get('IQA_METRIC'),
         swin_num=params.get('SWIN_NUM'),
@@ -783,5 +784,6 @@ def ensure_activation_cache(
         max_memory_gb=30.0,
         add_patch_mask_stats=True,
         include_pristine=needs_pristine_cache,
+        srground_include_sr_artifact=bool(params.get('SRGROUND_INCLUDE_SR_ARTIFACT', False)),
     )
     print('[run] Activation cache build completed.')

analysis/config.py

@@ -1,5 +1,5 @@
 """
-Конфигурация SAE-визуализаций.
+SAE visualization configuration.
 """
 
 from __future__ import annotations
@@ -12,14 +12,10 @@ 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 = {
-    'kadid10k': 'kadid10k',
-    'local_kadid': 'local_kadid',
-    'QGround': 'QGround',
-    'SRGround': 'SRGround',
-}
+_DATASET_IMAGE_SUBDIRS = {dataset: dataset for dataset in SUPPORTED_DATASETS}
 
 
 def dataset_images_root(datasets_root: str, dataset: str) -> str:
@@ -74,6 +70,7 @@ class SaeVisConfig:
     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
@@ -116,6 +113,7 @@ class SaeVisConfig:
             '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,
         }
 
 
@@ -256,10 +254,8 @@ def load_sae_vis_config(path: str | None = None) -> SaeVisConfig:
         raise ValueError('ACTIVATION_STEPS_TO_KEEP must contain only positive integers')
 
     dataset = str(_cfg_get(vis_cfg, 'DATASET', 'kadid10k')).strip()
-    supported_raw = _cfg_get(vis_cfg, 'SUPPORTED_DATASETS', list(SUPPORTED_DATASETS))
-    supported_datasets = tuple(str(value) for value in supported_raw)
-    if dataset not in supported_datasets:
-        raise ValueError(f'Unsupported DATASET={dataset!r}; expected one of {supported_datasets}')
+    if dataset not in SUPPORTED_DATASETS:
+        raise ValueError(f'Unsupported DATASET={dataset!r}; expected one of {SUPPORTED_DATASETS}')
 
     cache_dir = str(
         _cfg_get(
@@ -268,7 +264,7 @@ def load_sae_vis_config(path: str | None = None) -> SaeVisConfig:
             os.path.join(os.path.dirname(os.path.dirname(sae_checkpoint_path)), 'cache/'),
         )
     )
-    dataset_cache_configs = build_dataset_cache_paths(cache_dir, supported_datasets)
+    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:
@@ -301,6 +297,7 @@ def load_sae_vis_config(path: str | None = None) -> SaeVisConfig:
         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)),
@@ -324,7 +321,7 @@ def load_sae_vis_config(path: str | None = None) -> SaeVisConfig:
         SHOW_PROGRESS_BARS=bool(_cfg_get(vis_cfg, 'SHOW_PROGRESS_BARS', True)),
         CACHE_DIR=cache_dir,
         DATASET=dataset,
-        SUPPORTED_DATASETS=supported_datasets,
+        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),

analysis/datasets.py

@@ -1,5 +1,5 @@
 """
-Датасеты и вспомогательные структуры данных для KADID-10k. Взяты из репозитория PatchSAE.
+Datasets and helper data structures for KADID-10k. Taken from the PatchSAE repository.
 """
 
 import gzip
@@ -186,7 +186,20 @@ def _sr_artifact_labels(
     return np.where(sr_maps * sr_prom >= threshold, 6, 0).astype(np.uint8)
 
 
-def merge_srground_masks(annot_rd: np.ndarray, annot_sr: np.ndarray) -> np.ndarray:
+def merge_srground_masks(
+    annot_rd: np.ndarray | None,
+    annot_sr: np.ndarray | None,
+) -> np.ndarray:
+    """Merge real-distortion and SR-artifact label maps (0=background, 1..6=classes).
+
+    If ``annot_sr`` is None, returns ``annot_rd`` unchanged (no SR merge).
+    """
+    if annot_sr is None:
+        if annot_rd is None:
+            raise ValueError('merge_srground_masks requires annot_rd when annot_sr is None')
+        return annot_rd.astype(np.uint8, copy=False)
+    if annot_rd is None:
+        return annot_sr.astype(np.uint8, copy=False)
     return np.where(annot_sr == 6, 6, annot_rd).astype(np.uint8)
 
 
@@ -319,55 +332,114 @@ def srground_rows_for_image_paths(
 def srground_prominences_by_image_paths(
     image_paths: Sequence[str],
     *,
+    dataset_root: str | Path | None = None,
     datasets_root: str | None = None,
 ) -> dict[str, np.ndarray]:
     """Look up prominences for paths (absolute or SRGround-relative) via cached index."""
     if not image_paths:
         return {}
 
+    if datasets_root is None and dataset_root is not None:
+        datasets_root = parent_datasets_root(dataset_root)
     root = _resolve_datasets_root(datasets_root)
     index = srground_prominences_index(root)
     keys = {srground_image_key(path, datasets_root=root) for path in image_paths if path}
     return {key: index[key] for key in keys if key in index}
 
 
-def srground_mask_rgb_for_meta_row(
+def _meta_row_path(meta_row: pd.Series, column: str) -> str | None:
+    if column not in meta_row:
+        return None
+    value = meta_row.get(column)
+    if value is None or (isinstance(value, float) and np.isnan(value)):
+        return None
+    text = str(value).strip()
+    return text or None
+
+
+def parent_datasets_root(dataset_root: str | Path) -> str:
+    """Parent directory that contains dataset folders (e.g. ``Kadid`` for ``Kadid/SRGround``)."""
+    return str(Path(dataset_root).resolve().parent)
+
+
+def _resolve_dataset_file_path(rel_path: str, *, dataset_root: str | Path) -> Path:
+    path = Path(str(rel_path).strip())
+    if path.is_absolute():
+        return path
+    return Path(dataset_root) / path
+
+
+def infer_spatial_mask_dataset(meta_row: pd.Series | None) -> str | None:
+    """Return ``QGround`` or ``SRGround`` when meta row carries spatial mask fields."""
+    if meta_row is None:
+        return None
+    if _meta_row_path(meta_row, 'real_distortions_ann_path') or _meta_row_path(meta_row, 'sr_artifacts_ann_path'):
+        return 'SRGround'
+    if _meta_row_path(meta_row, 'mask_path'):
+        return 'QGround'
+    return None
+
+
+def qground_mask_rgb_for_meta_row(
     meta_row: pd.Series | None,
     *,
-    datasets_root: str | None = None,
-    include_sr_artifact: bool = True,
-    crop_size: int = 224,
-    sr_artifact_threshold: float = SRGROUND_SR_ARTIFACT_THRESHOLD,
-    prominences: np.ndarray | None = None,
+    dataset_root: str | Path,
 ) -> np.ndarray | None:
-    """Build an RGB annotation mask for one SRGround image (matches dashboard center crop)."""
+    """Load QGround RGB segmentation mask PNG referenced from cache metadata."""
     if meta_row is None:
         return None
 
-    root = _resolve_datasets_root(datasets_root)
+    mask_rel = _meta_row_path(meta_row, 'mask_path')
+    if not mask_rel:
+        return None
 
-    def _row_path(column: str) -> str | None:
-        if column not in meta_row:
-            return None
-        value = meta_row.get(column)
-        if value is None or (isinstance(value, float) and np.isnan(value)):
-            return None
-        text = str(value).strip()
-        return text or None
+    mask_path = _resolve_dataset_file_path(mask_rel, dataset_root=dataset_root)
+    if not mask_path.is_file():
+        return None
+    return np.asarray(Image.open(mask_path).convert('RGB'), dtype=np.uint8)
+
+
+def _resize_label_map_to_image(
+    mask_label: np.ndarray | None,
+    reference_size: tuple[int, int],
+) -> np.ndarray | None:
+    if mask_label is None:
+        return None
+    width, height = reference_size
+    if mask_label.shape[:2] == (height, width):
+        return mask_label.astype(np.uint8, copy=False)
+    mask_image = Image.fromarray(mask_label.astype(np.uint8), mode='L')
+    mask_image = mask_image.resize((width, height), resample=Image.NEAREST)
+    return np.asarray(mask_image, dtype=np.uint8)
+
+
+def srground_label_map_for_meta_row(
+    meta_row: pd.Series | None,
+    *,
+    dataset_root: str | Path,
+    sr_artifact_threshold: float = SRGROUND_SR_ARTIFACT_THRESHOLD,
+    prominences: np.ndarray | None = None,
+) -> np.ndarray | None:
+    """Build a full-frame SRGround label map from paths stored in cache metadata."""
+    if meta_row is None:
+        return None
 
     image_rel = _image_rel_from_meta_row(meta_row)
-    real_rel = _row_path('real_distortions_ann_path')
-    sr_rel = _row_path('sr_artifacts_ann_path')
-    if not real_rel and not (include_sr_artifact and sr_rel):
+    real_rel = _meta_row_path(meta_row, 'real_distortions_ann_path')
+    sr_rel = _meta_row_path(meta_row, 'sr_artifacts_ann_path')
+    if not real_rel and not sr_rel:
         return None
 
     if prominences is None and image_rel:
-        prominences = srground_prominences_by_image_paths([image_rel], datasets_root=root).get(image_rel)
+        prominences = srground_prominences_by_image_paths(
+            [image_rel],
+            dataset_root=dataset_root,
+        ).get(image_rel)
 
     annot_rd = None
     annot_sr = None
     if real_rel:
-        real_path = resolve_dataset_image_path('SRGround', real_rel, datasets_root=root)
+        real_path = _resolve_dataset_file_path(real_rel, dataset_root=dataset_root)
         if real_path.is_file():
             real_maps = _load_npy_gz(real_path)
             prom = prominences
@@ -375,8 +447,8 @@ def srground_mask_rgb_for_meta_row(
                 prom = np.ones(int(real_maps.shape[0]) + 1, dtype=np.float64)
             annot_rd = _real_distortion_labels(real_maps, prom)
 
-    if include_sr_artifact and sr_rel:
-        sr_path = resolve_dataset_image_path('SRGround', sr_rel, datasets_root=root)
+    if sr_rel:
+        sr_path = _resolve_dataset_file_path(sr_rel, dataset_root=dataset_root)
         if sr_path.is_file():
             sr_maps = _load_npy_gz(sr_path)
             prom = prominences
@@ -391,25 +463,97 @@ def srground_mask_rgb_for_meta_row(
     if annot_rd is None and annot_sr is None:
         return None
 
-    if annot_rd is None:
-        annot_rd = np.zeros_like(annot_sr, dtype=np.uint8)
-    if include_sr_artifact and annot_sr is not None:
-        merged = merge_srground_masks(annot_rd, annot_sr)
-    else:
-        merged = annot_rd.astype(np.uint8, copy=False)
-
+    reference_size = (annot_rd if annot_rd is not None else annot_sr).shape[1], (
+        annot_rd if annot_rd is not None else annot_sr
+    ).shape[0]
     if image_rel:
-        image_path = resolve_dataset_image_path('SRGround', image_rel, datasets_root=root)
+        image_path = _resolve_dataset_file_path(image_rel, dataset_root=dataset_root)
         if image_path.is_file():
             with Image.open(image_path) as image:
                 reference_size = image.size
-            merged = _center_crop_label_map(merged, crop_size, reference_size)
-        else:
-            merged = _center_crop_label_map(merged, crop_size, (merged.shape[1], merged.shape[0]))
-    else:
-        merged = _center_crop_label_map(merged, crop_size, (merged.shape[1], merged.shape[0]))
 
-    return label2rgb_srground(merged)
+    annot_rd = _resize_label_map_to_image(annot_rd, reference_size)
+    annot_sr = _resize_label_map_to_image(annot_sr, reference_size)
+    return merge_srground_masks(annot_rd, annot_sr)
+
+
+def annotation_mask_rgb_for_meta_row(
+    meta_row: pd.Series | None,
+    *,
+    dataset_root: str | Path,
+    dataset: str | None = None,
+    crop_size: int = 512,
+    full_frame: bool = True,
+    sr_artifact_threshold: float = SRGROUND_SR_ARTIFACT_THRESHOLD,
+    prominences: np.ndarray | None = None,
+) -> np.ndarray | None:
+    """RGB annotation mask for QGround / SRGround visualization (QGround-style full frame by default)."""
+    
+    if dataset_root is None:
+        from analysis.config import load_sae_vis_config
+        cfg = load_sae_vis_config()
+        datasets_root = str(cfg.DATASETS_ROOT)
+        dataset_root = Path(datasets_root) / dataset
+
+    dataset_name = dataset or infer_spatial_mask_dataset(meta_row)
+    if dataset_name == 'QGround':
+        return qground_mask_rgb_for_meta_row(meta_row, dataset_root=dataset_root)
+    if dataset_name == 'SRGround':
+        label_map = srground_label_map_for_meta_row(
+            meta_row,
+            dataset_root=dataset_root,
+            sr_artifact_threshold=sr_artifact_threshold,
+            prominences=prominences,
+        )
+        if label_map is None:
+            return None
+        if not full_frame:
+            image_rel = _image_rel_from_meta_row(meta_row)
+            reference_size = (label_map.shape[1], label_map.shape[0])
+            if image_rel:
+                image_path = _resolve_dataset_file_path(image_rel, dataset_root=dataset_root)
+                if image_path.is_file():
+                    with Image.open(image_path) as image:
+                        reference_size = image.size
+            label_map = _center_crop_label_map(label_map, crop_size, reference_size)
+        return label2rgb_srground(label_map)
+    return None
+
+
+def get_srground_rgb_mask(
+    meta_row: pd.Series | None,
+    *,
+    dataset_root: str | Path,
+    crop_size: int = 224,
+    full_frame: bool = True,
+    sr_artifact_threshold: float = SRGROUND_SR_ARTIFACT_THRESHOLD,
+) -> np.ndarray | None:
+    """RGB SRGround mask for one cache meta row (prominences from ``srground_train.json``)."""
+    return annotation_mask_rgb_for_meta_row(
+        meta_row,
+        dataset_root=dataset_root,
+        dataset='SRGround',
+        crop_size=crop_size,
+        full_frame=full_frame,
+        sr_artifact_threshold=sr_artifact_threshold,
+    )
+
+
+def srground_mask_rgb_for_meta_row(
+    meta_row: pd.Series | None,
+    *,
+    dataset_root: str | Path,
+    crop_size: int = 224,
+    sr_artifact_threshold: float = SRGROUND_SR_ARTIFACT_THRESHOLD,
+) -> np.ndarray | None:
+    """Center-cropped SRGround mask preview (alias of :func:`get_srground_rgb_mask`)."""
+    return get_srground_rgb_mask(
+        meta_row,
+        dataset_root=dataset_root,
+        crop_size=crop_size,
+        full_frame=False,
+        sr_artifact_threshold=sr_artifact_threshold,
+    )
 
 
 def _rgb2label_qground(mask_rgb: np.ndarray) -> np.ndarray:
@@ -443,7 +587,7 @@ def _to_relative_dataset_path(path: Path, root: Path) -> str:
 
 class Kadid10kDataset(Dataset):
     """
-    KADID-10k dataset. При семплинге применяется RandomCrop.
+    KADID-10k dataset. RandomCrop is applied during sampling.
     """
 
     def __init__(
@@ -595,11 +739,11 @@ class LocalKadidPresavedDataset(Dataset):
 
 def kadid_collate_fn(batch: List[dict]) -> dict:
     """
-    Collate для Kadid10kDataset.
+    Collate for Kadid10kDataset.
 
-    Возвращает:
+    Returns:
         images: Tensor (B, C, H, W)
-        + все остальные ключи как списки длины B
+        + all other keys as lists of length B
     """
     images = torch.stack([item["img"] for item in batch], dim=0)
     collated: dict = {"images": images}
@@ -799,10 +943,12 @@ class SRGroundSmallDataset(Dataset):
         allowed_methods: Optional[List[str]] = ['DiT4SR_x2'],
         crop_size: Optional[int] = None,
         sr_artifact_threshold: float = SRGROUND_SR_ARTIFACT_THRESHOLD,
+        include_sr_artifact: bool = False,
         transform=None,
     ):
         self.root = Path(root)
         self.sr_artifact_threshold = float(sr_artifact_threshold)
+        self.include_sr_artifact = bool(include_sr_artifact)
         self.json_path = self.root / 'srground_train.json'
         df = pd.read_json(self.json_path)
 
@@ -845,10 +991,15 @@ class SRGroundSmallDataset(Dataset):
         if real_ann_path.exists():
             annot_rd = _real_distortion_labels(_load_npy_gz(real_ann_path), prominences)
 
-        sr_path = sample.get('sr_artifacts_ann_path')
-        sr_ann_path = self._resolve_path(sr_path)
-        if sr_ann_path.exists():
-            annot_sr = _sr_artifact_labels(_load_npy_gz(sr_ann_path), prominences, threshold=self.sr_artifact_threshold)
+        if self.include_sr_artifact:
+            sr_path = sample.get('sr_artifacts_ann_path')
+            sr_ann_path = self._resolve_path(sr_path)
+            if sr_ann_path.exists():
+                annot_sr = _sr_artifact_labels(
+                    _load_npy_gz(sr_ann_path),
+                    prominences,
+                    threshold=self.sr_artifact_threshold,
+                )
 
         if annot_rd is None and annot_sr is None:
             return None, None
@@ -883,15 +1034,31 @@ class SRGroundSmallDataset(Dataset):
                 annot_sr = np.asarray(self.crop(mask_image), dtype=np.uint8)
 
         img_tensor = self.image_to_tensor(image)
+        mask_hw = (image.height, image.width)
 
-        mask_rd = torch.from_numpy(annot_rd.astype(np.float32)).unsqueeze(0)
-        mask_sr = torch.from_numpy(annot_sr.astype(np.float32)).unsqueeze(0)
-        
-        mask = torch.from_numpy(merge_srground_masks(annot_rd, annot_sr).astype(np.float32)).unsqueeze(0)
+        if annot_rd is None and annot_sr is None:
+            mask_label = np.zeros(mask_hw, dtype=np.uint8)
+        else:
+            mask_label = merge_srground_masks(annot_rd, annot_sr)
+        mask = torch.from_numpy(mask_label.astype(np.float32)).unsqueeze(0)
+        mask_rd = (
+            torch.from_numpy(annot_rd.astype(np.float32)).unsqueeze(0)
+            if annot_rd is not None
+            else torch.zeros((1, *mask_hw), dtype=torch.float32)
+        )
+        mask_sr = (
+            torch.from_numpy(annot_sr.astype(np.float32)).unsqueeze(0)
+            if annot_sr is not None
+            else torch.zeros((1, *mask_hw), dtype=torch.float32)
+        )
         mask_coverage = float((mask > 0).float().mean().item())
 
         real_ann_path = self._resolve_path(sample.get('real_distortions_ann_path'))
-        sr_ann_path = self._resolve_path(sample.get('sr_artifacts_ann_path'))
+        sr_ann_path = None
+        if self.include_sr_artifact:
+            candidate = self._resolve_path(sample.get('sr_artifacts_ann_path'))
+            if candidate.exists():
+                sr_ann_path = candidate
 
         return {
             'img': img_tensor,
@@ -903,7 +1070,9 @@ class SRGroundSmallDataset(Dataset):
             'mask_coverage': mask_coverage,
             'prominences': sample.get('prominences'),
             'has_markup': sample.get('has_markup', False),
-            'sr_artifacts_ann_path': _to_relative_dataset_path(sr_ann_path, self.root),
+            'sr_artifacts_ann_path': (
+                _to_relative_dataset_path(sr_ann_path, self.root) if sr_ann_path is not None else None
+            ),
             'real_distortions_ann_path': _to_relative_dataset_path(real_ann_path, self.root),
             'sample_id': sample.get('sample_id'),
             'distorted_img_path': _to_relative_dataset_path(img_path, self.root),
@@ -936,8 +1105,8 @@ def srground_collate_fn(batch: List[dict]) -> dict:
 
 class KadidPristineDataset(Dataset):
     """
-    KADID-10k pristine (reference) images dataset. 
-    Возвращает только оригинальные изображения без искажений с RandomCrop.
+    KADID-10k pristine (reference) images dataset.
+    Returns only original undistorted images with RandomCrop.
     """
 
     def __init__(
@@ -957,7 +1126,7 @@ class KadidPristineDataset(Dataset):
         else:
             self.transform = transform
 
-        # Находим все оригинальные изображения (формат I{number}.png без суффиксов)
+        # Find all original images (I{number}.png format without suffixes)
         images_dir = self.root / "images"
         pristine_pattern = re.compile(r'^I\d+\.png$')
 
@@ -983,13 +1152,13 @@ class KadidPristineDataset(Dataset):
 
         return {
             "img": img,
-            "mos": 5.0,  # максимальное качество для pristine
+            "mos": 5.0,  # maximum quality for pristine images
             "dist_type": "pristine",
             "dist_group": "pristine",
-            "dist_level": 0,  # уровень искажений = 0
+            "dist_level": 0,  # distortion level = 0
             "distorted_img_path": img_rel_path,
-            "original_img_path": img_rel_path,  # сам себе референс
-            "sample_id": img_path.stem,  # например "I01"
+            "original_img_path": img_rel_path,  # self-reference
+            "sample_id": img_path.stem,  # e.g. "I01"
         }
 
 
@@ -997,9 +1166,9 @@ class LocalKadidPristineDataset(Dataset):
     """
     Pristine KADID dataset.
 
-    Ожидает корневую директорию с index.csv.
-    Обязательные колонки: original_img_path
-    Опциональные: mask_path (если есть маска области без искажений)
+    Expects a root directory with index.csv.
+    Required columns: original_img_path
+    Optional: mask_path (if a distortion-free region mask is available)
     """
 
     def __init__(
@@ -1058,11 +1227,11 @@ class LocalKadidPristineDataset(Dataset):
 
 def kadid_pristine_collate_fn(batch: List[dict]) -> dict:
     """
-    Collate для KadidPristineDataset.
+    Collate for KadidPristineDataset.
 
-    Возвращает:
+    Returns:
         images: Tensor (B, C, H, W)
-        + все остальные ключи как списки длины B
+        + all other keys as lists of length B
     """
     images = torch.stack([item["img"] for item in batch], dim=0)
     collated: dict = {"images": images}
@@ -1077,18 +1246,18 @@ def kadid_pristine_collate_fn(batch: List[dict]) -> dict:
 
 def local_kadid_pristine_collate_fn(batch: List[dict]) -> dict:
     """
-    Collate для LocalKadidPristinePresavedDataset.
+    Collate for LocalKadidPristinePresavedDataset.
 
-    Возвращает:
+    Returns:
         images: Tensor (B, C, H, W)
-        masks: Tensor (B, 1, H, W) если has_masks=True
-        + остальные ключи как списки длины B
+        masks: Tensor (B, 1, H, W) if has_masks=True
+        + remaining keys as lists of length B
     """
     images = torch.stack([item["img"] for item in batch], dim=0)
 
     collated: dict = {"images": images}
 
-    # Если есть маски в батче
+    # If masks are present in the batch
     if "mask" in batch[0]:
         masks = torch.stack([item["mask"] for item in batch], dim=0)
         collated["masks"] = masks
@@ -1116,33 +1285,40 @@ def resolve_dataset_image_path(
     )
     dataset_root = Path(dataset_images_root(str(root), dataset))
     path = Path(path_from_meta)
-    return path if path.is_absolute() else (dataset_root / path)
+    if path.is_absolute():
+        parts = Path(path).parts
+        i = parts.index(dataset)
+        suffix = Path(*parts[i:])
+        path = dataset_root / suffix
+    else:
+        path = dataset_root / path
+    return path
 
 
 def dataset_image_paths(
     dataset: str,
-    kadid_images_path: str,
+    dataset_root: str,
     crop_size: int,
     min_distortion_level: int,
 ) -> List[str]:
     if dataset == 'kadid10k':
         ds = Kadid10kDataset(
-            kadid_images_path,
+            dataset_root,
             crop_size=crop_size,
             min_distortion_level=min_distortion_level,
         )
         return [_to_relative_dataset_path(Path(p), ds.root) for p in ds.images]
 
     if dataset == 'local_kadid':
-        ds = LocalKadidPresavedDataset(kadid_images_path, crop_size=crop_size)
+        ds = LocalKadidPresavedDataset(dataset_root, crop_size=crop_size)
         return [_to_relative_dataset_path(Path(str(p)), ds.root) for p in ds.df['distorted_img_path'].tolist()]
 
     if dataset == 'QGround':
-        ds = QGroundDataset(kadid_images_path, split='test', crop_size=crop_size)
+        ds = QGroundDataset(dataset_root, split='test', crop_size=crop_size)
         return [str(sample['image_rel']) for sample in ds.samples]
 
     if dataset == 'SRGround':
-        ds = SRGroundSmallDataset(kadid_images_path, json_path='srground_train.json', allowed_methods=['DiT4SR_x2'])
+        ds = SRGroundSmallDataset(dataset_root, json_path='srground_train.json', allowed_methods=['DiT4SR_x2'])
         return [_to_relative_dataset_path(Path(str(p)), ds.root) for p in ds.df['image_path'].tolist()]
 
     raise ValueError(f'Unsupported dataset: {dataset}')

analysis/features/feature_filters.py

@@ -85,57 +85,57 @@ def my_kruskal(
     col_data: np.ndarray,
     col_indices: np.ndarray,
     n_samples: int,
-    group_indices: list[np.ndarray],  # список массивов индексов для каждой группы
+    group_indices: list[np.ndarray],  # list of index arrays, one per group
 ) -> float:
     """
-    Kruskal-Wallis H-test для одной колонки CSC-матрицы.
+    Kruskal-Wallis H-test for a single CSC matrix column.
 
-    Параметры
+    Parameters
     ----------
-    col_data     : ненулевые значения колонки (csc.data[csc.indptr[j]:csc.indptr[j+1]])
-    col_indices  : строковые индексы ненулевых элементов (csc.indices[...])
-    n_samples    : полное число строк (n)
-    group_indices: group_indices[g] — массив строковых индексов группы g
-
-    Возвращает
-    ----------
-    p-value (float), или np.nan если тест неприменим
+    col_data     : nonzero column values (csc.data[csc.indptr[j]:csc.indptr[j+1]])
+    col_indices  : row indices of nonzero elements (csc.indices[...])
+    n_samples    : total number of rows (n)
+    group_indices: group_indices[g] — array of row indices for group g
+
+    Returns
+    -------
+    p-value (float), or np.nan if the test is not applicable
     """
     n = n_samples
     n_zeros = n - len(col_data)
 
-    # --- 1. Ранги ненулевых элементов среди ВСЕХ n значений ----------------
-    # Нули занимают позиции 1..n_zeros в общем порядке.
-    # Ненулевые элементы начинаются с позиции n_zeros + 1.
+    # --- 1. Ranks of nonzero elements among ALL n values -------------------
+    # Zeros occupy positions 1..n_zeros in the overall order.
+    # Nonzero elements start at position n_zeros + 1.
 
-    # Сортируем ненулевые значения, получаем ранги внутри них
+    # Sort nonzero values and compute ranks within them
     nonzero_vals = col_data.astype(np.float64)
     order = np.argsort(nonzero_vals, kind='stable')
-    # Обрабатываем ties: каждому уникальному значению — средний ранг
+    # Handle ties: assign each unique value its mean rank
     ranks_local = np.empty(len(nonzero_vals), dtype=np.float64)
-    # временно назначаем 1-based ранги среди ненулевых
+    # Temporarily assign 1-based ranks among nonzero values
     ranks_local[order] = np.arange(1, len(nonzero_vals) + 1, dtype=np.float64)
-    # Разрешаем ties усреднением
+    # Resolve ties by averaging
     sorted_vals = nonzero_vals[order]
     unique_vals, inverse, counts = np.unique(sorted_vals, return_inverse=True, return_counts=True)
-    # средний ранг для каждого уникального значения (1-based среди ненулевых)
+    # mean rank for each unique value (1-based among nonzero values)
     cum = np.concatenate([[0], np.cumsum(counts)])
     mean_local_ranks = (cum[:-1] + cum[1:] + 1) / 2.0  # (first + last + 1) / 2
     ranks_local = mean_local_ranks[inverse]
 
-    # Сдвигаем ранги ненулевых элементов на n_zeros (они все правее нулей)
-    # Но нужно учесть ties между нулями и ненулевыми.
-    global_ranks_nonzero = ranks_local + n_zeros  # сдвиг на количество нулей
+    # Shift nonzero ranks by n_zeros (they all lie to the right of zeros)
+    # Ties between zeros and nonzero values must be accounted for.
+    global_ranks_nonzero = ranks_local + n_zeros  # shift by the number of zeros
 
-    # Средний ранг нулей: нули занимают ранги 1..n_zeros
+    # Mean rank of zeros: zeros occupy ranks 1..n_zeros
     zero_mean_rank = (n_zeros + 1) / 2.0 if n_zeros > 0 else 0.0
 
-    # --- 2. Для каждой группы считаем сумму рангов -------------------------
+    # --- 2. Sum ranks for each group ----------------------------------------
     n_groups = len(group_indices)
     if n_groups < 2:
         return np.nan
 
-    # Маска ненулевых: строим dict index->global_rank для быстрого lookup
+    # Nonzero mask: build index->global_rank dict for fast lookup
     rank_map = dict(zip(col_indices, global_ranks_nonzero))
 
     H_num = 0.0
@@ -149,7 +149,7 @@ def my_kruskal(
         if ng == 0:
             continue
 
-        # Ранги элементов группы
+        # Ranks of group elements
         r_sum = 0.0
         for i in g_idx:
             r_sum += rank_map.get(int(i), zero_mean_rank)
@@ -161,9 +161,9 @@ def my_kruskal(
     if len(group_sizes) < 2:
         return np.nan
 
-    n_t = n_total_valid  # должно равняться n
+    n_t = n_total_valid  # should equal n
 
-    # --- 3. H-статистика (стандартная формула) ------------------------------
+    # --- 3. H-statistic (standard formula) ----------------------------------
     # H = 12 / (n*(n+1)) * sum(R_i^2 / n_i) - 3*(n+1)
     H = 0.0
     for ng, rs in zip(group_sizes, group_rank_sums):
@@ -171,13 +171,13 @@ def my_kruskal(
 
     H = (12.0 / (n_t * (n_t + 1))) * H - 3.0 * (n_t + 1)
 
-    # --- 4. Поправка на ties ------------------------------------------------
-    # T = sum(t^3 - t) для каждой группы ties, C = 1 - T / (n^3 - n)
-    # Ties среди ненулевых:
+    # --- 4. Tie correction --------------------------------------------------
+    # T = sum(t^3 - t) for each tie group, C = 1 - T / (n^3 - n)
+    # Ties among nonzero values:
     tie_correction = 1.0
     if len(unique_vals) < len(nonzero_vals):
         T = np.sum(counts ** 3 - counts, dtype=np.float64)
-        # Добавляем ties нулей:
+        # Add ties among zeros:
         if n_zeros > 1:
             T += n_zeros ** 3 - n_zeros
         denom = float(n_t) ** 3 - n_t
@@ -190,7 +190,7 @@ def my_kruskal(
     if H < 0:
         H = 0.0
 
-    # --- 5. p-value из chi2 с df = n_groups - 1 ----------------------------
+    # --- 5. p-value from chi2 with df = n_groups - 1 ------------------------
     df = len(group_sizes) - 1
     p_value = float(stats.chi2.sf(H, df))
     return p_value
@@ -314,7 +314,7 @@ class KruskalWallisFeatureFilter(BaseFeatureFilter):
             except Exception as exc:
                 print(f'[cache][feature_filter] invalid cache {cache_path}: {exc}. Recomputing...')
 
-        # Предвычисляем индексы строк для каждой группы — один раз
+        # Precompute row indices for each group once
         group_indices = []
         for group in unique_groups:
             idx = np.where(group_values == group)[0]
@@ -331,7 +331,7 @@ class KruskalWallisFeatureFilter(BaseFeatureFilter):
             col_data = codes_csc.data[start:stop]
             col_indices = codes_csc.indices[start:stop]
 
-            # Пропускаем константные колонки
+            # Skip constant columns
             n_nonzero = stop - start
             all_same = (
                 n_nonzero == 0
@@ -339,9 +339,9 @@ class KruskalWallisFeatureFilter(BaseFeatureFilter):
                 or (n_nonzero < n_samples and n_nonzero == 0)
             )
             if all_same and n_nonzero == 0:
-                continue  # все нули — константа
+                continue  # all zeros — constant
             if n_nonzero > 0 and n_nonzero == n_samples and np.all(col_data == col_data[0]):
-                continue  # все одинаковые ненулевые
+                continue  # all identical nonzero values
 
             p_value = my_kruskal(col_data, col_indices, n_samples, group_indices)
 

analysis/features/feature_matrix.py

@@ -3,10 +3,17 @@ from __future__ import annotations
 from typing import List, Sequence
 
 import numpy as np
+import scipy.sparse as sp
 
 from analysis.features.feature_indexing import FeatureMatrix
 
 
+def _axis0_to_dense(values) -> np.ndarray:
+    if hasattr(values, 'toarray'):
+        return np.asarray(values.toarray()).ravel().astype(np.float32)
+    return np.asarray(values).ravel().astype(np.float32)
+
+
 def _zscore_columns(x: np.ndarray) -> np.ndarray:
     mu = x.mean(axis=0, keepdims=True)
     sigma = x.std(axis=0, keepdims=True)
@@ -40,7 +47,7 @@ def build_image_feature_matrix(
 
     `image_row_indices` is a sequence where each element is an array of row indices
     (patch indices) belonging to that image.
-    Aggregation modes: 'max', 'sum', 'mean_acts' (mean over positive activations) or 'mean'.
+    Aggregation modes: 'max', 'sum', 'mean_acts' (mean over positive activations only).
     """
     return _zscore_columns(
         build_image_feature_matrix_raw(
@@ -52,6 +59,96 @@ def build_image_feature_matrix(
     )
 
 
+def _to_dense_image_features(values) -> np.ndarray:
+    if sp.issparse(values):
+        return np.asarray(values.toarray(), dtype=np.float32)
+    return np.asarray(values, dtype=np.float32)
+
+
+def _image_group_matrix(image_idx_arr: np.ndarray, n_images_used: int) -> sp.csr_matrix:
+    """Map each patch row to one image column (n_patches, n_images)."""
+    n_patches = int(len(image_idx_arr))
+    rows = np.arange(n_patches, dtype=np.int32)
+    cols = np.asarray(image_idx_arr, dtype=np.int32)
+    return sp.csr_matrix(
+        (np.ones(n_patches, dtype=np.float32), (rows, cols)),
+        shape=(n_patches, int(n_images_used)),
+    )
+
+
+def _grouped_sum(codes, group: sp.csr_matrix) -> np.ndarray:
+    return _to_dense_image_features(group.T @ codes)
+
+
+def _grouped_mean_acts(codes, group: sp.csr_matrix) -> np.ndarray:
+    if sp.issparse(codes):
+        positive = codes.multiply(codes > 0)
+        pos_count = group.T @ (codes > 0).astype(np.float32)
+    else:
+        codes_arr = np.asarray(codes, dtype=np.float32)
+        positive = codes_arr * (codes_arr > 0)
+        pos_count = group.T @ (codes_arr > 0).astype(np.float32)
+
+    pos_sum = _to_dense_image_features(group.T @ positive)
+    pos_count = _to_dense_image_features(pos_count)
+    out = np.zeros_like(pos_sum, dtype=np.float32)
+    nonzero = pos_count > 0
+    out[nonzero] = pos_sum[nonzero] / pos_count[nonzero]
+    return out
+
+
+def _grouped_max(codes, image_idx_arr: np.ndarray, n_images_used: int, n_features_total: int) -> np.ndarray:
+    x = np.zeros((int(n_images_used), int(n_features_total)), dtype=np.float32)
+    for img_i in range(int(n_images_used)):
+        patch_mask = image_idx_arr == img_i
+        if not np.any(patch_mask):
+            continue
+        chunk = codes[patch_mask, :]
+        x[img_i] = _axis0_to_dense(chunk.max(axis=0))
+    return x
+
+
+def _aggregate_codes_by_image_idx(
+    codes,
+    image_idx_arr: np.ndarray,
+    n_images_used: int,
+    aggregation_mode: str,
+) -> np.ndarray:
+    """Aggregate patch-level codes per image using dense image_idx labels."""
+    mode = str(aggregation_mode)
+    if mode not in {'max', 'sum', 'mean_acts'}:
+        raise ValueError(f'Unknown aggregation mode: {mode!r}')
+
+    image_idx_arr = np.asarray(image_idx_arr, dtype=np.int32)
+    n_images_used = int(n_images_used)
+    n_features_total = int(codes.shape[1])
+
+    if mode == 'max':
+        return _grouped_max(codes, image_idx_arr, n_images_used, n_features_total)
+
+    group = _image_group_matrix(image_idx_arr, n_images_used)
+    if mode == 'sum':
+        return _grouped_sum(codes, group)
+    return _grouped_mean_acts(codes, group)
+
+
+def build_image_feature_matrix_from_image_idx(
+    codes,
+    image_idx_arr: np.ndarray,
+    n_images_used: int,
+    aggregation_mode: str = 'max',
+) -> np.ndarray:
+    """Build z-scored image-level matrix using per-patch ``image_idx`` labels."""
+    return _zscore_columns(
+        _aggregate_codes_by_image_idx(
+            codes,
+            image_idx_arr,
+            n_images_used,
+            aggregation_mode,
+        )
+    )
+
+
 def build_image_feature_matrix_raw(
     features: FeatureMatrix,
     image_row_indices: Sequence[Sequence[int]],
@@ -64,37 +161,19 @@ def build_image_feature_matrix_raw(
     skips the final z-score normalization so it can be used for paired deltas.
     """
     mode = str(aggregation_mode)
-    if mode not in {'max', 'sum', 'mean_acts', 'mean'}:
+    if mode not in {'max', 'sum', 'mean_acts'}:
         raise ValueError(f'Unknown aggregation mode: {mode!r}')
 
-    n_features_total = features.n_features
     codes_subset = features.codes
-    x = np.zeros((int(n_images_used), int(n_features_total)), dtype=np.float32)
-
+    image_idx_arr = np.empty(int(codes_subset.shape[0]), dtype=np.int32)
     for img_i, row_ids in enumerate(image_row_indices):
         if len(row_ids) == 0:
             continue
-        chunk = codes_subset[row_ids, :]
-        if mode == 'max':
-            x[img_i] = np.asarray(chunk.max(axis=0)).ravel().astype(np.float32)
-        elif mode == 'sum':
-            x[img_i] = np.asarray(chunk.sum(axis=0)).ravel().astype(np.float32)
-        elif mode == 'mean':
-            x[img_i] = np.asarray(chunk.mean(axis=0)).ravel().astype(np.float32)
-        else:  # mean_acts: mean over positive activations
-            if hasattr(chunk, 'multiply'):
-                positive_mask = chunk > 0
-                positive_chunk = chunk.multiply(positive_mask)
-                pos_sum = np.asarray(positive_chunk.sum(axis=0)).ravel().astype(np.float32)
-                pos_count = np.asarray(positive_mask.sum(axis=0)).ravel().astype(np.float32)
-            else:
-                chunk_arr = np.asarray(chunk, dtype=np.float32)
-                positive_mask = chunk_arr > 0
-                pos_sum = np.asarray((chunk_arr * positive_mask).sum(axis=0)).ravel().astype(np.float32)
-                pos_count = np.asarray(positive_mask.sum(axis=0)).ravel().astype(np.float32)
-            mean_pos = np.zeros(int(n_features_total), dtype=np.float32)
-            nonzero = pos_count > 0
-            mean_pos[nonzero] = pos_sum[nonzero] / pos_count[nonzero]
-            x[img_i] = mean_pos
+        image_idx_arr[list(row_ids)] = int(img_i)
 
-    return x
+    return _aggregate_codes_by_image_idx(
+        codes_subset,
+        image_idx_arr,
+        n_images_used,
+        mode,
+    )

analysis/features/feature_selectors.py

@@ -252,7 +252,7 @@ class TopKAbsCorrSelector(BaseFeatureSelector):
     name: str = 'topk_abs_corr'
     description: str = 'Top-k features by max absolute dist_type correlation.'
     corr_key: str = 'corr_type_df'
-    relation_label: str = 'высокая корреляция'
+    relation_label: str = 'high correlation'
 
     def compute_metric_tables(
         self,
@@ -316,7 +316,7 @@ class TopKMutualInfoSelector(BaseFeatureSelector):
     name: str = 'topk_mi'
     description: str = 'Top-k features by max dist_type mutual information.'
     mi_key: str = 'mi_type_df'
-    relation_label: str = 'высокая взаимная информация'
+    relation_label: str = 'high mutual information'
 
     def compute_metric_tables(
         self,
@@ -386,7 +386,7 @@ class TopKRocAucSelector(BaseFeatureSelector):
     name: str = 'topk_auc'
     description: str = 'Top-k features by one-vs-rest ROC-AUC at dist_type level.'
     auc_key: str = 'auc_type_df'
-    relation_label: str = 'высокая ROC-AUC различимость'
+    relation_label: str = 'high ROC-AUC separability'
 
     def compute_metric_tables(
         self,
@@ -456,7 +456,7 @@ class TopKPairedDeltaSelector(BaseFeatureSelector):
     name: str = 'paired_delta'
     description: str = 'Top-k features by max relative paired activation delta.'
     delta_key: str = 'paired_relative_dist_type_df'
-    relation_label: str = 'высокая парная дельта активаций'
+    relation_label: str = 'high paired activation delta'
 
     @staticmethod
     def _delta_mode_from_key(delta_key: str) -> str:
@@ -636,7 +636,7 @@ class TopKIoUSelector(BaseFeatureSelector):
 
     name: str = 'topk_iou_type'
     description: str = 'Top-k features by median IoU with distortion regions at dist_type level.'
-    relation_label: str = 'высокая пространственная локализация'
+    relation_label: str = 'high spatial localization'
 
     def compute_metric_tables(
         self,

analysis/features/feature_stats.py

@@ -1,5 +1,5 @@
 """
-Статистики по SAE-признакам и их визуализация (bar charts).
+SAE feature statistics and visualization (bar charts).
 """
 
 from typing import List, Optional, Tuple
@@ -13,19 +13,19 @@ from analysis.features.feature_indexing import FeatureMatrix
 
 def compute_feature_stats(features: FeatureMatrix) -> pd.DataFrame:
     """
-    Вычисляет статистики для каждого признака SAE.
+    Compute statistics for each SAE feature.
 
-    Параметры
+    Parameters
     ----------
     features : CSR activations; ``column_feature_ids[j]`` = global SAE id
 
-    Возвращает
-    ----------
-    DataFrame с колонками: feature_id (global SAE id), mean, frequency, max, mean_acts
-        mean      — средняя активация по всем патчам
-        frequency — доля патчей с ненулевой активацией
-        max       — максимальная активация среди всех патчей
-        mean_acts — средняя активация по ненулевым патчам
+    Returns
+    -------
+    DataFrame with columns: feature_id (global SAE id), mean, frequency, max, mean_acts
+        mean      — mean activation across all patches
+        frequency — fraction of patches with nonzero activation
+        max       — maximum activation among all patches
+        mean_acts — mean activation over nonzero patches
     """
     codes = features.codes
     mat = codes if codes.dtype == np.float32 else codes.astype(np.float32)
@@ -55,18 +55,18 @@ def get_top_features(
     min_mean_acts: Optional[float] = None,
 ) -> List[int]:
     """
-    Возвращает индексы top-K признаков по выбранному критерию.
+    Return top-K feature indices by the selected criterion.
 
-    Параметры
+    Parameters
     ----------
-    stats         : DataFrame из compute_feature_stats
-    top_k         : число топ-признаков
+    stats         : DataFrame from compute_feature_stats
+    top_k         : number of top features
     criterion     : 'mean' | 'frequency' | 'max' | 'mean_acts'
-    min_mean_acts : предварительный фильтр по mean_acts
+    min_mean_acts : preliminary filter on mean_acts
 
-    Возвращает
-    ----------
-    List[int] — feature_id в порядке убывания критерия
+    Returns
+    -------
+    List[int] — feature_id in descending order of the criterion
     """
     assert criterion in ('mean', 'frequency', 'max', 'mean_acts'), \
         f"criterion must be one of 'mean', 'frequency', 'max', 'mean_acts', got {criterion!r}"
@@ -84,8 +84,8 @@ def plot_top_features(
     min_mean_acts: Optional[float] = None,
 ) -> None:
     """
-    Bar-chart топ-K признаков по выбранному критерию.
-    Над каждым баром подписывается frequency (доля ненулевых патчей).
+    Bar chart of top-K features by the selected criterion.
+    Each bar is labeled with frequency (fraction of nonzero patches).
     """
     top_ids = get_top_features(stats, top_k=top_k, criterion=criterion,
                                min_mean_acts=min_mean_acts)

analysis/metrics/correlations.py

@@ -21,8 +21,8 @@ def _prepare_codes_and_categories(
     feature_names = [int(fid) for fid in features.column_feature_ids]
 
     if binarize:
-        threshold = 0.2  # Такой был в PatchSAE
-        print(f'[binarize] Бинаризация с порогом {threshold}')
+        threshold = 0.2  # Same value as in PatchSAE
+        print(f'[binarize] Binarizing with threshold {threshold}')
         if sp.issparse(codes):
             codes = codes.copy()
             codes.data = (codes.data > threshold).astype(np.float32)
@@ -74,7 +74,6 @@ def compute_distortion_correlations(
     cached = load_parquet_cache(cache_path, label='correlations')
     if cached is not None:
         return cached
-
     work_features = features.subset(global_feature_ids)
 
     work_codes, _, feature_names, unique_categories, cat_idx = _prepare_codes_and_categories(

analysis/models.py

@@ -15,16 +15,16 @@ _hook_handle = None
 
 
 def _make_hook(name: str, sequence_layout: str = 'blc'):
-    """Forward hook для ARNIQA/MANIQA/LIQE.
+    """Forward hook for ARNIQA/MANIQA/LIQE.
 
-    Поддерживаемые форматы выхода слоя:
+    Supported layer output formats:
     - (B, C, H, W) -> (B*H*W, C)
     - (B, L, C)    -> (B*L, C)   (MANIQA)
     - (L, B, C)    -> (B*L, C)   (LIQE/LIQE-MIX)
     """
     def hook(module, inp, out):
         if hasattr(out, 'last_hidden_state'):
-            out_detached = out.last_hidden_state.detach() # Пока костыль
+            out_detached = out.last_hidden_state.detach()  # Temporary workaround
         else:
             out_detached = out.detach()
         if out_detached.ndim == 4:
@@ -101,23 +101,23 @@ def load_sae(
     sae_config: Optional[Dict[str, Any]] = None,
     **overrides: Any,
 ) -> torch.nn.Module:
-    """Загружает SAE из чекпоинта, определяя архитектуру из JSON-конфига.
+    """Load SAE from a checkpoint, inferring architecture from the JSON config.
 
     Parameters
     ----------
     checkpoint_path : str
-        Путь к директории чекпоинта (output_dir/checkpoint-<step>/) или файлу весов.
+        Path to the checkpoint directory (output_dir/checkpoint-<step>/) or weights file.
     config_path : str, optional
-        Путь к sae_config.json. По умолчанию ищется в родительской директории
-        чекпоинта (т.е. output_dir/sae_config.json).
+        Path to sae_config.json. By default, looked up in the parent directory
+        of the checkpoint (i.e. output_dir/sae_config.json).
     device : str
-        Устройство для загрузки.
+        Device to load onto.
     dtype : torch.dtype
-        Тип данных для загрузки.
+        Data type for loading.
     sae_config : dict, optional
-        Уже загруженный конфиг.
+        Pre-loaded config.
     **overrides
-        Переопределяют отдельные поля конфига (например, mp_threshold=0.05).
+        Override individual config fields (e.g. mp_threshold=0.05).
     """
     import accelerate
     from train_code.sae import SAE, MatchingPursuitSAE
@@ -159,17 +159,17 @@ def load_iqa_model(
     swin_num: int = 2,
 ):
     """
-    Загружает IQA-модель и регистрирует hook на выбранный слой.
+    Load an IQA model and register a hook on the selected layer.
 
-    Поддерживаемые метрики:
-    - arniqa-kadid: hook на iqa.net.encoder[layer_num]
-    - maniqa: hook на iqa.net.swintransformer{swin_num}.layers[layer_num]
-    - liqe / liqe_mix: hook на iqa.net.clip_model.visual.transformer.resblocks[layer_num]
+    Supported metrics:
+    - arniqa-kadid: hook on iqa.net.encoder[layer_num]
+    - maniqa: hook on iqa.net.swintransformer{swin_num}.layers[layer_num]
+    - liqe / liqe_mix: hook on iqa.net.clip_model.visual.transformer.resblocks[layer_num]
 
-    Возвращает
-    ----------
-    iqa        : загруженная IQA-модель
-    layer_name : строковый ключ, под которым активации хранятся в _iqa_activations
+    Returns
+    -------
+    iqa        : loaded IQA model
+    layer_name : string key under which activations are stored in _iqa_activations
     """
     import pyiqa
 

analysis/utils.py

@@ -1,5 +1,5 @@
 """
-Общие вспомогательные утилиты пакета analysis.
+Shared helper utilities for the analysis package.
 """
 
 from typing import List
@@ -18,25 +18,25 @@ def get_top_images_for_feature(
     aggregation: str = 'mean_acts',
 ) -> List[int]:
     """
-    Возвращает индексы изображений, на которых признак feature_id активируется сильнее всего.
+    Return image indices where feature_id activates most strongly.
 
-    Активации патчей внутри одного изображения агрегируются в одно скалярное значение,
-    после чего изображения сортируются по убыванию.
+    Patch activations within each image are aggregated into a single scalar,
+    then images are sorted in descending order.
 
-    Параметры
+    Parameters
     ----------
     features    : CSR activations with global id per column
-    meta        : DataFrame с колонкой 'image_idx' (один патч — одна строка)
+    meta        : DataFrame with an 'image_idx' column (one patch per row)
     feature_id  : global SAE feature id
-    top_n       : число возвращаемых изображений
+    top_n       : number of images to return
     aggregation : 'mean_acts' | 'max' | 'sum'
-                  mean_acts — среднее по патчам с активацией > 0
-                  max       — максимальная активация среди патчей
-                  sum       — сумма всех активаций
+                  mean_acts — mean over patches with activation > 0
+                  max       — maximum activation among patches
+                  sum       — sum of all activations
 
-    Возвращает
-    ----------
-    List[int] — image_idx в порядке убывания агрегированной активации
+    Returns
+    -------
+    List[int] — image_idx in descending order of aggregated activation
     """
     assert aggregation in ('mean_acts', 'max', 'sum'), (
         f"aggregation must be 'mean_acts', 'max' or 'sum', got {aggregation!r}"
@@ -72,11 +72,11 @@ def get_top_images_for_feature_by_iou(
     dataset: str | None = None,
 ) -> List[int]:
     """
-    Возвращает индексы изображений с наибольшим IoU между бинарной картой
-    активаций признака и маской искажений для каждого изображения.
+    Return image indices with the highest IoU between the feature's binary
+    activation map and the distortion mask for each image.
 
-    Требует, чтобы в `meta` были колонки `image_idx` и либо `patch_mask_label`,
-    либо `patch_is_distorted` (см. iou_utils._load_patch_mask_for_group).
+    Requires `image_idx` and either `patch_mask_label` or `patch_is_distorted`
+    in `meta` (see iou_utils._load_patch_mask_for_group).
     """
     from analysis.metrics import iou_utils
 

analysis/viz/umap_utils.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import hashlib
 import json
+import os
 import pickle
 import re
 import time
@@ -9,10 +10,62 @@ from pathlib import Path
 from typing import Any, Callable, Dict, Hashable, Mapping, Optional
 
 import numpy as np
+import pandas as pd
 import umap
 
 _UMAP_CACHE: Dict[Hashable, Dict[str, Any]] = {}
 
+UMAP_IMAGE_LIMIT = 1000
+UMAP_IMAGE_RANDOM_STATE = 42
+
+
+def select_random_image_indices(
+    max_images: int,
+    *,
+    limit: int = UMAP_IMAGE_LIMIT,
+    random_state: int = UMAP_IMAGE_RANDOM_STATE,
+) -> np.ndarray:
+    """Pick up to ``limit`` random unique image_idx values from ``[0, max_images)``."""
+    max_images = int(max_images)
+    n_pick = min(int(limit), max_images)
+    if n_pick <= 0:
+        raise ValueError(f'max_images must be positive, got {max_images}')
+    rng = np.random.default_rng(int(random_state))
+    return np.sort(rng.choice(np.arange(max_images), size=n_pick, replace=False)).astype(np.int32)
+
+
+def build_kadid_index_lookup(meta_df: pd.DataFrame, kadid_ds: Any, n_images_used: int) -> np.ndarray:
+    """Map dense UMAP image index (0..n-1) to ``Kadid10kDataset`` positional index."""
+    lookup = np.arange(n_images_used, dtype=np.int32)
+    if kadid_ds is None:
+        return lookup
+
+    first_rows = meta_df.groupby('image_idx', sort=True).first()
+    kadid_paths = [str(p) for p in kadid_ds.images]
+    path_to_idx = {p: i for i, p in enumerate(kadid_paths)}
+    base_to_idx = {os.path.basename(p): i for i, p in enumerate(kadid_paths)}
+
+    for dense_i in range(n_images_used):
+        row = first_rows.iloc[dense_i]
+        kadid_idx = None
+        for col in ('distorted_img_path', 'image_path', 'original_img_path'):
+            if col not in row.index:
+                continue
+            val = row[col]
+            if val is None or (isinstance(val, float) and np.isnan(val)):
+                continue
+            path = str(val)
+            kadid_idx = path_to_idx.get(path)
+            if kadid_idx is None and not Path(path).is_absolute():
+                kadid_idx = path_to_idx.get(str(kadid_ds.root / Path(path)))
+            if kadid_idx is None:
+                kadid_idx = base_to_idx.get(os.path.basename(path))
+            if kadid_idx is not None:
+                break
+        if kadid_idx is not None:
+            lookup[dense_i] = int(kadid_idx)
+    return lookup
+
 
 def _json_default(value: object) -> object:
     if isinstance(value, (np.integer, np.floating)):
@@ -116,36 +169,40 @@ def compute_umap_from_features(features: np.ndarray, umap_params: Dict) -> Dict[
 
 
 def get_or_compute_umap_with_builder(
-    key: Hashable,
     build_features_fn: Callable[[], np.ndarray],
     umap_params: Dict,
-    cache: Optional[Dict[Hashable, Dict[str, Any]]] = None,
     *,
+    cache_id: str,
+    cache_signature: Mapping[str, object],
+    cache: Optional[Dict[Hashable, Dict[str, Any]]] = None,
     cache_dir: str | Path | None = None,
-    cache_id: str | None = None,
-    cache_signature: Mapping[str, object] | None = None,
 ) -> Dict[str, Any]:
     """Get or compute UMAP embedding using a builder function to produce features.
 
-    - ``key``: a hashable cache key (e.g. (feature_id, params_tuple)).
     - ``build_features_fn``: zero-arg callable that returns a np.ndarray of shape (n_samples, n_features).
     - ``umap_params``: parameters passed to ``umap.UMAP``.
+    - ``cache_id``: stable logical id used in cache filenames.
+    - ``cache_signature``: data fingerprint for cache lookup (paired with ``umap_params`` via
+      :func:`build_umap_cache_key` for both in-memory and on-disk cache).
     - ``cache``: optional dict to use for caching; if None, module-level cache is used.
-    - ``cache_dir``: optional directory for on-disk NPZ cache (like feature filters).
-    - ``cache_id``: stable logical id used in cache filenames (defaults to ``str(key)``).
-    - ``cache_signature``: optional data fingerprint for disk cache lookup before feature build.
+    - ``cache_dir``: optional directory for on-disk NPZ cache.
 
     Returns a dict with embedding and timing and other meta fields. Adds 'cache_hit' flag.
     """
+    key = build_umap_cache_key(
+        cache_id=str(cache_id),
+        umap_params=umap_params,
+        signature=cache_signature,
+    )
+    resolved_cache_id = _sanitize_cache_id(cache_id)
+
     cache = _UMAP_CACHE if cache is None else cache
     if key in cache:
         cached = dict(cache[key])
         cached["cache_hit"] = True
         return cached
 
-    resolved_cache_id = _sanitize_cache_id(cache_id if cache_id is not None else str(key))
-
-    if cache_dir is not None and cache_signature is not None:
+    if cache_dir is not None:
         disk_cache_path = _resolve_disk_cache_path(
             cache_dir=cache_dir,
             cache_id=resolved_cache_id,
@@ -169,7 +226,7 @@ def get_or_compute_umap_with_builder(
             cache_id=resolved_cache_id,
             umap_params=umap_params,
             signature=cache_signature,
-            x_shape=None if cache_signature is not None else x_shape,
+            x_shape=None,
         )
         loaded = _load_umap_from_disk(disk_cache_path)
         if loaded is not None:

analysis/viz/vis_heatmaps.py

@@ -1,5 +1,5 @@
 """
-Визуализация тепловых карт SAE-признаков поверх изображений KADID-10k.
+Visualization of SAE feature heatmaps overlaid on KADID-10k images.
 """
 
 import io
@@ -27,25 +27,25 @@ from analysis.datasets import (
     QGROUND_DISTORTION_TYPES,
     SRGROUND_DISTORTION_TYPES,
     SRGROUND_LEGEND_LABELS,
+    annotation_mask_rgb_for_meta_row,
     available_distortions,
     distortion_types_mapping,
-    _image_rel_from_meta_row,
-    srground_mask_rgb_for_meta_row,
-    srground_prominences_by_image_paths,
+    get_srground_rgb_mask,
+    infer_spatial_mask_dataset,
 )
 
 # Show per-image filenames (and distortion subtitles where enabled) under heatmap panels.
 SHOW_HEATMAP_IMAGE_TITLES = False
 
-MASK_LEGEND_FONTSIZE = 20
+MASK_LEGEND_FONTSIZE = 14
 MASK_LEGEND_HANDLE_HEIGHT = 1.15
 MASK_LEGEND_HANDLE_LENGTH = 1.55
 
 
 def _get_distortion_info(img_path: str):
     """
-    По имени файла KADID (I04_07_05.png) возвращает
-    (distortion_name, distortion_group) или (None, None) для оригиналов.
+    From a KADID filename (I04_07_05.png), return
+    (distortion_name, distortion_group) or (None, None) for originals.
     """
     name = Path(img_path).name
     m = _re.match(r'I\d+_(\d+)_(\d+)\.png$', name, _re.IGNORECASE)
@@ -59,8 +59,8 @@ def _get_distortion_info(img_path: str):
 
 def _get_original_path(distorted_path: str) -> Optional[str]:
     """
-    Из пути к искажённому изображению KADID (I04_07_05.png)
-    возвращает путь к оригиналу (I04.png).
+    From a KADID distorted image path (I04_07_05.png),
+    return the path to the original (I04.png).
     """
     p = Path(distorted_path)
     match = _re.match(r'(I\d+)_\d+_\d+\.png$', p.name, _re.IGNORECASE)
@@ -207,15 +207,15 @@ def _plot_heatmap_grid(
     title_feature_id: int | None = None,
 ) -> FigureSaveResult:
     """
-    Рисует сетку overlay-хитмапов для одного feature_id.
+    Draw a grid of overlay heatmaps for a single feature_id.
 
-    Параметры
+    Parameters
     ----------
-    imgs_list      : список тензоров изображений (C, H, W)
-    codes_tensor   : тензор активаций (B, spatial, spatial, inner_dim)
-    feature_id     : индекс признака SAE
-    image_paths    : пути к PNG-файлам (для заголовков)
-    img_size_inches: размер одной ячейки в дюймах
+    imgs_list      : list of image tensors (C, H, W)
+    codes_tensor   : activation tensor (B, spatial, spatial, inner_dim)
+    feature_id     : SAE feature index
+    image_paths    : paths to PNG files (for titles)
+    img_size_inches: size of one cell in inches
     """
     B = len(imgs_list)
     inner_concepts = codes_tensor.shape[-1]
@@ -268,15 +268,15 @@ def _plot_diff_grid(
     show_img: bool = True,
 ) -> FigureSaveResult:
     """
-    Рисует сетку |distorted − original| для одного feature_id.
+    Draw a grid of |distorted − original| for a single feature_id.
 
-    Параметры
+    Parameters
     ----------
-    imgs_list      : список тензоров искажённых изображений (C, H, W)
-    orig_tensors   : список тензоров оригиналов (или None, если файл не найден)
-    feature_id     : индекс признака SAE (для заголовка)
-    image_paths    : пути к PNG-файлам (для заголовков)
-    img_size_inches: размер одной ячейки в дюймах
+    imgs_list      : list of distorted image tensors (C, H, W)
+    orig_tensors   : list of original tensors (or None if the file is missing)
+    feature_id     : SAE feature index (for the title)
+    image_paths    : paths to PNG files (for titles)
+    img_size_inches: size of one cell in inches
     """
     B = len(imgs_list)
     fig_diff, axes = plt.subplots(1, B, figsize=(img_size_inches * B, img_size_inches * 1.25))
@@ -302,37 +302,54 @@ def _plot_diff_grid(
     return _finalize_figure(fig_diff, save_path=save_path, show_img=show_img)
 
 
-def _plot_qground_mask_grid(
-    mask_paths: List[Optional[str]],
+def annotation_mask_legend_handles(
+    dataset: str,
+    *,
+    mask_rgb_list: Optional[List[Optional[np.ndarray]]] = None,
+) -> list[Patch]:
+    if dataset == 'SRGround':
+        return srground_legend_handles(mask_rgb_list)
+    if dataset == 'QGround':
+        handles = [Patch(facecolor='black', edgecolor='none', label='background')]
+        for label_name, rgb in QGROUND_DISTORTION_TYPES.items():
+            color = tuple((np.asarray(rgb, dtype=np.float32) / 255.0).tolist())
+            handles.append(Patch(facecolor=color, edgecolor='none', label=label_name))
+        return handles
+    return [Patch(facecolor='black', edgecolor='none', label='background')]
+
+
+def _plot_annotation_mask_grid(
+    mask_rgb_list: List[Optional[np.ndarray]],
     feature_id: int,
     image_paths: List[str],
+    *,
+    dataset: str,
     crop_size: int,
     img_size_inches: float = 4.0,
     meta_rows: Optional[List[Optional[pd.Series]]] = None,
     save_path: Optional[str] = None,
     show_img: bool = True,
 ) -> FigureSaveResult:
-    B = len(mask_paths)
+    """Plot GT annotation masks (QGround PNG or SRGround synthesized RGB) in one row."""
+    B = len(mask_rgb_list)
     fig, axes = plt.subplots(1, B, figsize=(img_size_inches * B, img_size_inches * 1.35))
     if B == 1:
         axes = [axes]
     if meta_rows is None:
         meta_rows = [None] * B
 
-    for ax, mask_path, img_path, meta_row in zip(axes, mask_paths, image_paths, meta_rows):
-        if mask_path is not None and Path(mask_path).exists():
-            mask_rgb = np.asarray(Image.open(mask_path).convert('RGB'), dtype=np.uint8)
+    for ax, mask_rgb, img_path, meta_row in zip(axes, mask_rgb_list, image_paths, meta_rows):
+        if mask_rgb is not None:
             ax.imshow(mask_rgb)
             height, width = mask_rgb.shape[:2]
-            rect_width = float(crop_size)
-            rect_height = float(crop_size)
-            rect_x = (width - rect_width) / 2.0
-            rect_y = (height - rect_height) / 2.0
+            rect_size = float(crop_size)
+            rect_x = (width - rect_size) / 2.0
+            rect_y = (height - rect_size) / 2.0
             ax.add_patch(
                 Rectangle(
                     (rect_x, rect_y),
-                    rect_width,
-                    rect_height,
+                    rect_size,
+                    rect_size,
                     fill=False,
                     edgecolor='black',
                     linewidth=2.0,
@@ -341,21 +358,64 @@ def _plot_qground_mask_grid(
             )
             _set_image_subplot_title(ax, img_path, meta_row)
         else:
-            ax.text(0.5, 0.5, 'QGround mask not found', ha='center', va='center',
-                    transform=ax.transAxes, fontsize=9)
+            ax.text(
+                0.5,
+                0.5,
+                f'{dataset} annotation mask not found',
+                ha='center',
+                va='center',
+                transform=ax.transAxes,
+                fontsize=9,
+            )
         ax.axis('off')
 
-    handles = [Patch(facecolor='black', edgecolor='none', label='background')]
-    for label_name, rgb in QGROUND_DISTORTION_TYPES.items():
-        color = tuple((np.asarray(rgb, dtype=np.float32) / 255.0).tolist())
-        handles.append(Patch(facecolor=color, edgecolor='none', label=label_name))
-
-    _add_mask_legend(fig, handles)
-    fig.suptitle(f'QGround annotation mask (feature {feature_id})', fontsize=12, x=0.5, ha='center')
+    _add_mask_legend(
+        fig,
+        annotation_mask_legend_handles(dataset, mask_rgb_list=mask_rgb_list),
+    )
+    fig.suptitle(f'{dataset} annotation mask (feature {feature_id})', fontsize=12, x=0.5, ha='center')
     fig.tight_layout(rect=(0, 0.14, 1, 1))
     return _finalize_figure(fig, save_path=save_path, show_img=show_img)
 
 
+def _plot_qground_mask_grid(
+    mask_paths: List[Optional[str]],
+    feature_id: int,
+    image_paths: List[str],
+    crop_size: int,
+    img_size_inches: float = 4.0,
+    meta_rows: Optional[List[Optional[pd.Series]]] = None,
+    save_path: Optional[str] = None,
+    show_img: bool = True,
+    *,
+    dataset_root: str | None = None,
+) -> FigureSaveResult:
+    """Backward-compatible wrapper around :func:`_plot_annotation_mask_grid`."""
+    if meta_rows is None:
+        meta_rows = [None] * len(mask_paths)
+    if dataset_root is None:
+        raise ValueError('dataset_root is required for QGround mask visualization')
+    mask_rgb_list = [
+        annotation_mask_rgb_for_meta_row(
+            meta_row,
+            dataset_root=dataset_root,
+            dataset='QGround',
+        )
+        for meta_row in meta_rows
+    ]
+    return _plot_annotation_mask_grid(
+        mask_rgb_list,
+        feature_id,
+        image_paths,
+        dataset='QGround',
+        crop_size=crop_size,
+        img_size_inches=img_size_inches,
+        meta_rows=meta_rows,
+        save_path=save_path,
+        show_img=show_img,
+    )
+
+
 def _add_mask_legend(fig, handles: list[Patch], *, ncol: int = 3) -> None:
     fig.legend(
         handles=handles,
@@ -371,10 +431,23 @@ def _add_mask_legend(fig, handles: list[Patch], *, ncol: int = 3) -> None:
     )
 
 
-def srground_legend_handles(*, include_sr_artifact: bool = True) -> list[Patch]:
+def srground_legend_handles(
+    mask_rgb_list: Optional[List[Optional[np.ndarray]]] = None,
+) -> list[Patch]:
+    """Legend entries for distortion classes present in plotted SRGround RGB masks."""
+    present: set[str] | None = None
+    if mask_rgb_list is not None:
+        present = set()
+        for mask_rgb in mask_rgb_list:
+            if mask_rgb is None:
+                continue
+            for dist_name, rgb in SRGROUND_DISTORTION_TYPES.items():
+                if np.isclose(mask_rgb, rgb, rtol=0.2, atol=20).any():
+                    present.add(dist_name)
+
     handles = [Patch(facecolor='black', edgecolor='none', label='background')]
     for dist_name, rgb in SRGROUND_DISTORTION_TYPES.items():
-        if dist_name == 'sr_artifact' and not include_sr_artifact:
+        if present is not None and dist_name not in present:
             continue
         color = tuple((np.asarray(rgb, dtype=np.float32) / 255.0).tolist())
         label = SRGROUND_LEGEND_LABELS.get(dist_name, dist_name)
@@ -392,7 +465,6 @@ def _plot_overlay_and_srground_mask_rows(
     meta_rows: Optional[List[Optional[pd.Series]]],
     *,
     title_feature_id: int,
-    include_sr_artifact: bool,
     crop_size: int,
     save_path: Optional[str] = None,
     show_img: bool = False,
@@ -468,7 +540,7 @@ def _plot_overlay_and_srground_mask_rows(
 
     _add_mask_legend(
         fig,
-        srground_legend_handles(include_sr_artifact=include_sr_artifact),
+        srground_legend_handles(mask_rgb_list),
     )
     fig.suptitle(f'Feature {display_id}', fontsize=13, x=0.5, ha='center')
     fig.tight_layout(rect=(0, 0.16, 1, 0.96))
@@ -483,12 +555,12 @@ def render_top_feature_panel_srground(
     feature_id: int,
     *,
     patches_per_image: Optional[int] = None,
-    crop_size: int = 224,
+    crop_size: int = 512,
     img_size_inches: float = 3.5,
-    include_sr_artifact: bool = True,
-    datasets_root: str | None = None,
+    dataset_root: str,
 ) -> bytes | None:
     """Two-row panel: SAE heatmap overlays and SRGround annotation masks with legend."""
+    crop_size = 512 # TODO: fix later
     pil_images = [Image.open(p).convert('RGB') for p in image_paths]
     preprocess = transforms.Compose(
         [
@@ -511,20 +583,13 @@ def render_top_feature_panel_srground(
 
     meta_lookup = _build_meta_lookup(meta)
     meta_rows = [meta_lookup.get(int(img_idx)) for img_idx in image_indices]
-    image_rels = [_image_rel_from_meta_row(row) for row in meta_rows]
-    prom_map = srground_prominences_by_image_paths(
-        [rel for rel in image_rels if rel],
-        datasets_root=datasets_root,
-    )
     mask_rgb_list = [
-        srground_mask_rgb_for_meta_row(
+        get_srground_rgb_mask(
             row,
-            datasets_root=datasets_root,
-            include_sr_artifact=include_sr_artifact,
+            dataset_root=dataset_root,
             crop_size=crop_size,
-            prominences=prom_map.get(image_rel) if image_rel else None,
         )
-        for row, image_rel in zip(meta_rows, image_rels)
+        for row in meta_rows
     ]
 
     matrix_col = features.column_for(int(feature_id))
@@ -537,7 +602,6 @@ def render_top_feature_panel_srground(
         img_size_inches,
         meta_rows,
         title_feature_id=int(feature_id),
-        include_sr_artifact=include_sr_artifact,
         crop_size=crop_size,
         save_path=None,
         show_img=False,
@@ -554,39 +618,41 @@ def visualize_feature_heatmaps(
     patches_per_image: Optional[int] = None,
     crop_size: int = 224,
     img_size_inches: float = 4.0,
-    show_diff: bool = True,
+    show_mask: bool = True,
     save_dir: Optional[str] = None,
     file_prefix: str = '',
     show_img: bool = True,
+    dataset_root: Optional[str] = None,
+    dataset: Optional[str] = None,
 ) -> List[Dict[str, object]]:
     """
-    Визуализирует тепловые карты SAE-признаков из предвычисленных sparse активаций.
+    Visualize SAE feature heatmaps from precomputed sparse activations.
 
-    Для каждого feature_id отображается:
-    - overlay-хитмап поверх искажённого изображения
-    - (опционально) |distorted − original| попиксельно
+    For each feature_id:
+    - overlay heatmap on the distorted image
+    - (optional) second row: GT mask (QGround/SRGround) or |distorted − original|
 
-    Параметры
+    Parameters
     ----------
-    meta             : DataFrame с метаданными (image_idx, patch_idx, ...)
+    meta             : metadata DataFrame (image_idx, patch_idx, ...)
     features         : CSR activations with global id per column
-    image_indices    : индексы изображений (image_idx) для визуализации
-    image_paths      : пути к соответствующим PNG-файлам
+    image_indices    : image indices (image_idx) to visualize
+    image_paths      : paths to the corresponding PNG files
     feature_ids      : global SAE feature ids to visualize
-    patches_per_image: число патчей на изображение; если None — из patch_idx
-    crop_size        : размер кропа при кэшировании
-    img_size_inches  : размер одного изображения на фигуре (в дюймах)
-    show_diff        : показывать ли |distorted − original|
-    save_dir         : директория для сохранения PNG; если None — PNG в памяти (ключ ``bytes``)
-    file_prefix      : префикс имени файла (например, stage/agg)
-    show_img         : показывать ли фигуры inline
-
-    Возвращает
-    ----------
-    List[Dict[str, object]] — артефакты с ключом ``path`` (диск) или ``bytes`` (память)
+    patches_per_image: number of patches per image; if None — inferred from patch_idx
+    crop_size        : crop size used during caching
+    img_size_inches  : size of one image on the figure (in inches)
+    show_mask        : second row — GT mask or |distorted − original| if no mask is available
+    save_dir         : directory for saving PNGs; if None — PNGs in memory (``bytes`` key)
+    file_prefix      : filename prefix (e.g. stage/agg)
+    show_img         : whether to display figures inline
+
+    Returns
+    -------
+    List[Dict[str, object]] — artifacts with ``path`` (disk) or ``bytes`` (memory) key
     """
     assert len(image_indices) == len(image_paths), \
-        "image_indices и image_paths должны иметь одинаковую длину"
+        "image_indices and image_paths must have the same length"
 
     codes = features.codes
     inner_dim = codes.shape[1]
@@ -596,8 +662,8 @@ def visualize_feature_heatmaps(
 
     spatial = int(math.isqrt(patches_per_image))
     assert spatial * spatial == patches_per_image, (
-        f"Число патчей {patches_per_image} не является точным квадратом — "
-        f"проверьте crop_size или слой модели"
+        f"Number of patches {patches_per_image} is not a perfect square — "
+        f"check crop_size or model layer"
     )
 
     preprocess = transforms.Compose(
@@ -618,11 +684,13 @@ def visualize_feature_heatmaps(
         patches_per_image=patches_per_image,
         crop_size=crop_size,
         img_size_inches=img_size_inches,
-        show_diff=show_diff,
+        show_mask=show_mask,
         save_dir=save_dir,
         file_prefix=file_prefix,
         show_img=show_img,
         preprocess=preprocess,
+        dataset_root=dataset_root,
+        dataset=dataset,
     )
 
 
@@ -636,12 +704,14 @@ def visualize_feature_heatmaps_from_images(
     patches_per_image: Optional[int] = None,
     crop_size: int = 224,
     img_size_inches: float = 4.0,
-    show_diff: bool = True,
+    show_mask: bool = True,
     save_dir: Optional[str] = None,
     file_prefix: str = '',
     show_img: bool = True,
     *,
     preprocess: Optional[object] = None,
+    dataset_root: Optional[str] = None,
+    dataset: Optional[str] = None,
 ) -> List[Dict[str, str]]:
     """Same as `visualize_feature_heatmaps`, but accepts already opened PIL images.
 
@@ -659,8 +729,8 @@ def visualize_feature_heatmaps_from_images(
 
     spatial = int(math.isqrt(patches_per_image))
     assert spatial * spatial == patches_per_image, (
-        f"Число патчей {patches_per_image} не является точным квадратом — "
-        f"проверьте crop_size или слой модели"
+        f"Number of patches {patches_per_image} is not a perfect square — "
+        f"check crop_size or model layer"
     )
 
     if preprocess is None:
@@ -677,9 +747,10 @@ def visualize_feature_heatmaps_from_images(
     meta_rows = [meta_lookup.get(int(img_idx)) for img_idx in image_indices]
 
     orig_tensors: List[Optional[torch.Tensor]] = []
-    mask_paths: List[Optional[str]] = []
-    has_qground_masks = False
-    if show_diff:
+    mask_rgb_list: List[Optional[np.ndarray]] = []
+    has_annotation_masks = False
+    mask_dataset = dataset
+    if show_mask:
         for img_name, meta_row in zip(image_names, meta_rows):
             orig_path = _get_original_path(img_name) or _meta_path_value(meta_row, 'original_img_path')
             if orig_path and Path(orig_path).exists():
@@ -687,10 +758,21 @@ def visualize_feature_heatmaps_from_images(
             else:
                 orig_tensors.append(None)
 
-            mask_path = _meta_path_value(meta_row, 'mask_path')
-            if mask_path and Path(mask_path).exists():
-                has_qground_masks = True
-            mask_paths.append(mask_path)
+            row_dataset = mask_dataset or infer_spatial_mask_dataset(meta_row)
+            mask_rgb = None
+            if dataset_root is not None and row_dataset in {'QGround', 'SRGround'}:
+                mask_rgb = annotation_mask_rgb_for_meta_row(
+                    meta_row,
+                    dataset_root=dataset_root,
+                    dataset=row_dataset,
+                    crop_size=crop_size,
+                    full_frame=True,
+                )
+            if mask_rgb is not None:
+                has_annotation_masks = True
+            mask_rgb_list.append(mask_rgb)
+        if mask_dataset is None and meta_rows:
+            mask_dataset = infer_spatial_mask_dataset(meta_rows[0]) or 'QGround'
 
     image_idx_arr = meta['image_idx'].values
     codes_list = []
@@ -728,43 +810,44 @@ def visualize_feature_heatmaps_from_images(
             feature_id=int(global_feature_id),
             saved=saved_overlay,
         )
-        if show_diff and any(t is not None for t in orig_tensors):
-            diff_name = f'{safe_prefix}feature_{global_feature_id}_diff.png'
-            diff_path = str(save_root / diff_name) if save_root is not None else None
-            saved_diff = _plot_diff_grid(
-                imgs_list,
-                orig_tensors,
+        if show_mask and has_annotation_masks:
+            mask_name = f'{safe_prefix}feature_{global_feature_id}_mask.png'
+            mask_save_path = str(save_root / mask_name) if save_root is not None else None
+            saved_mask = _plot_annotation_mask_grid(
+                mask_rgb_list,
                 global_feature_id,
                 image_names,
-                img_size_inches,
+                dataset=str(mask_dataset or 'QGround'),
+                crop_size=crop_size,
                 meta_rows=meta_rows,
-                save_path=diff_path,
+                img_size_inches=img_size_inches,
+                save_path=mask_save_path,
                 show_img=show_img,
             )
             _append_heatmap_artifact(
                 artifacts,
-                kind='diff',
+                kind='mask',
                 feature_id=int(global_feature_id),
-                saved=saved_diff,
+                saved=saved_mask,
             )
-        elif show_diff and has_qground_masks:
-            mask_name = f'{safe_prefix}feature_{global_feature_id}_mask.png'
-            mask_path = str(save_root / mask_name) if save_root is not None else None
-            saved_mask = _plot_qground_mask_grid(
-                mask_paths,
+        elif show_mask and any(t is not None for t in orig_tensors):
+            diff_name = f'{safe_prefix}feature_{global_feature_id}_diff.png'
+            diff_path = str(save_root / diff_name) if save_root is not None else None
+            saved_diff = _plot_diff_grid(
+                imgs_list,
+                orig_tensors,
                 global_feature_id,
                 image_names,
+                img_size_inches,
                 meta_rows=meta_rows,
-                img_size_inches=img_size_inches,
-                crop_size=crop_size,
-                save_path=mask_path,
+                save_path=diff_path,
                 show_img=show_img,
             )
             _append_heatmap_artifact(
                 artifacts,
-                kind='mask',
+                kind='diff',
                 feature_id=int(global_feature_id),
-                saved=saved_mask,
+                saved=saved_diff,
             )
 
     return artifacts

analysis/viz/vis_scatter.py

@@ -1,6 +1,6 @@
 """
-Scatter-plot активаций SAE-признаков, часть функций тоже из PatchSAE:
-log(sparsity) × log(mean_acts), окрашенные по энтропии меток искажений.
+Scatter plot of SAE feature activations; some functions are also from PatchSAE:
+log(sparsity) × log(mean_acts), colored by distortion-label entropy.
 """
 
 from pathlib import Path
@@ -56,20 +56,20 @@ def get_top_k_patches(
     label_col: str = 'dist_group',
 ) -> Tuple[torch.Tensor, torch.Tensor, List[str]]:
     """
-    Для каждого признака SAE выбирает top-K патчей с наибольшими активациями.
+    For each SAE feature, select the top-K patches with the largest activations.
 
-    Параметры
+    Parameters
     ----------
-    codes     : CSR матрица (n_patches, n_features)
-    meta      : DataFrame с колонкой label_col
-    top_k     : число патчей; None — все ненулевые (с паддингом до max nnz)
-    label_col : 'dist_group' или 'dist_type'
+    codes     : CSR matrix (n_patches, n_features)
+    meta      : DataFrame with a label_col column
+    top_k     : number of patches; None — all nonzero (padded to max nnz)
+    label_col : 'dist_group' or 'dist_type'
 
-    Возвращает
-    ----------
+    Returns
+    -------
     top_val     : FloatTensor (n_features, effective_k)
     top_label   : LongTensor  (n_features, effective_k)
-    label_names : List[str] — расшифровка числовых меток
+    label_names : List[str] — mapping for numeric labels
     """
     n_patches, n_features = codes.shape
 
@@ -118,19 +118,19 @@ def prepare_scatter_stats(
     group_filter: Optional[str] = None,
 ) -> Tuple[Dict[str, torch.Tensor], List[str]]:
     """
-    Собирает словарь статистик, совместимый с get_stats_scatter_plot.
-
-    Параметры
-    ----------
-    codes        : CSR матрица (n_patches, n_features)
-    meta         : DataFrame с метаданными
-    top_k        : число top-патчей для расчёта энтропии
-    label_col    : 'dist_group' или 'dist_type'
-    group_filter : если задана — фильтрует строки по dist_group == group_filter
+    Build a stats dict compatible with get_stats_scatter_plot.
 
-    Возвращает
+    Parameters
     ----------
-    stats       : dict с ключами 'sparsity', 'mean_acts', 'top_entropy' — FloatTensors
+    codes        : CSR matrix (n_patches, n_features)
+    meta         : metadata DataFrame
+    top_k        : number of top patches for entropy computation
+    label_col    : 'dist_group' or 'dist_type'
+    group_filter : if set, filter rows by dist_group == group_filter
+
+    Returns
+    -------
+    stats       : dict with keys 'sparsity', 'mean_acts', 'top_entropy' — FloatTensors
     label_names : List[str]
     """
     if group_filter is not None:
@@ -198,10 +198,10 @@ def _compute_color_metric_values(color_metric: str,
                                  *,
                                  dataset: str) -> torch.Tensor:
     '''
-    Вычисляет вектор значений метрики для окрашивания точек на scatter plot.
-    Поддерживает 'entropy', 'roc_auc', 'iou', 'precision' и 'recall'.
-    Для табличных метрик возвращает вектор длины n_features,
-    агрегируя значения по категориям из label_col.
+    Compute a metric vector for coloring scatter-plot points.
+    Supports 'entropy', 'roc_auc', 'iou', 'precision', and 'recall'.
+    For tabular metrics, returns a vector of length n_features,
+    aggregating values across categories from label_col.
     '''
     color_metric = str(color_metric).strip().lower()
     if color_metric == 'entropy':

assets/style.css

@@ -712,16 +712,20 @@ body {
 }
 
 .feature-jump-input .Select-control,
+input.feature-jump-input,
 .feature-jump-input input {
     min-height: 42px;
 }
 
 /* Hide +/- steppers in feature id number input (browser-native spinners). */
+input.feature-jump-input[type="number"],
 .feature-jump-input input[type="number"] {
     -moz-appearance: textfield;
     appearance: textfield;
 }
 
+input.feature-jump-input[type="number"]::-webkit-outer-spin-button,
+input.feature-jump-input[type="number"]::-webkit-inner-spin-button,
 .feature-jump-input input[type="number"]::-webkit-outer-spin-button,
 .feature-jump-input input[type="number"]::-webkit-inner-spin-button {
     -webkit-appearance: none;

dashboard/__init__.py

@@ -1,3 +1,13 @@
 """Dash application helpers for xIQA."""
 
-from .model_catalog import MODEL_FAMILIES, ModelRecord, discover_models_for_metric, discover_supported_datasets
+from .model_catalog import (
+    MODEL_FAMILIES,
+    ModelRecord,
+    checkpoint_cache_dir,
+    checkpoint_dataset_cache_paths,
+    discover_models_for_metric,
+    discover_supported_datasets,
+    get_model_record,
+    require_model_record,
+    selector_cache_files_for_model,
+)

dashboard/image_utils.py

@@ -7,7 +7,7 @@ from typing import List
 from analysis.datasets import resolve_dataset_image_path
 from analysis.utils import get_top_images_for_feature, get_top_images_for_feature_by_iou
 from analysis.viz.vis_heatmaps import render_top_feature_panel_srground, visualize_feature_heatmaps
-from dashboard.model_catalog import discover_models_for_metric, load_and_filter_model_activations
+from dashboard.model_catalog import get_model_record, load_and_filter_model_activations
 
 ROOT = Path(__file__).resolve().parents[1]
 
@@ -52,12 +52,12 @@ def get_top_feature_overlays(
     """
     ranking_mode = str(ranking_mode or "iou")
 
-    record = next((r for r in discover_models_for_metric(metric) if r.model_key == model_key), None)
+    record = get_model_record(metric, model_key)
     if record is None:
         return []
 
     try:
-        filtered = load_and_filter_model_activations(Path(record.model_path), selection_dataset)
+        filtered = load_and_filter_model_activations(record.checkpoint_path, selection_dataset)
     except Exception as exc:
         print(f"Error occurred while loading activations for {selection_dataset!r}: {exc}")
         return []
@@ -92,7 +92,6 @@ def get_top_feature_overlays(
         return []
 
     first_rows = meta.groupby("image_idx", sort=False).first()
-    datasets_root = str(dataset_root) if dataset_root is not None else None
     image_paths = []
     for idx in top_image_idxs:
         try:
@@ -115,7 +114,7 @@ def get_top_feature_overlays(
         resolved = resolve_dataset_image_path(
             selection_dataset,
             img_path,
-            datasets_root=datasets_root,
+            datasets_root=str(dataset_root) if dataset_root is not None else None,
         )
         image_paths.append(str(resolved))
 
@@ -135,8 +134,7 @@ def get_top_feature_overlays(
             patches_per_image=None,
             crop_size=224,
             img_size_inches=3.5,
-            include_sr_artifact=INCLUDE_SRGROUND_SR_ARTIFACT,
-            datasets_root=str(dataset_root) if dataset_root is not None else None,
+            dataset_root=dataset_root,
         )
         if panel_bytes is None:
             return []
@@ -158,7 +156,7 @@ def get_top_feature_overlays(
         patches_per_image=None,
         crop_size=224,
         img_size_inches=3.5,
-        show_diff=False,
+        show_mask=False,
         save_dir=None,
         file_prefix="topmax",
         show_img=False,

dashboard/model_catalog.py

@@ -1,18 +1,25 @@
 from __future__ import annotations
 
 import ast
+import json
 import os
 import re
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Iterable, Mapping
+from typing import Any, Iterable, Mapping, Sequence
 
 import numpy as np
 from functools import lru_cache
 import pandas as pd
 import scipy.sparse as sp
 
-from analysis.cache_utils import load_cache
+from analysis.cache_utils import load_cache, zero_codes_outside_activation_steps
+from analysis.config import (
+    SUPPORTED_DATASETS,
+    DatasetCachePaths,
+    build_dataset_cache_paths,
+    load_sae_vis_config,
+)
 from analysis.features.feature_filters import (
     build_filter,
 )
@@ -21,7 +28,7 @@ from analysis.features.feature_indexing import FeatureMatrix
 
 PROJECT_ROOT = Path(__file__).resolve().parents[1]
 DEFAULT_MODEL_STORAGE = PROJECT_ROOT / "logs"
-MODEL_FAMILIES = ("MANIQA", "ARNIQA", "LIQE")
+MODEL_FAMILIES = ("ARNIQA", "MANIQA", "LIQE")
 METRIC_DESCRIPTIONS: dict[str, tuple[str, ...]] = {
     "MANIQA": (
         "No-reference IQA with two Swin Transformer stages: stage 1 (768-d) and stage 2 (384-d). "
@@ -41,8 +48,6 @@ METRIC_DESCRIPTIONS: dict[str, tuple[str, ...]] = {
         "l0 is the validation mean count of active SAE features per image (higher = less sparse).",
     ),
 }
-SUPPORTED_DATASETS = ("SRGround", "local_kadid", "kadid10k")
-
 _FAMILY_ROOT_CANDIDATES = {
     "MANIQA": ("maniqa_logs", "MANIQA_logs", "maniqa"),
     "ARNIQA": ("arniqa_logs", "ARNIQA_logs", "arniqa"),
@@ -63,7 +68,8 @@ _MODEL_DIR_PATTERN = re.compile(
 class ModelRecord:
     family: str
     model_name: str
-    model_path: str   
+    model_path: str
+    checkpoint_path: str
     experiment_num: int
     layer_num: int | None
     lambda_token: str | None
@@ -234,10 +240,27 @@ def discover_models_for_metric(metric: str, model_storage: Path | None = None) -
     )
 
 
+def _latest_checkpoint_dir(model_root: Path) -> Path | None:
+    candidates: list[tuple[int, Path]] = []
+    for entry in model_root.iterdir():
+        if not entry.is_dir() or not entry.name.startswith('checkpoint-'):
+            continue
+        try:
+            step = int(entry.name.removeprefix('checkpoint-'))
+        except ValueError:
+            continue
+        candidates.append((step, entry))
+    if not candidates:
+        return None
+    candidates.sort(key=lambda item: item[0])
+    return candidates[-1][1]
+
+
 @lru_cache(maxsize=32)
 def _discover_models_for_metric_cached(metric_name: str, storage_root_str: str) -> tuple[ModelRecord, ...]:
     storage_root = Path(storage_root_str)
     family_root = resolve_family_root(metric_name, model_storage=storage_root)
+
     if not family_root.exists():
         return ()
 
@@ -259,11 +282,16 @@ def _discover_models_for_metric_cached(metric_name: str, storage_root_str: str)
         except Exception:
             l0_loss_value = None
 
+        checkpoint_dir = _latest_checkpoint_dir(model_dir)
+        if checkpoint_dir is None:
+            continue
+
         records.append(
             ModelRecord(
                 family=metric_name,
                 model_name=str(model_dir.name),
-                model_path=str(model_dir),   # <-- FIX
+                model_path=str(model_dir),
+                checkpoint_path=str(checkpoint_dir.resolve()),
                 experiment_num=int(parsed_name["experiment_num"] or -1),
                 layer_num=parsed_name["layer_num"],
                 lambda_token=parsed_name["lambda_token"],
@@ -276,7 +304,33 @@ def _discover_models_for_metric_cached(metric_name: str, storage_root_str: str)
 
     records.sort(key=lambda r: (r.experiment_num, r.model_name), reverse=True)
 
-    return tuple(records)   # <-- FIX (immutable)
+    return tuple(records)  
+
+
+def get_model_record(
+    metric: str,
+    model_key: str,
+    model_storage: Path | None = None,
+) -> ModelRecord | None:
+    return next(
+        (
+            record
+            for record in discover_models_for_metric(metric, model_storage)
+            if record.model_key == model_key
+        ),
+        None,
+    )
+
+
+def require_model_record(
+    metric: str,
+    model_key: str,
+    model_storage: Path | None = None,
+) -> ModelRecord:
+    record = get_model_record(metric, model_key, model_storage)
+    if record is None:
+        raise FileNotFoundError(f'Model {model_key!r} not found for metric {metric!r}')
+    return record
 
 
 def build_model_option_label(record: ModelRecord) -> str:
@@ -297,25 +351,38 @@ def summarize_model_record(record: ModelRecord) -> str:
     return "\n • ".join(params_split)
 
 
-def discover_selector_cache_files(model_path: Path, dataset: str) -> dict[str, Path]:
+def checkpoint_cache_dir(checkpoint_path: str | Path) -> Path:
+    return Path(checkpoint_path) / 'cache'
+
+
+def checkpoint_dataset_cache_paths(checkpoint_path: str | Path, dataset: str) -> DatasetCachePaths:
+    dataset_name = str(dataset).strip()
+    if dataset_name not in SUPPORTED_DATASETS:
+        raise ValueError(
+            f'Unsupported dataset={dataset_name!r}; expected one of {SUPPORTED_DATASETS}'
+        )
+    cache_dir = checkpoint_cache_dir(checkpoint_path)
+    return build_dataset_cache_paths(str(cache_dir), SUPPORTED_DATASETS)[dataset_name]
+
+
+def selector_cache_files_for_model(checkpoint_path: Path | str, dataset: str) -> dict[str, Path]:
     dataset_name = str(dataset).strip()
     if not dataset_name:
         return {}
 
+    dataset_dir = checkpoint_cache_dir(checkpoint_path) / dataset_name
+    if not dataset_dir.is_dir():
+        return {}
+
+    patch_only = dataset_name.lower() in ('local_kadid', 'srground')
+    parquet_paths = sorted(dataset_dir.glob('*.parquet'))
+    patch_stems = {path.stem for path in parquet_paths if '_patch' in path.stem}
     cache_files: dict[str, Path] = {}
-    # For some datasets we only want patch-level metrics
-    patch_only = dataset_name.lower() in ("local_kadid", "srground")
-    for cache_root_name in ("cache", "activation_cache"):
-        for cache_dir in sorted(model_path.glob(f"**/{cache_root_name}/{dataset_name}")):
-            if not cache_dir.is_dir():
-                continue
-            for parquet_path in sorted(cache_dir.glob("*.parquet")):
-                stem = parquet_path.stem
-                if patch_only and "_patch" not in stem:
-                    # skip non-patch tables for these datasets
-                    continue
-                cache_files.setdefault(stem, parquet_path)
-    
+    for parquet_path in parquet_paths:
+        stem = parquet_path.stem
+        if patch_only and '_patch' not in stem and f'{stem}_patch' in patch_stems:
+            continue
+        cache_files[stem] = parquet_path
     return cache_files
 
 
@@ -326,7 +393,7 @@ def summarize_selector_cache_file(cache_path: Path) -> tuple[float | None, float
     aggregate those per-feature scores across all features. Returns
     (max, mean, variance, min) computed over per-feature scores.
     """
-    feature_table = load_selector_cache_feature_table(cache_path)
+    feature_table = load_selector_cache_feature_table(str(cache_path))
     if feature_table.empty:
         return None, None, None, None
 
@@ -338,7 +405,8 @@ def summarize_selector_cache_file(cache_path: Path) -> tuple[float | None, float
     return float(arr.max()), float(arr.mean()), float(arr.var()), float(arr.min())
 
 
-def load_selector_cache_feature_table(cache_path: Path) -> pd.DataFrame:
+@lru_cache(maxsize=256)
+def load_selector_cache_feature_table(cache_path_str: str) -> pd.DataFrame:
     """Return ordered per-feature selector scores with their feature ids.
 
     The returned frame contains:
@@ -346,18 +414,12 @@ def load_selector_cache_feature_table(cache_path: Path) -> pd.DataFrame:
     - feature_score: max absolute score for that feature column
     - feature_rank: 1-based rank after sorting by score descending, then feature id ascending
     """
-    # Delegate to cached loader keyed by the cache path string to avoid
-    # repeated expensive parquet reads when users rapidly switch features.
-    return _load_selector_cache_feature_table_cached(str(cache_path))
-
-
-@lru_cache(maxsize=256)
-def _load_selector_cache_feature_table_cached(cache_path_str: str) -> pd.DataFrame:
     columns = ["feature_id", "feature_score", "feature_rank"]
-    cache_path = Path(cache_path_str)
-    if not cache_path.exists():
+    path = Path(cache_path_str)
+    if not path.exists():
         return pd.DataFrame(columns=columns)
-    table = pd.read_parquet(cache_path)
+
+    table = pd.read_parquet(path)
     numeric = table.select_dtypes(include=["number"])
     if numeric.empty:
         return pd.DataFrame(columns=columns)
@@ -389,7 +451,7 @@ def load_selector_cache_values(cache_path: Path) -> np.ndarray:
 
     Each value is the max absolute value for one numeric feature column.
     """
-    feature_table = load_selector_cache_feature_table(cache_path)
+    feature_table = load_selector_cache_feature_table(str(cache_path))
     if feature_table.empty:
         return np.array([], dtype=float)
 
@@ -397,41 +459,47 @@ def load_selector_cache_values(cache_path: Path) -> np.ndarray:
     return values[np.isfinite(values)]
 
 
-def summarize_selector_cache(model_path: Path, dataset: str, exclude_prefixes: list[str] | None = None) -> list[SelectorCacheSummary]:
+def summarize_selector_cache(
+    checkpoint_path: Path | str,
+    dataset: str,
+    exclude_prefixes: list[str] | None = None,
+) -> list[SelectorCacheSummary]:
     if exclude_prefixes is None:
         exclude_prefixes = list(DEFAULT_SELECTOR_HIDE_PREFIXES)
-    return _summarize_selector_cache_cached(str(model_path), dataset, tuple(exclude_prefixes))
+    return _summarize_selector_cache_cached(str(checkpoint_path), dataset, tuple(exclude_prefixes))
 
 
 @lru_cache(maxsize=256)
 def _summarize_selector_cache_cached(
-    model_path_str: str,
+    checkpoint_path_str: str,
     dataset: str,
     exclude_prefixes: tuple[str, ...],
 ) -> tuple[SelectorCacheSummary, ...]:
 
-    cache_files = discover_selector_cache_files(Path(model_path_str), dataset)
+    cache_files = selector_cache_files_for_model(checkpoint_path_str, dataset)
 
     summaries: list[SelectorCacheSummary] = []
-
     for selector_name in sorted(cache_files):
         if any(selector_name.startswith(p) for p in exclude_prefixes):
             continue
 
         cache_path = cache_files[selector_name]
-    
+
         max_v, mean_v, var_v, min_v = summarize_selector_cache_file(cache_path)
 
         summaries.append(
+            SelectorCacheSummary(
                 selector_name=selector_name,
-                cache_path=str(cache_path),  
+                cache_path=str(cache_path),  # <-- FIX
                 max_value=max_v,
                 mean_value=mean_v,
                 variance_value=var_v,
                 min_value=min_v,
             )
+        )
     return tuple(summaries)
 
+
 def format_selector_cache_summary(summaries: list[SelectorCacheSummary]) -> list[str]:
     if not summaries:
         return ["No selector cache files were found for this model."]
@@ -451,54 +519,102 @@ def format_selector_cache_summary(summaries: list[SelectorCacheSummary]) -> list
     return lines
 
 
-def _discover_cache_dir_for_dataset(model_path: Path, dataset: str) -> Path | None:
-    dataset_name = str(dataset).strip()
-    if not dataset_name:
-        return None
+def _acts_meta_path(acts_cache_path: Path) -> Path:
+    stem = acts_cache_path.stem.removesuffix('.feather')
+    return acts_cache_path.parent / f'{stem}_meta.feather'
+
+
+@lru_cache(maxsize=1)
+def _dashboard_default_params():
+    return load_sae_vis_config()
+
+
+def _filters_to_cache_key(requested_filters: list[dict]) -> str:
+    return json.dumps(requested_filters, sort_keys=True, default=str)
+
+
+def _activation_steps_cache_key(activation_steps_to_keep: Sequence[int]) -> str:
+    return json.dumps([int(step) for step in activation_steps_to_keep])
+
+
+def clear_activation_load_cache() -> None:
+    """Drop in-process activation caches (e.g. after rebuilding cache files on disk)."""
+    _load_and_filter_model_activations_cached.cache_clear()
+
 
-    # Prefer the most specific/most-recent cache directory
-    candidates = sorted(model_path.glob(f"**/cache/{dataset_name}"), key=lambda p: str(p), reverse=True)
-    return candidates[0] if candidates else None
-
-
-def _resolve_activation_cache_base_path(cache_dir: Path) -> Path | None:
-    meta_paths = sorted(cache_dir.glob('*_meta.feather'), key=lambda p: p.name, reverse=True)
-    for meta_path in meta_paths:
-        base_name = meta_path.name.removesuffix('_meta.feather')
-        codes_path = meta_path.with_name(f'{base_name}_codes.npz')
-        if codes_path.exists():
-            return meta_path.with_name(base_name)
-    return None
-
-DEFAULT_FILTERS = [        {
-            'name': 'nonzero_max',
-            'params': {},
-        },
-        {
-            'name': 'kruskal_wallis',
-            'params': {
-                'alpha': 0.05,
-                'group_col': 'dist_type',
-                'min_group_size': 3,
-            },
-        }]
-
-def load_and_filter_model_activations(model_path: Path, 
-                                      dataset: str, 
-                                      requested_filters: list[dict] = DEFAULT_FILTERS,
-                                      min_distortion_level = 3,
-                                      max_distortion_level = 5) -> FilteredActivationCache:
-    cache_dir = _discover_cache_dir_for_dataset(model_path, dataset)
-    if cache_dir is None:
-        raise FileNotFoundError(f'No activation cache directory found for dataset={dataset!r} in {model_path}')
-
-    cache_dir = _discover_cache_dir_for_dataset(model_path, dataset)
-    if cache_dir is None:
-        raise FileNotFoundError(f'No activation cache directory found for dataset={dataset!r} in {model_path}')
-
-    cache_base_path = _resolve_activation_cache_base_path(cache_dir)
-    if cache_base_path is None:
-        raise FileNotFoundError(f'No activation cache files found in {cache_dir}')
+@lru_cache(maxsize=32)
+def _load_and_filter_model_activations_cached(
+    checkpoint_path: str,
+    dataset: str,
+    filters_key: str,
+    min_distortion_level: int,
+    max_distortion_level: int,
+    activation_steps_key: str,
+) -> FilteredActivationCache:
+    requested_filters = json.loads(filters_key)
+    activation_steps_to_keep = [int(step) for step in json.loads(activation_steps_key)]
+    return _load_and_filter_model_activations_impl(
+        checkpoint_path,
+        dataset,
+        requested_filters=requested_filters,
+        min_distortion_level=min_distortion_level,
+        max_distortion_level=max_distortion_level,
+        activation_steps_to_keep=activation_steps_to_keep,
+    )
+
+
+def load_and_filter_model_activations(
+    checkpoint_path: Path | str,
+    dataset: str,
+    *,
+    requested_filters: list[dict] | None = None,
+    min_distortion_level: int | None = None,
+    max_distortion_level: int | None = None,
+    activation_steps_to_keep: Sequence[int] | None = None,
+) -> FilteredActivationCache:
+    default_cfg = _dashboard_default_params()
+    filters = list(default_cfg.FEATURE_FILTERS if requested_filters is None else requested_filters)
+    min_level = int(
+        min_distortion_level
+        if min_distortion_level is not None
+        else default_cfg.KADID_MIN_DISTORTION_LEVEL
+    )
+    max_level = int(
+        max_distortion_level
+        if max_distortion_level is not None
+        else default_cfg.KADID_MAX_DISTORTION_LEVEL
+    )
+    keep_steps = list(
+        activation_steps_to_keep
+        if activation_steps_to_keep is not None
+        else default_cfg.ACTIVATION_STEPS_TO_KEEP
+    )
+    return _load_and_filter_model_activations_cached(
+        str(Path(checkpoint_path).resolve()),
+        str(dataset),
+        _filters_to_cache_key(filters),
+        min_level,
+        max_level,
+        _activation_steps_cache_key(keep_steps),
+    )
+
+
+def _load_and_filter_model_activations_impl(
+    checkpoint_path: str,
+    dataset: str,
+    *,
+    requested_filters: list[dict],
+    min_distortion_level: int,
+    max_distortion_level: int,
+    activation_steps_to_keep: Sequence[int],
+) -> FilteredActivationCache:
+    cache_paths = checkpoint_dataset_cache_paths(checkpoint_path, dataset)
+    acts_cache_path = Path(cache_paths.acts_cache_path)
+    meta_path = _acts_meta_path(acts_cache_path)
+    if not meta_path.exists():
+        raise FileNotFoundError(f'Activation cache not found: {meta_path}')
+
+    cache_base_path = acts_cache_path
 
     meta_df, codes_csr, activation_steps_csr = load_cache(
         str(cache_base_path),
@@ -506,8 +622,13 @@ def load_and_filter_model_activations(model_path: Path,
         min_distortion_level=min_distortion_level,
         max_distortion_level=max_distortion_level
     )
+    codes_csr = zero_codes_outside_activation_steps(
+        codes_csr,
+        activation_steps_csr,
+        activation_steps_to_keep,
+    )
 
-    feature_filter_cache_dir = cache_dir / '../feature_filter_cache'
+    feature_filter_cache_dir = Path(cache_paths.cache_dir) / 'feature_filter_cache'
     feature_filter_summary: list[dict[str, Any]] = []
 
     if not requested_filters:
@@ -578,7 +699,7 @@ def load_and_filter_model_activations(model_path: Path,
 
 
 @lru_cache(maxsize=256)
-def summarize_feature_filter_cache(model_path: Path, dataset: str) -> list[str]:
+def summarize_feature_filter_cache(checkpoint_path: Path | str, dataset: str) -> list[str]:
     """Return human-readable summary lines for feature filtering.
 
     The model cache is loaded first, then configured feature filters are applied
@@ -587,7 +708,7 @@ def summarize_feature_filter_cache(model_path: Path, dataset: str) -> list[str]:
     ``load_and_filter_model_activations()`` for downstream use.
     """
     try:
-        filtered_cache = load_and_filter_model_activations(Path(model_path), dataset)
+        filtered_cache = load_and_filter_model_activations(checkpoint_path, dataset)
     except Exception as exc:
         return [f'Feature filtering unavailable: {exc}']
 

dashboard/umap_service.py

@@ -11,14 +11,29 @@ import pandas as pd
 import plotly.graph_objects as go
 
 from analysis.cache_utils import load_pristine_cache, zero_codes_outside_activation_steps
-from analysis.config import dataset_images_root, load_sae_vis_config
+from analysis.config import dataset_images_root
 from analysis.datasets import Kadid10kDataset
-from analysis.features.feature_matrix import build_image_feature_matrix, build_image_feature_matrix_raw
+from analysis.features.feature_matrix import (
+    build_image_feature_matrix,
+    build_image_feature_matrix_from_image_idx,
+    build_image_feature_matrix_raw,
+)
 from analysis.features.feature_indexing import FeatureMatrix, global_id_set, global_to_column
 from analysis.metrics.correlations import compute_distortion_correlations
 from analysis.viz.umap_plot import build_umap_figure, make_thumb_b64 as _make_thumb_b64
-from analysis.viz.umap_utils import get_or_compute_umap_with_builder
-from dashboard.model_catalog import discover_models_for_metric, load_and_filter_model_activations
+from analysis.viz.umap_utils import (
+    UMAP_IMAGE_LIMIT,
+    UMAP_IMAGE_RANDOM_STATE,
+    build_kadid_index_lookup,
+    get_or_compute_umap_with_builder,
+    select_random_image_indices,
+)
+from dashboard.model_catalog import (
+    _dashboard_default_params,
+    get_model_record,
+    load_and_filter_model_activations,
+    require_model_record,
+)
 import scipy.sparse as sp
 from analysis.models import load_sae
 
@@ -42,7 +57,6 @@ CORR_GARBAGE_THRESHOLD = 0.20
 
 # Activation cache and image previews for dashboard UMAP (images + features modes).
 UMAP_DATASET = 'kadid10k'
-
 UMAP_PARAMS = dict(
     n_neighbors=15,
     min_dist=0.1,
@@ -62,6 +76,8 @@ class UmapContext:
     model_key: str
     features: FeatureMatrix
     image_row_indices: tuple[tuple[int, ...], ...]
+    image_idx_arr: np.ndarray
+    selected_image_idx: tuple[int, ...]
     n_images_used: int
     n_features_total: int
     dist_type_arr: np.ndarray
@@ -69,32 +85,36 @@ class UmapContext:
     dist_level_arr: np.ndarray
     mos_arr: np.ndarray
     umap_cache_dir: Path
+    dataset_cache_dir: Path
     kadid_ds: Any | None
     meta_df: pd.DataFrame
     use_reference_delta: bool
+    delta_codes: Any | None = None
     pristine_meta_df: pd.DataFrame | None = None
     pristine_features: FeatureMatrix | None = None
     kadid_image_idx: np.ndarray | None = None
 
 
-def _resolve_checkpoint_path(model_path: str | Path) -> Path:
-    model_root = Path(model_path)
-    if model_root.is_file():
-        return model_root
-
-    candidate_dirs: list[Path] = []
-    for pattern in ('checkpoint-*', 'checkpoints/checkpoint-*', 'checkpoint'):
-        candidate_dirs.extend(p for p in model_root.glob(pattern) if p.is_dir())
-
-    if not candidate_dirs:
-        return model_root
-
-    def _checkpoint_sort_key(path: Path) -> tuple[int, str]:
-        digits = ''.join(ch for ch in path.name if ch.isdigit())
-        numeric = int(digits) if digits else -1
-        return numeric, str(path)
-
-    return max(candidate_dirs, key=_checkpoint_sort_key)
+def _resolve_kadid_dataset(
+    dataset: str,
+    n_images_used: int,
+    kadid_images_path: str,
+    crop_size: int,
+    min_distortion_level: int,
+):
+    if 'kadid' not in str(dataset).lower():
+        return None
+    try:
+        kadid_ds = Kadid10kDataset(
+            kadid_images_path,
+            crop_size=crop_size,
+            min_distortion_level=min_distortion_level,
+        )
+        if n_images_used > len(kadid_ds):
+            return None
+        return kadid_ds
+    except Exception:
+        return None
 
 
 def empty_umap_figure(message: str) -> go.Figure:
@@ -137,39 +157,6 @@ def _compact_meta_image_indices(meta_df: pd.DataFrame) -> tuple[pd.DataFrame, in
     return meta_out, int(unique_ids.size)
 
 
-def _kadid_index_lookup(meta_df: pd.DataFrame, kadid_ds: Any, n_images_used: int) -> np.ndarray:
-    """Map dense UMAP image index (0..n-1) to ``Kadid10kDataset`` positional index."""
-    lookup = np.arange(n_images_used, dtype=np.int32)
-    if kadid_ds is None:
-        return lookup
-
-    first_rows = meta_df.groupby('image_idx', sort=True).first()
-    kadid_paths = [str(p) for p in kadid_ds.images]
-    path_to_idx = {p: i for i, p in enumerate(kadid_paths)}
-    base_to_idx = {os.path.basename(p): i for i, p in enumerate(kadid_paths)}
-
-    for dense_i in range(n_images_used):
-        row = first_rows.iloc[dense_i]
-        kadid_idx = None
-        for col in ('distorted_img_path', 'image_path', 'original_img_path'):
-            if col not in row.index:
-                continue
-            val = row[col]
-            if val is None or (isinstance(val, float) and np.isnan(val)):
-                continue
-            path = str(val)
-            kadid_idx = path_to_idx.get(path)
-            if kadid_idx is None and not Path(path).is_absolute():
-                kadid_idx = path_to_idx.get(str(kadid_ds.root / Path(path)))
-            if kadid_idx is None:
-                kadid_idx = base_to_idx.get(os.path.basename(path))
-            if kadid_idx is not None:
-                break
-        if kadid_idx is not None:
-            lookup[dense_i] = int(kadid_idx)
-    return lookup
-
-
 def _image_labels_from_meta(meta_df, n_images_used: int) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
     first_rows = meta_df.groupby('image_idx', sort=True).first()
     expected_index = np.arange(n_images_used, dtype=np.int32)
@@ -271,34 +258,6 @@ def _build_reference_delta_matrix(
     )
 
 
-def _resolve_kadid_dataset(
-    dataset: str,
-    n_images_used: int,
-    kadid_images_path: Optional[str],
-    crop_size: Optional[int],
-    min_distortion_level: Optional[int],
-):
-    if 'kadid' not in str(dataset).lower():
-        return None
-    try:
-        if kadid_images_path is None or crop_size is None or min_distortion_level is None:
-            cfg = load_sae_vis_config()
-            kadid_images_path = kadid_images_path or cfg.KADID_IMAGES_PATH
-            crop_size = crop_size or cfg.CROP_SIZE
-            min_distortion_level = min_distortion_level or cfg.KADID_MIN_DISTORTION_LEVEL
-
-        kadid_ds = Kadid10kDataset(
-            kadid_images_path,
-            crop_size=crop_size,
-            min_distortion_level=min_distortion_level,
-        )
-        if n_images_used > len(kadid_ds):
-            return None
-        return kadid_ds
-    except Exception:
-        return None
-
-
 @lru_cache(maxsize=32)
 def load_umap_context(
     metric: str,
@@ -312,54 +271,48 @@ def load_umap_context(
     crop_size: Optional[int] = None,
     use_reference_delta: bool = True,
 ) -> UmapContext:
-    record = next((r for r in discover_models_for_metric(metric) if r.model_key == model_key), None)
-    if record is None:
-        raise FileNotFoundError(f'Model {model_key!r} not found for metric {metric!r}')
-
-    # Load default config only for missing parameters
-    cfg: Optional[object] = None
-    need_cfg = any(
-        v is None
-        for v in (
-            feature_filters,
-            kadid_min_dist_level,
-            kadid_max_dist_level,
-            activation_steps_to_keep,
-            kadid_images_path,
-            crop_size,
-        )
-    )
-    if need_cfg:
-        cfg = load_sae_vis_config()
+    record = require_model_record(metric, model_key)
 
-    requested_filters = list(feature_filters if feature_filters is not None else cfg.FEATURE_FILTERS)
-    min_level = kadid_min_dist_level if kadid_min_dist_level is not None else cfg.KADID_MIN_DISTORTION_LEVEL
-    max_level = kadid_max_dist_level if kadid_max_dist_level is not None else cfg.KADID_MAX_DISTORTION_LEVEL
+    default_cfg = _dashboard_default_params()
+    keep_steps = list(
+        activation_steps_to_keep
+        if activation_steps_to_keep is not None
+        else default_cfg.ACTIVATION_STEPS_TO_KEEP
+    )
 
     filtered = load_and_filter_model_activations(
-        Path(record.model_path),
+        record.checkpoint_path,
         dataset,
-        requested_filters=requested_filters,
-        min_distortion_level=min_level,
-        max_distortion_level=max_level,
+        requested_filters=feature_filters,
+        min_distortion_level=kadid_min_dist_level,
+        max_distortion_level=kadid_max_dist_level,
+        activation_steps_to_keep=keep_steps,
     )
 
-    keep_steps = activation_steps_to_keep if activation_steps_to_keep is not None else (cfg.ACTIVATION_STEPS_TO_KEEP if cfg is not None else [])
-    codes_csr = zero_codes_outside_activation_steps(
-        filtered.features.codes,
-        filtered.activation_steps_csr,
-        keep_steps,
-    )
-    features = FeatureMatrix(
-        codes=codes_csr,
-        column_feature_ids=filtered.features.column_feature_ids,
-    )
+    features = filtered.features
+    codes_csr = features.codes
 
     meta_df = filtered.meta_df
     if 'image_idx' not in meta_df.columns:
         raise ValueError('Activation cache metadata is missing required column "image_idx"')
 
-    meta_df, n_images_used = _compact_meta_image_indices(meta_df)
+    image_idx_arr = meta_df['image_idx'].to_numpy(dtype=np.int64)
+    max_images = int(image_idx_arr.max()) + 1
+    selected_image_idx = select_random_image_indices(
+        max_images,
+        limit=UMAP_IMAGE_LIMIT,
+        random_state=UMAP_IMAGE_RANDOM_STATE,
+    )
+    row_mask = np.isin(image_idx_arr, selected_image_idx)
+    dense_map = {int(orig): i for i, orig in enumerate(selected_image_idx.tolist())}
+    meta_df = meta_df.loc[row_mask].copy()
+    meta_df['image_idx'] = meta_df['image_idx'].map(lambda v: dense_map[int(v)]).astype(np.int32)
+    codes_csr = codes_csr[row_mask]
+    features = FeatureMatrix(
+        codes=codes_csr,
+        column_feature_ids=features.column_feature_ids,
+    )
+    n_images_used = int(selected_image_idx.size)
     image_idx_arr = meta_df['image_idx'].to_numpy(dtype=np.int32)
     n_features_total = features.n_features
 
@@ -370,12 +323,12 @@ def load_umap_context(
     image_row_indices = _build_image_row_indices(image_idx_arr, n_images_used)
 
     dist_type_arr, dist_group_arr, dist_level_arr, mos_arr = _image_labels_from_meta(meta_df, n_images_used)
-    cache_dir = Path(record.model_path)
-    discovered = sorted(cache_dir.glob(f'**/cache/{dataset}'), key=lambda p: str(p), reverse=True)
-    umap_cache_dir = (discovered[0] / 'umap_cache') if discovered else (cache_dir / 'cache' / dataset / 'umap_cache')
+    dataset_cache_dir = Path(record.checkpoint_path) / 'cache' / dataset
+    umap_cache_dir = Path(record.checkpoint_path) / 'cache' / 'umap_cache'
 
     pristine_meta_df: pd.DataFrame | None = None
     pristine_features: FeatureMatrix | None = None
+    delta_codes: Any | None = None
     if use_reference_delta:
         try:
             pristine_meta_df, pristine_codes, pristine_steps = load_pristine_cache(str(filtered.cache_base_path), return_activation_steps=True)
@@ -391,9 +344,14 @@ def load_umap_context(
                 raise ValueError(
                     f'Pristine and distorted feature dimensions differ: {pristine_features.n_features} != {n_features_total}'
                 )
+            pristine_row_indices = _build_reference_row_indices(meta_df, pristine_meta_df)
+            delta_codes = features.codes - pristine_features.codes[pristine_row_indices]
         except Exception as exc:
             print(f'[umap] reference delta unavailable, falling back to raw features: {exc}')
             use_reference_delta = False
+            pristine_meta_df = None
+            pristine_features = None
+            delta_codes = None
 
     kadid_ds = _kadid_dataset_for_context_cached(
         metric,
@@ -401,7 +359,7 @@ def load_umap_context(
         dataset,
         n_images_used,
     )
-    kadid_image_idx = _kadid_index_lookup(meta_df, kadid_ds, n_images_used)
+    kadid_image_idx = build_kadid_index_lookup(meta_df, kadid_ds, n_images_used)
 
     return UmapContext(
         metric=metric,
@@ -409,6 +367,8 @@ def load_umap_context(
         model_key=model_key,
         features=features,
         image_row_indices=image_row_indices,
+        image_idx_arr=image_idx_arr,
+        selected_image_idx=tuple(int(v) for v in selected_image_idx.tolist()),
         n_images_used=n_images_used,
         n_features_total=n_features_total,
         dist_type_arr=dist_type_arr,
@@ -416,9 +376,11 @@ def load_umap_context(
         dist_level_arr=dist_level_arr,
         mos_arr=mos_arr,
         umap_cache_dir=umap_cache_dir,
+        dataset_cache_dir=dataset_cache_dir,
         kadid_ds=kadid_ds,
         meta_df=meta_df,
         use_reference_delta=use_reference_delta,
+        delta_codes=delta_codes,
         pristine_meta_df=pristine_meta_df,
         pristine_features=pristine_features,
         kadid_image_idx=kadid_image_idx,
@@ -434,15 +396,18 @@ def _kadid_dataset_for_context_cached(
 ) -> Any:
     if 'kadid' not in str(dataset).lower():
         return None
+    record = get_model_record(metric, model_key)
+    if record is None:
+        return None
     try:
-        cfg = load_sae_vis_config()
-        images_root = dataset_images_root(cfg.DATASETS_ROOT, dataset)
+        default_cfg = _dashboard_default_params()
+        images_root = dataset_images_root(default_cfg.DATASETS_ROOT, dataset)
         return _resolve_kadid_dataset(
             dataset,
             n_images_used,
             images_root,
-            cfg.CROP_SIZE,
-            cfg.KADID_MIN_DISTORTION_LEVEL,
+            default_cfg.CROP_SIZE,
+            default_cfg.KADID_MIN_DISTORTION_LEVEL,
         )
     except Exception:
         return None
@@ -464,11 +429,13 @@ def _labels_for_color_mode(context: UmapContext, color_mode: str) -> np.ndarray:
 
 
 @lru_cache(maxsize=16)
-def _resolve_feature_atoms_cached(metric: str, model_key: str) -> tuple[np.ndarray, tuple[int, ...]]:
-    record = next((r for r in discover_models_for_metric(metric) if r.model_key == model_key), None)
-    if record is None:
-        raise FileNotFoundError(f'Model {model_key!r} not found for metric {metric!r}')
-    checkpoint_path = _resolve_checkpoint_path(record.model_path)
+def _resolve_feature_atoms_cached(
+    metric: str,
+    model_key: str,
+    dataset: str,
+) -> tuple[np.ndarray, tuple[int, ...]]:
+    record = require_model_record(metric, model_key)
+    checkpoint_path = Path(record.checkpoint_path)
     sae = load_sae(checkpoint_path=str(checkpoint_path), device='cpu')
     try:
         decoder_weight = sae.decoder.weight
@@ -513,13 +480,17 @@ def _feature_stats_for_ids(
     return means, stds, nonzeros
 
 
+def _feature_umap_cache_signature(metric: str, model_key: str, dataset: str) -> dict[str, object]:
+    """Cache fingerprint for decoder-atom UMAP (matches notebook 07)."""
+    atoms, _aligned = _resolve_feature_atoms_cached(metric, model_key, dataset)
+    return {
+        'n_feature_atoms': int(atoms.shape[0]),
+        'atom_dim': int(atoms.shape[1]),
+    }
+
+
 def _feature_stats(context: UmapContext) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
-    return _feature_stats_cached(
-        context.metric,
-        context.dataset,
-        context.model_key,
-        bool(context.use_reference_delta),
-    )
+    return _feature_stats_cached(context.metric, context.dataset, context.model_key)
 
 
 @lru_cache(maxsize=32)
@@ -527,9 +498,8 @@ def _feature_stats_cached(
     metric: str,
     dataset: str,
     model_key: str,
-    use_reference_delta: bool,
 ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
-    context = load_umap_context(metric, dataset, model_key, use_reference_delta=use_reference_delta)
+    context = load_umap_context(metric, dataset, model_key, use_reference_delta=False)
     return _feature_stats_all_columns(context)
 
 
@@ -544,12 +514,10 @@ def _feature_marker_sizes(mean_feats: np.ndarray) -> np.ndarray:
 
 
 def _feature_similarity_cache_path(context: UmapContext, source: str = 'type') -> str:
-    cfg = load_sae_vis_config()
-    dataset_cfg = cfg.DATASET_CACHE_CONFIGS[context.dataset]
     if source == 'type':
-        return dataset_cfg.corr_type_cache_path
+        return str(context.dataset_cache_dir / 'corr_type.parquet')
     if source == 'group':
-        return dataset_cfg.corr_group_cache_path
+        return str(context.dataset_cache_dir / 'corr_group.parquet')
     raise ValueError(f"source must be 'type' or 'group', got {source!r}")
 
 
@@ -646,24 +614,28 @@ def _feature_top_corr_labels(
     return np.asarray(top_labels, dtype=object)
 
 
-def _cache_key(context: UmapContext, aggregation_mode: str) -> Tuple:
-    return (
-        context.metric,
-        context.dataset,
-        context.model_key,
-        str(aggregation_mode),
-        bool(context.use_reference_delta),
-        tuple(int(fid) for fid in context.features.column_feature_ids),
-        tuple(sorted(UMAP_PARAMS.items())),
-    )
+def _umap_codes_for_signature(context: UmapContext):
+    if context.use_reference_delta and context.delta_codes is not None:
+        return context.delta_codes
+    return context.features.codes
+
+
+def _images_umap_cache_id(context: UmapContext, aggregation_mode: str) -> str:
+    if context.use_reference_delta and context.delta_codes is not None:
+        return f'grid_delta_agg_{aggregation_mode}'
+    return f'grid_raw_agg_{aggregation_mode}'
 
 
 def _umap_cache_signature(context: UmapContext) -> dict[str, object]:
-    signature = {
-        'codes_shape': tuple(int(v) for v in context.features.codes.shape),
+    codes_for_sig = _umap_codes_for_signature(context)
+    signature: dict[str, object] = {
+        'codes_shape': tuple(int(v) for v in codes_for_sig.shape),
         'n_images_used': int(context.n_images_used),
         'n_features_total': int(context.n_features_total),
         'use_reference_delta': bool(context.use_reference_delta),
+        'image_selection': 'random',
+        'random_state': int(UMAP_IMAGE_RANDOM_STATE),
+        'selected_image_idx': [int(v) for v in context.selected_image_idx],
         'feature_ids': [int(fid) for fid in context.features.column_feature_ids],
     }
     if context.pristine_features is not None:
@@ -671,32 +643,21 @@ def _umap_cache_signature(context: UmapContext) -> dict[str, object]:
     return signature
 
 
-def get_or_compute_umap(context: UmapContext, aggregation_mode: str) -> dict[str, object]:
-    key = _cache_key(context, aggregation_mode)
-    cache_id = f'{context.model_key}_{context.dataset}_agg_{aggregation_mode}'
-
-    def _build_features() -> np.ndarray:
-        if context.use_reference_delta and context.pristine_features is not None and context.pristine_meta_df is not None:
-            return _build_reference_delta_matrix(
-                context.features,
-                context.image_row_indices,
-                context.meta_df,
-                context.pristine_features,
-                context.pristine_meta_df,
-                context.n_images_used,
-                aggregation_mode,
-            )
+def _build_images_umap_matrix(context: UmapContext, aggregation_mode: str) -> np.ndarray:
+    codes = _umap_codes_for_signature(context)
+    return build_image_feature_matrix_from_image_idx(
+        codes,
+        context.image_idx_arr,
+        context.n_images_used,
+        aggregation_mode,
+    )
 
-        return build_image_feature_matrix(
-            context.features,
-            context.image_row_indices,
-            context.n_images_used,
-            aggregation_mode,
-        )
+
+def get_or_compute_umap(context: UmapContext, aggregation_mode: str) -> dict[str, object]:
+    cache_id = _images_umap_cache_id(context, aggregation_mode)
 
     return get_or_compute_umap_with_builder(
-        key,
-        _build_features,
+        lambda: _build_images_umap_matrix(context, aggregation_mode),
         UMAP_PARAMS,
         cache_dir=context.umap_cache_dir,
         cache_id=cache_id,
@@ -704,35 +665,26 @@ def get_or_compute_umap(context: UmapContext, aggregation_mode: str) -> dict[str
     )
 
 
-def get_or_compute_feature_umap(context: UmapContext) -> dict[str, object]:
-    key = (
-        context.metric,
-        context.dataset,
-        context.model_key,
-        'features',
-        tuple(sorted(UMAP_PARAMS.items())),
-    )
-    cache_id = f"{context.model_key}_{context.dataset}_features_v3"
+def get_or_compute_feature_umap(
+    metric: str,
+    model_key: str,
+    dataset: str,
+    umap_cache_dir: Path,
+) -> dict[str, object]:
+    cache_id = 'features_grid'
 
     def _build_atoms() -> np.ndarray:
-        atoms, _aligned = _resolve_feature_atoms_cached(context.metric, context.model_key)
+        atoms, _aligned = _resolve_feature_atoms_cached(metric, model_key, dataset)
         return atoms
 
-    def _build_signature() -> dict[str, object]:
-        sig = _umap_cache_signature(context)
-        _atoms, aligned = _resolve_feature_atoms_cached(context.metric, context.model_key)
-        sig['n_feature_atoms'] = len(aligned)
-        return sig
-
     info = get_or_compute_umap_with_builder(
-        key,
         _build_atoms,
         UMAP_PARAMS,
-        cache_dir=context.umap_cache_dir,
+        cache_dir=umap_cache_dir,
         cache_id=cache_id,
-        cache_signature=_build_signature(),
+        cache_signature=_feature_umap_cache_signature(metric, model_key, dataset),
     )
-    _atoms, aligned = _resolve_feature_atoms_cached(context.metric, context.model_key)
+    _atoms, aligned = _resolve_feature_atoms_cached(metric, model_key, dataset)
     expected_n = len(aligned)
     emb = np.asarray(info['embedding_2d'])
     if emb.shape[0] != expected_n:
@@ -760,10 +712,10 @@ def build_feature_umap_figure(
     of the code (including `context`, `feature_ids`, `unique_categories`, and
     cache-related fields returned by the UMAP computation).
     """
-    context = load_umap_context(metric, dataset, model_key)
-    _atoms, aligned_feature_ids = _resolve_feature_atoms_cached(context.metric, context.model_key)
+    context = load_umap_context(metric, dataset, model_key, use_reference_delta=False)
+    _atoms, aligned_feature_ids = _resolve_feature_atoms_cached(context.metric, context.model_key, dataset)
     aligned_feature_ids = list(aligned_feature_ids)
-    info = get_or_compute_feature_umap(context)
+    info = get_or_compute_feature_umap(context.metric, context.model_key, dataset, context.umap_cache_dir)
     emb = np.asarray(info['embedding_2d'])
     n_points = min(int(emb.shape[0]), len(aligned_feature_ids))
     emb = emb[:n_points]
@@ -872,7 +824,7 @@ def build_feature_umap_figure(
 
 @lru_cache(maxsize=32)
 def activation_feature_id_set(metric: str, dataset: str, model_key: str) -> frozenset[int]:
-    context = load_umap_context(metric, dataset, model_key)
+    context = load_umap_context(metric, dataset, model_key, use_reference_delta=False)
     return global_id_set(context.features.column_feature_ids)
 
 
@@ -973,10 +925,8 @@ def make_hover_thumb_b64(context: UmapContext, img_idx: int) -> str | None:
 
 
 def format_feature_umap_status(info: dict[str, object], color_mode: str) -> str:
-    context = info.get('context')
-    feature_mode = 'delta(ref)' if getattr(context, 'use_reference_delta', False) else 'raw'
     corr_source = str(info.get('feature_corr_source', DEFAULT_FEATURE_CORR_SOURCE))
-    return f'mode={feature_mode} | color=corr({corr_source}) | X={info.get("x_shape")}'
+    return f'features=decoder_atoms | color=corr({corr_source}) | X={info.get("x_shape")}'
 
 
 def format_image_color_mode(color_mode: str) -> str:

dashboard/user_image_service.py

@@ -19,7 +19,7 @@ from analysis.models import load_iqa_model, load_sae, read_sae_config
 import analysis.models as _analysis_models
 from analysis.viz.vis_heatmaps import visualize_feature_heatmaps, visualize_feature_heatmaps_from_images
 from dashboard.image_utils import png_bytes_to_data_url
-from dashboard.model_catalog import discover_models_for_metric
+from dashboard.model_catalog import get_model_record, require_model_record
 
 
 PROJECT_ROOT = Path(__file__).resolve().parents[1]
@@ -38,24 +38,6 @@ class FeatureMapRequest:
     upload_filename: str | None = None
 
 
-def _resolve_checkpoint_path(model_path: str | Path) -> Path:
-    model_root = Path(model_path)
-    if model_root.is_file():
-        return model_root
-
-    candidate_dirs: list[Path] = []
-    for pattern in ("checkpoint-*", "checkpoints/checkpoint-*", "checkpoint"):
-        candidate_dirs.extend(p for p in model_root.glob(pattern) if p.is_dir())
-    if not candidate_dirs:
-        return model_root
-
-    def _checkpoint_sort_key(path: Path) -> tuple[int, str]:
-        digits = "".join(ch for ch in path.name if ch.isdigit())
-        numeric = int(digits) if digits else -1
-        return numeric, str(path)
-
-    return max(candidate_dirs, key=_checkpoint_sort_key)
-
 
 def _assets_url(path: Path) -> str:
     rel = path.relative_to(ASSETS_DIR)
@@ -213,12 +195,9 @@ def list_baseline_examples(model_key: str) -> list[str]:
 
 
 @lru_cache(maxsize=16)
-def _load_models(metric: str, model_key: str, device: str):
-    record = next((r for r in discover_models_for_metric(metric) if r.model_key == model_key), None)
-    if record is None:
-        raise FileNotFoundError(f"Model {model_key!r} not found for metric {metric!r}")
-
-    checkpoint_path = _resolve_checkpoint_path(record.model_path)
+def _load_models(metric: str, model_key: str, dataset: str, device: str):
+    record = require_model_record(metric, model_key)
+    checkpoint_path = Path(record.checkpoint_path)
     sae_cfg = read_sae_config(str(checkpoint_path))
 
     layer_num = int(sae_cfg.get("layer_num", 3))
@@ -282,10 +261,11 @@ def _codes_for_single_image(
     *,
     metric: str,
     model_key: str,
+    dataset: str,
     image_path: Path,
 ) -> tuple[pd.DataFrame, FeatureMatrix, int]:
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    models = _load_models(metric, model_key, device)
+    models = _load_models(metric, model_key, dataset, device)
 
     iqa = models["iqa"]
     layer_name = str(models["layer_name"])
@@ -333,6 +313,7 @@ def _resolve_overlay_activations(
 ) -> tuple[pd.DataFrame, FeatureMatrix, int]:
     metric = str(request.metric).strip().upper()
     model_key = str(request.model_key).strip()
+    dataset = str(request.dataset).strip()
 
     meta_df: pd.DataFrame | None = None
     features: FeatureMatrix | None = None
@@ -346,7 +327,7 @@ def _resolve_overlay_activations(
     if meta_df is None or features is None or crop_size is None:
         if upload_image is not None:
             device = "cuda" if torch.cuda.is_available() else "cpu"
-            models = _load_models(metric, model_key, device)
+            models = _load_models(metric, model_key, dataset, device)
             iqa = models["iqa"]
             layer_name = str(models["layer_name"])
             sae = models["sae"]
@@ -383,6 +364,7 @@ def _resolve_overlay_activations(
             meta_df, features, crop_size = _codes_for_single_image(
                 metric=metric,
                 model_key=model_key,
+                dataset=dataset,
                 image_path=image_path,
             )
         if request.baseline_image_path and image_path is not None and str(image_path).startswith(str(ASSETS_DIR)):
@@ -436,7 +418,7 @@ def render_feature_overlay_png(request: FeatureMapRequest) -> bytes | None:
         patches_per_image=None,
         crop_size=crop_size,
         img_size_inches=4.5,
-        show_diff=False,
+        show_mask=False,
         save_dir=None,
         show_img=False,
     )

dashboard/viz_helpers.py

@@ -8,8 +8,8 @@ from dash import dcc, html
 
 from dashboard.model_catalog import (
     MODEL_FAMILIES,
-    discover_models_for_metric,
     discover_supported_datasets,
+    get_model_record,
     load_selector_cache_feature_table,
     summarize_selector_cache,
 )
@@ -51,17 +51,17 @@ def feature_rows_for_selector(
     if not selector_name:
         return []
 
-    record = next((r for r in discover_models_for_metric(metric) if r.model_key == model_key), None)
+    record = get_model_record(metric, model_key)
     if record is None:
         return []
 
-    summaries = summarize_selector_cache(record.model_path, dataset)
+    summaries = summarize_selector_cache(record.checkpoint_path, dataset)
     summary_by_name = {summary.selector_name: summary for summary in summaries}
     selected_summary = summary_by_name.get(selector_name, summaries[0] if summaries else None)
     if selected_summary is None:
         return []
 
-    feature_table = load_selector_cache_feature_table(selected_summary.cache_path)
+    feature_table = load_selector_cache_feature_table(str(selected_summary.cache_path))
     if feature_table.empty:
         return []
 
@@ -86,14 +86,14 @@ def _feature_comparison_rows_cached(
     feature_id: int,
     active_selector_name: str | None,
 ) -> list[dict[str, object]]:
-    record = next((r for r in discover_models_for_metric(metric) if r.model_key == model_key), None)
+    record = get_model_record(metric, model_key)
     if record is None:
         return []
 
-    summaries = summarize_selector_cache(record.model_path, dataset)
+    summaries = summarize_selector_cache(record.checkpoint_path, dataset)
     comparison_rows: list[dict[str, object]] = []
     for summary in summaries:
-        feature_table = load_selector_cache_feature_table(summary.cache_path)
+        feature_table = load_selector_cache_feature_table(str(summary.cache_path))
         if feature_table.empty:
             feature_score = None
             feature_rank = None
@@ -242,14 +242,12 @@ def select_feature_id(
 
 def feature_unavailable_message(
     metric: str,
-    dataset: str,
     model_key: str,
     feature_id: int,
     *,
-    activation_dataset: str | None = None,
+    activation_dataset: str,
 ) -> str | None:
-    cache_dataset = activation_dataset if activation_dataset is not None else dataset
-    if feature_in_activation_cache(metric, cache_dataset, model_key, int(feature_id)):
+    if feature_in_activation_cache(metric, activation_dataset, model_key, int(feature_id)):
         return None
     return FILTERED_FEATURE_MESSAGE
 
@@ -258,6 +256,60 @@ def feature_empty_state(message: str) -> html.Div:
     return html.Div(message, className="feature-empty")
 
 
+def build_feature_detail_content(
+    metric: str,
+    dataset: str,
+    model_key: str,
+    feature_id: int | None,
+    selector_name: str | None = None,
+) -> list[object] | html.Div:
+    record = get_model_record(metric, model_key)
+    if record is None:
+        return html.Div("No selector cache files were found for this model.", className="feature-empty")
+
+    if not selector_name:
+        summaries = summarize_selector_cache(record.checkpoint_path, dataset)
+        if not summaries:
+            return html.Div("No selector cache files were found for this model.", className="feature-empty")
+        selector_name = summaries[0].selector_name
+
+    rows = feature_rows_for_selector(metric, dataset, model_key, selector_name)
+    if not rows:
+        return html.Div("No selector cache files were found for this model.", className="feature-empty")
+
+    selected_ids = feature_ids(rows)
+    if not selected_ids:
+        return html.Div("No selector cache files were found for this model.", className="feature-empty")
+
+    try:
+        selected_feature_id = int(feature_id)
+    except (TypeError, ValueError):
+        selected_feature_id = selected_ids[0]
+
+    unavailable = feature_unavailable_message(
+        metric,
+        model_key,
+        selected_feature_id,
+        activation_dataset=UMAP_DATASET,
+    )
+    if unavailable is not None:
+        return feature_empty_state(unavailable)
+
+    comparison_rows = feature_comparison_rows(
+        metric,
+        dataset,
+        model_key,
+        selected_feature_id,
+        selector_name,
+    )
+    return feature_detail_children(
+        selector_name or "selector",
+        rows,
+        selected_feature_id,
+        comparison_rows,
+    )
+
+
 def feature_detail_children(
     selector_name: str,
     feature_rows: list[dict[str, object]],

default_configs/06_config.json

@@ -4,6 +4,7 @@
   "DATASET_ROOT": "/home_2/28m_gov@lab.graphicon.ru/SAE/Kadid",
   "CACHE_DIR": "/home_2/28m_gov@lab.graphicon.ru/SAE/xIQA/logs/arniqa_logs/ARNIQA_layer5_lambda5e-4_scale1_exp37/cache",
   "IMAGE_LIMIT": 1000,
+  "IMAGE_RANDOM_STATE": 42,
   "AGGREGATION_MODES": ["mean", "max", "sum"],
   "DEFAULT_AGGREGATION_MODE": "mean",
   "UMAP_PARAMS": {

default_configs/sae_vis_config.json

@@ -1,7 +1,7 @@
 {
-  "SAE_CHECKPOINT_PATH": "/data/logs/ARNIQA_layer3_lambda1e-5_scale1_exp2/checkpoints/checkpoint-200200",
-  "DATASETS_ROOT": "/data/Kadid",
-  "CACHE_DIR": "/data/logs/ARNIQA_layer3_lambda1e-5_scale1_exp2/cache",
+  "SAE_CHECKPOINT_PATH": "/home/28m_gov@lab.graphicon.ru/SAE/xIQA/logs/ARNIQA_layer3_lambda1e-5_scale1_exp2/checkpoint-200200",
+  "DATASETS_ROOT": "/home/28m_gov@lab.graphicon.ru/SAE/Kadid",
+  "CACHE_DIR": "/home/28m_gov@lab.graphicon.ru/SAE/xIQA/logs/ARNIQA_layer3_lambda1e-5_scale1_exp2/checkpoint-200200/cache",
   "DATASET": "kadid10k",
   "CROP_SIZE": 224,
   "DOWNSCALE_FACTOR": 2,
@@ -9,10 +9,6 @@
   "NUM_WORKERS": 4,
   "DEVICE": "cuda",
   "ACTIVATION_STEPS_TO_KEEP": [],
-  "SUPPORTED_DATASETS": [
-    "kadid10k",
-    "local_kadid"
-  ],
   "CORR_TOP_K": 30,
   "HEATMAP_CRITERION": "max",
   "N_TOP_FEATURES_HEATMAPS": 3,

pages/feature_maps.py

@@ -12,7 +12,7 @@ from dashboard.user_image_service import (
     list_baseline_example_records,
     render_feature_overlay,
 )
-from dashboard.viz_helpers import visualization_params
+from dashboard.viz_helpers import build_feature_detail_content, visualization_params
 
 
 ROOT = Path(__file__).resolve().parents[1]
@@ -103,6 +103,25 @@ layout = html.Div(
                                     debounce=True,
                                     className="feature-jump-input",
                                 ),
+                                html.Div(
+                                    dcc.Loading(
+                                        id="fm-feature-detail-loading",
+                                        type="default",
+                                        color="#2563eb",
+                                        className="loading-block feature-detail-body-loading",
+                                        children=html.Div(
+                                            id="fm-feature-detail-body",
+                                            className="feature-detail-card",
+                                            children=[
+                                                html.Div(
+                                                    "Enter a feature id to see statistics.",
+                                                    className="feature-empty",
+                                                ),
+                                            ],
+                                        ),
+                                    ),
+                                    style={"marginTop": "16px"},
+                                ),
                                 html.Div(
                                     id="fm-status",
                                     className="home-status",
@@ -290,3 +309,28 @@ def update_overlay(
         return "", "Overlay unavailable."
     return overlay_url, ""
 
+
+@callback(
+    Output("fm-feature-detail-body", "children"),
+    Input("_pages_location", "search"),
+    Input("fm-feature-id", "value"),
+)
+def render_fm_feature_detail(search: str | None, feature_id_value):
+    params = visualization_params(search)
+    if params is None:
+        return html.Div("Missing page parameters.", className="feature-empty")
+
+    metric, dataset, model_key = params
+    if not model_key:
+        return html.Div("Select a model on the Home page first.", className="feature-empty")
+
+    try:
+        feature_id = int(feature_id_value) if feature_id_value is not None else None
+    except (TypeError, ValueError):
+        return html.Div("Enter a valid feature id.", className="feature-empty")
+
+    if feature_id is None:
+        return html.Div("Enter a feature id to see statistics.", className="feature-empty")
+
+    return build_feature_detail_content(metric, dataset, model_key, feature_id)
+

pages/visualizations.py

@@ -9,7 +9,7 @@ from dash import Input, Output, State, callback, clientside_callback, ctx, dcc,
 from dashboard.image_utils import get_top_feature_overlays, overlay_artifact_data_url
 from dashboard.layout import page_back_nav
 from dashboard.model_catalog import (
-    discover_models_for_metric,
+    get_model_record,
     summarize_model_record,
     summarize_selector_cache,
 )
@@ -31,10 +31,9 @@ from dashboard.umap_service import (
     format_umap_status,
 )
 from dashboard.viz_helpers import (
+    build_feature_detail_content,
     cached_hover_meta_line,
     cached_hover_thumb,
-    feature_comparison_rows,
-    feature_detail_children,
     feature_empty_state,
     feature_ids,
     feature_rows_for_selector,
@@ -506,9 +505,9 @@ layout = html.Div(
                                                 ),
                                                 dcc.Input(
                                                     id="feature-jump-input",
-                                                    type="number",
-                                                    min=0,
-                                                    step=1,
+                                                    type="text",
+                                                    inputMode="numeric",
+                                                    pattern="[0-9]*",
                                                     debounce=True,
                                                     placeholder="Jump to feature id",
                                                     className="feature-jump-input",
@@ -607,7 +606,7 @@ def render_model_info(search: str | None):
         return html.Div("Missing visualization parameters.", className="feature-empty")
 
     metric, selection_dataset, model_key = params
-    record = next((r for r in discover_models_for_metric(metric) if r.model_key == model_key), None)
+    record = get_model_record(metric, model_key)
 
     children = [
         html.Div([html.Span("Metric: "), html.Strong(metric)]),
@@ -624,7 +623,7 @@ def render_model_info(search: str | None):
     try:
         from dashboard.model_catalog import summarize_feature_filter_cache
 
-        filter_lines = summarize_feature_filter_cache(record.model_path, selection_dataset)
+        filter_lines = summarize_feature_filter_cache(record.checkpoint_path, selection_dataset)
     except Exception:
         filter_lines = ["Feature filter summary unavailable"]
 
@@ -653,11 +652,11 @@ def init_feature_selector(search: str | None):
         return [], None, True
 
     metric, selection_dataset, model_key = params
-    record = next((r for r in discover_models_for_metric(metric) if r.model_key == model_key), None)
+    record = get_model_record(metric, model_key)
     if record is None:
         return [], None, True
 
-    summaries = summarize_selector_cache(record.model_path, selection_dataset)
+    summaries = summarize_selector_cache(record.checkpoint_path, selection_dataset)
     if not summaries:
         return [], None, True
 
@@ -918,10 +917,9 @@ def display_umap_hover_content(hover_data, search: str | None, umap_mode: str |
 
         params = visualization_params(search)
         if params is not None:
-            metric, selection_dataset, model_key = params
+            metric, _selection_dataset, model_key = params
             unavailable = feature_unavailable_message(
                 metric,
-                selection_dataset,
                 model_key,
                 fid,
                 activation_dataset=UMAP_DATASET,
@@ -1053,13 +1051,11 @@ def update_selected_feature(
         requested_feature_id=requested_feature_id,
     )
     if selected_feature_id is not None:
-        activation_dataset = UMAP_DATASET if trigger_id == "umap-graph" else None
         unavailable = feature_unavailable_message(
             metric,
-            selection_dataset,
             model_key,
             int(selected_feature_id),
-            activation_dataset=activation_dataset,
+            activation_dataset=UMAP_DATASET,
         )
         if unavailable is not None:
             status = unavailable
@@ -1082,36 +1078,13 @@ def render_feature_detail(
         return html.Div("No selector cache files were found for this model.", className="feature-empty")
 
     metric, selection_dataset, model_key = params
-    rows = feature_rows_for_selector(metric, selection_dataset, model_key, selector_name)
-    if not rows:
-        return html.Div("No selector cache files were found for this model.", className="feature-empty")
-
-    selected_ids = feature_ids(rows)
-    if not selected_ids:
-        return html.Div("No selector cache files were found for this model.", className="feature-empty")
-
-    try:
-        selected_feature_id = int(selected_feature_id)
-    except (TypeError, ValueError):
-        selected_feature_id = selected_ids[0]
-
-    unavailable = feature_unavailable_message(metric, selection_dataset, model_key, selected_feature_id)
-    if unavailable is not None:
-        return feature_empty_state(unavailable)
-
-    comparison_rows = feature_comparison_rows(
+    return build_feature_detail_content(
         metric,
         selection_dataset,
         model_key,
         selected_feature_id,
         selector_name,
     )
-    return feature_detail_children(
-        selector_name or "selector",
-        rows,
-        selected_feature_id,
-        comparison_rows,
-    )
 
 
 @callback(
@@ -1137,7 +1110,12 @@ def render_feature_top_images(
     except Exception:
         return _empty_feature_image_outputs("No feature selected.")
 
-    unavailable = feature_unavailable_message(metric, selection_dataset, model_key, feature_id)
+    unavailable = feature_unavailable_message(
+        metric,
+        model_key,
+        feature_id,
+        activation_dataset=UMAP_DATASET,
+    )
     if unavailable is not None:
         return _empty_feature_image_outputs(unavailable)
 

train_code/benjdataset.py

@@ -8,14 +8,14 @@ import torch
 from torch.utils.data import Dataset, DataLoader
 from torchvision import transforms  # type:ignore
 
-# искажения из ARNIQA для аугментации
+# ARNIQA distortions for augmentation
 from train_code.aug_utils.utils_data import distort_images
 from utils.utils_data import resize_crop
 
 
 def collate_fn(data):
-    # Надо подумать насчёт того, правильно ли было выносить ресайз кроп в гетайтем NOTE
-    # Но у них коде вроде так же
+    # TODO: reconsider whether resize/crop should live in __getitem__ instead
+    # Their code appears to do it the same way
     # images = torch.stack([resize_crop(example['image'], crop_size=CROP_SIZE,
     #                      downscale_factor=DOWNSCALE_FACTOR) for example in data])
     # print(data.keys())
@@ -61,8 +61,8 @@ class DeepFeaturesDataset(Dataset):
             raw_image = resize_crop(raw_image, crop_size=self.crop_size, downscale_factor=self.downscale_factor)
         if random.random() > self.prestine_prob:
             try:
-                # Для того, чтобы при обучении SAE метрика видела искажённые изображения, использую пайплайн искажений из статьи ARNIQA.
-                # Код в aug_utils взят из репо ARNIQA
+                # During SAE training, the metric should see distorted images; use the ARNIQA paper distortion pipeline.
+                # aug_utils code is taken from the ARNIQA repo
                 raw_image = distort_images(raw_image.clone())[0]
             except Exception as e:
                 print(f'[Warning] Degradation operation failed due to {e}!')
@@ -82,8 +82,8 @@ def create_dataloaders(dataset_name: str,
                        token: str,
                        transform: Optional[Callable] = None) -> Tuple[DataLoader, DataLoader]:
     # ds = load_dataset("timm/imagenet-1k-wds", split='train', cache_dir='D:/adapter_experiment/hf_cache', data_files='**/*-validation-*.tar')
-    # тут грузится большой датасет (несколько тб), можно заменить на любой другой источник данных. Закомментил, чтобы нечаянно не начать их грузить
-    # Используем небольшой датасет для тестирования
+    # The full dataset is huge (several TB); swap in any other data source. Commented out to avoid loading it by accident.
+    # Use a small dataset for testing
 
     ds = load_dataset(dataset_name, split=split, streaming=True, cache_dir='./datasets/')
     # ds_data = datasets.Dataset.from_list(ds_head)

train_code/timmdataset.py

@@ -6,13 +6,13 @@ import torch
 from torch.utils.data import Dataset, DataLoader
 from torchvision import transforms  # type:ignore
 
-# искажения из ARNIQA для аугментации
+# ARNIQA distortions for augmentation
 from train_code.aug_utils.utils_data import distort_images
 from utils.utils_data import resize_crop
 
 def collate_fn(data):
-    # Надо подумать насчёт того, правильно ли было выносить ресайз кроп в гетайтем NOTE
-    # Но у них коде вроде так же
+    # TODO: reconsider whether resize/crop should live in __getitem__ instead
+    # Their code appears to do it the same way
     # images = torch.stack([resize_crop(example['image'], crop_size=CROP_SIZE,
     #                      downscale_factor=DOWNSCALE_FACTOR) for example in data])
     # print(data.keys())
@@ -58,8 +58,8 @@ class DeepFeaturesDataset(Dataset):
             raw_image = resize_crop(raw_image, crop_size=self.crop_size, downscale_factor=self.downscale_factor)
         if random.random() > self.prestine_prob:
             try:
-                # Для того, чтобы при обучении SAE метрика видела искажённые изображения, использую пайплайн искажений из статьи ARNIQA.
-                # Код в aug_utils взят из репо ARNIQA
+                # During SAE training, the metric should see distorted images; use the ARNIQA paper distortion pipeline.
+                # aug_utils code is taken from the ARNIQA repo
                 raw_image = distort_images(raw_image.clone())[0]
             except Exception as e:
                 print(f'[Warning] Degradation operation failed due to {e}!')
@@ -79,15 +79,15 @@ def create_dataloaders(dataset_name: str,
                        token: str,
                        transform: Optional[Callable] = None) -> Tuple[DataLoader, DataLoader]:
     # ds = load_dataset("timm/imagenet-1k-wds", split='train', cache_dir='D:/adapter_experiment/hf_cache', data_files='**/*-validation-*.tar')
-    # тут грузится большой датасет (несколько тб), можно заменить на любой другой источник данных. Закомментил, чтобы нечаянно не начать их грузить
-    # Используем небольшой датасет для тестирования
+    # The full dataset is huge (several TB); swap in any other data source. Commented out to avoid loading it by accident.
+    # Use a small dataset for testing
 
     # ds = load_dataset(dataset_name, split=split, streaming=True, cache_dir='./datasets/')
     # ds_head = list(ds.take(10))
     # ds_data = datasets.Dataset.from_list(ds_head)
     # dset = DeepFeaturesDataset(ds_data, prestine_prob = prestine_prob)
 
-    # загружаем только первые 1000 изображений для тестирования
+    # Load only the first 1000 images for testing
     ds = load_dataset(dataset_name, split=split, cache_dir='/mgov/datasets/', token=token)
     # ds = ds.select(range(10000)) # !!!!!!!!!!!
     dset = DeepFeaturesDataset(

train_code/train_script_arniqa.py

@@ -78,7 +78,7 @@ def parse_args():
         help="Initialization of SAE weights with fixed norm.",
     )
     # parser.add_argument(
-    #     "--sae_input_dim", # Будет определяться автоматически
+    #     "--sae_input_dim", # Determined automatically
     #     type=int,
     #     default=256,
     #     help="Dimentionality of activations.",
@@ -311,8 +311,8 @@ def get_layer_output_size(model: InferenceModel, device: Union[str, torch.device
 
 def set_hook(model: InferenceModel, layer_num: int, layer_name: str) -> Tuple[torch.utils.hooks.RemovableHandle, dict]:
 
-    # Навешивание хука на нужный слой модели IQA для получения активаций
-    # подробнее про хуки: https://web.stanford.edu/~nanbhas/blog/forward-hooks-pytorch/
+    # Register a hook on the target IQA layer to capture activations
+    # More on hooks: https://web.stanford.edu/~nanbhas/blog/forward-hooks-pytorch/
     activations = {}
 
     def getActivation(name):
@@ -331,7 +331,7 @@ def create_model(device: Union[str, torch.device], iqa_weight_dtype: torch.dtype
     return model
 
 # create_sae, setup_lambda_scaling, setup_optimizer_and_scheduler, extract_activations,
-# compute_lambda, train_step, log_train_metrics, save_checkpoint импортируются из train_utils.py
+# compute_lambda, train_step, log_train_metrics, save_checkpoint are imported from train_utils.py
 
 
 def save_sae_config(
@@ -340,19 +340,19 @@ def save_sae_config(
     inner_dim: int,
     output_dir: Optional[str] = None,
 ) -> None:
-    """Сохраняет гиперпараметры SAE в JSON-файл рядом с чекпоинтами.
-    Формат достаточен для воспроизведения архитектуры модели без дополнительных аргументов.
+    """Save SAE hyperparameters to a JSON file next to checkpoints.
+    The format is sufficient to reproduce the model architecture without extra arguments.
     """
     cfg = {
         "iqa_metric": "arniqa-kadid",
         "layer_num": args.layer_num,
         "scaling_factor": args.scaling_factor,
-        # архитектура
+        # architecture
         "sae_type": args.sae_type,
         "sae_input_dim": sae_input_dim,
         "inner_dim": inner_dim,
         "weight_norm_init": args.weight_norm_init,
-        # MatchingPursuitSAE-специфичные параметры
+        # MatchingPursuitSAE-specific parameters
         "mp_threshold": args.mp_threshold,
         "mp_normalize": bool(args.mp_normalize),
     }
@@ -364,7 +364,7 @@ def save_sae_config(
     print(f"SAE config saved to {path}")
 
 
-# validate_epoch, train_epoch, fix_seed импортируются из train_utils.py
+# validate_epoch, train_epoch, fix_seed are imported from train_utils.py
 
 
 def print_hyperparams(

train_code/train_script_liqe.py

@@ -361,7 +361,7 @@ def create_model(
 
 
 # create_sae, setup_lambda_scaling, setup_optimizer_and_scheduler, extract_activations,
-# compute_lambda, train_step, log_train_metrics, save_checkpoint импортируются из train_utils.py
+# compute_lambda, train_step, log_train_metrics, save_checkpoint are imported from train_utils.py
 
 
 def save_sae_config(

train_code/train_script_maniqa.py

@@ -75,12 +75,12 @@ def get_layer_output_size(
 
 
 def set_hook(model, layer_name: str, swin_num: int = 2, layer_num: int = 1) -> Tuple[Any, dict]:
-    """Навешивает hook на выбранный BasicLayer SwinTransformer модели MANIQA.
+    """Register a hook on the selected MANIQA SwinTransformer BasicLayer.
 
-    swin_num: 1 (swintransformer1, dim=768) или 2 (swintransformer2, dim=384).
-    layer_num: индекс BasicLayer внутри выбранного SwinTransformer (0 или 1).
-    Выход слоя: (B, L, C) — последовательность токенов.
-    Hook флэттит до (B*L, C), чтобы SAE мог работать с отдельными токенами.
+    swin_num: 1 (swintransformer1, dim=768) or 2 (swintransformer2, dim=384).
+    layer_num: BasicLayer index inside the selected SwinTransformer (0 or 1).
+    Layer output: (B, L, C) token sequence.
+    The hook flattens to (B*L, C) so the SAE can process individual tokens.
     """
     activations: dict = {}
 
@@ -99,7 +99,7 @@ def save_sae_config(
     inner_dim: int,
     output_dir: Optional[str] = None,
 ) -> None:
-    """Сохраняет гиперпараметры SAE в JSON-файл рядом с чекпойнтами."""
+    """Save SAE hyperparameters to a JSON file next to checkpoints."""
     cfg = {
         "iqa_metric": "maniqa",
         "swin_num": args.swin_num,
@@ -494,7 +494,7 @@ def main():
         save_sae_config(args, sae_input_dim, inner_dim)
 
     dataloader_train, dataloader_test = create_dataloaders(
-        DATASET_NAME, 'train', 0, args.crop_size, 1, args.train_frac, bs, num_workers, collate_fn, token=args.hf_token) # Downscale=1, потому что так в оригинале, но возможно стоит добавить.
+        DATASET_NAME, 'train', 0, args.crop_size, 1, args.train_frac, bs, num_workers, collate_fn, token=args.hf_token) # Downscale=1 as in the original; may be worth making configurable.
 
     autoenc_model = create_sae(
         sae_input_dim, inner_dim, weight_norm_init, device, weight_dtype,
@@ -503,7 +503,7 @@ def main():
         mp_normalize=bool(args.mp_normalize),
     )
 
-    # Выбираем функцию потерь в зависимости от типа SAE
+    # Choose the loss function based on SAE type
     _loss_base = get_loss_fn(args.sae_type)
     _aux_alpha = args.aux_alpha
 

train_code/train_utils.py

@@ -1,4 +1,4 @@
-"""Общие утилиты для обучения SAE.
+"""Shared utilities for SAE training.
 """
 
 import os
@@ -52,7 +52,7 @@ def fix_seed(seed: int) -> None:
     torch.backends.cudnn.deterministic = True
     torch.backends.cudnn.benchmark = False
 
-# Чекпоинты 
+# Checkpoints
 
 def _extract_checkpoint_step(name: str) -> Optional[int]:
     prefix = "checkpoint-"
@@ -117,7 +117,7 @@ def create_sae(
     autoenc_model.train()
     return autoenc_model
 
-# Оптимизаторы и scheduler'ы
+# Optimizers and schedulers
 
 def create_optimizer(
     model: torch.nn.Module,
@@ -203,7 +203,7 @@ def compute_lambda(
 def calc_loss_reconstr(x: torch.Tensor, x_rec: torch.Tensor) -> torch.Tensor:
     return torch.nn.functional.mse_loss(x_rec, x)
 
-# Извлечение активаций 
+# Activation extraction
 
 def extract_activations(
     model: torch.nn.Module,
@@ -222,7 +222,7 @@ def extract_activations(
         data *= scaling_factor
     return data
 
-# Обучение
+# Training
 
 def train_step(
     autoenc_model: torch.nn.Module,
@@ -240,7 +240,7 @@ def train_step(
     real_batch_size = data.shape[0]
     optimizer.zero_grad()
 
-    # Перемешиваем батч и семплируем sample_frac долю пространственных позиций
+    # Shuffle the batch and sample sample_frac of spatial positions
     perm = torch.randperm(real_batch_size, device=data.device)
     data = data[perm[: int(len(perm) * sample_frac)]]
     chunk_size = max(1, real_batch_size // grad_accum_steps)