File size: 3,429 Bytes
de1e3fc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | """Generate embeddings for a dataset's images with the *active* embedder, with terminal progress.
Use this instead of the UI "Embed" button for large datasets / the real (CPU) HF model, where the
synchronous HTTP endpoint would time out. Idempotent: skips pictures that already have an embedding
for the active model, so it's safe to re-run / resume.
Usage:
python -m scripts.embed_dataset --dataset-id 3 # one dataset
python -m scripts.embed_dataset --all # every dataset
python -m scripts.embed_dataset --all --limit 5 # quick smoke (first 5 pictures each)
"""
from __future__ import annotations
import argparse
import time
from sqlalchemy import select
from app.db import SessionLocal, engine
from app.ml import get_embedder
from app.models import Base, Dataset, Picture
from app.services.datasets import _known_ids, _picture_ids, _unknown_ids
from app.services.images import embed_picture
BATCH = 100
def run(dataset_id: int | None, do_all: bool, limit: int | None) -> None:
Base.metadata.create_all(bind=engine)
db = SessionLocal()
try:
embedder = get_embedder()
print(f"Active embedder: {embedder.name}/{embedder.version} (dim {embedder.dim})")
if do_all:
datasets = db.execute(select(Dataset).order_by(Dataset.id)).scalars().all()
elif dataset_id is not None:
d = db.get(Dataset, dataset_id)
datasets = [d] if d else []
if not d:
print(f"Dataset #{dataset_id} not found.")
else:
print("Pass --dataset-id N or --all.")
return
for ds in datasets:
pic_ids = _picture_ids(db, _known_ids(db, ds.id), _unknown_ids(db, ds.id))
if limit:
pic_ids = pic_ids[:limit]
print(f"\nDataset #{ds.id} {ds.name!r}: {len(pic_ids)} picture(s) to consider")
embedded = skipped = errors = 0
start = time.time()
for i, pid in enumerate(pic_ids, 1):
pic = db.get(Picture, pid)
if pic is None:
continue
try:
if embed_picture(db, pic, skip_if_exists=True):
embedded += 1
else:
skipped += 1
except Exception as exc: # noqa: BLE001
errors += 1
print(f" ! picture {pid}: {exc}")
if i % BATCH == 0:
db.commit()
rate = i / max(time.time() - start, 1e-6)
print(f" {i}/{len(pic_ids)} embedded={embedded} skipped={skipped} "
f"errors={errors} ({rate:.1f} img/s)")
db.commit()
print(f" done: embedded={embedded} skipped={skipped} errors={errors} "
f"in {time.time() - start:.0f}s")
finally:
db.close()
def main() -> None:
parser = argparse.ArgumentParser(description="Embed a dataset's images with the active embedder.")
parser.add_argument("--dataset-id", type=int, default=None)
parser.add_argument("--all", action="store_true", help="Embed every dataset")
parser.add_argument("--limit", type=int, default=None, help="Only the first N pictures (smoke test)")
args = parser.parse_args()
run(args.dataset_id, args.all, args.limit)
if __name__ == "__main__":
main()
|