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>
This commit is contained in:
@@ -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"]
|
||||
@@ -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
|
||||
)
|
||||
@@ -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
|
||||
@@ -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"
|
||||
@@ -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
|
||||
@@ -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]
|
||||
Reference in New Issue
Block a user