dashboard/image_utils.py

from __future__ import annotations

import base64
from pathlib import Path
from typing import List

from analysis.config import dataset_images_root
from analysis.datasets import parent_datasets_root, 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 analysis.viz.vis_heatmaps_plotly import empty_heatmap_figure, set_heatmap_overlay_visible
from dashboard.model_catalog import (
    _dashboard_default_params,
    get_model_record,
    load_and_filter_model_activations,
)
from log_config import get_logger

logger = get_logger(__name__)

ROOT = Path(__file__).resolve().parents[1]

# Page 1 (Visualizations): include SR-artifact regions on the SRGround mask row.
INCLUDE_SRGROUND_SR_ARTIFACT = True
# Show filename captions under top-image cards in the dashboard.
SHOW_TOP_IMAGE_CAPTIONS = False


def png_bytes_to_data_url(png: bytes) -> str:
    encoded = base64.b64encode(png).decode("ascii")
    return f"data:image/png;base64,{encoded}"


def overlay_artifact_data_url(artifact: dict[str, object]) -> str | None:
    raw = artifact.get("bytes")
    if isinstance(raw, bytes):
        return png_bytes_to_data_url(raw)
    path = artifact.get("path")
    if path is None:
        return None
    try:
        rel = Path(str(path)).relative_to(ROOT / "assets")
        return f"/assets/{rel.as_posix()}"
    except Exception:
        return f"file://{path}"


def resolve_dashboard_dataset_root(
    selection_dataset: str,
    dataset_root: Path | str | None = None,
) -> Path:
    if dataset_root is not None:
        return Path(dataset_root)
    default_cfg = _dashboard_default_params()
    return Path(dataset_images_root(default_cfg.DATASETS_ROOT, selection_dataset))


def overlay_artifact_figure(
    artifact: dict[str, object] | None,
    *,
    show_heatmap_overlay: bool = True,
) -> dict:
    if not artifact:
        return empty_heatmap_figure("Top images overlay is not available yet.").to_plotly_json()
    figure = artifact.get("figure")
    if isinstance(figure, dict):
        return set_heatmap_overlay_visible(figure, show_heatmap_overlay)
    return empty_heatmap_figure("Top images overlay is not available yet.").to_plotly_json()


def overlay_show_heatmap_enabled(show_value: list[str] | None) -> bool:
    if not show_value:
        return False
    return "show" in show_value


def get_top_feature_overlays(
    metric: str,
    selection_dataset: str,
    model_key: str,
    feature_id: int,
    top_n: int = 5,
    ranking_mode: str = "iou",
    dataset_root: Path | None = None,
) -> List[dict[str, object]]:
    """Return overlay artifacts for top-N images (PNG bytes in memory, no disk cache).

    Uses activation cache for ``selection_dataset`` (from the home-page selector),
    not the UMAP kadid10k cache.
    """
    ranking_mode = str(ranking_mode or "iou")
    resolved_dataset_root = resolve_dashboard_dataset_root(selection_dataset, dataset_root)
    datasets_root_parent = parent_datasets_root(resolved_dataset_root)

    record = get_model_record(metric, model_key)
    if record is None:
        return []

    try:
        filtered = load_and_filter_model_activations(record.checkpoint_path, selection_dataset)
    except Exception as exc:
        logger.error('Error occurred while loading activations for %r: %s', selection_dataset, exc)
        return []

    meta = filtered.meta_df
    features = filtered.features

    if meta is None or features.codes is None:
        logger.error('Meta or codes data is missing in the loaded model activations.')
        return []

    if int(feature_id) not in features.column_feature_ids:
        return []

    if ranking_mode == "iou":
        top_image_idxs = get_top_images_for_feature_by_iou(
            features,
            meta,
            int(feature_id),
            top_n=top_n,
            dataset=selection_dataset,
        )
    else:
        top_image_idxs = get_top_images_for_feature(
            features,
            meta,
            int(feature_id),
            top_n=top_n,
            aggregation="max",
        )
    if not top_image_idxs:
        return []

    first_rows = meta.groupby("image_idx", sort=False).first()
    image_paths = []
    for idx in top_image_idxs:
        try:
            row = first_rows.loc[int(idx)]
        except Exception:
            row = None
        img_path = None
        if row is not None:
            for col in ("distorted_img_path", "image_path", "original_img_path"):
                if col in row and row[col] is not None:
                    img_path = str(row[col])
                    break
        if img_path is None and row is not None:
            for v in row.values:
                if isinstance(v, str) and v.lower().endswith((".png", ".jpg", ".jpeg")):
                    img_path = v
                    break
        if img_path is None:
            continue
        resolved = resolve_dataset_image_path(
            selection_dataset,
            img_path,
            datasets_root=datasets_root_parent,
        )
        image_paths.append(str(resolved))

    if not image_paths:
        return []

    image_indices = top_image_idxs[: len(image_paths)]
    caption = f"topmax_feature_{int(feature_id)}_overlay.png"

    if str(selection_dataset) == "SRGround":
        panel_figure = render_top_feature_panel_srground(
            meta=meta,
            features=features,
            image_indices=image_indices,
            image_paths=image_paths,
            feature_id=int(feature_id),
            patches_per_image=None,
            crop_size=224,
            img_size_inches=3.5,
            dataset_root=str(resolved_dataset_root),
            backend="plotly",
        )
        if panel_figure is None or not isinstance(panel_figure, dict):
            return []
        return [
            {
                "kind": "overlay",
                "feature_id": str(feature_id),
                "figure": panel_figure,
                "caption": caption if SHOW_TOP_IMAGE_CAPTIONS else "",
            }
        ]

    artifacts = visualize_feature_heatmaps(
        meta=meta,
        features=features,
        image_indices=image_indices,
        image_paths=image_paths,
        feature_ids=[int(feature_id)],
        patches_per_image=None,
        crop_size=224,
        img_size_inches=3.5,
        show_mask=False,
        save_dir=None,
        file_prefix="topmax",
        show_img=False,
        backend="plotly",
    )
    overlays: list[dict[str, object]] = []
    for item in artifacts:
        if item.get("kind") != "overlay":
            continue
        figure = item.get("figure")
        if not isinstance(figure, dict):
            continue
        overlays.append(
            {
                "kind": "overlay",
                "feature_id": str(feature_id),
                "figure": figure,
                "caption": caption if SHOW_TOP_IMAGE_CAPTIONS else "",
            }
        )
    return overlays