Build an offline visual defect search system with CLIP and VectorAI DB

uv init defect-search && cd defect-search
uv add torch transformers pillow actian-vectorai-client matplotlib

uv run huggingface-cli download openai/clip-vit-base-patch32

services:
  vectorai:
    image: actian/vectorai:latest
    ports:
      - "6573:6573"   # REST API
      - "6574:6574"   # gRPC API, used by the Python SDK
      - "6575:6575"   # local web UI
    volumes:
      - ./vectorai_data:/var/lib/actian-vectorai
    environment:
      - ACTIAN_VECTORAI_ACCEPT_EULA=YES
    restart: unless-stopped

docker-compose up -d

# smoke_test.py
from importlib import import_module

VectorAIClient = import_module("actian_vectorai").VectorAIClient

with VectorAIClient("localhost:6574") as client:
    print(client.health_check())

uv run smoke_test.py

from __future__ import annotations

from pathlib import Path
from typing import Sequence

import torch
from PIL import Image
from transformers import CLIPModel, CLIPProcessor

MODEL_ID = "openai/clip-vit-base-patch32"
EMBED_DIM = 512


def _features(out) -> torch.Tensor:
    """
    transformers 4.x returns a Tensor from get_image_features and
    get_text_features. transformers 5.x returns a
    BaseModelOutputWithPooling whose pooler_output holds the
    projected vector. Handle both so this does not break on upgrade.
    """
    return out if isinstance(out, torch.Tensor) else out.pooler_output


class ClipEncoder:
    def __init__(self, model_id: str = MODEL_ID, device: str | None = None) -> None:
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.model = CLIPModel.from_pretrained(model_id).to(self.device).eval()
        self.processor = CLIPProcessor.from_pretrained(model_id)

    @torch.inference_mode()
    def encode_images(self, paths: Sequence[str | Path]) -> torch.Tensor:
        images = [Image.open(p).convert("RGB") for p in paths]
        inputs = self.processor(images=images, return_tensors="pt").to(self.device)
        vecs = _features(self.model.get_image_features(**inputs))
        return torch.nn.functional.normalize(vecs, dim=-1).cpu()

    @torch.inference_mode()
    def encode_text(self, texts: Sequence[str]) -> torch.Tensor:
        inputs = self.processor(
            text=list(texts), return_tensors="pt", padding=True, truncation=True
        ).to(self.device)
        vecs = _features(self.model.get_text_features(**inputs))
        return torch.nn.functional.normalize(vecs, dim=-1).cpu()

# ingest.py
from __future__ import annotations

import argparse
import time
from pathlib import Path

from actian_vectorai import (
    VectorAIClient, PointStruct, VectorParams, Distance, CollectionExistsError,
)

from clip_encoder import ClipEncoder, EMBED_DIM

IMAGE_SUFFIXES = {".png", ".jpg", ".jpeg", ".bmp"}


def scan_corpus(root: Path) -> list[dict]:
    """Read metadata straight out of the directory layout."""
    records: list[dict] = []
    for path in sorted(root.rglob("*")):
        if path.suffix.lower() not in IMAGE_SUFFIXES:
            continue
        if "ground_truth" in path.parts:      # segmentation masks, not photos
            continue
        try:
            rel = path.relative_to(root)
            part_category, _split, defect_type = rel.parts[0], rel.parts[1], rel.parts[2]
        except (ValueError, IndexError):
            continue
        records.append(
            {
                "image_path": str(path),
                "part_category": part_category,
                "defect_type": defect_type,
                "is_defect": defect_type != "good",
            }
        )
    return records


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", default="./mvtec_ad")
    ap.add_argument("--collection", default="defect_search")
    ap.add_argument("--url", default="localhost:6574")
    ap.add_argument("--batch-size", type=int, default=32)
    args = ap.parse_args()

    data_root = Path(args.data)
    if not data_root.is_dir():
        raise SystemExit(
            f"Dataset directory not found: {data_root}\n"
            "Download and extract MVTec AD, then pass its directory with --data."
        )

    records = scan_corpus(data_root)
    if not records:
        raise SystemExit(
            f"No dataset images found under {data_root}\n"
            "Expected paths like <data>/<category>/test/<defect>/*.png."
        )
    print(f"Found {len(records)} images")

    encoder = ClipEncoder()
    print(f"CLIP running on {encoder.device}")

    with VectorAIClient(args.url) as client:
        try:
            client.collections.create(
                args.collection,
                vectors_config=VectorParams(size=EMBED_DIM, distance=Distance.Cosine),
            )
            print(f"Created collection '{args.collection}'")
        except CollectionExistsError:
            print(f"Collection '{args.collection}' already exists, appending")

        started = time.perf_counter()
        done = 0

        for start in range(0, len(records), args.batch_size):
            batch = records[start : start + args.batch_size]
            vectors = encoder.encode_images([r["image_path"] for r in batch])

            client.points.upsert(
                args.collection,
                points=[
                    PointStruct(id=start + offset, vector=vec.tolist(), payload=record)
                    for offset, (record, vec) in enumerate(zip(batch, vectors))
                ],
            )

            done += len(batch)
            elapsed = time.perf_counter() - started
            rate = done / elapsed if elapsed else 0.0
            remaining = (len(records) - done) / rate if rate else 0.0
            print(
                f"\r{done}/{len(records)} images "
                f"| {rate:.1f} img/s | ~{remaining/60:.1f} min left",
                end="", flush=True,
            )

        print(f"\nIndexed {done} images in {(time.perf_counter()-started)/60:.1f} min")
        print(f"Collection now holds {client.points.count(args.collection)} points")


if __name__ == "__main__":
    main()

uv run ingest.py --data ./mvtec_ad

docker compose restart
uv run query.py --text "cracked surface"

# query.py
from __future__ import annotations

import argparse
import time

from actian_vectorai import VectorAIClient, FilterBuilder, Field

from clip_encoder import ClipEncoder

COLLECTION = "defect_search"   # must match the name ingest.py created


def build_filter(part_category: str | None = None, defect_only: bool = False):
    """Metadata filter, applied inside the search rather than after it."""
    builder = FilterBuilder()
    if part_category:
        builder = builder.must(Field("part_category").eq(part_category))
    if defect_only:
        builder = builder.must(Field("is_defect").eq(True))
    return None if builder.is_empty() else builder.build()


def search_by_text(client, encoder, text, part_category=None, defect_only=False, limit=5):
    vector = encoder.encode_text([text])[0].tolist()
    return client.points.search(
        COLLECTION, vector, limit=limit,
        filter=build_filter(part_category, defect_only), with_payload=True,
    )


def search_by_image(client, encoder, image_path, part_category=None, defect_only=False, limit=5):
    vector = encoder.encode_images([image_path])[0].tolist()
    return client.points.search(
        COLLECTION, vector, limit=limit,
        filter=build_filter(part_category, defect_only), with_payload=True,
    )


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--text")
    ap.add_argument("--image")
    ap.add_argument("--part-category")
    ap.add_argument("--defect-only", action="store_true")
    ap.add_argument("--limit", type=int, default=5)
    args = ap.parse_args()
    if not args.text and not args.image:
        raise SystemExit("Pass --text or --image")

    encoder = ClipEncoder()
    with VectorAIClient("localhost:6574") as client:
        t0 = time.perf_counter()
        if args.text:
            results = search_by_text(client, encoder, args.text,
                                     args.part_category, args.defect_only, args.limit)
        else:
            results = search_by_image(client, encoder, args.image,
                                      args.part_category, args.defect_only, args.limit)
        total = (time.perf_counter() - t0) * 1000

    print(f"  score  {'category':<12} {'defect':<16} path")
    for r in results:
        p = r.payload
        print(f"{r.score:7.4f}  {p['part_category']:<12} {p['defect_type']:<16} {p['image_path']}")
    print(f"\ntotal {total:.1f} ms")


if __name__ == "__main__":
    main()

uv run query.py --text "surface scratch on metal component" --part-category metal_nut
uv run query.py --image ./line_capture_0412.png --defect-only

# benchmark.py
def measure_recall_at_k(client, encoder, collection, test_records, k=5) -> float:
    hits = 0
    for record in test_records:
        vec = encoder.encode_images([record["image_path"]])[0].tolist()
        results = client.points.search(
            collection, vector=vec, limit=k + 1, with_payload=True
        )
        neighbours = [
            r for r in results if r.payload["image_path"] != record["image_path"]
        ][:k]
        if any(r.payload["defect_type"] == record["defect_type"] for r in neighbours):
            hits += 1
    return hits / len(test_records) if test_records else 0.0

uv run benchmark.py --out benchmark.csv
uv run make_chart.py --csv benchmark.csv --hardware "Ryzen 7 5800X, 32GB RAM, CPU-only"