backend/scripts/load_dataset.py

"""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()