"""Batch loader: turn a folder of images + a CSV into Dataset-linked DB records. Reuses the existing image pipeline (``process_and_store_picture``) for embedding/breed generation — nothing is reimplemented here. One CSV row == one image == one KnownDog/UnknownDog record (matching ``prepare_test_data.py``). Owners/finders are de-duped by email and linked to the dataset. Individual image failures are logged and skipped (the load never aborts); total errors are reported. """ from __future__ import annotations import csv import logging import random from dataclasses import dataclass, field from datetime import date from pathlib import Path from typing import Callable from sqlalchemy import select from sqlalchemy.orm import Session from ..models import Case, Dataset, KnownDog, UnknownDog, User from ..models.base import ( CaseStatus, CaseType, DatasetType, DogSize, KnownDogStatus, SubjectType, UnknownDogStatus, UserRole, ) from .images import ImageValidationError, process_and_store_picture logger = logging.getLogger("pawtrace.batch_loader") ProgressCb = Callable[[dict], None] @dataclass class LoadResult: dataset_id: int dogs_loaded: int = 0 images_processed: int = 0 cases_created: int = 0 users_created: int = 0 errors: list[str] = field(default_factory=list) matching: dict | None = None def as_dict(self) -> dict: return { "dataset_id": self.dataset_id, "dogs_loaded": self.dogs_loaded, "images_processed": self.images_processed, "cases_created": self.cases_created, "users_created": self.users_created, "error_count": len(self.errors), "errors": self.errors, "matching": self.matching, } def _coerce_size(value: str | None) -> DogSize | None: if not value: return None try: return DogSize(value.strip().lower()) except ValueError: return None def _read_csv(csv_path: Path) -> list[dict]: with csv_path.open(newline="", encoding="utf-8") as fh: return list(csv.DictReader(fh)) def _get_or_create_user( db: Session, dataset_id: int, *, name: str, email: str, phone: str | None, zip_code: str, role: UserRole, created_counter: list[int], ) -> User: email = (email or "").strip().lower() user = None if email: user = db.execute(select(User).where(User.email == email)).scalar_one_or_none() if user is None: # Synthesize a unique email if the CSV omitted one (anonymous finder). if not email: email = f"anon-{dataset_id}-{created_counter[0]}@example.invalid" user = User( name=name or "Unknown", email=email, phone=phone or None, zip=zip_code or "00000", role=role, dataset_id=dataset_id, ) db.add(user) db.flush() created_counter[0] += 1 return user def _resolve_image(root: Path, folder: str, image_file: str) -> Path | None: candidate = root / folder / image_file if candidate.is_file(): return candidate return None def load_dataset( db: Session, *, folder: str | Path, dataset_type: str, name: str, description: str | None, csv_path: str | Path, mark_lost: bool = False, mark_lost_pct: int = 100, run_matching: bool = False, skip_embeddings: bool = False, group_by_folder: bool = True, seed: int = 42, progress: ProgressCb | None = None, ) -> LoadResult: root = Path(folder) csv_file = Path(csv_path) if not csv_file.is_file(): raise FileNotFoundError(f"CSV not found: {csv_file}") dtype = DatasetType(dataset_type) is_known = dtype in (DatasetType.known, DatasetType.test_known) dataset = Dataset( name=name, type=dtype, description=description, source_path=str(root), dog_count=0, ) db.add(dataset) db.flush() rows = _read_csv(csv_file) total = len(rows) rng = random.Random(seed) result = LoadResult(dataset_id=dataset.id) created_counter = [0] # Group rows so every image of one identity becomes ONE dog with multiple pictures. The CSV # `folder` column is the identity key (each DogFaceNet subfolder == one dog). With # group_by_folder=False each row becomes its own single-picture dog (legacy behavior). if group_by_folder: grouped: dict[str, list[dict]] = {} for row in rows: grouped.setdefault((row.get("folder") or "").strip(), []).append(row) groups = list(grouped.items()) else: groups = [((row.get("folder") or "").strip(), [row]) for row in rows] done_rows = 0 def _emit() -> None: if progress: progress({"processed": done_rows, "total": total, "dataset_id": dataset.id}) for folder_name, frows in groups: try: if is_known: _load_known_folder(db, dataset, folder_name, frows, root, rng, mark_lost, mark_lost_pct, result, created_counter, generate=not skip_embeddings) else: _load_unknown_folder(db, dataset, folder_name, frows, root, result, created_counter, generate=not skip_embeddings) except Exception as exc: # noqa: BLE001 - skip-and-continue per spec msg = f"folder {folder_name!r}: {exc}" logger.warning(msg, exc_info=True) result.errors.append(msg) done_rows += len(frows) _emit() dataset.dog_count = result.dogs_loaded db.commit() if run_matching: from .datasets import match_dataset result.matching = match_dataset(db, dataset) return result def _load_known_folder(db, dataset, folder_name, frows, root, rng, mark_lost, mark_lost_pct, result, created_counter, *, generate: bool = True) -> None: """Create one KnownDog for the folder, attaching every readable image as a picture.""" meta = frows[0] dog: KnownDog | None = None owner: User | None = None first_pic = True for row in frows: image_file = (row.get("image_file") or "").strip() img_path = _resolve_image(root, folder_name, image_file) if img_path is None: result.errors.append(f"folder {folder_name!r}: image not found {image_file!r}") continue try: data = img_path.read_bytes() if dog is None: owner = _get_or_create_user( db, dataset.id, name=meta.get("owner_name", ""), email=meta.get("owner_email", ""), phone=meta.get("owner_phone"), zip_code=meta.get("zip", ""), role=UserRole.owner, created_counter=created_counter, ) dog = KnownDog( owner_id=owner.id, dataset_id=dataset.id, name=meta.get("dog_name") or "Unnamed", breed=meta.get("breed") or None, color=meta.get("color") or None, size=_coerce_size(meta.get("size")), description=meta.get("description") or "", status=KnownDogStatus.home, ) db.add(dog) db.flush() result.dogs_loaded += 1 result.users_created = created_counter[0] process_and_store_picture( db, subject_type=SubjectType.known, subject_id=dog.id, data=data, is_primary=first_pic, generate_embedding=generate, generate_breed=generate, ) first_pic = False result.images_processed += 1 except ImageValidationError as exc: result.errors.append(f"folder {folder_name!r} image {image_file!r}: invalid ({exc})") except Exception as exc: # noqa: BLE001 result.errors.append(f"folder {folder_name!r} image {image_file!r}: {exc}") if dog is not None and mark_lost and rng.random() < (mark_lost_pct / 100.0): zip_code = meta.get("zip") or (owner.zip if owner else "") dog.status = KnownDogStatus.lost dog.last_known_zip = zip_code db.add( Case( person_id=owner.id if owner else None, known_dog_id=dog.id, type=CaseType.lost, event_zip=zip_code, event_date=date.today(), search_radius_miles=0, status=CaseStatus.open, ) ) result.cases_created += 1 def _load_unknown_folder(db, dataset, folder_name, frows, root, result, created_counter, *, generate: bool = True) -> None: """Create one UnknownDog + one found Case for the folder, attaching every readable image.""" meta = frows[0] dog: UnknownDog | None = None finder: User | None = None first_pic = True for row in frows: image_file = (row.get("image_file") or "").strip() img_path = _resolve_image(root, folder_name, image_file) if img_path is None: result.errors.append(f"folder {folder_name!r}: image not found {image_file!r}") continue try: data = img_path.read_bytes() if dog is None: finder = _get_or_create_user( db, dataset.id, name=meta.get("finder_name", ""), email=meta.get("finder_email", ""), phone=meta.get("finder_phone"), zip_code=meta.get("found_zip", ""), role=UserRole.finder, created_counter=created_counter, ) dog = UnknownDog( dataset_id=dataset.id, description=meta.get("description") or "", color=meta.get("color") or None, size=_coerce_size(meta.get("size")), current_zip=meta.get("found_zip") or "", current_location_detail=meta.get("current_location") or None, status=UnknownDogStatus.at_shelter, ) db.add(dog) db.flush() result.dogs_loaded += 1 result.users_created = created_counter[0] process_and_store_picture( db, subject_type=SubjectType.unknown, subject_id=dog.id, data=data, is_primary=first_pic, generate_embedding=generate, generate_breed=generate, ) first_pic = False result.images_processed += 1 except ImageValidationError as exc: result.errors.append(f"folder {folder_name!r} image {image_file!r}: invalid ({exc})") except Exception as exc: # noqa: BLE001 result.errors.append(f"folder {folder_name!r} image {image_file!r}: {exc}") if dog is not None and finder is not None: found_zip = meta.get("found_zip") or "" db.add( Case( person_id=finder.id, finder_name=meta.get("finder_name") or finder.name, finder_email=meta.get("finder_email") or finder.email, finder_phone=meta.get("finder_phone"), unknown_dog_id=dog.id, type=CaseType.found, event_zip=found_zip, event_date=date.today(), current_location=meta.get("current_location") or None, search_radius_miles=0, status=CaseStatus.open, ) ) result.cases_created += 1