router/ai/scripts/inspect_quickdraw_embeddings_pgvector.py

from __future__ import annotations

import argparse
import json
from pathlib import Path

from PIL import Image, ImageDraw
import psycopg
from psycopg import sql


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


def save_grid(label: str, image_paths: list[Path], out_path: Path) -> None:
    if not image_paths:
        return
    thumbs: list[Image.Image] = []
    for path in image_paths:
        thumbs.append(Image.open(path).convert("RGB").resize((128, 128), Image.Resampling.BILINEAR))

    cols = 4
    rows = (len(thumbs) + cols - 1) // cols
    canvas = Image.new("RGB", (cols * 128, rows * 128 + 30), (255, 255, 255))
    draw = ImageDraw.Draw(canvas)
    draw.text((8, 6), f"{label} pgvector top-{len(image_paths)}", fill=(0, 0, 0))

    for idx, thumb in enumerate(thumbs):
        r = idx // cols
        c = idx % cols
        canvas.paste(thumb, (c * 128, r * 128 + 30))

    out_path.parent.mkdir(parents=True, exist_ok=True)
    canvas.save(out_path)


def parse_labels(raw: str) -> list[str]:
    return [item.strip() for item in raw.split(",") if item.strip()]


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("--embedding-table", default="quickdraw_embeddings")
    parser.add_argument("--prototype-table", default="quickdraw_prototypes")
    parser.add_argument("--embedding-source", default="quickdraw_train_v1")
    parser.add_argument("--prototype-source", default="quickdraw_avg_v1")
    parser.add_argument("--labels", required=True, help="Comma-separated labels")
    parser.add_argument("--top-k", type=int, default=12)
    parser.add_argument("--out-dir", default="data/quickdraw/pgvector_inspect")
    args = parser.parse_args()

    labels = parse_labels(args.labels)
    if not labels:
        raise RuntimeError("labels must not be empty")

    data_root = Path(args.data_root)
    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    db_url = normalize_pg_url(args.db_url)
    summary: dict[str, dict[str, object]] = {}

    query = sql.SQL(
        """
        SELECT e.image_path, 1 - (e.embedding <=> p.embedding) AS similarity
        FROM {embedding_table} e
        JOIN {prototype_table} p
          ON p.label = %s
         AND p.source = %s
        WHERE e.label = %s
          AND e.source = %s
        ORDER BY e.embedding <=> p.embedding
        LIMIT %s
        """
    ).format(
        embedding_table=sql.Identifier(args.embedding_table),
        prototype_table=sql.Identifier(args.prototype_table),
    )

    with psycopg.connect(db_url) as conn:
        with conn.cursor() as cur:
            for label in labels:
                cur.execute(
                    query,
                    (
                        label,
                        args.prototype_source,
                        label,
                        args.embedding_source,
                        args.top_k,
                    ),
                )
                rows = cur.fetchall()
                if not rows:
                    print(f"[warn] {label}: no rows found")
                    summary[label] = {"top_k": 0, "items": [], "output": None}
                    continue

                items: list[dict[str, object]] = []
                found_paths: list[Path] = []
                missing_count = 0
                for image_path, similarity in rows:
                    rel = Path(str(image_path))
                    abs_path = data_root / rel
                    exists = abs_path.exists()
                    if exists:
                        found_paths.append(abs_path)
                    else:
                        missing_count += 1
                    items.append(
                        {
                            "image_path": str(rel).replace("\\", "/"),
                            "absolute_path": str(abs_path),
                            "similarity": float(similarity),
                            "exists": exists,
                        }
                    )

                out_png = out_dir / f"{label}_top{len(found_paths)}.png"
                if found_paths:
                    save_grid(label, found_paths, out_png)
                    print(f"[done] {label}: wrote {out_png} (missing={missing_count})")
                else:
                    print(f"[warn] {label}: all top-k files missing under {data_root}")

                summary[label] = {
                    "top_k_requested": args.top_k,
                    "top_k_found_files": len(found_paths),
                    "missing_files": missing_count,
                    "output": str(out_png) if found_paths else None,
                    "items": items,
                }

    summary_path = out_dir / "summary.json"
    summary_path.write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"[done] summary: {summary_path}")


if __name__ == "__main__":
    main()