Mini-Nav/mini-nav/benchmarks/tasks/retrieval.py

"""Retrieval task for benchmark evaluation (Recall@K)."""

from typing import TYPE_CHECKING, Any, cast

import lancedb
import numpy as np
import pyarrow as pa
import torch
import torch.nn.functional as F
from benchmarks.base import BaseBenchmarkTask
from benchmarks.tasks.registry import RegisterTask
from compressors.model_loader import get_dino_dim, load_dino_model, load_hash_compressor
from configs import cfg_manager
from rich.progress import track
from torch import nn
from torch.utils.data import DataLoader
from transformers import BitImageProcessor
from utils.feature_extractor import extract_batch_features, infer_vector_dim

if TYPE_CHECKING:
    from compressors.hash_compressor import HashCompressor


class RetrievalEncoder(nn.Module):
    """Benchmark encoder for DINO and optional hash compression."""

    def __init__(
        self,
        dino: nn.Module,
        compressor: "HashCompressor | None" = None,
    ) -> None:
        """Initialize retrieval encoder.

        Args:
            dino: DINO backbone used for feature extraction.
            compressor: Optional hash compressor for recall evaluation.
        """
        super().__init__()
        self.dino: nn.Module = dino
        self.compressor: HashCompressor | None = compressor

    def forward(self, inputs: Any) -> torch.Tensor:
        """Encode processor inputs into benchmark vectors.

        Args:
            inputs: Batched processor outputs.

        Returns:
            Float tensor used for LanceDB insertion and retrieval.
        """
        outputs = self.dino(**inputs)
        tokens = outputs.last_hidden_state

        if self.compressor is None:
            features = tokens.mean(dim=1)
            return F.normalize(features, dim=-1)

        bits = self.compressor.encode(tokens)
        return bits.to(dtype=torch.float32)


def _build_eval_schema(vector_dim: int) -> pa.Schema:
    """Build PyArrow schema for evaluation database table.

    Args:
        vector_dim: Feature vector dimension.

    Returns:
        PyArrow schema with id, label, and vector fields.
    """
    return pa.schema(
        [
            pa.field("id", pa.int32()),
            pa.field("label", pa.int32()),
            pa.field("vector", pa.list_(pa.float32(), vector_dim)),
        ]
    )


def _establish_eval_database(
    processor: BitImageProcessor,
    model: nn.Module,
    table: lancedb.table.Table,
    dataloader: DataLoader[Any],
    label_column: str = "label",
) -> None:
    """Extract features from training images and store them in a database table.

    Args:
        processor: Image preprocessor.
        model: Feature extraction model.
        table: LanceDB table to store features.
        dataloader: DataLoader for the training dataset.
        label_column: Column name for labels in the batch dict.
    """
    # Extract all features using the utility function
    all_features = extract_batch_features(processor, model, dataloader)
    config = cfg_manager.get()

    # Store features to database
    global_idx = 0
    for batch in track(dataloader, description="Storing eval database"):
        labels = batch[label_column]
        labels_list = labels.tolist()
        batch_size = len(labels_list)

        table.add(
            [
                {
                    "id": global_idx + j,
                    "label": labels_list[j],
                    "vector": all_features[global_idx + j].detach().cpu().numpy(),
                }
                for j in range(batch_size)
            ]
        )
        global_idx += batch_size


def _evaluate_recall(
    processor: BitImageProcessor,
    model: nn.Module,
    table: lancedb.table.Table,
    dataloader: DataLoader[Any],
    top_k: int,
    label_column: str = "label",
) -> dict[str, Any]:
    """Evaluate Recall@K by searching the database for each test image.

    Args:
        processor: Image preprocessor.
        model: Feature extraction model.
        table: LanceDB table to search against.
        dataloader: DataLoader for the test dataset.
        top_k: Number of top results to retrieve.
        label_column: Column name for labels in the batch dict.

    Returns:
        A dict with keys:
            - correct: Number of correct predictions
            - total: Total number of test samples
            - query_labels: True labels for each query (length N)
            - topk_ids: IDs of top-K retrieved items (N x K)
            - topk_labels: Labels of top-K retrieved items (N x K)
    """
    # Extract all features using the utility function
    all_features = extract_batch_features(processor, model, dataloader)
    config = cfg_manager.get()

    correct = 0
    total = 0
    feature_idx = 0
    query_labels: list[int] = []
    topk_ids: list[list[int]] = []
    topk_labels: list[list[int]] = []

    for batch in track(dataloader, description=f"Evaluating Recall@{top_k}"):
        labels = batch[label_column]
        labels_list = labels.tolist()

        for j in range(len(labels_list)):
            feature = all_features[feature_idx + j].tolist()
            true_label = labels_list[j]

            results = (
                table.search(feature)
                .select(["id", "label", "_distance"])
                .limit(top_k)
                .to_polars()
            )

            retrieved_ids = results["id"].to_list()
            retrieved_labels = results["label"].to_list()
            if true_label in retrieved_labels:
                correct += 1
            total += 1

            query_labels.append(true_label)
            topk_ids.append(retrieved_ids)
            topk_labels.append(retrieved_labels)

        feature_idx += len(labels_list)

    return {
        "correct": correct,
        "total": total,
        "query_labels": query_labels,
        "topk_ids": topk_ids,
        "topk_labels": topk_labels,
    }


def _build_confusion_matrix(
    query_labels: list[int],
    topk_labels: list[list[int]],
    num_classes: int,
) -> np.ndarray:
    """Build row-normalized confusion matrix from Top-1 predictions.

    For each query, the true label is compared against the first
    retrieved label (Top-1). The resulting matrix is row-normalized
    so that each row sums to 1.0, representing "given true class X,
    what fraction of queries were predicted as each class Y".

    Args:
        query_labels: True labels for each query (length N).
        topk_labels: Top-K retrieved labels for each query (N x K).
        num_classes: Total number of label classes.

    Returns:
        Row-normalized confusion matrix as (num_classes, num_classes)
        float64 array. Rows with zero queries remain all-zero.
    """
    cm = np.zeros((num_classes, num_classes), dtype=np.int64)
    for y_true, top_labels in zip(query_labels, topk_labels):
        y_pred = top_labels[0]  # Top-1
        cm[int(y_true), int(y_pred)] += 1

    # Row normalize: each row sums to 1.0
    row_sums = cm.sum(axis=1, keepdims=True)
    cm_norm = np.divide(
        cm.astype(np.float64),
        row_sums,
        out=np.zeros_like(cm, dtype=np.float64),
        where=row_sums > 0,
    )
    return cm_norm


@RegisterTask("retrieval")
class RetrievalTask(BaseBenchmarkTask):
    """Retrieval evaluation task (Recall@K)."""

    def __init__(
        self,
        top_k: int = 10,
        top_k_list: list[int] | None = None,
        dino_model: str = "facebook/dinov2-large",
        compression_dim: int = 512,
        compressor_path: str | None = None,
    ):
        """Initialize retrieval task.

        Args:
            top_k: Maximum K for retrieval search.
                Deprecated — prefer ``top_k_list``.
            top_k_list: List of K values to evaluate (all derived from
                a single max-K search). When not provided, ``[top_k]``
                is used as fallback.
            dino_model: DINO model name used for feature extraction.
            compression_dim: Output dimension of the hash compressor.
            compressor_path: Optional path to trained hash compressor weights.
        """
        if top_k_list is not None:
            top_k_list = sorted(set(top_k_list))
        else:
            top_k_list = [top_k]
        super().__init__(top_k_list=top_k_list)
        self.top_k_list: list[int] = top_k_list
        self.max_k: int = max(self.top_k_list)
        self.dino_model = dino_model
        self.compression_dim = compression_dim
        self.compressor_path = compressor_path
        self._processor: BitImageProcessor | None = None
        self._model: nn.Module | None = None
        self._model_name = "hash_compressor" if compressor_path else "dinov2"

    def prepare_benchmark(
        self,
        model: Any,
        processor: Any,
        model_name: str = "model",
    ) -> tuple[nn.Module, BitImageProcessor, str]:
        """Resolve benchmark resources for this task.

        Args:
            model: Optional pre-built model from the caller.
            processor: Optional pre-built processor from the caller.
            model_name: Fallback table model name.

        Returns:
            Tuple of benchmark model, processor, and resolved model name.
        """
        if model is not None and processor is not None:
            return (
                cast(nn.Module, model),
                cast(BitImageProcessor, processor),
                model_name,
            )

        self._ensure_resources_loaded()
        return (
            cast(nn.Module, self._model),
            cast(BitImageProcessor, self._processor),
            self._model_name,
        )

    def _ensure_resources_loaded(self) -> None:
        """Lazy-load retrieval benchmark resources."""
        if self._processor is not None and self._model is not None:
            return

        processor, dino = load_dino_model(self.dino_model)

        compressor = None
        if self.compressor_path is not None:
            compressor = load_hash_compressor(
                input_dim=get_dino_dim(self.dino_model),
                hash_bits=self.compression_dim,
                compressor_path=self.compressor_path,
            )
            compressor.eval()

        self._processor = processor
        self._model = RetrievalEncoder(dino=dino, compressor=compressor)
        self._model.eval()

    def build_database(
        self,
        model: Any,
        processor: Any,
        train_dataset: Any,
        table: lancedb.table.Table,
        batch_size: int,
        label_column: str = "label",
    ) -> None:
        """Build the evaluation database from training data.

        Args:
            model: Feature extraction model.
            processor: Image preprocessor.
            train_dataset: Training dataset.
            table: LanceDB table to store features.
            batch_size: Batch size for DataLoader.
            label_column: Column name for labels in the batch dict.
        """
        # Get a sample image to infer vector dimension
        sample = train_dataset[0]
        sample_image = sample["img"]

        vector_dim = infer_vector_dim(processor, model, sample_image)
        expected_schema = _build_eval_schema(vector_dim)

        # Check schema compatibility
        if table.schema != expected_schema:
            raise ValueError(
                f"Table schema mismatch. Expected: {expected_schema}, "
                f"Got: {table.schema}"
            )

        # Build database
        train_loader = DataLoader(
            train_dataset.with_format("torch"),
            batch_size=batch_size,
            shuffle=False,
            num_workers=4,
        )
        _establish_eval_database(processor, model, table, train_loader, label_column)

    def evaluate(
        self,
        model: Any,
        processor: Any,
        test_dataset: Any,
        table: lancedb.table.Table,
        batch_size: int,
        label_column: str = "label",
    ) -> dict[str, Any]:
        """Evaluate the model on the test dataset.

        Args:
            model: Feature extraction model.
            processor: Image preprocessor.
            test_dataset: Test dataset.
            table: LanceDB table to search against.
            batch_size: Batch size for DataLoader.
            label_column: Column name for labels in the batch dict.

        Returns:
            Dictionary containing evaluation results with keys:
                - recalls: Dict of ``{"recall@K": value}`` for each K
                - total: Total number of test samples
                - top_k_list: List of K values evaluated
                - num_classes: Number of label classes
                - confusion_matrix: Row-normalized confusion matrix (Top-1)
                - query_labels: True labels for each query sample
                - topk_labels: Top-K retrieved labels per query sample
        """
        test_loader = DataLoader(
            test_dataset.with_format("torch"),
            batch_size=batch_size,
            shuffle=False,
            num_workers=4,
        )
        eval_data = _evaluate_recall(
            processor, model, table, test_loader, self.max_k, label_column
        )

        total = eval_data["total"]
        query_labels = eval_data["query_labels"]
        topk_labels = eval_data["topk_labels"]

        # Compute Recall@k for each k in top_k_list
        recalls: dict[str, float] = {}
        for k in self.top_k_list:
            hits = sum(
                1
                for true, preds in zip(query_labels, topk_labels)
                if true in preds[:k]
            )
            recalls[f"recall@{k}"] = hits / total if total > 0 else 0.0

        # Infer number of classes from the data
        all_labels = query_labels + [
            label for labels in topk_labels for label in labels
        ]
        num_classes = max(all_labels) + 1 if all_labels else 1

        # Confusion matrix from Top-1 only
        cm = _build_confusion_matrix(query_labels, topk_labels, num_classes)

        return {
            "recalls": recalls,
            "total": total,
            "top_k_list": self.top_k_list,
            "num_classes": num_classes,
            "confusion_matrix": cm.tolist(),
            "query_labels": query_labels,
            "topk_labels": topk_labels,
        }