Implement class masking using the post-processing framework

ami/main/admin.py

@@ -5,13 +5,17 @@
 from django.db.models.query import QuerySet
 from django.http.request import HttpRequest
 from django.template.defaultfilters import filesizeformat
+from django.urls import reverse
 from django.utils.formats import number_format
+from django.utils.html import format_html
 from guardian.admin import GuardedModelAdmin
 
 import ami.utils
 from ami import tasks
 from ami.jobs.models import Job
+from ami.ml.models.algorithm import Algorithm
 from ami.ml.models.project_pipeline_config import ProjectPipelineConfig
+from ami.ml.post_processing.class_masking import update_single_occurrence
 from ami.ml.tasks import remove_duplicate_classifications
 
 from .models import (
@@ -288,20 +292,29 @@ class ClassificationInline(admin.TabularInline):
     model = Classification
     extra = 0
     fields = (
+        "classification_link",
         "taxon",
         "algorithm",
         "timestamp",
         "terminal",
         "created_at",
     )
     readonly_fields = (
+        "classification_link",
         "taxon",
         "algorithm",
         "timestamp",
         "terminal",
         "created_at",
     )
 
+    @admin.display(description="Classification")
+    def classification_link(self, obj: Classification) -> str:
+        if obj.pk:
+            url = reverse("admin:main_classification_change", args=[obj.pk])
+            return format_html('<a href="{}">{}</a>', url, f"Classification #{obj.pk}")
+        return "-"
+
     def get_queryset(self, request: HttpRequest) -> QuerySet[Any]:
         qs = super().get_queryset(request)
         return qs.select_related("taxon", "algorithm", "detection")
@@ -311,20 +324,29 @@ class DetectionInline(admin.TabularInline):
     model = Detection
     extra = 0
     fields = (
+        "detection_link",
         "detection_algorithm",
         "source_image",
         "timestamp",
         "created_at",
         "occurrence",
     )
     readonly_fields = (
+        "detection_link",
         "detection_algorithm",
         "source_image",
         "timestamp",
         "created_at",
         "occurrence",
     )
 
+    @admin.display(description="ID")
+    def detection_link(self, obj):
+        if obj.pk:
+            url = reverse("admin:main_detection_change", args=[obj.pk])
+            return format_html('<a href="{}">{}</a>', url, obj.pk)
+        return "-"
+
 
 @admin.register(Detection)
 class DetectionAdmin(admin.ModelAdmin[Detection]):
@@ -382,7 +404,7 @@ class OccurrenceAdmin(admin.ModelAdmin[Occurrence]):
         "determination__rank",
         "created_at",
     )
-    search_fields = ("determination__name", "determination__search_names")
+    search_fields = ("id", "determination__name", "determination__search_names")
 
     def get_queryset(self, request: HttpRequest) -> QuerySet[Any]:
         qs = super().get_queryset(request)
@@ -404,11 +426,60 @@ def get_queryset(self, request: HttpRequest) -> QuerySet[Any]:
     def detections_count(self, obj) -> int:
         return obj.detections_count
 
+    @admin.action(description="Update occurrence with Newfoundland species taxa list")
+    def update_with_newfoundland_species(self, request: HttpRequest, queryset: QuerySet[Occurrence]) -> None:
+        """
+        Update selected occurrences using the 'Newfoundland species' taxa list
+        and 'Quebec & Vermont Species Classifier - Apr 2024' algorithm.
+        """
+        try:
+            # Get the taxa list by name
+            taxa_list = TaxaList.objects.get(name="Newfoundland Species")
+        except TaxaList.DoesNotExist:
+            self.message_user(
+                request,
+                "Error: TaxaList 'Newfoundland species' not found.",
+                level="error",
+            )
+            return
+
+        try:
+            # Get the algorithm by name
+            algorithm = Algorithm.objects.get(name="Quebec & Vermont Species Classifier - Apr 2024")
+        except Algorithm.DoesNotExist:
+            self.message_user(
+                request,
+                "Error: Algorithm 'Quebec & Vermont Species Classifier - Apr 2024' not found.",
+                level="error",
+            )
+            return
+
+        # Process each occurrence
+        count = 0
+        for occurrence in queryset:
+            try:
+                update_single_occurrence(
+                    occurrence=occurrence,
+                    algorithm=algorithm,
+                    taxa_list=taxa_list,
+                )
+                count += 1
+            except Exception as e:
+                self.message_user(
+                    request,
+                    f"Error processing occurrence {occurrence.pk}: {str(e)}",
+                    level="error",
+                )
+
+        self.message_user(request, f"Successfully updated {count} occurrence(s).")
+
     ordering = ("-created_at",)
 
     # Add classifications as inline
     inlines = [DetectionInline]
 
+    actions = [update_with_newfoundland_species]
+
 
 @admin.register(Classification)
 class ClassificationAdmin(admin.ModelAdmin[Classification]):
@@ -432,6 +503,8 @@ class ClassificationAdmin(admin.ModelAdmin[Classification]):
         "taxon__rank",
     )
 
+    autocomplete_fields = ("taxon",)
+
     def get_queryset(self, request: HttpRequest) -> QuerySet[Any]:
         qs = super().get_queryset(request)
         return qs.select_related(
@@ -662,10 +735,32 @@ def run_small_size_filter(self, request: HttpRequest, queryset: QuerySet[SourceI
 
         self.message_user(request, f"Queued Small Size Filter for {queryset.count()} collection(s). Jobs: {jobs}")
 
+    @admin.action(description="Run Rank Rollup post-processing task (async)")
+    def run_rank_rollup(self, request: HttpRequest, queryset: QuerySet[SourceImageCollection]) -> None:
+        """Trigger the Rank Rollup post-processing job asynchronously."""
+        jobs = []
+        DEFAULT_THRESHOLDS = {"species": 0.8, "genus": 0.6, "family": 0.4}
+
+        for collection in queryset:
+            job = Job.objects.create(
+                name=f"Post-processing: RankRollup on Collection {collection.pk}",
+                project=collection.project,
+                job_type_key="post_processing",
+                params={
+                    "task": "rank_rollup",
+                    "config": {"source_image_collection_id": collection.pk, "thresholds": DEFAULT_THRESHOLDS},
+                },
+            )
+            job.enqueue()
+            jobs.append(job.pk)
+
+        self.message_user(request, f"Queued Rank Rollup for {queryset.count()} collection(s). Jobs: {jobs}")
+
     actions = [
         populate_collection,
         populate_collection_async,
         run_small_size_filter,
+        run_rank_rollup,
     ]
 
     # Hide images many-to-many field from form. This would list all source images in the database.

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(
+            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 +0,0 @@
-from . import small_size_filter  # noqa: F401