Compare commits

...

2 Commits

Author SHA1 Message Date
pluto 2df137ee2a feat: add recognition engine with CLIP matching and Ollama vision LLM
- OllamaClient: analyze_image, compare_images, is_available
- Recognition pipeline: CLIP embedding -> cosine similarity -> threshold
  routing (known >0.85, unknown <0.60, ambiguous -> Ollama fallback)
- Vision dispatcher: get_description() supports haiku/ollama/both backends
- 9 embedding serialization tests

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-07 17:17:49 -06:00
pluto 74e12eca60 feat: add CLIP embedding service sidecar for GPU machine
FastAPI service that loads ViT-L-14 model and exposes /embed,
/similarity, and /batch-similarity endpoints. Runs on RTX 3090
with NVIDIA GPU reservation. Includes 9 cosine similarity tests.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-07 17:17:42 -06:00
10 changed files with 951 additions and 0 deletions
+132
View File
@@ -0,0 +1,132 @@
"""Ollama vision LLM client for frigate-notify v2."""
import base64
import logging
import aiohttp
logger = logging.getLogger(__name__)
class OllamaClient:
"""Client for Ollama vision models (e.g., llava)."""
def __init__(self, base_url: str, model: str):
self.base_url = base_url.rstrip("/")
self.model = model
async def analyze_image(self, image_bytes: bytes, prompt: str) -> str | None:
"""Send an image to Ollama for analysis.
Args:
image_bytes: Raw image bytes (JPEG).
prompt: Text prompt for the vision model.
Returns:
Response text or None on failure.
"""
b64 = base64.b64encode(image_bytes).decode("utf-8")
payload = {
"model": self.model,
"prompt": prompt,
"images": [b64],
"stream": False,
"options": {"num_predict": 200},
}
try:
async with aiohttp.ClientSession() as session:
async with session.post(
f"{self.base_url}/api/generate",
json=payload,
timeout=aiohttp.ClientTimeout(total=30),
) as resp:
if resp.status == 200:
data = await resp.json()
text = data.get("response", "").strip()
logger.info(
"Ollama analysis complete: %d chars", len(text)
)
return text if text else None
body = await resp.text()
logger.error(
"Ollama analyze_image failed %s: %s", resp.status, body
)
except Exception:
logger.exception("Ollama analyze_image request failed")
return None
async def compare_images(
self,
image_bytes: bytes,
reference_bytes: bytes,
object_name: str,
) -> bool | None:
"""Compare two images to determine if they show the same object/person.
Args:
image_bytes: The new event image bytes.
reference_bytes: The reference (known) image bytes.
object_name: Name of the known object for context.
Returns:
True if same, False if different, None on failure/inconclusive.
"""
b64_new = base64.b64encode(image_bytes).decode("utf-8")
b64_ref = base64.b64encode(reference_bytes).decode("utf-8")
prompt = (
f"Look at these two images. The first image is from a security camera. "
f"The second image is a reference photo of '{object_name}'. "
f"Are these the same person/vehicle/animal? "
f"Answer with ONLY 'yes' or 'no'."
)
payload = {
"model": self.model,
"prompt": prompt,
"images": [b64_new, b64_ref],
"stream": False,
"options": {"num_predict": 200},
}
try:
async with aiohttp.ClientSession() as session:
async with session.post(
f"{self.base_url}/api/generate",
json=payload,
timeout=aiohttp.ClientTimeout(total=30),
) as resp:
if resp.status == 200:
data = await resp.json()
text = data.get("response", "").strip().lower()
logger.info("Ollama compare result: %s", text)
if "yes" in text.split()[:3]:
return True
if "no" in text.split()[:3]:
return False
logger.warning(
"Ollama compare inconclusive: %s", text
)
return None
body = await resp.text()
logger.error(
"Ollama compare_images failed %s: %s",
resp.status,
body,
)
except Exception:
logger.exception("Ollama compare_images request failed")
return None
async def is_available(self) -> bool:
"""Check if Ollama is reachable and responding.
Returns:
True if Ollama responds with HTTP 200, False otherwise.
"""
try:
async with aiohttp.ClientSession() as session:
async with session.get(
f"{self.base_url}/api/tags",
timeout=aiohttp.ClientTimeout(total=3),
) as resp:
return resp.status == 200
except Exception:
return False
+232
View File
@@ -0,0 +1,232 @@
"""Recognition pipeline using CLIP embeddings with Ollama vision fallback."""
import base64
import logging
import struct
from typing import Optional
import aiohttp
import numpy as np
from config import config
from database import get_connection
from ollama_client import OllamaClient
logger = logging.getLogger(__name__)
def embedding_to_bytes(embedding: list[float]) -> bytes:
"""Pack a float32 embedding list into bytes for SQLite storage.
Args:
embedding: List of float values (typically 768-dimensional).
Returns:
Packed bytes (len(embedding) * 4 bytes).
"""
return struct.pack(f"{len(embedding)}f", *embedding)
def bytes_to_embedding(data: bytes) -> np.ndarray:
"""Unpack bytes back into a float32 numpy array.
Args:
data: Packed bytes from embedding_to_bytes().
Returns:
Numpy array of float32 values.
"""
count = len(data) // 4
values = struct.unpack(f"{count}f", data)
return np.array(values, dtype=np.float32)
async def get_clip_embedding(image_bytes: bytes) -> list[float] | None:
"""Get CLIP embedding for an image from the clip-service sidecar.
Args:
image_bytes: Raw image bytes (JPEG).
Returns:
List of float embedding values, or None on failure.
"""
b64 = base64.b64encode(image_bytes).decode("utf-8")
try:
async with aiohttp.ClientSession() as session:
async with session.post(
f"{config.CLIP_SERVICE_URL}/embed",
json={"image": b64},
timeout=aiohttp.ClientTimeout(total=10),
) as resp:
if resp.status == 200:
data = await resp.json()
embedding = data.get("embedding")
if embedding:
logger.info(
"CLIP embedding received: %d dimensions",
len(embedding),
)
return embedding
logger.warning("CLIP response missing embedding field")
else:
body = await resp.text()
logger.error(
"CLIP embed failed %s: %s", resp.status, body
)
except Exception:
logger.exception("CLIP embedding request failed")
return None
async def find_best_match(
query_embedding: list[float],
) -> tuple[int | None, float, str | None]:
"""Find the best matching known object for a query embedding.
Loads all embeddings from the database, computes cosine similarity,
groups by known_object_id and takes the max score per object.
Args:
query_embedding: The query embedding to match against.
Returns:
Tuple of (best_object_id, best_score, best_name).
Returns (None, 0.0, None) if no embeddings exist.
"""
conn = get_connection()
cursor = conn.execute(
"""
SELECT e.id, e.known_object_id, e.embedding, ko.name
FROM embeddings e
JOIN known_objects ko ON e.known_object_id = ko.id
"""
)
rows = cursor.fetchall()
if not rows:
return (None, 0.0, None)
query_vec = np.array(query_embedding, dtype=np.float32)
query_norm = np.linalg.norm(query_vec)
if query_norm == 0:
return (None, 0.0, None)
query_vec = query_vec / query_norm
# Group scores by known_object_id, keeping max per object
best_per_object: dict[int, tuple[float, str]] = {}
for row in rows:
stored_vec = bytes_to_embedding(row["embedding"])
stored_norm = np.linalg.norm(stored_vec)
if stored_norm == 0:
continue
stored_vec = stored_vec / stored_norm
similarity = float(np.dot(query_vec, stored_vec))
obj_id = row["known_object_id"]
obj_name = row["name"]
if obj_id not in best_per_object or similarity > best_per_object[obj_id][0]:
best_per_object[obj_id] = (similarity, obj_name)
if not best_per_object:
return (None, 0.0, None)
best_id = max(best_per_object, key=lambda k: best_per_object[k][0])
best_score, best_name = best_per_object[best_id]
return (best_id, best_score, best_name)
async def recognize_event(
image_bytes: bytes,
ollama_client: Optional[OllamaClient] = None,
) -> tuple[str, int | None, float, str | None]:
"""Full recognition pipeline for a Frigate event image.
Steps:
1. Get CLIP embedding for the image.
2. Find the best matching known object.
3. High confidence (>= threshold_high): return "known".
4. Low confidence (< threshold_low): return "unknown".
5. Ambiguous range: use Ollama vision LLM to confirm/deny.
6. If LLM unavailable or inconclusive: return "ambiguous".
7. If no embedding could be obtained: return "unreviewed".
Args:
image_bytes: Raw image bytes (JPEG).
ollama_client: Optional OllamaClient for ambiguous case fallback.
Returns:
Tuple of (match_status, object_id, score, object_name).
match_status is one of: "known", "unknown", "ambiguous", "unreviewed".
"""
# Step 1: Get CLIP embedding
embedding = await get_clip_embedding(image_bytes)
if embedding is None:
logger.warning("Could not get CLIP embedding, returning unreviewed")
return ("unreviewed", None, 0.0, None)
# Step 2: Find best match
best_id, best_score, best_name = await find_best_match(embedding)
# No embeddings in DB at all
if best_id is None:
return ("unknown", None, 0.0, None)
high_threshold = config.RECOGNITION_THRESHOLD_HIGH
low_threshold = config.RECOGNITION_THRESHOLD_LOW
# Step 3: High confidence match
if best_score >= high_threshold:
logger.info(
"High confidence match: %s (%.3f)", best_name, best_score
)
return ("known", best_id, best_score, best_name)
# Step 4: Low confidence - clearly not a match
if best_score < low_threshold:
logger.info("Low confidence, no match (%.3f)", best_score)
return ("unknown", None, best_score, None)
# Step 5: Ambiguous range - try Ollama vision LLM
if ollama_client is not None:
logger.info(
"Ambiguous score %.3f for %s, trying Ollama comparison",
best_score,
best_name,
)
# Get a reference thumbnail from the embeddings table
conn = get_connection()
cursor = conn.execute(
"""
SELECT thumbnail FROM embeddings
WHERE known_object_id = ? AND thumbnail IS NOT NULL
LIMIT 1
""",
(best_id,),
)
ref_row = cursor.fetchone()
if ref_row and ref_row["thumbnail"] is not None:
result = await ollama_client.compare_images(
image_bytes, ref_row["thumbnail"], best_name
)
if result is True:
logger.info(
"Ollama confirmed match: %s (%.3f)",
best_name,
best_score,
)
return ("known", best_id, best_score, best_name)
elif result is False:
logger.info(
"Ollama denied match: %s (%.3f)", best_name, best_score
)
return ("unknown", None, best_score, None)
# result is None: inconclusive, fall through to ambiguous
# Step 6: LLM unavailable or inconclusive
logger.info(
"Ambiguous match: %s (%.3f)", best_name, best_score
)
return ("ambiguous", best_id, best_score, best_name)
+67
View File
@@ -1,7 +1,11 @@
import base64
import logging
from typing import Optional
import anthropic
from ollama_client import OllamaClient
logger = logging.getLogger(__name__)
PROMPT = (
@@ -40,3 +44,66 @@ def analyze_snapshot(client: anthropic.Anthropic, image_bytes: bytes) -> str | N
except Exception:
logger.exception("Claude Vision API failed")
return None
async def analyze_snapshot_ollama(
ollama_client: OllamaClient, image_bytes: bytes
) -> str | None:
"""Send a snapshot to Ollama vision model for description.
Args:
ollama_client: Configured OllamaClient instance.
image_bytes: Raw image bytes (JPEG).
Returns:
Description string or None on failure.
"""
return await ollama_client.analyze_image(image_bytes, PROMPT)
async def get_description(
image_bytes: bytes,
vision_backend: str,
claude_client: Optional[anthropic.Anthropic] = None,
ollama_client: Optional[OllamaClient] = None,
) -> str | None:
"""Get a scene description using the configured vision backend.
Args:
image_bytes: Raw image bytes (JPEG).
vision_backend: One of "haiku", "ollama", or "both".
"both" tries Ollama first, falls back to Haiku.
claude_client: Anthropic client (required for "haiku" and "both").
ollama_client: OllamaClient instance (required for "ollama" and "both").
Returns:
Description string or None if all backends fail.
"""
if vision_backend == "ollama":
if ollama_client is None:
logger.error("Ollama backend selected but no client provided")
return None
return await analyze_snapshot_ollama(ollama_client, image_bytes)
if vision_backend == "haiku":
if claude_client is None:
logger.error("Haiku backend selected but no client provided")
return None
return analyze_snapshot(claude_client, image_bytes)
if vision_backend == "both":
# Try Ollama first, fall back to Haiku
if ollama_client is not None:
result = await analyze_snapshot_ollama(ollama_client, image_bytes)
if result is not None:
return result
logger.info("Ollama failed, falling back to Haiku")
if claude_client is not None:
return analyze_snapshot(claude_client, image_bytes)
logger.error("Both backend selected but no clients provided")
return None
logger.error("Unknown vision backend: %s", vision_backend)
return None
+15
View File
@@ -0,0 +1,15 @@
FROM pytorch/pytorch:2.5.1-cuda12.4-cudnn9-runtime
RUN apt-get update && apt-get install -y --no-install-recommends curl \
&& rm -rf /var/lib/apt/lists/*
WORKDIR /app
COPY app/requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
COPY app/ .
EXPOSE 5200
CMD ["python", "-u", "-m", "uvicorn", "main:app", "--host", "0.0.0.0", "--port", "5200"]
+194
View File
@@ -0,0 +1,194 @@
"""CLIP embedding service for frigate-notify.
Loads OpenAI ViT-L/14 CLIP model and exposes endpoints for generating
image embeddings and computing cosine similarity between embeddings.
Designed to run on a GPU machine (RTX 3090) as a sidecar service.
"""
import base64
import io
import logging
import time
from contextlib import asynccontextmanager
import numpy as np
import open_clip
import torch
from fastapi import FastAPI, HTTPException
from PIL import Image
from pydantic import BaseModel
logger = logging.getLogger("clip-service")
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
# --- Pydantic models ---
class EmbedRequest(BaseModel):
image: str # base64-encoded image
class EmbedResponse(BaseModel):
embedding: list[float]
inference_ms: float
class SimilarityRequest(BaseModel):
embedding1: list[float]
embedding2: list[float]
class SimilarityResponse(BaseModel):
score: float
class BatchSimilarityRequest(BaseModel):
query_embedding: list[float]
candidate_embeddings: list[list[float]]
class BatchSimilarityResponse(BaseModel):
scores: list[float]
best_index: int
best_score: float
# --- Globals ---
model = None
preprocess = None
device = None
model_name = "ViT-L-14"
pretrained = "openai"
def load_model():
"""Load the CLIP model onto the best available device."""
global model, preprocess, device
if torch.cuda.is_available():
device = torch.device("cuda")
logger.info("CUDA available — using GPU: %s", torch.cuda.get_device_name(0))
else:
device = torch.device("cpu")
logger.warning("CUDA not available — falling back to CPU")
logger.info("Loading CLIP model %s (pretrained=%s)...", model_name, pretrained)
t0 = time.monotonic()
model, _, preprocess = open_clip.create_model_and_transforms(
model_name, pretrained=pretrained, device=device
)
model.eval()
elapsed = time.monotonic() - t0
logger.info("Model loaded in %.1fs on %s", elapsed, device)
@asynccontextmanager
async def lifespan(app: FastAPI):
load_model()
yield
app = FastAPI(title="CLIP Embedding Service", version="1.0.0", lifespan=lifespan)
# --- Utility functions ---
def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
"""Compute cosine similarity between two vectors."""
dot = np.dot(a, b)
norm_a = np.linalg.norm(a)
norm_b = np.linalg.norm(b)
if norm_a == 0 or norm_b == 0:
return 0.0
return float(dot / (norm_a * norm_b))
# --- Endpoints ---
@app.get("/health")
async def health():
return {
"status": "healthy",
"device": str(device) if device else "not_loaded",
"model": f"{model_name}/{pretrained}",
}
@app.post("/embed", response_model=EmbedResponse)
async def embed(request: EmbedRequest):
if model is None or preprocess is None:
raise HTTPException(status_code=503, detail="Model not loaded")
# Decode base64 image
try:
image_bytes = base64.b64decode(request.image)
image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
except Exception as e:
raise HTTPException(status_code=400, detail=f"Invalid image data: {e}")
# Preprocess and encode
try:
image_tensor = preprocess(image).unsqueeze(0).to(device)
t0 = time.monotonic()
with torch.no_grad():
embedding = model.encode_image(image_tensor)
inference_ms = (time.monotonic() - t0) * 1000
# Normalize the embedding
embedding = embedding / embedding.norm(dim=-1, keepdim=True)
embedding_list = embedding.squeeze(0).cpu().numpy().tolist()
except Exception as e:
raise HTTPException(status_code=500, detail=f"Inference failed: {e}")
return EmbedResponse(embedding=embedding_list, inference_ms=round(inference_ms, 2))
@app.post("/similarity", response_model=SimilarityResponse)
async def similarity(request: SimilarityRequest):
a = np.array(request.embedding1, dtype=np.float32)
b = np.array(request.embedding2, dtype=np.float32)
if a.shape != b.shape:
raise HTTPException(
status_code=400,
detail=f"Embedding dimensions must match: {a.shape} vs {b.shape}",
)
score = cosine_similarity(a, b)
return SimilarityResponse(score=round(score, 6))
@app.post("/batch-similarity", response_model=BatchSimilarityResponse)
async def batch_similarity(request: BatchSimilarityRequest):
query = np.array(request.query_embedding, dtype=np.float32)
if not request.candidate_embeddings:
raise HTTPException(status_code=400, detail="No candidate embeddings provided")
candidates = np.array(request.candidate_embeddings, dtype=np.float32)
if query.shape[0] != candidates.shape[1]:
raise HTTPException(
status_code=400,
detail=f"Dimension mismatch: query={query.shape[0]}, candidates={candidates.shape[1]}",
)
# Compute cosine similarity for all candidates at once
query_norm = query / (np.linalg.norm(query) + 1e-10)
candidate_norms = candidates / (
np.linalg.norm(candidates, axis=1, keepdims=True) + 1e-10
)
scores = candidate_norms @ query_norm
scores_list = [round(float(s), 6) for s in scores]
best_index = int(np.argmax(scores))
best_score = scores_list[best_index]
return BatchSimilarityResponse(
scores=scores_list, best_index=best_index, best_score=best_score
)
+7
View File
@@ -0,0 +1,7 @@
fastapi>=0.115.0
uvicorn[standard]>=0.34.0
open-clip-torch>=2.29.0
torch>=2.5.0
Pillow>=10.0.0
numpy>=1.26.0
python-multipart>=0.0.18
+27
View File
@@ -0,0 +1,27 @@
services:
clip-service:
build: .
container_name: clip-service
restart: unless-stopped
ports:
- "5200:5200"
deploy:
resources:
reservations:
devices:
- driver: nvidia
count: 1
capabilities: [gpu]
mem_limit: 4G
memswap_limit: 4G
healthcheck:
test: ["CMD-SHELL", "curl -sf http://localhost:5200/health || exit 1"]
interval: 30s
timeout: 10s
retries: 3
start_period: 120s
logging:
driver: "json-file"
options:
max-size: "50m"
max-file: "3"
+65
View File
@@ -0,0 +1,65 @@
"""Test configuration: mock heavy dependencies that aren't needed for unit tests."""
import sys
from unittest.mock import MagicMock
# Mock modules that are not installed in the test environment.
# The tests only exercise pure numpy math, not model inference.
for mod_name in [
"torch",
"open_clip",
"fastapi",
"fastapi.responses",
"pydantic",
"PIL",
"PIL.Image",
"uvicorn",
]:
if mod_name not in sys.modules:
sys.modules[mod_name] = MagicMock()
# Provide a real BaseModel stand-in so Pydantic models can be defined
# (they are referenced at import time in main.py).
import types
pydantic_mock = sys.modules["pydantic"]
class _FakeBaseModel:
def __init_subclass__(cls, **kwargs):
pass
def __init__(self, **kwargs):
for k, v in kwargs.items():
setattr(self, k, v)
pydantic_mock.BaseModel = _FakeBaseModel
# Provide a fake FastAPI and HTTPException so the module loads
fastapi_mock = sys.modules["fastapi"]
class _FakeFastAPI:
def __init__(self, **kwargs):
pass
def get(self, *a, **kw):
def decorator(fn):
return fn
return decorator
def post(self, *a, **kw):
def decorator(fn):
return fn
return decorator
class _FakeHTTPException(Exception):
def __init__(self, status_code=500, detail=""):
self.status_code = status_code
self.detail = detail
fastapi_mock.FastAPI = _FakeFastAPI
fastapi_mock.HTTPException = _FakeHTTPException
+130
View File
@@ -0,0 +1,130 @@
"""Tests for CLIP service math utilities.
These tests exercise the cosine similarity logic using pure numpy --
no GPU, no model loading, no network calls required.
"""
import numpy as np
import pytest
# Import the cosine_similarity function directly from the app module.
# We add the app directory to sys.path so the import works without
# installing the package.
import sys
import os
sys.path.insert(
0, os.path.join(os.path.dirname(__file__), "..", "app")
)
from main import cosine_similarity
class TestCosineSimilarity:
"""Test the cosine_similarity helper function."""
def test_identical_vectors(self):
"""Identical vectors should have cosine similarity of 1.0."""
v = np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32)
score = cosine_similarity(v, v)
assert score == pytest.approx(1.0, abs=1e-6)
def test_identical_unit_vectors(self):
"""Identical unit vectors should also yield 1.0."""
v = np.array([0.0, 1.0, 0.0], dtype=np.float32)
score = cosine_similarity(v, v)
assert score == pytest.approx(1.0, abs=1e-6)
def test_orthogonal_vectors(self):
"""Orthogonal vectors should have cosine similarity of 0.0."""
a = np.array([1.0, 0.0, 0.0], dtype=np.float32)
b = np.array([0.0, 1.0, 0.0], dtype=np.float32)
score = cosine_similarity(a, b)
assert score == pytest.approx(0.0, abs=1e-6)
def test_opposite_vectors(self):
"""Opposite vectors should have cosine similarity of -1.0."""
a = np.array([1.0, 2.0, 3.0], dtype=np.float32)
b = np.array([-1.0, -2.0, -3.0], dtype=np.float32)
score = cosine_similarity(a, b)
assert score == pytest.approx(-1.0, abs=1e-6)
def test_zero_vector_returns_zero(self):
"""A zero vector should return 0.0 (guarded against division by zero)."""
a = np.array([0.0, 0.0, 0.0], dtype=np.float32)
b = np.array([1.0, 2.0, 3.0], dtype=np.float32)
score = cosine_similarity(a, b)
assert score == pytest.approx(0.0, abs=1e-6)
def test_high_dimensional(self):
"""Test with 768-dimensional vectors (CLIP ViT-L/14 output size)."""
rng = np.random.RandomState(42)
a = rng.randn(768).astype(np.float32)
# Similarity with itself should be 1.0
score = cosine_similarity(a, a)
assert score == pytest.approx(1.0, abs=1e-5)
class TestBatchSimilarity:
"""Test batch similarity logic (matching the endpoint's numpy math)."""
def _batch_similarity(
self, query: np.ndarray, candidates: np.ndarray
) -> tuple[list[float], int, float]:
"""Replicate the batch-similarity endpoint logic."""
query_norm = query / (np.linalg.norm(query) + 1e-10)
candidate_norms = candidates / (
np.linalg.norm(candidates, axis=1, keepdims=True) + 1e-10
)
scores = candidate_norms @ query_norm
scores_list = [float(s) for s in scores]
best_index = int(np.argmax(scores))
best_score = scores_list[best_index]
return scores_list, best_index, best_score
def test_best_match_index(self):
"""The candidate most similar to the query should be identified."""
query = np.array([1.0, 0.0, 0.0], dtype=np.float32)
candidates = np.array(
[
[0.0, 1.0, 0.0], # orthogonal
[0.9, 0.1, 0.0], # close match
[1.0, 0.0, 0.0], # exact match
[0.5, 0.5, 0.0], # partial match
],
dtype=np.float32,
)
scores, best_idx, best_score = self._batch_similarity(query, candidates)
assert best_idx == 2
assert best_score == pytest.approx(1.0, abs=1e-5)
assert len(scores) == 4
# Orthogonal candidate should score ~0
assert scores[0] == pytest.approx(0.0, abs=1e-5)
def test_single_candidate(self):
"""Batch with a single candidate should work correctly."""
query = np.array([1.0, 1.0], dtype=np.float32)
candidates = np.array([[1.0, 1.0]], dtype=np.float32)
scores, best_idx, best_score = self._batch_similarity(query, candidates)
assert best_idx == 0
assert best_score == pytest.approx(1.0, abs=1e-5)
def test_scores_ordering(self):
"""Scores should reflect relative similarity to the query."""
query = np.array([1.0, 0.0], dtype=np.float32)
candidates = np.array(
[
[0.5, 0.866], # ~60 degrees
[0.866, 0.5], # ~30 degrees
[1.0, 0.0], # 0 degrees (exact)
],
dtype=np.float32,
)
scores, best_idx, _ = self._batch_similarity(query, candidates)
# Exact match should be best
assert best_idx == 2
# 30 degrees should score higher than 60 degrees
assert scores[1] > scores[0]
+82
View File
@@ -0,0 +1,82 @@
"""Tests for the recognition module — embedding serialization helpers."""
import numpy as np
import pytest
from recognition import bytes_to_embedding, embedding_to_bytes
class TestEmbeddingToBytes:
"""Tests for embedding_to_bytes()."""
def test_roundtrip_preserves_values(self):
"""Packing and unpacking should return the same values."""
original = [0.1, 0.2, 0.3, -0.5, 1.0, 0.0]
packed = embedding_to_bytes(original)
recovered = bytes_to_embedding(packed)
np.testing.assert_allclose(recovered, original, rtol=1e-6)
def test_correct_byte_length_768(self):
"""A 768-dim embedding should produce 768 * 4 = 3072 bytes."""
embedding = [0.0] * 768
packed = embedding_to_bytes(embedding)
assert len(packed) == 768 * 4
def test_correct_byte_length_arbitrary(self):
"""Byte length should always be len(embedding) * 4."""
for dim in [1, 10, 512, 1024]:
embedding = [1.0] * dim
packed = embedding_to_bytes(embedding)
assert len(packed) == dim * 4
def test_empty_embedding(self):
"""An empty embedding should produce 0 bytes."""
packed = embedding_to_bytes([])
assert len(packed) == 0
recovered = bytes_to_embedding(packed)
assert len(recovered) == 0
class TestBytesToEmbedding:
"""Tests for bytes_to_embedding()."""
def test_correct_shape(self):
"""Should return a 1-D array with the correct number of elements."""
embedding = [0.5] * 768
packed = embedding_to_bytes(embedding)
recovered = bytes_to_embedding(packed)
assert recovered.shape == (768,)
def test_correct_dtype(self):
"""Should return float32 numpy array."""
embedding = [0.1, 0.2, 0.3]
packed = embedding_to_bytes(embedding)
recovered = bytes_to_embedding(packed)
assert recovered.dtype == np.float32
def test_negative_values(self):
"""Negative values should survive roundtrip."""
original = [-1.0, -0.5, 0.0, 0.5, 1.0]
packed = embedding_to_bytes(original)
recovered = bytes_to_embedding(packed)
np.testing.assert_allclose(recovered, original, rtol=1e-6)
def test_large_values(self):
"""Large float values should survive roundtrip within float32 precision."""
original = [1e10, -1e10, 3.14159, 2.71828]
packed = embedding_to_bytes(original)
recovered = bytes_to_embedding(packed)
np.testing.assert_allclose(recovered, original, rtol=1e-5)
def test_realistic_clip_embedding(self):
"""Simulate a realistic CLIP embedding (768-dim, normalized)."""
rng = np.random.default_rng(42)
vec = rng.standard_normal(768).astype(np.float32)
vec = vec / np.linalg.norm(vec) # normalize
packed = embedding_to_bytes(vec.tolist())
recovered = bytes_to_embedding(packed)
assert recovered.shape == (768,)
assert recovered.dtype == np.float32
np.testing.assert_allclose(recovered, vec, rtol=1e-6)