PawTrace / backend /app /services /datasets.py
Elliott Duke
PawTrace: HF demo image with dog dataset + demo photos
6213763
Raw
History Blame Contribute Delete
35.5 kB
"""Dataset stats, transactional purge, embed-all, and match-all (admin data management).
Pictures link to dogs by convention (subject_type/subject_id, no FK), and cases/matches link to
dogs/users by FK, so a clean purge deletes children before parents in FK-dependency order, all
inside one transaction (everything or nothing).
"""
from __future__ import annotations
import logging
from datetime import date
from sqlalchemy import delete, func, select
from sqlalchemy.orm import Session
from ..models import (
BreedPrediction,
Case,
Dataset,
Embedding,
KnownDog,
Match,
Notification,
Picture,
UnknownDog,
User,
)
from ..models.base import SubjectType
from ..storage import get_storage
from .images import embed_and_breed_picture
from .matching import run_matching_for_case
logger = logging.getLogger("pawtrace.datasets")
# ---- membership helpers ----
def _known_ids(db: Session, dataset_id: int) -> list[int]:
return list(
db.execute(select(KnownDog.id).where(KnownDog.dataset_id == dataset_id)).scalars()
)
def _unknown_ids(db: Session, dataset_id: int) -> list[int]:
return list(
db.execute(select(UnknownDog.id).where(UnknownDog.dataset_id == dataset_id)).scalars()
)
def _user_ids(db: Session, dataset_id: int) -> list[int]:
return list(db.execute(select(User.id).where(User.dataset_id == dataset_id)).scalars())
def _case_ids(db: Session, known_ids, unknown_ids, user_ids) -> list[int]:
conds = []
if known_ids:
conds.append(Case.known_dog_id.in_(known_ids))
if unknown_ids:
conds.append(Case.unknown_dog_id.in_(unknown_ids))
if user_ids:
conds.append(Case.person_id.in_(user_ids))
if not conds:
return []
from sqlalchemy import or_
return list(db.execute(select(Case.id).where(or_(*conds))).scalars())
def _picture_ids(db: Session, known_ids, unknown_ids) -> list[int]:
from sqlalchemy import and_, or_
conds = []
if known_ids:
conds.append(and_(Picture.subject_type == SubjectType.known, Picture.subject_id.in_(known_ids)))
if unknown_ids:
conds.append(and_(Picture.subject_type == SubjectType.unknown, Picture.subject_id.in_(unknown_ids)))
if not conds:
return []
return list(db.execute(select(Picture.id).where(or_(*conds))).scalars())
# ---- stats ----
def dataset_stats(db: Session, dataset: Dataset) -> dict:
known_ids = _known_ids(db, dataset.id)
unknown_ids = _unknown_ids(db, dataset.id)
user_ids = _user_ids(db, dataset.id)
case_ids = _case_ids(db, known_ids, unknown_ids, user_ids)
pic_ids = _picture_ids(db, known_ids, unknown_ids)
match_count = 0
if case_ids:
match_count = db.execute(
select(func.count()).select_from(Match).where(Match.case_id.in_(case_ids))
).scalar_one()
embedded = 0
if pic_ids:
embedded = db.execute(
select(func.count(func.distinct(Embedding.picture_id))).where(
Embedding.picture_id.in_(pic_ids)
)
).scalar_one()
return {
"known_dog_count": len(known_ids),
"unknown_dog_count": len(unknown_ids),
"dog_count": len(known_ids) + len(unknown_ids),
"user_count": len(user_ids),
"case_count": len(case_ids),
"match_count": match_count,
"picture_count": len(pic_ids),
"embedded_picture_count": embedded,
# No persistent per-load error log yet; load-time errors are reported via the job status.
"error_count": 0,
}
def list_dogs(db: Session, dataset: Dataset, limit: int, offset: int) -> dict:
"""Paginated list of the dataset's dogs (known + unknown) for the detail view."""
rows: list[dict] = []
for dog in db.execute(
select(KnownDog).where(KnownDog.dataset_id == dataset.id).order_by(KnownDog.id)
).scalars():
rows.append(
{"kind": "known", "id": dog.id, "name": dog.name, "status": dog.status.value,
"zip": dog.last_known_zip, "thumb_url": _thumb_url(db, SubjectType.known, dog.id)}
)
for dog in db.execute(
select(UnknownDog).where(UnknownDog.dataset_id == dataset.id).order_by(UnknownDog.id)
).scalars():
rows.append(
{"kind": "unknown", "id": dog.id, "name": dog.description or "Found dog",
"status": dog.status.value, "zip": dog.current_zip,
"thumb_url": _thumb_url(db, SubjectType.unknown, dog.id)}
)
total = len(rows)
return {"items": rows[offset : offset + limit], "total": total, "limit": limit, "offset": offset}
def _picture_count(db: Session, subject_type: SubjectType, subject_id: int) -> int:
return db.execute(
select(func.count()).select_from(Picture).where(
Picture.subject_type == subject_type, Picture.subject_id == subject_id
)
).scalar_one()
# ---- breed predictions (HF softmax labels) ----
def active_breed_model_in_db(db: Session) -> tuple[str, str] | None:
"""The breed model to read predictions from: the most-used one, preferring a non-mock model.
Avoids loading the HF model just to read its name; everything here is pure DB.
"""
rows = db.execute(
select(BreedPrediction.model_name, BreedPrediction.model_version, func.count())
.group_by(BreedPrediction.model_name, BreedPrediction.model_version)
.order_by(func.count().desc())
).all()
if not rows:
return None
non_mock = [r for r in rows if not str(r[0]).startswith("mock")]
chosen = non_mock[0] if non_mock else rows[0]
return chosen[0], chosen[1]
def active_embed_model_in_db(db: Session) -> tuple[str, str] | None:
"""The embedding model to compare against: the most-used one, preferring a non-mock model.
Lets matching pick the right precomputed vectors WITHOUT loading the (HF) embedder — matching
only needs stored embeddings, so it must never depend on the model being loadable/online.
"""
rows = db.execute(
select(Embedding.model_name, Embedding.model_version, func.count())
.group_by(Embedding.model_name, Embedding.model_version)
.order_by(func.count().desc())
).all()
if not rows:
return None
non_mock = [r for r in rows if not str(r[0]).startswith("mock")]
chosen = non_mock[0] if non_mock else rows[0]
return chosen[0], chosen[1]
def aggregated_breeds(db, subject_type, subject_id, model, limit=3) -> list[str]:
"""A dog's estimated breeds aggregated across ALL its photos, blending frequency and confidence.
Each breed is scored by the SUM of its confidence over every photo it appears in — so a breed
earns weight both from showing up often (frequency) and from the model being sure (confidence).
A breed seen weakly in many photos can beat one seen strongly in a single photo, and vice-versa.
Ties break by how many photos ranked it #1 (top-1 votes), then by best rank. For a single-photo
dog this reduces to plain rank order. Display only — matching's breed gate uses a separate union
set and is unaffected.
"""
if not model:
return []
name, version = model
rows = db.execute(
select(BreedPrediction.label, BreedPrediction.rank, BreedPrediction.score)
.join(Picture, BreedPrediction.picture_id == Picture.id)
.where(
Picture.subject_type == subject_type, Picture.subject_id == subject_id,
BreedPrediction.model_name == name, BreedPrediction.model_version == version,
)
).all()
if not rows:
return []
score_sum: dict[str, float] = {} # frequency-weighted confidence (the combined score)
votes: dict[str, int] = {} # top-1 appearances (tiebreak)
best_rank: dict[str, int] = {}
for label, rank, score in rows:
score_sum[label] = score_sum.get(label, 0.0) + score
if rank == 0:
votes[label] = votes.get(label, 0) + 1
best_rank[label] = min(best_rank.get(label, rank), rank)
labels = list(score_sum)
labels.sort(
key=lambda l: (score_sum[l], votes.get(l, 0), -best_rank[l]),
reverse=True,
)
return labels[:limit]
def _breed_match_ids(db, subject_type, breed, k, model) -> set[int]:
"""Dog ids whose pictures predict ``breed`` within their top-``k`` (rank < k) for ``model``."""
if not model or not breed:
return set()
name, version = model
rows = db.execute(
select(Picture.subject_id)
.join(BreedPrediction, BreedPrediction.picture_id == Picture.id)
.where(
Picture.subject_type == subject_type,
BreedPrediction.model_name == name, BreedPrediction.model_version == version,
BreedPrediction.label == breed, BreedPrediction.rank < k,
)
.distinct()
).scalars().all()
return set(rows)
def list_breeds(db: Session) -> list[dict]:
"""Distinct predicted breed labels (active model) + how many predictions each has."""
model = active_breed_model_in_db(db)
if not model:
return []
name, version = model
rows = db.execute(
select(BreedPrediction.label, func.count())
.where(BreedPrediction.model_name == name, BreedPrediction.model_version == version)
.group_by(BreedPrediction.label)
.order_by(BreedPrediction.label)
).all()
return [{"label": r[0], "count": r[1]} for r in rows]
def _known_profile(db, dog, ds_names, breed_model=None) -> dict:
return {
"kind": "known", "id": dog.id, "name": dog.name, "breed": dog.breed, "color": dog.color,
"size": dog.size.value if dog.size else None, "status": dog.status.value,
"zip": dog.last_known_zip, "dataset_id": dog.dataset_id,
"dataset_name": ds_names.get(dog.dataset_id) if dog.dataset_id else None,
"picture_count": _picture_count(db, SubjectType.known, dog.id),
"thumb_url": _thumb_url(db, SubjectType.known, dog.id),
"predicted_breeds": aggregated_breeds(db, SubjectType.known, dog.id, breed_model),
}
def _unknown_profile(db, dog, ds_names, breed_model=None) -> dict:
return {
"kind": "unknown", "id": dog.id, "name": dog.description or "Found dog",
"breed": dog.est_breed, "color": dog.color,
"size": dog.size.value if dog.size else None, "status": dog.status.value,
"zip": dog.current_zip, "dataset_id": dog.dataset_id,
"dataset_name": ds_names.get(dog.dataset_id) if dog.dataset_id else None,
"picture_count": _picture_count(db, SubjectType.unknown, dog.id),
"thumb_url": _thumb_url(db, SubjectType.unknown, dog.id),
"predicted_breeds": aggregated_breeds(db, SubjectType.unknown, dog.id, breed_model),
}
def dog_profile(db: Session, subject_type: SubjectType, dog_id: int) -> dict | None:
"""Single dog profile dict (same shape as the browse list), or None if missing."""
ds_names = {row[0]: row[1] for row in db.execute(select(Dataset.id, Dataset.name)).all()}
breed_model = active_breed_model_in_db(db)
if subject_type == SubjectType.known:
dog = db.get(KnownDog, dog_id)
return _known_profile(db, dog, ds_names, breed_model) if dog else None
dog = db.get(UnknownDog, dog_id)
return _unknown_profile(db, dog, ds_names, breed_model) if dog else None
def list_all_dogs(
db: Session, kind: str, limit: int, offset: int,
breed: str | None = None, breed_k: int = 10,
sort: str = "newest", zip_prefix: str | None = None,
added_from: date | None = None, added_to: date | None = None,
) -> dict:
"""Paginated browse of every dog, optionally filtered by breed / ZIP / added-date and sorted.
``kind`` is ``known`` | ``unknown`` | ``all``. For ``all`` the two pools are merged and sorted
together so ordering is global. ``sort`` is ``newest`` (default) or ``oldest`` by added date.
``zip_prefix`` matches the dog's ZIP by prefix (e.g. ``770`` = all Houston). ``added_from`` /
``added_to`` bound the created-at date (inclusive). When ``breed`` is given, only dogs whose
pictures predict that breed within their top-``breed_k`` are included (breed_k=1 == top match).
Tab counts (``known`` / ``unknown``) reflect every active filter, so they match what's listed.
"""
ds_names = {row[0]: row[1] for row in db.execute(select(Dataset.id, Dataset.name)).all()}
breed_model = active_breed_model_in_db(db)
known_match = unknown_match = None
if breed:
known_match = _breed_match_ids(db, SubjectType.known, breed, breed_k, breed_model)
unknown_match = _breed_match_ids(db, SubjectType.unknown, breed, breed_k, breed_model)
def _refs(model, subject_type, zip_col, match_set) -> list[tuple]:
q = select(model.id, model.created_at)
if match_set is not None:
q = q.where(model.id.in_(match_set))
if zip_prefix:
q = q.where(zip_col.like(f"{zip_prefix}%"))
if added_from is not None:
q = q.where(func.date(model.created_at) >= added_from)
if added_to is not None:
q = q.where(func.date(model.created_at) <= added_to)
return [(ca, i, subject_type) for i, ca in db.execute(q).all()]
known_refs = _refs(KnownDog, SubjectType.known, KnownDog.last_known_zip, known_match)
unknown_refs = _refs(UnknownDog, SubjectType.unknown, UnknownDog.current_zip, unknown_match)
known_count, unknown_count = len(known_refs), len(unknown_refs)
if kind == "known":
refs = known_refs
elif kind == "unknown":
refs = unknown_refs
else:
refs = known_refs + unknown_refs
# Sort by (created_at, id); newest-first unless oldest requested.
refs.sort(key=lambda r: (r[0], r[1]), reverse=(sort != "oldest"))
total = len(refs)
items: list[dict] = []
for _created_at, dog_id, subject_type in refs[offset : offset + limit]:
if subject_type == SubjectType.known:
dog = db.get(KnownDog, dog_id)
if dog:
items.append(_known_profile(db, dog, ds_names, breed_model))
else:
dog = db.get(UnknownDog, dog_id)
if dog:
items.append(_unknown_profile(db, dog, ds_names, breed_model))
return {
"items": items, "total": total, "limit": limit, "offset": offset,
"counts": {"known": known_count, "unknown": unknown_count},
}
def _case_dog(db: Session, case: Case) -> dict | None:
"""Compact view of the dog a case is about (for the admin all-cases list)."""
if case.known_dog_id:
d = db.get(KnownDog, case.known_dog_id)
if d:
return {"kind": "known", "id": d.id, "name": d.name, "status": d.status.value,
"zip": d.last_known_zip, "thumb_url": _thumb_url(db, SubjectType.known, d.id)}
if case.unknown_dog_id:
d = db.get(UnknownDog, case.unknown_dog_id)
if d:
return {"kind": "unknown", "id": d.id, "name": d.description or "Found dog",
"status": d.status.value, "zip": d.current_zip,
"thumb_url": _thumb_url(db, SubjectType.unknown, d.id)}
return None
def _case_row(db: Session, case: Case) -> dict:
person = None
if case.person_id:
u = db.get(User, case.person_id)
if u:
person = {"id": u.id, "name": u.name, "email": u.email}
match_count = db.execute(
select(func.count()).select_from(Match).where(Match.case_id == case.id)
).scalar_one()
return {
"id": case.id, "type": case.type.value, "status": case.status.value,
"event_zip": case.event_zip, "event_date": case.event_date.isoformat(),
"created_at": case.created_at.isoformat(),
"dog": _case_dog(db, case), "person": person,
"finder_name": case.finder_name, "finder_email": case.finder_email,
"match_count": match_count,
}
def list_all_cases(
db: Session, *, case_type: str | None = None, status: str | None = None,
limit: int = 30, offset: int = 0,
) -> dict:
"""Paginated view of ALL cases (admin), newest first, optionally filtered by type/status."""
from ..models.base import CaseStatus, CaseType
conds = []
if case_type in ("lost", "found"):
conds.append(Case.type == CaseType(case_type))
if status in ("open", "matched", "resolved", "closed"):
conds.append(Case.status == CaseStatus(status))
total = db.execute(select(func.count()).select_from(Case).where(*conds)).scalar_one()
rows = db.execute(
select(Case).where(*conds).order_by(Case.created_at.desc()).limit(limit).offset(offset)
).scalars().all()
lost = db.execute(select(func.count()).select_from(Case).where(Case.type == CaseType.lost)).scalar_one()
found = db.execute(select(func.count()).select_from(Case).where(Case.type == CaseType.found)).scalar_one()
return {
"items": [_case_row(db, c) for c in rows], "total": total, "limit": limit, "offset": offset,
"counts": {"lost": lost, "found": found},
}
def list_owners(db: Session, *, q: str | None = None, limit: int = 30, offset: int = 0) -> dict:
"""Paginated owners (newest first) + the dogs they own; optional name/email search (admin)."""
from sqlalchemy import or_
from ..models.base import UserRole
conds = [User.role == UserRole.owner]
if q and q.strip():
like = f"%{q.strip()}%"
conds.append(or_(User.name.ilike(like), User.email.ilike(like)))
total = db.execute(select(func.count()).select_from(User).where(*conds)).scalar_one()
owners = db.execute(
select(User).where(*conds)
.order_by(User.created_at.desc(), User.id.desc())
.limit(limit).offset(offset)
).scalars().all()
items = []
for o in owners:
dogs = db.execute(
select(KnownDog).where(KnownDog.owner_id == o.id).order_by(KnownDog.id)
).scalars().all()
items.append({
"id": o.id, "name": o.name, "email": o.email, "zip": o.zip,
"dog_count": len(dogs),
"dogs": [{"id": d.id, "name": d.name, "status": d.status.value,
"zip": d.last_known_zip, "thumb_url": _thumb_url(db, SubjectType.known, d.id)}
for d in dogs],
})
return {"items": items, "total": total, "limit": limit, "offset": offset}
def owner_detail(db: Session, user_id: int) -> dict | None:
"""One owner's profile plus ALL their dogs (full AdminDog shape) and cases they reported.
Returns None if the user doesn't exist or isn't an owner (the admin Owners view is owner-only).
"""
from ..models.base import UserRole
o = db.get(User, user_id)
if not o or o.role != UserRole.owner:
return None
ds_names = {row[0]: row[1] for row in db.execute(select(Dataset.id, Dataset.name)).all()}
breed_model = active_breed_model_in_db(db)
dogs = db.execute(
select(KnownDog)
.where(KnownDog.owner_id == o.id)
.order_by(KnownDog.created_at.desc(), KnownDog.id.desc())
).scalars().all()
cases = db.execute(
select(Case).where(Case.person_id == o.id).order_by(Case.created_at.desc())
).scalars().all()
return {
"owner": {
"id": o.id, "name": o.name, "email": o.email, "zip": o.zip,
"phone": o.phone, "dog_count": len(dogs),
},
"dogs": [_known_profile(db, d, ds_names, breed_model) for d in dogs],
"cases": [_case_row(db, c) for c in cases],
}
def _thumb_url(db: Session, subject_type: SubjectType, subject_id: int) -> str | None:
pic = db.execute(
select(Picture)
.where(Picture.subject_type == subject_type, Picture.subject_id == subject_id)
.order_by(Picture.is_primary.desc(), Picture.id)
).scalars().first()
if not pic:
return None
key = pic.thumb_path or pic.file_path
return f"/media/{key}" if key else None
def case_count(db: Session, dataset: Dataset) -> int:
known_ids = _known_ids(db, dataset.id)
unknown_ids = _unknown_ids(db, dataset.id)
user_ids = _user_ids(db, dataset.id)
return len(_case_ids(db, known_ids, unknown_ids, user_ids))
# ---- transactional purge ----
def purge_dataset(db: Session, dataset: Dataset) -> dict[str, int]:
"""Delete every record linked to ``dataset`` then the dataset itself, in one transaction.
Returns per-type deleted counts. Children are removed before parents so FK constraints hold.
"""
from sqlalchemy import or_
known_ids = _known_ids(db, dataset.id)
unknown_ids = _unknown_ids(db, dataset.id)
user_ids = _user_ids(db, dataset.id)
case_ids = _case_ids(db, known_ids, unknown_ids, user_ids)
pic_ids = _picture_ids(db, known_ids, unknown_ids)
# Collect the storage keys NOW, before the picture rows are deleted, so we can remove the files
# after the DB transaction commits (a rolled-back purge never deletes anything on disk).
media_keys: list[str] = []
if pic_ids:
for file_path, thumb_path in db.execute(
select(Picture.file_path, Picture.thumb_path).where(Picture.id.in_(pic_ids))
).all():
if file_path:
media_keys.append(file_path)
if thumb_path:
media_keys.append(thumb_path)
# Matches to remove: any whose case OR candidate_case belongs to this dataset, PLUS any whose
# CANDIDATE is one of this dataset's dogs — a cross-dataset match (its case lives elsewhere)
# would otherwise be left pointing at a deleted dog.
from sqlalchemy import and_
match_conds = []
if case_ids:
match_conds.append(Match.case_id.in_(case_ids))
match_conds.append(Match.candidate_case_id.in_(case_ids))
if known_ids:
match_conds.append(
and_(Match.candidate_type == SubjectType.known, Match.candidate_id.in_(known_ids))
)
if unknown_ids:
match_conds.append(
and_(Match.candidate_type == SubjectType.unknown, Match.candidate_id.in_(unknown_ids))
)
match_ids: list[int] = (
list(db.execute(select(Match.id).where(or_(*match_conds))).scalars()) if match_conds else []
)
counts: dict[str, int] = {}
def _del(stmt) -> int:
return db.execute(stmt).rowcount or 0
try:
# notifications reference users/cases/matches
notif_conds = []
if user_ids:
notif_conds.append(Notification.user_id.in_(user_ids))
if case_ids:
notif_conds.append(Notification.case_id.in_(case_ids))
if match_ids:
notif_conds.append(Notification.match_id.in_(match_ids))
counts["notifications"] = _del(delete(Notification).where(or_(*notif_conds))) if notif_conds else 0
# breed_predictions + embeddings reference pictures
if pic_ids:
counts["breed_predictions"] = _del(
delete(BreedPrediction).where(BreedPrediction.picture_id.in_(pic_ids))
)
counts["embeddings"] = _del(delete(Embedding).where(Embedding.picture_id.in_(pic_ids)))
else:
counts["breed_predictions"] = 0
counts["embeddings"] = 0
# matches reference cases (and users via reviewed_by)
counts["matches"] = _del(delete(Match).where(Match.id.in_(match_ids))) if match_ids else 0
# null any remaining match.reviewed_by pointing at users we're about to delete
if user_ids:
db.execute(
Match.__table__.update()
.where(Match.reviewed_by.in_(user_ids))
.values(reviewed_by=None)
)
counts["pictures"] = _del(delete(Picture).where(Picture.id.in_(pic_ids))) if pic_ids else 0
counts["cases"] = _del(delete(Case).where(Case.id.in_(case_ids))) if case_ids else 0
counts["unknown_dogs"] = (
_del(delete(UnknownDog).where(UnknownDog.id.in_(unknown_ids))) if unknown_ids else 0
)
counts["known_dogs"] = (
_del(delete(KnownDog).where(KnownDog.id.in_(known_ids))) if known_ids else 0
)
# A user can be SHARED with another dataset (e.g. test CSVs that reuse an owner/finder email):
# it still owns a dog or a case that survives this purge, so deleting it would violate a FK.
# Only delete users with no surviving references; DETACH the rest (null dataset_id) so the
# dataset row can be removed without dangling them on a now-deleted dataset.
retained: set[int] = set()
if user_ids:
owners = db.execute(
select(KnownDog.owner_id).where(KnownDog.owner_id.in_(user_ids)).distinct()
).scalars()
persons = db.execute(
select(Case.person_id).where(Case.person_id.in_(user_ids)).distinct()
).scalars()
retained = {u for u in owners} | {u for u in persons if u is not None}
safe_users = [u for u in user_ids if u not in retained]
counts["users"] = _del(delete(User).where(User.id.in_(safe_users))) if safe_users else 0
if retained:
db.execute(
User.__table__.update().where(User.id.in_(retained)).values(dataset_id=None)
)
counts["users_detached"] = len(retained)
counts["datasets"] = _del(delete(Dataset).where(Dataset.id == dataset.id))
db.commit()
except Exception:
db.rollback()
raise
# DB is committed and consistent; now remove the files. Best-effort — a failure here only
# leaves orphans (recoverable with scripts/prune_media.py), never inconsistent data.
counts["media_files"] = _delete_media_keys(media_keys)
return counts
def _delete_media_keys(keys: list[str]) -> int:
# Dedupe: thumb_path may equal file_path (full-image thumbnails), so each file appears once.
unique = list(dict.fromkeys(keys))
if not unique:
return 0
storage = get_storage()
removed = 0
for key in unique:
try:
storage.delete(key)
removed += 1
except Exception: # noqa: BLE001
logger.warning("Could not delete media %s during purge", key, exc_info=True)
return removed
# ---- single-record delete (admin convenience / re-demoing) ----
def delete_case(db: Session, case_id: int) -> dict[str, int]:
"""Delete one case + its matches + related notifications (leaves the dog itself)."""
from sqlalchemy import or_
match_ids = list(
db.execute(
select(Match.id).where(or_(Match.case_id == case_id, Match.candidate_case_id == case_id))
).scalars()
)
counts: dict[str, int] = {}
ncond = [Notification.case_id == case_id]
if match_ids:
ncond.append(Notification.match_id.in_(match_ids))
counts["notifications"] = db.execute(delete(Notification).where(or_(*ncond))).rowcount or 0
counts["matches"] = (
db.execute(delete(Match).where(Match.id.in_(match_ids))).rowcount or 0 if match_ids else 0
)
counts["cases"] = db.execute(delete(Case).where(Case.id == case_id)).rowcount or 0
db.commit()
return counts
def delete_dog(db: Session, subject_type: SubjectType, dog_id: int) -> dict[str, int]:
"""Delete one dog and everything attached: pictures (+ media), embeddings, breed predictions,
its cases, any matches referencing it, and related notifications. Owners are left intact."""
from sqlalchemy import and_, or_
pic_rows = db.execute(
select(Picture.id, Picture.file_path, Picture.thumb_path).where(
Picture.subject_type == subject_type, Picture.subject_id == dog_id
)
).all()
pic_ids = [r[0] for r in pic_rows]
media_keys = [k for r in pic_rows for k in (r[1], r[2]) if k]
if subject_type == SubjectType.known:
case_ids = list(db.execute(select(Case.id).where(Case.known_dog_id == dog_id)).scalars())
else:
case_ids = list(db.execute(select(Case.id).where(Case.unknown_dog_id == dog_id)).scalars())
mconds = [and_(Match.candidate_type == subject_type, Match.candidate_id == dog_id)]
if case_ids:
mconds += [Match.case_id.in_(case_ids), Match.candidate_case_id.in_(case_ids)]
match_ids = list(db.execute(select(Match.id).where(or_(*mconds))).scalars())
counts: dict[str, int] = {}
ncond = []
if case_ids:
ncond.append(Notification.case_id.in_(case_ids))
if match_ids:
ncond.append(Notification.match_id.in_(match_ids))
counts["notifications"] = db.execute(delete(Notification).where(or_(*ncond))).rowcount or 0 if ncond else 0
if pic_ids:
counts["breed_predictions"] = db.execute(
delete(BreedPrediction).where(BreedPrediction.picture_id.in_(pic_ids))
).rowcount or 0
counts["embeddings"] = db.execute(
delete(Embedding).where(Embedding.picture_id.in_(pic_ids))
).rowcount or 0
counts["matches"] = (
db.execute(delete(Match).where(Match.id.in_(match_ids))).rowcount or 0 if match_ids else 0
)
counts["pictures"] = (
db.execute(delete(Picture).where(Picture.id.in_(pic_ids))).rowcount or 0 if pic_ids else 0
)
counts["cases"] = (
db.execute(delete(Case).where(Case.id.in_(case_ids))).rowcount or 0 if case_ids else 0
)
subj = KnownDog if subject_type == SubjectType.known else UnknownDog
counts["dogs"] = db.execute(delete(subj).where(subj.id == dog_id)).rowcount or 0
db.commit()
counts["media_files"] = _delete_media_keys(media_keys)
return counts
def delete_person(db: Session, user_id: int) -> dict[str, int]:
"""Delete a person (owner) and everything of theirs: the dogs they own (+ pictures/media/
embeddings/breeds), the cases they filed, and matches on those. Admins are refused elsewhere."""
from sqlalchemy import and_, or_
known_ids = list(db.execute(select(KnownDog.id).where(KnownDog.owner_id == user_id)).scalars())
pic_rows = (
db.execute(
select(Picture.id, Picture.file_path, Picture.thumb_path).where(
Picture.subject_type == SubjectType.known, Picture.subject_id.in_(known_ids)
)
).all()
if known_ids else []
)
pic_ids = [r[0] for r in pic_rows]
media_keys = [k for r in pic_rows for k in (r[1], r[2]) if k]
cconds = [Case.person_id == user_id]
if known_ids:
cconds.append(Case.known_dog_id.in_(known_ids))
case_ids = list(db.execute(select(Case.id).where(or_(*cconds))).scalars())
mconds = []
if case_ids:
mconds += [Match.case_id.in_(case_ids), Match.candidate_case_id.in_(case_ids)]
if known_ids:
mconds.append(and_(Match.candidate_type == SubjectType.known, Match.candidate_id.in_(known_ids)))
match_ids = list(db.execute(select(Match.id).where(or_(*mconds))).scalars()) if mconds else []
counts: dict[str, int] = {}
nconds = [Notification.user_id == user_id]
if case_ids:
nconds.append(Notification.case_id.in_(case_ids))
if match_ids:
nconds.append(Notification.match_id.in_(match_ids))
counts["notifications"] = db.execute(delete(Notification).where(or_(*nconds))).rowcount or 0
if pic_ids:
counts["breed_predictions"] = db.execute(
delete(BreedPrediction).where(BreedPrediction.picture_id.in_(pic_ids))
).rowcount or 0
counts["embeddings"] = db.execute(
delete(Embedding).where(Embedding.picture_id.in_(pic_ids))
).rowcount or 0
counts["matches"] = (
db.execute(delete(Match).where(Match.id.in_(match_ids))).rowcount or 0 if match_ids else 0
)
db.execute(Match.__table__.update().where(Match.reviewed_by == user_id).values(reviewed_by=None))
counts["pictures"] = (
db.execute(delete(Picture).where(Picture.id.in_(pic_ids))).rowcount or 0 if pic_ids else 0
)
counts["cases"] = (
db.execute(delete(Case).where(Case.id.in_(case_ids))).rowcount or 0 if case_ids else 0
)
counts["known_dogs"] = (
db.execute(delete(KnownDog).where(KnownDog.id.in_(known_ids))).rowcount or 0 if known_ids else 0
)
counts["users"] = db.execute(delete(User).where(User.id == user_id)).rowcount or 0
db.commit()
counts["media_files"] = _delete_media_keys(media_keys)
return counts
# ---- embed-all ----
def embed_all(db: Session, dataset: Dataset, *, progress=None, commit_every: int = 25) -> dict:
"""Generate embeddings AND breed predictions for every dataset picture lacking them.
Runs the model ONCE per image when the embedder and breed classifier share an HF model (spec
§9.1, "same model, two roles"); otherwise uses the two models independently. Idempotent /
best-effort. ``embedded`` and ``breeds`` count newly-written rows; ``skipped`` counts pictures
that already had both.
Commits (and reports ``progress``) every ``commit_every`` pictures instead of once at the end, so
a large encode run streams live progress to the UI job poller and doesn't hold one giant
transaction / unbounded session. Safe because encoding is additive and idempotent — a crash mid
run just leaves the already-committed embeddings (re-running skips them). ``progress`` is called
with ``{"processed", "total", "dataset_id"}``.
"""
dataset_id = dataset.id
known_ids = _known_ids(db, dataset_id)
unknown_ids = _unknown_ids(db, dataset_id)
pic_ids = _picture_ids(db, known_ids, unknown_ids)
total = len(pic_ids)
embedded = breeds = skipped = errors = 0
def _emit(done: int) -> None:
if progress:
progress({"processed": done, "total": total, "dataset_id": dataset_id})
_emit(0)
for i, pid in enumerate(pic_ids, start=1):
picture = db.get(Picture, pid)
if picture is not None:
try:
emb, breed = embed_and_breed_picture(db, picture, skip_if_exists=True)
if emb:
embedded += 1
if breed:
breeds += 1
if not emb and not breed:
skipped += 1
except Exception: # noqa: BLE001
errors += 1
logger.warning("Embed+breed failed for picture %s", pid, exc_info=True)
if i % commit_every == 0:
db.commit()
db.expunge_all() # keep the session small over a long run
_emit(i)
db.commit()
_emit(total)
return {
"total_pictures": total,
"embedded": embedded,
"breeds": breeds,
"skipped": skipped,
"errors": errors,
}
# ---- match-all ----
def match_dataset(db: Session, dataset: Dataset, candidate_dataset_id: int | None = None) -> dict:
"""Run matching for every case belonging to ``dataset`` (optionally scoped to a 2nd dataset)."""
known_ids = _known_ids(db, dataset.id)
unknown_ids = _unknown_ids(db, dataset.id)
case_ids = _case_ids(db, known_ids, unknown_ids, [])
cases_matched = 0
matches_created = 0
for cid in case_ids:
case = db.get(Case, cid)
if case is None:
continue
matches = run_matching_for_case(db, case, candidate_dataset_id=candidate_dataset_id)
cases_matched += 1
matches_created += len(matches)
db.commit()
return {
"cases_processed": cases_matched,
"matches_created": matches_created,
"candidate_dataset_id": candidate_dataset_id,
}