feat(post-processing): implement class masking + rank rollup on the admin framework

Rework of #999 onto the #1289 post-processing framework. The branch was cut from
an old main with its own hand-rolled admin action; ClassMaskingTask and
RankRollupTask now subclass BasePostProcessingTask with pydantic config schemas
and are triggered through make_post_processing_action (collection scope on
SourceImageCollection, single-occurrence scope on Occurrence for class masking).

Correctness fixes from the review threads:
- Class masking selects the top class from an -inf-masked softmax, so a class
  excluded by the taxa list can never win even when it had the highest logit;
  raises when the taxa list excludes every class in the category map. Stored
  logits stay raw (JSON-safe) and the mask is captured in scores (excluded -> 0).
- The masked-output Algorithm is one per (source algorithm, taxa list) and its
  category map is persisted (previously set in memory only, so masked
  classifications referenced a null map).
- applied_to is populated on new masked classifications (the provenance the API
  exposes was left blank).
- Rank rollup preloads category-map labels in two queries and select_relates the
  per-row relations instead of dereferencing category_map per classification.

Surfaces provenance in the API: applied_to is added to the Classification
serializers, and applied_to__algorithm is prefetched in the occurrence
list/detail prefetch and the classification viewset to avoid an N+1 on render.

Tests: pydantic config validation, admin trigger for both scopes, the masking
maths (including the excluded-class guarantee and the all-excluded error),
ClassMaskingTask.run() end to end for both scopes, and rank rollup. 20 new tests;
full post_processing suite and occurrence query-count tests pass.

Co-Authored-By: Claude <noreply@anthropic.com>

Author: mihow

ami/main/admin.py

@@ -14,7 +14,11 @@
 from ami.jobs.models import Job
 from ami.ml.models.project_pipeline_config import ProjectPipelineConfig
 from ami.ml.post_processing.admin.actions import make_post_processing_action
+from ami.ml.post_processing.admin.class_masking_form import ClassMaskingActionForm
+from ami.ml.post_processing.admin.rank_rollup_form import RankRollupActionForm
 from ami.ml.post_processing.admin.small_size_filter_form import SmallSizeFilterActionForm
+from ami.ml.post_processing.class_masking import ClassMaskingTask
+from ami.ml.post_processing.rank_rollup import RankRollupTask
 from ami.ml.post_processing.small_size_filter import SmallSizeFilterTask
 from ami.ml.tasks import remove_duplicate_classifications
 
@@ -528,6 +532,12 @@ def detections_count(self, obj) -> int:
         scope_resolver=lambda occurrence: {"occurrence_id": occurrence.pk},
         name_resolver=lambda task_cls, occurrence: (f"Post-processing: {task_cls.name} on Occurrence {occurrence.pk}"),
     )
+    run_class_masking = make_post_processing_action(
+        ClassMaskingTask,
+        ClassMaskingActionForm,
+        scope_resolver=lambda occurrence: {"occurrence_id": occurrence.pk},
+        name_resolver=lambda task_cls, occurrence: (f"Post-processing: {task_cls.name} on Occurrence {occurrence.pk}"),
+    )
 
     @admin.action(description="Recompute determination from current classifications and identifications")
     def recompute_determination(self, request: HttpRequest, queryset: QuerySet[Any]) -> None:
@@ -544,7 +554,7 @@ def recompute_determination(self, request: HttpRequest, queryset: QuerySet[Any])
             count += 1
         self.message_user(request, f"Recomputed determination for {count} occurrence(s).")
 
-    actions = [run_small_size_filter, recompute_determination]
+    actions = [run_small_size_filter, run_class_masking, recompute_determination]
 
     # Order by -id (the indexed primary key) rather than -created_at, which has no
     # index and would force a full sort of the table to find the newest page. id
@@ -841,11 +851,29 @@ def populate_collection_async(self, request: HttpRequest, queryset: QuerySet[Sou
             f"Post-processing: {task_cls.name} on Capture Set {collection.pk}"
         ),
     )
+    run_class_masking = make_post_processing_action(
+        ClassMaskingTask,
+        ClassMaskingActionForm,
+        scope_resolver=lambda collection: {"source_image_collection_id": collection.pk},
+        name_resolver=lambda task_cls, collection: (
+            f"Post-processing: {task_cls.name} on Capture Set {collection.pk}"
+        ),
+    )
+    run_rank_rollup = make_post_processing_action(
+        RankRollupTask,
+        RankRollupActionForm,
+        scope_resolver=lambda collection: {"source_image_collection_id": collection.pk},
+        name_resolver=lambda task_cls, collection: (
+            f"Post-processing: {task_cls.name} on Capture Set {collection.pk}"
+        ),
+    )
 
     actions = [
         populate_collection,
         populate_collection_async,
         run_small_size_filter,
+        run_class_masking,
+        run_rank_rollup,
     ]
 
     # Hide images many-to-many field from form. This would list all source images in the database.

ami/main/api/serializers.py

@@ -944,10 +944,26 @@ class ClassificationPredictionItemSerializer(serializers.Serializer):
     logit = serializers.FloatField(read_only=True)
 
 
+class ClassificationAppliedToSerializer(serializers.ModelSerializer):
+    """Lightweight nested representation of the parent classification this was derived from.
+
+    Post-processing tasks (class masking, rank rollup) record provenance via
+    ``Classification.applied_to``; this exposes just enough to show what a result
+    was derived from without recursing back into the full classification.
+    """
+
+    algorithm = AlgorithmSerializer(read_only=True)
+
+    class Meta:
+        model = Classification
+        fields = ["id", "created_at", "algorithm"]
+
+
 class ClassificationSerializer(DefaultSerializer):
     taxon = TaxonNestedSerializer(read_only=True)
     algorithm = AlgorithmSerializer(read_only=True)
     top_n = ClassificationPredictionItemSerializer(many=True, read_only=True)
+    applied_to = ClassificationAppliedToSerializer(read_only=True)
 
     class Meta:
         model = Classification
@@ -960,6 +976,7 @@ class Meta:
             "scores",
             "logits",
             "top_n",
+            "applied_to",
             "created_at",
             "updated_at",
         ]
@@ -982,6 +999,8 @@ class Meta(ClassificationSerializer.Meta):
 
 
 class ClassificationListSerializer(DefaultSerializer):
+    applied_to = ClassificationAppliedToSerializer(read_only=True)
+
     class Meta:
         model = Classification
         fields = [
@@ -990,6 +1009,7 @@ class Meta:
             "taxon",
             "score",
             "algorithm",
+            "applied_to",
             "created_at",
             "updated_at",
         ]
@@ -1009,6 +1029,7 @@ class Meta:
             "score",
             "terminal",
             "algorithm",
+            "applied_to",
             "created_at",
         ]
 

ami/main/api/views.py

@@ -2053,7 +2053,7 @@ class ClassificationViewSet(DefaultViewSet, ProjectMixin):
     """
 
     require_project_for_list = True  # Unfiltered list scans are too expensive on this table
-    queryset = Classification.objects.all().select_related("taxon", "algorithm")  # , "detection")
+    queryset = Classification.objects.all().select_related("taxon", "algorithm", "applied_to__algorithm")
     serializer_class = ClassificationSerializer
     filterset_fields = [
         # Docs about slow loading API browser because of large choice fields

ami/main/models_future/occurrence.py

@@ -58,7 +58,10 @@ def _detections_prefetch(*, ordering: tuple[str, ...], with_source_image: bool)
     qs = Detection.objects.prefetch_related(
         Prefetch(
             "classifications",
-            queryset=Classification.objects.select_related("taxon", "algorithm"),
+            # applied_to__algorithm: post-processed classifications (class masking,
+            # rank rollup) serialize their provenance parent; pull it here so the
+            # nested applied_to render doesn't issue a query per classification.
+            queryset=Classification.objects.select_related("taxon", "algorithm", "applied_to__algorithm"),
         )
     ).order_by(*ordering)
     if with_source_image:

ami/ml/management/commands/run_class_masking.py

@@ -0,0 +1,83 @@
+from django.core.management.base import BaseCommand, CommandError
+
+from ami.main.models import SourceImageCollection, TaxaList
+from ami.ml.models.algorithm import Algorithm
+from ami.ml.post_processing.class_masking import ClassMaskingTask
+
+
+class Command(BaseCommand):
+    help = (
+        "Run class masking post-processing on a source image collection. "
+        "Masks classifier logits for species not in the given taxa list and recalculates softmax scores."
+    )
+
+    def add_arguments(self, parser):
+        parser.add_argument("--collection-id", type=int, required=True, help="SourceImageCollection ID to process")
+        parser.add_argument("--taxa-list-id", type=int, required=True, help="TaxaList ID to use as the species mask")
+        parser.add_argument(
+            "--algorithm-id", type=int, required=True, help="Algorithm ID whose classifications to mask"
+        )
+        parser.add_argument("--dry-run", action="store_true", help="Show what would be done without making changes")
+
+    def handle(self, *args, **options):
+        collection_id = options["collection_id"]
+        taxa_list_id = options["taxa_list_id"]
+        algorithm_id = options["algorithm_id"]
+        dry_run = options["dry_run"]
+
+        # Validate inputs
+        try:
+            collection = SourceImageCollection.objects.get(pk=collection_id)
+        except SourceImageCollection.DoesNotExist:
+            raise CommandError(f"SourceImageCollection {collection_id} does not exist.")
+
+        try:
+            taxa_list = TaxaList.objects.get(pk=taxa_list_id)
+        except TaxaList.DoesNotExist:
+            raise CommandError(f"TaxaList {taxa_list_id} does not exist.")
+
+        try:
+            algorithm = Algorithm.objects.get(pk=algorithm_id)
+        except Algorithm.DoesNotExist:
+            raise CommandError(f"Algorithm {algorithm_id} does not exist.")
+
+        if not algorithm.category_map:
+            raise CommandError(f"Algorithm '{algorithm.name}' does not have a category map.")
+
+        from ami.main.models import Classification
+
+        classification_count = (
+            Classification.objects.filter(
+                detection__source_image__collections=collection,
+                terminal=True,
+                algorithm=algorithm,
+                scores__isnull=False,
+            )
+            .distinct()
+            .count()
+        )
+
+        taxa_count = taxa_list.taxa.count()
+
+        self.stdout.write(
+            f"Collection:     {collection.name} (id={collection.pk})\n"
+            f"Taxa list:      {taxa_list.name} (id={taxa_list.pk}, {taxa_count} taxa)\n"
+            f"Algorithm:      {algorithm.name} (id={algorithm.pk})\n"
+            f"Classifications to process: {classification_count}"
+        )
+
+        if classification_count == 0:
+            raise CommandError("No terminal classifications with scores found for this collection/algorithm.")
+
+        if dry_run:
+            self.stdout.write(self.style.WARNING("Dry run — no changes made."))
+            return
+
+        self.stdout.write("Running class masking...")
+        task = ClassMaskingTask(
+            source_image_collection_id=collection_id,
+            taxa_list_id=taxa_list_id,
+            algorithm_id=algorithm_id,
+        )
+        task.run()
+        self.stdout.write(self.style.SUCCESS("Class masking completed."))

ami/ml/post_processing/__init__.py

@@ -1 +1,3 @@
+from . import class_masking  # noqa: F401
+from . import rank_rollup  # noqa: F401
 from . import small_size_filter  # noqa: F401

ami/ml/post_processing/admin/class_masking_form.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+from django import forms
+
+from ami.main.models import TaxaList
+from ami.ml.models import Algorithm
+from ami.ml.models.algorithm import AlgorithmTaskType
+from ami.ml.post_processing.admin.forms import BasePostProcessingActionForm
+
+
+class ClassMaskingActionForm(BasePostProcessingActionForm):
+    """Knobs surfaced when an admin triggers Class masking.
+
+    The operator picks the source classifier and the taxa list to keep; the
+    scope (which collection or occurrence) is supplied by the admin entry point,
+    not the form. Selections are model instances, so ``to_config`` hands the
+    schema their primary keys (``ClassMaskingConfig`` expects ``*_id`` ints).
+    """
+
+    algorithm_id = forms.ModelChoiceField(
+        queryset=Algorithm.objects.filter(task_type=AlgorithmTaskType.CLASSIFICATION.value).order_by("name"),
+        label="Source classifier",
+        help_text="The classification algorithm whose terminal predictions will be re-scored.",
+    )
+    taxa_list_id = forms.ModelChoiceField(
+        queryset=TaxaList.objects.all().order_by("name"),
+        label="Taxa list to keep",
+        help_text=(
+            "Classes whose taxon is not in this list are masked out; each "
+            "classification's softmax is renormalised over the classes that remain."
+        ),
+    )
+
+    def to_config(self) -> dict:
+        return {
+            "algorithm_id": self.cleaned_data["algorithm_id"].pk,
+            "taxa_list_id": self.cleaned_data["taxa_list_id"].pk,
+        }

ami/ml/post_processing/admin/rank_rollup_form.py

@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+from ami.ml.post_processing.admin.forms import BasePostProcessingActionForm
+
+
+class RankRollupActionForm(BasePostProcessingActionForm):
+    """Knob form for Rank rollup.
+
+    Rank rollup runs with the per-rank score thresholds and rollup order defined
+    on ``RankRollupConfig``. There are no per-run knobs yet, so the form only
+    confirms the selected capture set(s); the empty ``cleaned_data`` lets the
+    schema apply its defaults. Threshold overrides can be added here later.
+    """