"""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, }