| """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: |
| |
| 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] |
|
|
| |
| |
| |
| 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: |
| 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: |
| 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: |
| 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 |
|
|
