feat(tracking): revive occurrence tracking as a post-processing task#1272
feat(tracking): revive occurrence tracking as a post-processing task#1272
Conversation
Re-introduces the tracking feature originally developed by @mohamedelabbas1996 on feat/restore-tracking (PR #863), ported to land on top of current main and integrated with the generic post-processing framework introduced in #954: - Adds pgvector extension and Classification.features_2048 (2048-d embedding from the model backbone). Pipeline.save_results writes classification_resp.features into this column. - Adds Detection.next_detection self-FK (related_name=previous_detection) for storing the tracking chain. - Implements tracking as a BasePostProcessingTask (key="tracking") rather than a bespoke JobType. Greedy lowest-cost matching across consecutive captures (cosine similarity over features + IoU + box ratio + distance), followed by chain-based occurrence reassignment. - Adds a SourceImageCollection admin action to enqueue a PostProcessingJob with task=tracking. - Adds a unit test that reconstructs ground-truth occurrence groups from mocked per-occurrence feature vectors. Embedding-storage strategy (heavy outputs vs Classification row size) is intentionally left as a follow-up; this port keeps features_2048 on Classification for v1. Sister PR: RolnickLab/ami-data-companion#77 (produces `features` on ClassificationResponse). Supersedes #863. Co-authored-by: Mohamed Elabbas <hack1996man@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>
✅ Deploy Preview for antenna-preview canceled.
|
✅ Deploy Preview for antenna-ssec canceled.
|
There was a problem hiding this comment.
Pull request overview
Revives occurrence tracking by reintroducing the detection-linking + occurrence-reassignment algorithm as a generic post-processing task (TrackingTask) executed via PostProcessingJob, and adds support for storing backbone feature embeddings used by the matcher.
Changes:
- Add
Classification.features_2048(pgvector) andDetection.next_detectionto persist embeddings and tracking chains. - Implement and register
TrackingTaskas a post-processing task, plus an admin action to enqueue it. - Extend the processing-service schema/pipeline persistence path to accept and store
ClassificationResponse.features.
Reviewed changes
Copilot reviewed 11 out of 12 changed files in this pull request and generated 7 comments.
Show a summary per file
| File | Description |
|---|---|
| requirements/base.txt | Adds pgvector Python dependency for vector storage. |
| ami/ml/schemas.py | Adds optional features field to ClassificationResponse. |
| ami/ml/post_processing/tracking_task.py | New tracking implementation as a post-processing task (pairing + chain walking + occurrence reassignment). |
| ami/ml/post_processing/tests/test_tracking_task.py | Adds regression test ensuring tracking reconstructs occurrence groupings. |
| ami/ml/post_processing/registry.py | Registers TrackingTask in the post-processing registry. |
| ami/ml/models/pipeline.py | Persists ClassificationResponse.features into Classification.features_2048 on create and duplicate-update paths. |
| ami/main/models.py | Adds features_2048 to Classification and next_detection to Detection. |
| ami/main/migrations/0084_add_pgvector_extension.py | Creates the vector Postgres extension. |
| ami/main/migrations/0085_classification_features_2048.py | Adds the features_2048 vector column. |
| ami/main/migrations/0086_detection_next_detection.py | Adds the next_detection self OneToOne relation. |
| ami/main/admin.py | Adds admin action to enqueue tracking as a post-processing job. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| d.save() | ||
|
|
||
| occurrence.save() | ||
|
|
There was a problem hiding this comment.
The current implementation does per-row save() calls when reassigning occurrences / links. For Detection, save() triggers update_calculated_fields() (and may cause an extra save), making this much heavier than needed. Consider updating the occurrence field via bulk_update (or QuerySet.update) and, when saving individual rows is required, pass update_fields=[...] and update_calculated_fields=False to avoid the extra work.
| d.save() | |
| occurrence.save() | |
| Detection.objects.bulk_update(chain, ["occurrence"]) |
There was a problem hiding this comment.
Claude says: Deferring this one. Looked at Detection.update_calculated_fields (ami/main/models.py:2903) — it only sets timestamp from source_image when missing. For chain elements, timestamp is already populated, so the per-row save() is essentially a single UPDATE with no second pass. Switching to bulk_update here would skip update_calculated_fields entirely; cheap to do but I'd rather not bundle it with this set of fixes. Will track as a perf follow-up if profiling shows it matters.
| def event_fully_processed(event: Event, logger: logging.Logger, algorithm: Algorithm) -> bool: | ||
| total = event.captures.count() | ||
| processed = ( | ||
| event.captures.filter( | ||
| detections__classifications__features_2048__isnull=False, | ||
| detections__classifications__algorithm=algorithm, | ||
| ) | ||
| .distinct() | ||
| .count() | ||
| ) |
There was a problem hiding this comment.
event_fully_processed() treats an event as fully processed only if every capture has at least one detection classification with features_2048 for the chosen algorithm. Events commonly include captures with zero detections (or placeholder/null-bbox detections), which would make processed < total forever even when processing is complete. Consider counting only captures with at least one real detection (non-null bbox) and requiring those to have embeddings, or using a more direct “pipeline finished” signal if available.
There was a problem hiding this comment.
Claude says: Acknowledged but not changing in this PR. require_completely_processed_session defaults to False and is currently advisory; event_fully_processed is only invoked when an operator opts in. Agreed the "every capture" check is overly strict (captures with zero real detections will block forever). Filed as a follow-up: either count only captures with non-null-bbox detections, or wire this gate to a more direct "pipeline finished" signal once one exists.
📝 WalkthroughWalkthroughAdds occurrence-tracking: pgvector support and a 2048-d feature field, a TrackingTask post-processing implementation that links detections across captures, admin actions to enqueue tracking jobs, migrations and Docker support for pgvector, and tests covering tracking logic and admin flows. (49 words) Changes
Sequence DiagramsequenceDiagram
actor Admin
participant DjangoAdmin as "Django Admin"
participant JobQueue as "Job / DB"
participant TrackingTask as "TrackingTask"
participant Event as "Event / Captures"
participant Algo as "Feature Algorithm"
participant Detection as "Detection objects"
participant Similarity as "Cost calculator"
participant Occurrence as "Occurrence store"
Admin->>DjangoAdmin: Trigger run_tracking action / confirm
DjangoAdmin->>JobQueue: Create post_processing Job(s) with params.task="tracking"
JobQueue->>TrackingTask: Worker runs TrackingTask(job.params)
TrackingTask->>Event: Resolve scoped events (filter by project / event_ids)
TrackingTask->>Algo: Select algorithm most consistent for event
Algo-->>TrackingTask: Provide chosen algorithm id / availability
loop per event
TrackingTask->>Event: Load ordered source captures
loop per adjacent capture pair
TrackingTask->>Detection: Load detections for pair
Detection->>Algo: Obtain features_2048 for detections
Algo-->>TrackingTask: Return embeddings
TrackingTask->>Similarity: Compute embedding + geometric costs
Similarity-->>TrackingTask: Return cost matrix
TrackingTask->>Detection: Assign next_detection links (greedy)
end
TrackingTask->>Occurrence: Consolidate chains -> keeper occurrence(s)
end
TrackingTask-->>JobQueue: Persist updates, log stats, finish
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Review rate limit: 0/1 reviews remaining, refill in 60 minutes.Comment |
There was a problem hiding this comment.
Actionable comments posted: 5
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@ami/ml/post_processing/tracking_task.py`:
- Around line 134-140: The code currently returns early when all detections in a
discovered chain share the same occurrence_id, preventing re-materialization of
chains; remove that early continue and always materialize each chain into an
occurrence: in TrackingTask.run (the loop over variable chain) ensure you create
a new Occurrence (or reuse a freshly created one) for every chain found, update
each Detection.occurrence/occurrence_id in the chain to point to that
Occurrence, save the Detection rows, and handle singleton chains the same way so
stale groupings are replaced rather than left untouched.
- Around line 76-89: The query in get_most_common_algorithm_for_event can pick
up records with algorithm_id=None; change the Classification queryset to exclude
NULL algorithms (e.g., add .filter(algorithm_id__isnull=False) or
.exclude(algorithm_id__isnull=True)) so the annotated most_common never has
algorithm_id=None, then keep the existing lookup of
Algorithm.objects.get(id=most_common["algorithm_id"]); this ensures you skip
NULL algorithm groups instead of calling Algorithm.objects.get(id=None) which
raises.
- Around line 189-197: The nested loop repeatedly calls get_feature_vector(det,
algorithm) and get_feature_vector(nxt, algorithm), causing redundant queries;
instead precompute and cache feature vectors for all entries in
current_detections and next_detections (e.g., dicts keyed by detection id or
object) before the inner loop, filter out None vectors, then iterate using the
cached vectors when computing cost via total_cost(det_vec, nxt_vec, det.bbox,
nxt.bbox, diag); update references to use the cached det_vec/nxt_vec and remove
duplicate calls to get_feature_vector.
- Around line 291-306: The run() loop ignores
TrackingParams.feature_extraction_algorithm_id and always calls
get_most_common_algorithm_for_event(event); change the algorithm selection in
tracking_task.run() to first check the params from self._params() for
feature_extraction_algorithm_id and, if present, resolve that specific algorithm
(e.g., lookup by id) and use it, otherwise fall back to
get_most_common_algorithm_for_event(event); also update the log message around
the algorithm choice to indicate when the override from TrackingParams is being
applied.
In `@ami/ml/schemas.py`:
- Around line 136-141: features is accepted as any-length list but must be
exactly 2048 floats; add validation so incorrect-length vectors fail fast.
Update the schema by either changing the field type to a fixed-length list
(e.g., use pydantic.conlist(float, min_items=2048, max_items=2048) for the
features field) or add a `@pydantic.validator/field_validator` for "features" that
asserts features is None or len(features) == 2048 and raises a clear ValueError
(mentioning "features must be length 2048") so invalid payloads are rejected
before DB persistence.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 8503bb0f-6df8-4345-a40d-a698dddb48f8
📒 Files selected for processing (12)
ami/main/admin.pyami/main/migrations/0084_add_pgvector_extension.pyami/main/migrations/0085_classification_features_2048.pyami/main/migrations/0086_detection_next_detection.pyami/main/models.pyami/ml/models/pipeline.pyami/ml/post_processing/registry.pyami/ml/post_processing/tests/__init__.pyami/ml/post_processing/tests/test_tracking_task.pyami/ml/post_processing/tracking_task.pyami/ml/schemas.pyrequirements/base.txt
CodeRabbit + Copilot review fixes for the tracking post-processing task: - Materialize singleton chains so unlinked detections still get an occurrence; previously len(chain)<=1 was skipped, leaving detections orphaned. Skip rebuild only when the chain is already coherent (single occurrence, all detections assigned). - Cache feature vectors per image before the nested pair_detections loop. Cuts get_feature_vector() calls from O(m*n) to O(m+n). - Honor TrackingParams.feature_extraction_algorithm_id when set, falling back to most-common-algorithm detection only when unset. - Filter out NULL algorithms in get_most_common_algorithm_for_event to avoid Algorithm.objects.get(id=None). - Use try/except Detection.DoesNotExist for reverse OneToOne previous_detection access; clearer than relying on the RelatedObjectDoesNotExist/AttributeError inheritance trick. - Validate ClassificationResponse.features length == 2048 in the pydantic schema so wrong-length vectors fail at the boundary, not at DB save time. - Make the pgvector extension migration's reverse a no-op; dropping a shared extension on rollback is brittle in some environments. - Install postgresql-16-pgvector in the local/CI postgres image so CREATE EXTENSION succeeds when migrations run. Deferred to follow-up (noted in code @todo): full re-tracking that splits a previously-merged occurrence into multiple chains. Today, if every subchain still references the old occurrence_id, the existing occurrence stays in place. Addressing this needs a deeper design discussion. Co-Authored-By: Claude <noreply@anthropic.com>
…tomic per-event Locks v1 to fresh-data only and tightens correctness so the deferred re-track / split-detection question can wait for v2. Behavior changes: - New TrackingParams.require_fresh_event (default True) + event_is_fresh() precondition. v1 only operates on events where every detection has its own auto-created occurrence (1:1) and no chain links exist yet. Other events are skipped with a log explaining why. This sidesteps the Identification.occurrence CASCADE risk by construction. - Merge-into-first occurrence rebuild: keep chain[0]'s existing occurrence as the keeper, fold other detections in, delete now-empty siblings. Replaces the "delete-all-and-recreate" path. v1's fresh invariant guarantees absorbed siblings have no Identifications; v2 will need to reassign Identification.occurrence to the keeper instead of relying on this invariant. - Per-event transaction.atomic() boundary. A crash mid-event rolls back that event's chain links + occurrence consolidation only. - Width/height: skip transition (continue) instead of aborting the whole event (return). Single bad image no longer kills tracking elsewhere. - Multi-algo events: skip with warning naming the candidate algorithms, not silent majority pick. Operator must pass an explicit feature_extraction_algorithm_id to disambiguate. Replaces get_most_common_algorithm_for_event with get_unique_feature_algorithm_for_event. - Determinism: pair_detections candidate sort uses (cost, det.pk, nxt.pk) so tied costs produce reproducible pairings across runs. - TrackingParams.cost_threshold gets a docstring flagging that it was calibrated against synthetic features; needs per-dataset tuning before use on real backbone embeddings. Test updated to reflect the v1 fresh-data semantics: only chain links are wiped before re-running tracking; occurrences stay so event_is_fresh passes. Co-Authored-By: Claude <noreply@anthropic.com>
`update_occurrence_determination` only set `new_score` when the best prediction's taxon differed from the current one. Tracking can fold a new detection into an existing occurrence whose top species classification scores higher than the keeper's; in that case the taxon is unchanged but the score should still refresh to reflect the strongest evidence across all detections in the chain. Surfaced during PR #1272 E2E testing on a 363-capture event. Of the 13 multi-detection occurrences formed, 7 had stale `determination_score` values: the keeper's original score from its first detection rather than the highest-scoring same-taxon classification across the chain. Example: occ 634532 (Agriphila vulgivagellus, 41 detections) reported score 0.203 instead of 0.554. Add a regression test in `TestOccurrenceDeterminationScoreRefresh`. Co-Authored-By: Claude <noreply@anthropic.com>
Add `_resolve_events()` to TrackingTask that resolves the scope from either `config["event_ids"]` (priority) or the existing collection path. The new path lets EventAdmin trigger tracking on selected events directly without materializing a SourceImageCollection. When a job is attached, events are filtered to `job.project` so cross-project IDs from a misclicked or malicious POST cannot enter a single job. Missing or cross-project IDs are warned, not errored, so partial selections still make progress. `_resolve_collection()` now returns None instead of raising when no collection is set; the combined error message lives in `_resolve_events()`. The `unknown` allow-list in `_params()` is extended via a `_SCOPE_CONFIG_KEYS` constant. Backward compatible. The collection path remains; will be retired once the admin action is migrated to pass event_ids inline. Co-Authored-By: Claude <noreply@anthropic.com>
Add TrackingActionForm (django.forms.Form) as the source of truth for the admin-trigger UI: labels, help-text, validation rules, and the cleaned_data shape that becomes Job.params['config']. Form fields mirror TrackingParams (cost_threshold, skip_if_human_identifications, require_fresh_event) plus an optional feature_extraction_algorithm_id override whose dropdown is scoped to algorithms that produced features_2048 on the selected events. Scoping keeps the choice query bounded on production-sized DBs and prevents leaking algorithms from other projects into the dropdown. `to_config()` drops the algo override when blank so TrackingTask._params() falls through to per-event auto-detection rather than logging an unknown-key warning. Template (admin/main/tracking_confirmation.html) follows Django's stock delete_selected confirmation pattern, iterates form fields directly, and preserves the action checkbox / queryset round-trip. Form smoke-tested in shell — 4 fields render, defaults match DEFAULT_TRACKING_PARAMS, valid POST round-trips correctly, blank algo override is omitted from the resulting config. Co-Authored-By: Claude <noreply@anthropic.com>
Add `run_tracking_on_events` to EventAdmin. The action:
1. Renders an intermediate confirmation page (admin/main/tracking_confirmation.html)
with the TrackingActionForm so the operator can adjust cost_threshold,
skip_if_human_identifications, require_fresh_event, and (optionally) the
feature-extraction-algorithm override before queueing.
2. Partitions selected events by project and queues one Job per project, each
carrying its slice of event_ids in `params['config']`. Avoids Job-spam when
the operator picks "all events on this deployment" but still keeps the
project FK on each Job correct.
3. Skips events with `project_id=None` (legacy data without a project FK,
would otherwise hit `jobs_job.project_id NOT NULL` IntegrityError) and
surfaces the count + IDs as a WARNING-level admin message.
4. Gates the action to superusers. Tracking changes determinations across an
event; the broader "let any project admin trigger" decision is deferred.
Drop the duplicate `<h1>` from the template — Django admin's base.html already
renders `{{ title }}` via the `content_title` block.
Browser-tested end-to-end: filled form, submitted, Job 1550 created with
correct project FK, event_ids in config, status SUCCESS.
Co-Authored-By: Claude <noreply@anthropic.com>
…n branch from task SourceImageCollectionAdmin.run_tracking now mirrors the EventAdmin action: shows the TrackingActionForm on an intermediate confirmation page, then queues one Job per collection with event_ids computed inline from the collection's captures. TrackingTask.run no longer reads source_image_collection_id from config — collections are flattened to event_ids by the trigger. Job.source_image_collection FK is still set for traceability/UI but the task itself only knows about events. This collapses two scope-resolution paths into one. Tested end-to-end: queueing tracking on capture set 175 → Job 1551 SUCCESS with params event_ids=[2702]. Co-Authored-By: Claude <noreply@anthropic.com>
TrackingTask._resolve_events: event_ids resolution, cross-project drop, ValueError on missing config. Admin actions: EventAdmin renders intermediate page without confirm, queues per-project jobs with config passthrough; SourceImageCollectionAdmin flattens collection to event_ids on the resulting Job. Co-Authored-By: Claude <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
ami/main/models.py (1)
3381-3401:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winHandle zero-valued scores when refreshing
determination_score.Line 3398 uses a truthy check (
if new_score and ...), so a valid0.0score won’t be persisted.Fix falsy-score edge case
- if new_score and new_score != occurrence.determination_score: + if new_score is not None and new_score != occurrence.determination_score: logger.debug(f"Changing det. score of {occurrence} from {occurrence.determination_score} to {new_score}") occurrence.determination_score = new_score needs_update = True🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ami/main/models.py` around lines 3381 - 3401, The truthy check for new_score drops valid 0.0 values so determination_score won't be updated; change the conditional that updates occurrence.determination_score to use an explicit None check (e.g., "if new_score is not None and new_score != occurrence.determination_score") so zero-valued scores are persisted, referencing variables new_score, occurrence.determination_score and the update block that sets occurrence.determination_score and needs_update.
🧹 Nitpick comments (3)
ami/main/models.py (1)
2793-2800: ⚡ Quick winAdd a DB guard to prevent self-referential detection links.
Line 2793 introduces chain pointers, but there’s no DB-level protection against
next_detection_id == id. A self-loop can break traversal logic and create hard-to-debug chain corruption.Suggested model constraint
class Detection(BaseModel): @@ class Meta: ordering = [ "frame_num", "timestamp", ] + constraints = [ + models.CheckConstraint( + check=~models.Q(pk=models.F("next_detection_id")), + name="detection_next_not_self", + ), + ]🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ami/main/models.py` around lines 2793 - 2800, Add a DB-level CheckConstraint to the model that defines the next_detection OneToOneField to forbid self-references (i.e., next_detection_id == id). In the model that declares next_detection (the class containing the next_detection = models.OneToOneField(...) field), add a CheckConstraint such as one using condition=~Q(id=F('next_detection')) with a clear name like "prevent_self_next_detection" and include it in the model's Meta.constraints, then create and apply a migration to enforce it at the database level.ami/main/admin.py (1)
248-346: ⚡ Quick winExtract the shared tracking-action flow before these paths drift further.
These two actions are now mostly copy/paste, and they have already diverged in small ways (
run_trackingsortsevent_ids,run_tracking_on_eventsdoes not). Pull the confirmation-page context andJobpayload construction into one helper so future changes to tracking fields or guardrails stay aligned across both entry points.Also applies to: 771-867
ami/main/test_admin.py (1)
17-118: ⚡ Quick winAdd regression coverage for the permission and confirmation guardrails.
The new tests only exercise superuser success paths. Please add assertions that a non-superuser POST does not create
Jobs for either action, and mirror the collection action's intermediate confirmation-page flow as well; those branches are now part of the feature contract.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ami/main/test_admin.py` around lines 17 - 118, Add assertions and new test methods to cover permission and confirmation guardrails: (1) in _AdminTrackingCase or a new test class, create a regular non-superuser, force_login them and POST the same payloads to admin:main_event_changelist (use TestEventAdminTrackingAction._post_action) and admin:main_sourceimagecollection_changelist, then assert Job.objects.filter(job_type_key="post_processing").count() remains 0 and no Job created for the target SourceImageCollection; (2) mirror TestEventAdminTrackingAction.test_renders_intermediate_page_without_confirm for the collection action by POSTing without "confirm" to the view name "admin:main_sourceimagecollection_changelist" and assert a 200 response containing the intermediate page text (e.g. "Run Occurrence Tracking" or "Tracking parameters") and that no Job was created. Ensure tests reference Job, SourceImageCollection, TestCollectionAdminTrackingAction.test_creates_job_with_event_ids_from_collection, and the admin action names to find the code paths.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@ami/ml/post_processing/tracking_task.py`:
- Around line 265-274: A detection with bbox=None can reach the pairing loop and
crash total_cost when indexing det.bbox or nxt.bbox; before calling total_cost
inside the nested loops iterating current_detections and next_detections (using
current_vectors/next_vectors and det.pk/nxt.pk), add a guard that skips the pair
if either det.bbox or nxt.bbox is None (or otherwise invalid/empty), so cost is
only computed when both bounding boxes are present and valid before comparing
against cost_threshold.
- Around line 225-230: The current loop catches all exceptions when deleting
Occurrence objects (occ_id) which can mask DB/integrity failures and allow the
surrounding event transaction to commit; update the code in tracking_task.py
around the loop that iterates old_occ_ids - {keeper.pk} (where
Occurrence.objects.filter(id=occ_id).delete() is called) to not swallow
errors—either remove the try/except entirely so exceptions propagate, or after
logging with logger.error re-raise the exception (raise) so the transaction will
roll back; keep references to occ_id, old_occ_ids and keeper.pk so the logic and
logging remain clear.
---
Outside diff comments:
In `@ami/main/models.py`:
- Around line 3381-3401: The truthy check for new_score drops valid 0.0 values
so determination_score won't be updated; change the conditional that updates
occurrence.determination_score to use an explicit None check (e.g., "if
new_score is not None and new_score != occurrence.determination_score") so
zero-valued scores are persisted, referencing variables new_score,
occurrence.determination_score and the update block that sets
occurrence.determination_score and needs_update.
---
Nitpick comments:
In `@ami/main/models.py`:
- Around line 2793-2800: Add a DB-level CheckConstraint to the model that
defines the next_detection OneToOneField to forbid self-references (i.e.,
next_detection_id == id). In the model that declares next_detection (the class
containing the next_detection = models.OneToOneField(...) field), add a
CheckConstraint such as one using condition=~Q(id=F('next_detection')) with a
clear name like "prevent_self_next_detection" and include it in the model's
Meta.constraints, then create and apply a migration to enforce it at the
database level.
In `@ami/main/test_admin.py`:
- Around line 17-118: Add assertions and new test methods to cover permission
and confirmation guardrails: (1) in _AdminTrackingCase or a new test class,
create a regular non-superuser, force_login them and POST the same payloads to
admin:main_event_changelist (use TestEventAdminTrackingAction._post_action) and
admin:main_sourceimagecollection_changelist, then assert
Job.objects.filter(job_type_key="post_processing").count() remains 0 and no Job
created for the target SourceImageCollection; (2) mirror
TestEventAdminTrackingAction.test_renders_intermediate_page_without_confirm for
the collection action by POSTing without "confirm" to the view name
"admin:main_sourceimagecollection_changelist" and assert a 200 response
containing the intermediate page text (e.g. "Run Occurrence Tracking" or
"Tracking parameters") and that no Job was created. Ensure tests reference Job,
SourceImageCollection,
TestCollectionAdminTrackingAction.test_creates_job_with_event_ids_from_collection,
and the admin action names to find the code paths.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: af1982cf-be83-45d5-b51f-ee88e427df7e
📒 Files selected for processing (11)
ami/main/admin.pyami/main/migrations/0084_add_pgvector_extension.pyami/main/models.pyami/main/test_admin.pyami/main/tests.pyami/ml/post_processing/admin_forms.pyami/ml/post_processing/tests/test_tracking_task.pyami/ml/post_processing/tracking_task.pyami/ml/schemas.pyami/templates/admin/main/tracking_confirmation.htmlcompose/local/postgres/Dockerfile
✅ Files skipped from review due to trivial changes (1)
- compose/local/postgres/Dockerfile
🚧 Files skipped from review as they are similar to previous changes (1)
- ami/ml/schemas.py
| for occ_id in old_occ_ids - {keeper.pk}: | ||
| try: | ||
| Occurrence.objects.filter(id=occ_id).delete() | ||
| merged += 1 | ||
| except Exception as e: | ||
| logger.error(f"Failed to delete occurrence {occ_id}: {e}") |
There was a problem hiding this comment.
Do not swallow occurrence-delete failures inside the event transaction.
Catching Exception here can hide integrity/DB failures and still commit partial reassignment state for the event.
💡 Suggested fix
for occ_id in old_occ_ids - {keeper.pk}:
- try:
- Occurrence.objects.filter(id=occ_id).delete()
- merged += 1
- except Exception as e:
- logger.error(f"Failed to delete occurrence {occ_id}: {e}")
+ deleted, _ = Occurrence.objects.filter(id=occ_id).delete()
+ if deleted:
+ merged += 1🧰 Tools
🪛 Ruff (0.15.12)
[warning] 229-229: Do not catch blind exception: Exception
(BLE001)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@ami/ml/post_processing/tracking_task.py` around lines 225 - 230, The current
loop catches all exceptions when deleting Occurrence objects (occ_id) which can
mask DB/integrity failures and allow the surrounding event transaction to
commit; update the code in tracking_task.py around the loop that iterates
old_occ_ids - {keeper.pk} (where Occurrence.objects.filter(id=occ_id).delete()
is called) to not swallow errors—either remove the try/except entirely so
exceptions propagate, or after logging with logger.error re-raise the exception
(raise) so the transaction will roll back; keep references to occ_id,
old_occ_ids and keeper.pk so the logic and logging remain clear.
| for det in current_detections: | ||
| det_vec = current_vectors[det.pk] | ||
| if det_vec is None: | ||
| continue | ||
| for nxt in next_detections: | ||
| nxt_vec = next_vectors[nxt.pk] | ||
| if nxt_vec is None: | ||
| continue | ||
| cost = total_cost(det_vec, nxt_vec, det.bbox, nxt.bbox, diag) | ||
| if cost < cost_threshold: |
There was a problem hiding this comment.
Guard null bounding boxes before cost computation.
total_cost(...) assumes both boxes are indexable. If a detection with bbox=None reaches this loop (e.g., non-fresh runs), Line 273 can crash the event run.
💡 Suggested fix
for det in current_detections:
+ if det.bbox is None:
+ continue
det_vec = current_vectors[det.pk]
if det_vec is None:
continue
for nxt in next_detections:
+ if nxt.bbox is None:
+ continue
nxt_vec = next_vectors[nxt.pk]
if nxt_vec is None:
continue
cost = total_cost(det_vec, nxt_vec, det.bbox, nxt.bbox, diag)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@ami/ml/post_processing/tracking_task.py` around lines 265 - 274, A detection
with bbox=None can reach the pairing loop and crash total_cost when indexing
det.bbox or nxt.bbox; before calling total_cost inside the nested loops
iterating current_detections and next_detections (using
current_vectors/next_vectors and det.pk/nxt.pk), add a guard that skips the pair
if either det.bbox or nxt.bbox is None (or otherwise invalid/empty), so cost is
only computed when both bounding boxes are present and valid before comparing
against cost_threshold.
Adds the shared infrastructure that PRs #999 (class masking) and #1272 (tracking) each grew independently. Uses SmallSizeFilterTask as the migration consumer to prove the pattern without carving up another open PR. Changes: - BasePostProcessingTask gains required pydantic config_schema class attr; validates Job.params['config'] at task __init__ via self.config_schema(**config). self.config is now a BaseModel instance (was: freeform dict). - SmallSizeFilterTask: adds SmallSizeFilterConfig schema (size_threshold + source_image_collection_id, validators for 0<x<1, extra=forbid). Existing default 0.0008 preserved. - New admin form base BasePostProcessingActionForm with to_config() contract, plus SmallSizeFilterActionForm exposing size_threshold knob (currently unreachable — admin trigger hardcoded empty config). - Parameterized confirmation template at admin/post_processing/confirmation.html with overridable {% block intro %} and a _form_fieldset.html partial. - run_small_size_filter admin action rewrites onto new pattern: intermediate confirmation page on first POST, validate + enqueue on confirm. Per-collection Job creation preserved (each Job gets correct project FK). - 17 tests covering schema contract, form validation, intermediate page render, multi-collection partitioning by project FK. Pydantic v1 syntax throughout (repo pins pydantic<2.0). Memory note: container is v1, .dict() / .__fields__ are correct here. Out of scope (explicitly deferred): - Project-partitioning helper (defer to whichever multi-scope adopter lands first, likely #1272's tracking which partitions events across projects) - REST API trigger surface (eventual primary; admin not future-primary) - Rank rollup, class masking, tracking tasks themselves — stay in their PRs Design doc: docs/claude/planning/2026-05-01-post-processing-admin-scaffolding-design.md Co-Authored-By: Claude <noreply@anthropic.com>
Summary
Revives the occurrence-tracking feature originally built by @mohamedelabbas1996 in #863, ported to land on top of current
mainand integrated with the generic post-processing framework introduced in #954.The tracking algorithm and the data model (chains of detections via
Detection.next_detection, occurrence reassignment from chains, feature-space cost function) are unchanged in spirit and substantially in code. What changed is how it plugs in: instead of being a bespokeJobType, tracking is now aBasePostProcessingTask(key="tracking") executed byPostProcessingJob, matching the pattern established by Small Size Filter.This PR also re-introduces the
Classification.features_2048column (2048-d embedding from the model backbone) and the matchingfeaturesfield on the processing-serviceClassificationResponseschema. Embeddings are produced by the sister PR in ami-data-companion: RolnickLab/ami-data-companion#77.Scope: v1 fresh-data only
This PR explicitly targets v1: tracking on fresh data. A "fresh event" = the state immediately after the processing pipeline runs, where every detection has its own auto-created occurrence (1:1) and no chain links exist yet. v1 collapses those 1:1 mappings into chain-based occurrences using a merge-into-first strategy (the earliest existing occurrence in the chain is kept, others are absorbed into it).
Re-tracking previously-tracked data (and the related correctness questions around splitting merged occurrences) is out of scope and lands in v2 (incremental append/prepend — see notes below). v1 enforces this with an
event_is_fresh()precondition: events whose detections are already partially or fully consolidated are skipped with a log line.Related
What's preserved from #863 (Mohamed's work)
Quoting from #863, all of which still applies here:
The cost function (
1 - cos(features) + 1 - IoU + 1 - box_ratio + distance/diag), the greedy-lowest-cost assignment with claim sets, the chain walk that creates one Occurrence per chain, and the human-ID skip — all carried over.What changed for the framework integration
#954 landed a generic post-processing framework in the meantime, so the revival adapts tracking to fit that pattern rather than re-introducing the bespoke
TrackingJobsubclass:ami/ml/post_processing/tracking_task.py—TrackingTask(BasePostProcessingTask)withkey="tracking",name="Occurrence Tracking".TrackingParams(cost threshold, skip-on-human-IDs, fresh-event guard, etc.) reads fromself.config. Source image collection resolves fromjob.source_image_collection, with asource_image_collection_idconfig fallback.ami/ml/post_processing/registry.py— registers the task alongsideSmallSizeFilterTask.ami/main/admin.py— adds a "Run Occurrence Tracking" admin action onSourceImageCollectionthat enqueues aPostProcessingJobwithparams={"task": "tracking", "config": dataclasses.asdict(DEFAULT_TRACKING_PARAMS)}. Mirrors the Small Size Filter action.ami/jobs/models.py— no changes needed. The earlier branch added aTrackingJobJobType and a job-type-key migration; both are obsolete now thatPostProcessingJobdispatches generically.Schema additions
Three migrations on top of
0083_dedupe_taxalist_names:0084_add_pgvector_extension.py—CREATE EXTENSION IF NOT EXISTS vector;(reverse is a no-op since the extension can be shared)0085_classification_features_2048.py—Classification.features_2048(pgvector.django.VectorField(dimensions=2048, null=True))0086_detection_next_detection.py—Detection.next_detection(OneToOneField("self", related_name="previous_detection", null=True))Adds
pgvector==0.3.6torequirements/base.txt. The local/CI Postgres image (compose/local/postgres/Dockerfile) installspostgresql-16-pgvectorsoCREATE EXTENSIONsucceeds.ami/ml/schemas.py::ClassificationResponsegains an optionalfeatures: list[float] | Nonewith a pydantic v1@validatorenforcing length == 2048 (so wrong-length payloads fail at the boundary, not at DB save time).pipeline.save_resultswrites it intoClassification.features_2048on both create and the existing duplicate-update path.v1 hardening (post-review-pass)
After the initial review surfaced several silent-bug risks, the v1 implementation was tightened. The relevant guards:
event_is_fresh()precondition (default on viarequire_fresh_event=True). Tracking only runs on events where every detection has an occurrence and every occurrence has exactly one detection. Anything else is skipped with a log explaining why. This sidesteps the "re-track destroys identifications via CASCADE" risk by construction (fresh occurrences have no identifications).Identification.occurrenceinstead of relying on this invariant.transaction.atomic(). Chain materialization for a single event is atomic; a crash mid-event rolls back that event only, leaving the rest of the job unaffected.feature_extraction_algorithm_idis unset, the event is skipped with a warning naming the candidates. Operator must pass an explicit override to disambiguate. (Fresh events from a single pipeline run should never hit this case.)(cost, det.pk, nxt.pk)as the secondary key so tied costs produce reproducible pairings across runs.TrackingParams.cost_threshold: the default is calibrated against synthetic features in tests, not real backbone embeddings. Tune per dataset before relying on it.How to test
docker compose up -dand run migrations:docker compose run --rm django python manage.py migrate/processresponse carriesfeatureson classifications).Classification.features_2048rows. Each detection will get an auto-created occurrence (1:1).Detection.next_detectionlinks chain detections in temporal order.Deployment notes
pgvectorextension must be installed on the Postgres server. Migration0084runsCREATE EXTENSION IF NOT EXISTS vector;, which requires the extension package to be available (the standardpgvector/pgvectorPostgres image, the official Postgrespostgresql-XX-pgvectorapt package, or RDS'vectorextension). The local/CI Postgres image is built to include it.pgvector==0.3.6.features(i.e. the ADC Web UI setup table views #77 build); existing rows will simply havefeatures_2048 IS NULLand be skipped by tracking, which is the desired behaviour.v2 plan: incremental tracking (append/prepend)
The v1 design above is the right batch primitive but the wrong UX: users won't see merged occurrences until the entire event has been processed and the post-processing job has run. v2 closes that latency gap by tracking incrementally as captures arrive, so chain consolidation happens in near-real-time.
Architecture readiness
The v1 building blocks are deliberately structured to support v2 with minimal rework:
pair_detections(cur_image, next_image, …)already operates on two adjacent images and is idempotent under theDetection.next_detectionOneToOne uniqueness constraint. Calling it twice for the same transition is safe.assign_occurrences_from_detection_chainswalks chains by followingnext_detection, so it works regardless of whether the chain was built incrementally or in one batch pass. The merge-into-first keeper rule handles prepend: when a new detection is added to the head of an existing chain, the existing keeper survives and the new detection's auto-created singleton occurrence is absorbed.event_is_fresh()is the v1 gate; v2 will replace it with a tighter "incremental-safe" check (allow events with multi-detection occurrences as long as their identifications are intact).What v2 needs to add
pipeline.save_resultsenqueues a tracking task scoped to the just-processed image's adjacent transitions (image-1 ↔ image,image ↔ image+1). Reuses the existing job framework, no new signal plumbing.post_saveonDetectionschedules a debounced per-event tracking pass. Simpler scope but heavier signal management.Identification.occurrenceto the keeper before deleting the absorbed occurrence. v1 sidesteps this by refusing non-fresh events; v2 cannot.i+1is processed before imagei, prepending must work cleanly. The current chain walk already handles this (walks from "no previous" detection), but the v2 trigger needs to re-pair thei ↔ i+1transition onceiarrives.What stays unchanged
Classification.features_2048).Detection.next_detectionschema and the OneToOne uniqueness it provides.Out of scope / follow-ups
features_2048directly onClassificationfor v1. Moving heavy outputs (embeddings, logits) to a siblingClassificationOutput(or a genericAlgorithmOutput) table — or derivingscoresfromlogits— is intentionally deferred to a separate change.Algorithm.task_type. The base post-processing class hardcodestask_type=POST_PROCESSINGon the auto-created Algorithm row;AlgorithmTaskType.TRACKINGalready exists and would be more accurate. Trivial follow-up.require_completely_processed_sessiondefaults toFalsefor now, matching the late-stage state of Add support for occurrence tracking #863. Re-enable once we have a reliable signal that an event is fully processed.event_fully_processed()over-strictness. Treats events where any capture has zero real detections as "not fully processed" forever. Gated behind the default-off feature flag above; will need rework when that flag flips on.bulk_updatefor chain detection reassignment. Per-rowsave()is fine for the current detection volumes, butbulk_update(["occurrence"])is the obvious perf follow-up if profiling shows it matters.Checklist
ami/ml/post_processing/tests/test_tracking_task.py— ground-truth-occurrence-reproduction test ported from Add support for occurrence tracking #863, updated for the v1 fresh-event invariant).features_2048. Without it the migrations and admin action are inert but harmless.Summary by CodeRabbit
New Features
Migrations
Tests
Chores
E2E test results (local, 2026-04-29)
End-to-end validated against ADC #77 (
feat/add-classification-features-to-response) running in pull-mode (NATS) against this branch. Both runs used pipelinequebec_vermont_moths_2023.Run 1 — sparse intervals (5 min): 10 captures, single event.
features_2048populated.Run 2 — dense intervals (~2s active): 363 captures, full-night session.
features_2048(rest filtered out by binary moth/non-moth before species classifier ran).Occurrence.duration()updates correctly on multi-detection occurrences in both runs.Detection.next_detectionlinkage and the fresh-event guard behaved as designed.Bug found during testing (separate, not blocking)
Occurrence.last_appearance_timestampreturnsnullon the detail endpoint (GET /api/v2/occurrences/{id}/) for multi-det occurrences. The annotation that powers it is only applied on the list queryset (with_timestamps()), not on the detail view.first_appearance_timestamp,duration, andduration_labelall populate correctly. Pre-existing — not introduced by this PR.What should be tested on staging
Pull-mode (NATS) deployments need two configuration steps verified before a real-data run:
Worker-side feature toggle. Until Propagate PipelineRequestConfigParameters through pull-mode (NATS) tasks #1275 is fixed, the processing-service worker must be started with
AMI_INCLUDE_FEATURES=truein its environment.Pipeline.default_config = {"include_features": True}andProjectPipelineConfig.configare silently ignored on the pull-mode path becausePipelineProcessingTaskdoes not carry aconfigfield. ConfirmClassification.features_2048is populated on a fresh ML run before triggering tracking.pgvectorin the Celery worker image. Thepgvector==0.3.6Python module must be installed in the Celery worker image, not just the Django web image. If only the web image has it,bulk_createofClassificationrows can silently drop thefeatures_2048value while the rest of the row saves successfully (observed locally before the worker image was rebuilt). The compose file's celeryworker service must rebuild from the samerequirements/base.txtafter this PR lands.Suggested staging exercise:
requirements/base.txt.manage.py migrate(createsvectorextension, addsfeatures_2048andnext_detectioncolumns).features_2048will be null).Classification.objects.filter(features_2048__isnull=False).count()is non-zero on the affected detections.next_detectionlinks chain detections in temporal order.Edge cases worth probing:
feature_extraction_algorithm_id) — should skip with a warning, not silently pick a majority.width/height— single transition skipped, event continues.(cost, det.pk, nxt.pk)).