router/ai/scripts/sync_quickdraw_embeddings_pgvector.py

from __future__ import annotations

import argparse
import random
from pathlib import Path

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

PROFILE_LABELS = {
    "stabilize": ["cup", "face", "cloud", "ice_cream", "t_shirt"],
    "animals": ["dog", "cat", "rabbit", "horse", "cow", "character", "butterfly"],
}


def normalize_pg_url(url: str) -> str:
    if url.startswith("postgresql+psycopg://"):
        return "postgresql://" + url.split("postgresql+psycopg://", 1)[1]
    return url


def to_vector_literal(vec: torch.Tensor) -> str:
    vals = vec.detach().cpu().tolist()
    return "[" + ",".join(f"{float(v):.8f}" for v in vals) + "]"


def list_label_images(
    data_root: Path,
    split: str,
    label: str,
    max_per_label: int,
    seed: int,
) -> list[Path]:
    label_dir = data_root / split / label
    if not label_dir.exists():
        raise FileNotFoundError(f"Label directory not found: {label_dir}")

    files = sorted(label_dir.glob("*.png"))
    if not files:
        raise RuntimeError(f"No images found for label '{label}' at {label_dir}")

    if max_per_label > 0 and len(files) > max_per_label:
        rng = random.Random(seed)
        files = rng.sample(files, max_per_label)
        files.sort()
    return files


def ensure_table(cur: psycopg.Cursor, table: str, dim: int) -> None:
    cur.execute("CREATE EXTENSION IF NOT EXISTS vector;")
    cur.execute(
        f"""
        CREATE TABLE IF NOT EXISTS {table} (
          id bigserial PRIMARY KEY,
          label text NOT NULL,
          image_path text NOT NULL,
          source text NOT NULL,
          embedding vector({dim}) NOT NULL,
          created_at timestamptz NOT NULL DEFAULT now(),
          UNIQUE(image_path, source)
        );
        """
    )
    cur.execute(
        f"""
        CREATE INDEX IF NOT EXISTS {table}_embedding_idx
        ON {table} USING ivfflat (embedding vector_cosine_ops)
        WITH (lists = 100);
        """
    )
    cur.execute(
        f"""
        CREATE UNIQUE INDEX IF NOT EXISTS {table}_path_source_uniq
        ON {table} (image_path, source);
        """
    )


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--db-url", required=True, help="PostgreSQL URL")
    parser.add_argument("--data-root", default="data/quickdraw/images")
    parser.add_argument("--split", default="train", choices=["train", "val", "test"])
    parser.add_argument("--profile", default="stabilize", choices=["stabilize", "animals"])
    parser.add_argument("--labels", default="", help="Comma-separated labels")
    parser.add_argument("--max-per-label", type=int, default=2400)
    parser.add_argument("--batch-size", type=int, default=64)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--model", default="openai/clip-vit-base-patch32")
    parser.add_argument("--table", default="quickdraw_embeddings")
    parser.add_argument("--source", default="quickdraw_train_v1")
    args = parser.parse_args()

    random.seed(args.seed)
    torch.manual_seed(args.seed)

    if args.labels.strip():
        labels = [s.strip() for s in args.labels.split(",") if s.strip()]
    else:
        labels = PROFILE_LABELS[args.profile]
    if not labels:
        raise RuntimeError("labels must not be empty")

    data_root = Path(args.data_root)
    db_url = normalize_pg_url(args.db_url)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = CLIPModel.from_pretrained(args.model)
    model.to(device)
    model.eval()
    processor = CLIPProcessor.from_pretrained(args.model)

    total = 0
    with psycopg.connect(db_url) as conn:
        with conn.cursor() as cur:
            table_ready = False
            upsert_sql = (
                f"INSERT INTO {args.table} (label, image_path, source, embedding) "
                f"VALUES (%s, %s, %s, %s::vector) "
                f"ON CONFLICT (image_path, source) DO UPDATE SET "
                f"label = EXCLUDED.label, embedding = EXCLUDED.embedding"
            )

            for label in labels:
                paths = list_label_images(
                    data_root=data_root,
                    split=args.split,
                    label=label,
                    max_per_label=args.max_per_label,
                    seed=args.seed,
                )

                inserted_label = 0
                for i in range(0, len(paths), args.batch_size):
                    batch_paths = paths[i : i + args.batch_size]
                    images = [Image.open(path).convert("RGB") for path in batch_paths]
                    inputs = processor(images=images, return_tensors="pt")
                    pixel_values = inputs["pixel_values"].to(device)

                    with torch.no_grad():
                        vision_outputs = model.vision_model(pixel_values=pixel_values)
                        pooled = vision_outputs.pooler_output
                        feats = model.visual_projection(pooled)
                        feats = feats / feats.norm(dim=-1, keepdim=True).clamp_min(1e-6)
                    feats = feats.cpu()

                    if not table_ready:
                        ensure_table(cur, args.table, int(feats.shape[1]))
                        table_ready = True

                    rows: list[tuple[str, str, str, str]] = []
                    for path, feat in zip(batch_paths, feats):
                        rel = path.relative_to(data_root).as_posix()
                        rows.append((label, rel, args.source, to_vector_literal(feat)))
                    cur.executemany(upsert_sql, rows)

                    inserted_label += len(rows)
                    total += len(rows)
                    print(f"[batch] {label}: {min(i + args.batch_size, len(paths))}/{len(paths)}")

                print(f"[done] {label}: upserted {inserted_label} rows")
            conn.commit()

    print(f"[done] total upserted: {total} rows into {args.table} (source={args.source})")


if __name__ == "__main__":
    main()