Semantic local image indexing and search on windows

README.md

imgsearch

Local semantic image search with SigLIP 2. Index a folder once, search it with plain text.

Install

​bash pip install torch transformers pillow numpy sentencepiece ​

First run downloads the model (~1.7 GB) from Hugging Face. CUDA is used if available, otherwise CPU.

Index

​bash python imgsearch.py index ~/Pictures ​

Rescans are incremental — only new or modified files are embedded. Entries whose files have disappeared are listed and removed only if you confirm (-y skips the prompt). To sweep the whole database, not just one folder:

​bash python imgsearch.py prune ​

Search

​bash python imgsearch.py search "cat on a keyboard" python imgsearch.py search -i # interactive, keeps the model loaded ​

Options: --top N (default 20), --uri for file:// output, --no-template to embed the query verbatim instead of wrapping it in "This is a photo of …".

Classify

One-shot classification of every indexed image into a fixed set of categories, using a categories.json file:

​json [ { "category": "Portraits & Selfies", "description": "This is a photo of an individual person, featuring a close-up face, self-portrait, or headshot." }, { "category": "Group & Social Gatherings", "description": "This is a photo of multiple people gathered together, such as a family event, wedding, party, or social celebration." }, { "category": "Pets & Domestic Animals", "description": "This is a photo of a domestic pet or companion animal, such as a dog, cat, or indoor pet." } ] ​

Each description is embedded verbatim (no extra template wrapping — write it as a full caption), and every image gets the category with the highest similarity:

​bash python imgsearch.py classify -c categories.json ​

Results are written to a classifications table in images.db (path, category, score) and re-running with an updated categories.json overwrites prior assignments. Add -o results.csv to also export a CSV, --uri to use file:// paths there.

List

List all images assigned to a category (requires classify to have run), sorted by score descending:

​bash python imgsearch.py list --category "Portraits & Selfies" ​

Options: -c/--category (exact name, required), --scores to prefix each line with its similarity score, --uri for file:// output, -q to suppress the trailing count.

Known limitation

Paths are stored exactly as given, so index . and index /full/path create separate rows for the same file.

imgsearch.py

#!/usr/bin/env python3
"""Local semantic image search with SigLIP 2. Index once, search whenever."""
import argparse, csv, json, os, sqlite3, sys
from collections import Counter
from pathlib import Path
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
from PIL import Image, ImageFile
from transformers import AutoModel, AutoProcessor

ImageFile.LOAD_TRUNCATED_IMAGES = True  # tolerate slightly corrupted files
MODEL_ID = "google/siglip2-so400m-patch14-384"
EXTS = {".jpg", ".jpeg", ".png", ".gif", ".bmp", ".webp", ".tiff", ".tif"}
DB = "images.db"
DIM = 1152
PRUNE_PREVIEW = 100  # max entries listed before collapsing to a count


def _as_tensor(out):
    """get_*_features may return a tensor or an output object depending on the version."""
    if isinstance(out, torch.Tensor):
        return out
    for attr in ("pooler_output", "image_embeds", "text_embeds", "last_hidden_state"):
        v = getattr(out, attr, None)
        if v is not None:
            return v
    raise TypeError(f"Unexpected output type: {type(out)}")


def get_model():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = AutoModel.from_pretrained(MODEL_ID).to(device).eval()
    proc = AutoProcessor.from_pretrained(MODEL_ID)
    return model, proc, device


def db_connect():
    con = sqlite3.connect(DB)
    con.execute("PRAGMA journal_mode=WAL")
    con.execute("CREATE TABLE IF NOT EXISTS images("
                "path TEXT PRIMARY KEY, mtime REAL, emb BLOB)")
    con.execute("CREATE TABLE IF NOT EXISTS classifications("
                "path TEXT PRIMARY KEY, category TEXT, score REAL)")
    return con


def _under(path, root):
    """True if abspath `path` lies inside abspath `root`."""
    return path == root or path.startswith(root + os.sep)


def confirm_prune(con, stale, assume_yes=False):
    """List index entries with no file on disk, delete only on explicit y."""
    if not stale:
        return
    print(f"\n{len(stale)} indexed files no longer exist on disk:")
    for p in stale[:PRUNE_PREVIEW]:
        print(f"  {p}")
    if len(stale) > PRUNE_PREVIEW:
        print(f"  ... and {len(stale) - PRUNE_PREVIEW} more")

    if assume_yes:
        ans = "y"
    else:
        try:
            ans = input(f"Remove these {len(stale)} from the index? [y/N] ").strip().lower()
        except (EOFError, KeyboardInterrupt):
            print()
            ans = ""

    if ans in {"y", "yes"}:
        con.executemany("DELETE FROM images WHERE path=?", [(p,) for p in stale])
        con.commit()
        print(f"Removed {len(stale)} from the index.")
    else:
        print("Kept — nothing deleted.")


class ImageDataset(Dataset):
    """Reads and decodes images in worker processes to keep the GPU saturated."""
    def __init__(self, paths, proc):
        self.paths = paths
        self.proc = proc

    def __len__(self):
        return len(self.paths)

    def __getitem__(self, i):
        p = self.paths[i]
        try:
            img = Image.open(p).convert("RGB")
            pixel = self.proc(images=img, return_tensors="pt")["pixel_values"][0]
            return p, pixel
        except Exception as e:
            print(f"  skipping {p}: {e}", file=sys.stderr)
            return p, None


def _collate(batch):
    paths = [b[0] for b in batch if b[1] is not None]
    tensors = [b[1] for b in batch if b[1] is not None]
    if not tensors:
        return [], None
    return paths, torch.stack(tensors)


@torch.no_grad()
def cmd_index(args):
    root = Path(args.root)
    if not root.is_dir():
        sys.exit(f"Not a directory: {root}")

    disk = {}
    for p in root.rglob("*"):
        if p.suffix.lower() in EXTS:
            try:
                disk[str(p)] = p.stat().st_mtime
            except OSError:
                pass  # broken symlink, or vanished mid-scan
    print(f"Found {len(disk)} images under {root}")

    con = db_connect()
    have = {r[0]: r[1] for r in con.execute("SELECT path, mtime FROM images")}

    # index -> disk: rows under this root whose file is gone (moved or deleted)
    root_abs = os.path.abspath(root)
    on_disk = {os.path.abspath(p) for p in disk}
    have_abs = {p: os.path.abspath(p) for p in have}
    stale = sorted(p for p, a in have_abs.items()
                   if _under(a, root_abs) and a not in on_disk)
    confirm_prune(con, stale, args.yes)

    # disk -> index: new or modified files (mtime reused from the scan above)
    todo = [p for p, m in disk.items() if have.get(p) != m]
    print(f"{len(todo)} new/changed to index ({len(disk) - len(todo)} already current)")
    if not todo:
        return

    model, proc, device = get_model()
    print(f"Device: {device} | batch={args.batch} | workers={args.workers}")

    ds = ImageDataset(todo, proc)
    loader = DataLoader(ds, batch_size=args.batch, num_workers=args.workers,
                        collate_fn=_collate, pin_memory=(device == "cuda"),
                        persistent_workers=(args.workers > 0))

    done = 0
    for paths, pixels in loader:
        if pixels is not None:
            pixels = pixels.to(device, non_blocking=True)
            with torch.autocast(device, dtype=torch.float16, enabled=(device == "cuda")):
                feats = model.get_image_features(pixel_values=pixels)
            feats = torch.nn.functional.normalize(_as_tensor(feats).float(), dim=-1)
            embs = feats.cpu().numpy().astype("float32")
            rows = [(pp, disk[pp], embs[j].tobytes()) for j, pp in enumerate(paths)]
            con.executemany(
                "INSERT OR REPLACE INTO images(path, mtime, emb) VALUES(?,?,?)", rows)
            con.commit()  # commits per batch → safe to interrupt and resume
        done += len(paths)
        print(f"  {min(done, len(todo))}/{len(todo)}", end="\r", flush=True)

    total = con.execute("SELECT COUNT(*) FROM images").fetchone()[0]
    print(f"\nDone. Total in database: {total}")


def cmd_prune(args):
    con = db_connect()
    paths = [r[0] for r in con.execute("SELECT path FROM images")]
    stale = sorted(p for p in paths if not os.path.exists(p))
    if not stale:
        print(f"Index is clean: all {len(paths)} files present.")
        return
    confirm_prune(con, stale, args.yes)


@torch.no_grad()
def cmd_search(args):
    con = db_connect()
    rows = con.execute("SELECT path, emb FROM images").fetchall()
    if not rows:
        print("The index is empty. Run `index` first.")
        return
    if not args.interactive and not args.query:
        print("Give a query, or use -i for interactive mode.", file=sys.stderr)
        return
    paths = [r[0] for r in rows]
    mat = np.frombuffer(b"".join(r[1] for r in rows),
                        dtype="float32").reshape(len(rows), DIM)
    model, proc, device = get_model()

    def run(query, top, no_template):
        text = query if no_template else f"This is a photo of {query}."
        inputs = proc(text=[text], padding="max_length", max_length=64,
                      truncation=True, return_tensors="pt").to(device)
        q = _as_tensor(model.get_text_features(**inputs))
        q = torch.nn.functional.normalize(q, dim=-1).cpu().numpy().astype("float32")[0]
        scores = mat @ q
        for rank, idx in enumerate(np.argsort(-scores)[:top], 1):
            p = Path(paths[idx]).as_uri() if args.uri else f'"{paths[idx]}"'
            print(f"{rank:2d}. {scores[idx]:.3f}  {p}")

    if not args.interactive:
        run(args.query, args.top, args.no_template)
        return

    print(f"{len(paths)} images indexed. Ctrl-D to quit.\n")
    while True:
        try:
            toks = input("$: ").split()
        except (EOFError, KeyboardInterrupt):
            print()
            return
        top, no_tpl = args.top, args.no_template
        try:
            if "--no-template" in toks:
                toks.remove("--no-template")
                no_tpl = True
            if "--top" in toks:
                i = toks.index("--top")
                top = int(toks.pop(i + 1))
                toks.pop(i)
        except (IndexError, ValueError):
            print("  usage: <query> [--top N] [--no-template]\n", file=sys.stderr)
            continue
        if toks:
            run(" ".join(toks), top, no_tpl)
            print()


def load_categories(path):
    with open(path, "r", encoding="utf-8") as f:
        data = json.load(f)
    if not data:
        sys.exit(f"No categories found in {path}")
    cats = [d["category"] for d in data]
    descs = [d["description"] for d in data]
    if len(set(cats)) != len(cats):
        dupes = [c for c, n in Counter(cats).items() if n > 1]
        sys.exit(f"Duplicate category names: {dupes}")
    return cats, descs


@torch.no_grad()
def cmd_classify(args):
    con = db_connect()
    rows = con.execute("SELECT path, emb FROM images").fetchall()
    if not rows:
        print("The index is empty. Run `index` first.")
        return

    cats, descs = load_categories(args.categories)
    paths = [r[0] for r in rows]
    mat = np.frombuffer(b"".join(r[1] for r in rows),
                        dtype="float32").reshape(len(rows), DIM)

    model, proc, device = get_model()
    print(f"Classifying {len(paths)} images into {len(cats)} categories...")

    # descriptions are used verbatim — they're already phrased as full captions
    inputs = proc(text=descs, padding="max_length", max_length=64,
                  truncation=True, return_tensors="pt").to(device)
    q = _as_tensor(model.get_text_features(**inputs))
    q = torch.nn.functional.normalize(q, dim=-1).cpu().numpy().astype("float32")  # (C, DIM)

    scores = mat @ q.T  # (N, C)
    best_idx = np.argmax(scores, axis=1)
    best_score = scores[np.arange(len(paths)), best_idx]

    out_rows = [(paths[i], cats[best_idx[i]], float(best_score[i]))
                for i in range(len(paths))]
    con.executemany(
        "INSERT OR REPLACE INTO classifications(path, category, score) VALUES(?,?,?)",
        out_rows)
    con.commit()

    counts = Counter(cats[i] for i in best_idx)
    print(f"\nResults (also saved to the `classifications` table in {DB}):\n")
    for cat in cats:
        print(f"  {counts.get(cat, 0):5d}  {cat}")

    if args.output:
        with open(args.output, "w", newline="", encoding="utf-8") as f:
            w = csv.writer(f)
            w.writerow(["path", "category", "score"])
            for path, cat, score in out_rows:
                p = Path(path).as_uri() if args.uri else path
                w.writerow([p, cat, f"{score:.4f}"])
        print(f"\nWrote {args.output}")


def cmd_list(args):
    con = db_connect()
    rows = con.execute(
        "SELECT path, score FROM classifications WHERE category=? ORDER BY score DESC",
        (args.category,)).fetchall()
    if not rows:
        known = [r[0] for r in con.execute(
            "SELECT DISTINCT category FROM classifications ORDER BY category")]
        print(f'No images classified as "{args.category}".', file=sys.stderr)
        if known:
            print("Known categories:", file=sys.stderr)
            for c in known:
                print(f"  {c}", file=sys.stderr)
        return

    for path, score in rows:
        p = Path(path).as_uri() if args.uri else path
        if args.scores:
            print(f"{score:.3f}  {p}")
        else:
            print(p)

    if not args.quiet:
        print(f"\n{len(rows)} image(s) in \"{args.category}\"", file=sys.stderr)


def main():
    ap = argparse.ArgumentParser(description="Local SigLIP 2 image search")
    sub = ap.add_subparsers(dest="cmd", required=True)

    pi = sub.add_parser("index", help="index a folder recursively")
    pi.add_argument("root")
    pi.add_argument("--batch", type=int, default=64)
    pi.add_argument("--workers", type=int, default=6,
                    help="read/decode workers for disk I/O (0 = in the main process)")
    pi.add_argument("-y", "--yes", action="store_true",
                    help="remove missing entries without asking")
    pi.set_defaults(func=cmd_index)

    pp = sub.add_parser("prune", help="drop index entries whose files are gone")
    pp.add_argument("-y", "--yes", action="store_true",
                    help="remove missing entries without asking")
    pp.set_defaults(func=cmd_prune)

    ps = sub.add_parser("search", help="search with text")
    ps.add_argument("query", nargs="?")
    ps.add_argument("--uri", action="store_true",
                    help="print file:// URIs instead of quoted paths")
    ps.add_argument("--top", type=int, default=20)
    ps.add_argument("-i", "--interactive", action="store_true",
                    help="interactive prompt; keeps model and index in memory")
    ps.add_argument("--no-template", action="store_true",
                    help="embed the raw query without the photo template")
    ps.set_defaults(func=cmd_search)

    pc = sub.add_parser("classify", help="one-shot classify all indexed images into categories")
    pc.add_argument("-c", "--categories", required=True,
                    help='JSON file: [{"category": "...", "description": "..."}, ...]')
    pc.add_argument("-o", "--output", help="also write results to a CSV file")
    pc.add_argument("--uri", action="store_true",
                    help="use file:// URIs in the CSV output")
    pc.set_defaults(func=cmd_classify)

    pl = sub.add_parser("list", help="list images in a classified category")
    pl.add_argument("-c", "--category", required=True, help="exact category name")
    pl.add_argument("--uri", action="store_true", help="print file:// URIs")
    pl.add_argument("--scores", action="store_true", help="prefix each path with its score")
    pl.add_argument("-q", "--quiet", action="store_true", help="suppress the summary line")
    pl.set_defaults(func=cmd_list)

    args = ap.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()