diff --git a/ami/ml/post_processing/class_masking.py b/ami/ml/post_processing/class_masking.py index ffc77f346..f7c18d224 100644 --- a/ami/ml/post_processing/class_masking.py +++ b/ami/ml/post_processing/class_masking.py @@ -263,19 +263,31 @@ def _get_or_create_masking_algorithm( return algorithm def _scoped_classifications( - self, config: ClassMaskingConfig, source_algorithm: Algorithm + self, config: ClassMaskingConfig, source_algorithm: Algorithm, masking_algorithm: Algorithm ) -> tuple[QuerySet[Classification], str]: """Resolve the terminal classifications to re-score from the config's scope. ``config_schema`` guarantees exactly one scope id is set, so the single ``else`` branch is sound. + + Sources already re-scored by ``masking_algorithm`` are excluded via the + ``applied_to`` lineage so the mask is idempotent: a source is masked at most + once per masking algorithm. This keeps re-runs safe — finishing a partially + completed run (e.g. one the health-check reaper revoked) processes only the + remainder, and a source that became terminal again (after a dedup or + re-classification pass) is not masked a second time. The guard is on the + lineage rather than the terminal flag, which is why it survives that churn. """ - base = Classification.objects.filter( - terminal=True, - algorithm=source_algorithm, - scores__isnull=False, - logits__isnull=False, - ).select_related("detection", "detection__occurrence") + base = ( + Classification.objects.filter( + terminal=True, + algorithm=source_algorithm, + scores__isnull=False, + logits__isnull=False, + ) + .exclude(derived_classifications__algorithm=masking_algorithm) + .select_related("detection", "detection__occurrence") + ) if config.occurrence_id is not None: if not Occurrence.objects.filter(pk=config.occurrence_id).exists(): @@ -312,7 +324,7 @@ def run(self) -> None: masking_algorithm = self._get_or_create_masking_algorithm( source_algorithm, taxa_list, reweight=config.reweight ) - classifications, scope_desc = self._scoped_classifications(config, source_algorithm) + classifications, scope_desc = self._scoped_classifications(config, source_algorithm, masking_algorithm) self.logger.info(f"Applying class masking on {scope_desc} using taxa list {taxa_list.pk}") def _on_batch(m: dict) -> None: diff --git a/ami/ml/post_processing/tests/test_class_masking.py b/ami/ml/post_processing/tests/test_class_masking.py index b0578272c..458018ca7 100644 --- a/ami/ml/post_processing/tests/test_class_masking.py +++ b/ami/ml/post_processing/tests/test_class_masking.py @@ -253,6 +253,53 @@ def test_task_run_collection_scope_persists_masking_algorithm(self): occ.refresh_from_db() self.assertEqual(occ.determination, self.species_taxa[1], "Occurrence determination follows the masked result") + def test_rerun_does_not_duplicate_masked_classifications(self): + """Re-running the same mask must not create a second masked classification for + a source already re-scored, even if that source became terminal again in between. + + Idempotency is keyed on the ``applied_to`` lineage — a source is masked at most + once per masking algorithm — not on the terminal flag. This makes it safe to + finish a partially completed run (e.g. one the health-check reaper revoked) or + to re-run after a re-classification / dedup pass re-terminalized a source. + """ + logits = [0.5, 3.0, 3.5] # excluded index 2 is top; index 1 is the in-list winner + taxa_list = TaxaList.objects.create(name="Idempotency list") + taxa_list.taxa.set(self.species_taxa[:2]) + + det, _ = self._detection_with_occurrence() + original = self._create_classification_with_logits(det, self.species_taxa[2], _softmax(logits), logits) + + def run(): + ClassMaskingTask( + source_image_collection_id=self.collection.pk, + taxa_list_id=taxa_list.pk, + algorithm_id=self.algorithm.pk, + ).run() + + run() + masking_algo = Algorithm.objects.get( + key=f"{self.algorithm.key}_filtered_by_taxa_list_{taxa_list.pk}_reweighted" + ) + self.assertEqual( + Classification.objects.filter(algorithm=masking_algo, applied_to=original).count(), + 1, + "First run masks the source exactly once", + ) + + # Simulate the source becoming terminal again (a dedup or re-classification pass) + # while its masked child still exists. The terminal filter alone would re-select + # it; the applied_to guard must still skip it. + original.refresh_from_db() + original.terminal = True + original.save(update_fields=["terminal"]) + + run() + self.assertEqual( + Classification.objects.filter(algorithm=masking_algo, applied_to=original).count(), + 1, + "Re-run must not create a duplicate masked classification for an already-masked source", + ) + def test_reweight_modes_get_distinct_masking_algorithms(self): """The reweight mode is part of the masking algorithm's identity. diff --git a/ami/ml/tests.py b/ami/ml/tests.py index cb88ee6fd..dd3cc1b68 100644 --- a/ami/ml/tests.py +++ b/ami/ml/tests.py @@ -2037,8 +2037,10 @@ def test_deployment_counts_refresh_after_save_results(self): class TestAlgorithmViewSetProjectFilter(APITestCase): """ - The algorithm list endpoint is scoped to algorithms belonging to - pipelines enabled for the active project. + The algorithm list endpoint is scoped to algorithms relevant to the active + project: those belonging to an enabled pipeline, plus any that produced + classifications in the project (e.g. post-processing algorithms like class + masking, which are created standalone with no pipeline). """ def setUp(self): @@ -2103,3 +2105,33 @@ def test_detail_endpoint_unscoped_even_with_project_id(self): response = self.client.get(url) self.assertEqual(response.status_code, 200) self.assertEqual(response.json()["name"], "Algo Disabled") + + def _classify_in_project(self, algorithm, project): + """Give ``algorithm`` a terminal classification whose capture is in ``project``.""" + source_image = SourceImage.objects.create(project=project) + detection = Detection.objects.create(source_image=source_image) + return Classification.objects.create( + detection=detection, + algorithm=algorithm, + timestamp=datetime.datetime.now(datetime.timezone.utc), + ) + + def test_lists_post_processing_algorithm_with_classifications_in_project(self): + """A post-processing algorithm has no pipeline but produces determinations in + the project, so the list must include it — otherwise the user cannot filter + occurrences by the masked result.""" + masked_algo = Algorithm.objects.create(name="Class Masked Classifier", version=1) + self._classify_in_project(masked_algo, self.project) + + names = self._list_algorithm_names(project_id=self.project.pk) + self.assertIn("Class Masked Classifier", names) + self.assertIn("Algo Enabled", names, "Enabled-pipeline algorithms still appear") + self.assertNotIn("Algo Disabled", names, "A disabled pipeline with no classifications stays hidden") + + def test_classifications_in_other_project_do_not_leak(self): + """An algorithm whose classifications live in another project must not appear.""" + other_masked_algo = Algorithm.objects.create(name="Other Project Masked", version=1) + self._classify_in_project(other_masked_algo, self.other_project) + + names = self._list_algorithm_names(project_id=self.project.pk) + self.assertNotIn("Other Project Masked", names) diff --git a/ami/ml/views.py b/ami/ml/views.py index 7de502f4a..47a4c33b3 100644 --- a/ami/ml/views.py +++ b/ami/ml/views.py @@ -1,7 +1,7 @@ import logging from django.db import transaction -from django.db.models import Prefetch +from django.db.models import Prefetch, Q from django.db.models.query import QuerySet from django.utils.text import slugify from drf_spectacular.utils import extend_schema @@ -15,7 +15,7 @@ from ami.base.views import ProjectMixin from ami.main.api.schemas import project_id_doc_param from ami.main.api.views import DefaultViewSet -from ami.main.models import Project, SourceImage +from ami.main.models import Classification, Project, SourceImage from ami.ml.schemas import PipelineRegistrationResponse from .models.algorithm import Algorithm, AlgorithmCategoryMap @@ -56,14 +56,25 @@ class AlgorithmViewSet(DefaultViewSet, ProjectMixin): def get_queryset(self) -> QuerySet["Algorithm"]: qs: QuerySet["Algorithm"] = super().get_queryset() qs = qs.with_category_count() # type: ignore[union-attr] # Custom queryset method - # Only scope list by project. Detail stays unscoped so links from historical + # Only scope the list by project. Detail stays unscoped so links from historical # classifications whose pipeline is no longer enabled still resolve. if getattr(self, "action", None) == "list": project = self.get_active_project() if project: + # An algorithm is relevant to the project if it is configured via an + # enabled pipeline OR it produced classifications in the project. + # Post-processing algorithms (e.g. class masking) are created standalone + # with no pipeline, so the pipeline join alone would hide them even though + # they own determinations the user needs to filter occurrences by. + classified_in_project = Classification.objects.filter(detection__source_image__project=project).values( + "algorithm" + ) qs = qs.filter( - pipelines__project_pipeline_configs__project=project, - pipelines__project_pipeline_configs__enabled=True, + Q( + pipelines__project_pipeline_configs__project=project, + pipelines__project_pipeline_configs__enabled=True, + ) + | Q(pk__in=classified_in_project) ).distinct() return qs