File size: 2,652 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 | """CLI for the batch loader (spec §4).
Usage:
python -m scripts.load_dataset --folder PATH --type known --name "My set" \
--csv PATH/known_dogs.csv [--description ...] [--mark-lost] [--mark-lost-pct 50] \
[--run-matching] [--seed 42]
"""
from __future__ import annotations
import argparse
from app.db import SessionLocal, engine
from app.models import Base
from app.services.batch_loader import load_dataset
def main() -> None:
parser = argparse.ArgumentParser(description="Batch-load a dataset of dog images.")
parser.add_argument("--folder", required=True)
parser.add_argument("--type", required=True, choices=["known", "unknown", "test_known", "test_found"])
parser.add_argument("--name", required=True)
parser.add_argument("--csv", required=True)
parser.add_argument("--description", default=None)
parser.add_argument("--mark-lost", action="store_true")
parser.add_argument("--mark-lost-pct", type=int, default=100)
parser.add_argument("--run-matching", action="store_true")
parser.add_argument("--skip-embeddings", action="store_true",
help="Store images only; generate embeddings later via admin embed-all")
parser.add_argument("--one-dog-per-image", action="store_true",
help="Legacy: one dog record per image instead of grouping a folder's "
"images into a single dog with multiple pictures")
parser.add_argument("--seed", type=int, default=42)
args = parser.parse_args()
Base.metadata.create_all(bind=engine)
db = SessionLocal()
try:
result = load_dataset(
db,
folder=args.folder,
dataset_type=args.type,
name=args.name,
description=args.description,
csv_path=args.csv,
mark_lost=args.mark_lost,
mark_lost_pct=args.mark_lost_pct,
run_matching=args.run_matching,
skip_embeddings=args.skip_embeddings,
group_by_folder=not args.one_dog_per_image,
seed=args.seed,
)
finally:
db.close()
print(f"Dataset #{result.dataset_id} loaded:")
print(f" dogs loaded: {result.dogs_loaded}")
print(f" images processed: {result.images_processed}")
print(f" cases created: {result.cases_created}")
print(f" users created: {result.users_created}")
print(f" errors: {len(result.errors)}")
if result.matching:
print(f" matching: {result.matching}")
for err in result.errors[:20]:
print(f" ! {err}")
if __name__ == "__main__":
main()
|