Spaces:
Sleeping
Sleeping
File size: 5,366 Bytes
f343f06 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 | """
Retrieval engine combining feature extraction and FAISS search.
Provides high-level API for image retrieval.
"""
import torch
import time
from PIL import Image
from typing import List, Dict, Optional, Tuple
from dataclasses import dataclass
from ..features.extractor import FeatureExtractor
from .index import FAISSIndex
@dataclass
class RetrievalResult:
"""Result of a single retrieval query."""
indices: List[int]
scores: List[float]
query_time_ms: float
modality: str
class RetrievalEngine:
"""
High-level retrieval engine.
Combines feature extraction with FAISS search for fast image retrieval.
"""
def __init__(
self,
feature_extractor: Optional[FeatureExtractor] = None,
index: Optional[FAISSIndex] = None,
device: Optional[str] = None
):
"""
Initialize retrieval engine.
Args:
feature_extractor: Feature extractor (creates new if None)
index: FAISS index (creates new if None)
device: Device to use
"""
self.feature_extractor = feature_extractor or FeatureExtractor(device=device)
self.index = index or FAISSIndex(embed_dim=self.feature_extractor.embed_dim)
# Timing statistics
self._query_times: List[float] = []
def build_index(
self,
embeddings: torch.Tensor,
save_path: Optional[str] = None
) -> None:
"""
Build index from pre-computed embeddings.
Args:
embeddings: Gallery embeddings, shape (N, embed_dim)
save_path: Optional path to save index
"""
self.index.build(embeddings)
if save_path:
self.index.save(save_path)
def query(
self,
image: Image.Image,
modality: str = "optical",
k: int = 5
) -> RetrievalResult:
"""
Query with a single image.
Args:
image: Query image
modality: Image modality
k: Number of results
Returns:
RetrievalResult with indices, scores, and timing
"""
start_time = time.perf_counter()
# Extract features
query_embedding = self.feature_extractor.extract_features(
image, modality=modality, normalize=True
)
# Search
scores, indices = self.index.search(query_embedding, k=k)
elapsed_ms = (time.perf_counter() - start_time) * 1000
self._query_times.append(elapsed_ms)
return RetrievalResult(
indices=indices[0].tolist(),
scores=scores[0].tolist(),
query_time_ms=elapsed_ms,
modality=modality
)
def batch_query(
self,
images: List[Image.Image],
modality: str = "optical",
k: int = 5
) -> List[RetrievalResult]:
"""
Query with multiple images.
Args:
images: List of query images
modality: Image modality
k: Number of results
Returns:
List of RetrievalResult
"""
results = []
for image in images:
result = self.query(image, modality=modality, k=k)
results.append(result)
return results
def get_timing_stats(self) -> Dict[str, float]:
"""
Get timing statistics.
Returns:
Dict with mean, median, p95, p99 query times
"""
if not self._query_times:
return {"mean": 0, "median": 0, "p95": 0, "p99": 0}
times = sorted(self._query_times)
n = len(times)
return {
"mean": sum(times) / n,
"median": times[n // 2],
"p95": times[int(n * 0.95)] if n >= 20 else times[-1],
"p99": times[int(n * 0.99)] if n >= 100 else times[-1],
}
@property
def _query_times(self) -> List[float]:
"""Query times list (lazy init)."""
if not hasattr(self, '_query_times_list'):
self._query_times_list = []
return self._query_times_list
# Self-check
if __name__ == "__main__":
print("Testing RetrievalEngine...")
# Create dummy data
n_gallery = 50
embed_dim = 768
# Build index
embeddings = torch.randn(n_gallery, embed_dim)
embeddings = torch.nn.functional.normalize(embeddings, dim=1)
# Initialize engine (without model for testing)
engine = RetrievalEngine.__new__(RetrievalEngine)
engine.index = FAISSIndex(embed_dim)
engine._query_times_list = []
# Build index
engine.build_index(embeddings)
print(f"Index built with {engine.index.size} embeddings")
# Simulate query timing
for _ in range(10):
start = time.perf_counter()
query = torch.randn(embed_dim)
query = torch.nn.functional.normalize(query, dim=0)
scores, indices = engine.index.search(query, k=5)
elapsed = (time.perf_counter() - start) * 1000
engine._query_times_list.append(elapsed)
# Get stats
stats = engine.get_timing_stats()
print(f"Timing stats: {stats}")
print("\nRetrievalEngine test passed!")
|