Implement class masking using the post-processing framework

ami/jobs/migrations/0018_alter_job_job_type_key.py

@@ -0,0 +1,29 @@
+# Generated by Django 4.2.10 on 2025-09-30 12:25
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("jobs", "0017_alter_job_logs_alter_job_progress"),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name="job",
+            name="job_type_key",
+            field=models.CharField(
+                choices=[
+                    ("ml", "ML pipeline"),
+                    ("populate_captures_collection", "Populate captures collection"),
+                    ("data_storage_sync", "Data storage sync"),
+                    ("unknown", "Unknown"),
+                    ("data_export", "Data Export"),
+                    ("post_processing", "Post Processing"),
+                ],
+                default="unknown",
+                max_length=255,
+                verbose_name="Job Type",
+            ),
+        ),
+    ]

ami/jobs/models.py

@@ -645,6 +645,40 @@ def run(cls, job: "Job"):
         job.update_status(JobState.SUCCESS, save=True)
 
 
+class PostProcessingJob(JobType):
+    name = "Post Processing"
+    key = "post_processing"
+
+    @classmethod
+    def run(cls, job: "Job"):
+        import ami.ml.post_processing  # noqa F401
+        from ami.ml.post_processing.base import get_postprocessing_task
+
+        job.progress.add_stage(cls.name, key=cls.key)
+        job.update_status(JobState.STARTED)
+        job.started_at = datetime.datetime.now()
+        job.save()
+
+        params = job.params or {}
+        task_key: str = params.get("task", "")
+        config = params.get("config", {})
+        job.logger.info(f"Post-processing task: {task_key} with params: {config}")
+        # Get the registered task class
+        task_cls = get_postprocessing_task(task_key)
+        if not task_cls:
+            raise ValueError(f"Unknown post-processing task '{task_key}'")
+
+        # Instantiate the task with job context and config
+        task = task_cls(job=job, **config)
+
+        # Run the task
+        task.run()
+        job.progress.update_stage(cls.key, status=JobState.SUCCESS, progress=1)
+        job.finished_at = datetime.datetime.now()
+        job.update_status(JobState.SUCCESS)
+        job.save()
+
+
 class UnknownJobType(JobType):
     name = "Unknown"
     key = "unknown"
@@ -654,7 +688,14 @@ def run(cls, job: "Job"):
         raise ValueError(f"Unknown job type '{job.job_type()}'")
 
 
-VALID_JOB_TYPES = [MLJob, SourceImageCollectionPopulateJob, DataStorageSyncJob, UnknownJobType, DataExportJob]
+VALID_JOB_TYPES = [
+    MLJob,
+    SourceImageCollectionPopulateJob,
+    DataStorageSyncJob,
+    UnknownJobType,
+    DataExportJob,
+    PostProcessingJob,
+]
 
 
 def get_job_type_by_key(key: str) -> type[JobType] | None:

ami/main/admin.py

@@ -10,6 +10,7 @@
 
 import ami.utils
 from ami import tasks
+from ami.jobs.models import Job
 from ami.ml.models.project_pipeline_config import ProjectPipelineConfig
 from ami.ml.tasks import remove_duplicate_classifications
 
@@ -619,7 +620,54 @@ def populate_collection_async(self, request: HttpRequest, queryset: QuerySet[Sou
             f"Populating {len(queued_tasks)} collection(s) background tasks: {queued_tasks}.",
         )
 
-    actions = [populate_collection, populate_collection_async]
+    @admin.action(description="Run Small Size Filter post-processing task (async)")
+    def run_small_size_filter(self, request: HttpRequest, queryset: QuerySet[SourceImageCollection]) -> None:
+        jobs = []
+        for collection in queryset:
+            job = Job.objects.create(
+                name=f"Post-processing: SmallSizeFilter on Collection {collection.pk}",
+                project=collection.project,
+                job_type_key="post_processing",
+                params={
+                    "task": "small_size_filter",
+                    "config": {
+                        "size_threshold": 0.01,  # default threshold
+                        "source_image_collection_id": collection.pk,
+                    },
+                },
+            )
+            job.enqueue()
+            jobs.append(job.pk)
+
+        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.
     exclude = ("images",)

ami/main/migrations/0075_alter_project_options.py

@@ -0,0 +1,59 @@
+# Generated by Django 4.2.10 on 2025-10-14 05:01
+
+from django.db import migrations
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("main", "0074_taxon_cover_image_credit_taxon_cover_image_url_and_more"),
+    ]
+
+    operations = [
+        migrations.AlterModelOptions(
+            name="project",
+            options={
+                "ordering": ["-priority", "created_at"],
+                "permissions": [
+                    ("create_identification", "Can create identifications"),
+                    ("update_identification", "Can update identifications"),
+                    ("delete_identification", "Can delete identifications"),
+                    ("create_job", "Can create a job"),
+                    ("update_job", "Can update a job"),
+                    ("run_ml_job", "Can run/retry/cancel ML jobs"),
+                    ("run_populate_captures_collection_job", "Can run/retry/cancel Populate Collection jobs"),
+                    ("run_data_storage_sync_job", "Can run/retry/cancel Data Storage Sync jobs"),
+                    ("run_data_export_job", "Can run/retry/cancel Data Export jobs"),
+                    ("run_single_image_ml_job", "Can process a single capture"),
+                    ("run_post_processing_job", "Can run/retry/cancel Post-Processing jobs"),
+                    ("delete_job", "Can delete a job"),
+                    ("create_deployment", "Can create a deployment"),
+                    ("delete_deployment", "Can delete a deployment"),
+                    ("update_deployment", "Can update a deployment"),
+                    ("sync_deployment", "Can sync images to a deployment"),
+                    ("create_sourceimagecollection", "Can create a collection"),
+                    ("update_sourceimagecollection", "Can update a collection"),
+                    ("delete_sourceimagecollection", "Can delete a collection"),
+                    ("populate_sourceimagecollection", "Can populate a collection"),
+                    ("create_sourceimage", "Can create a source image"),
+                    ("update_sourceimage", "Can update a source image"),
+                    ("delete_sourceimage", "Can delete a source image"),
+                    ("star_sourceimage", "Can star a source image"),
+                    ("create_sourceimageupload", "Can create a source image upload"),
+                    ("update_sourceimageupload", "Can update a source image upload"),
+                    ("delete_sourceimageupload", "Can delete a source image upload"),
+                    ("create_s3storagesource", "Can create storage"),
+                    ("delete_s3storagesource", "Can delete storage"),
+                    ("update_s3storagesource", "Can update storage"),
+                    ("test_s3storagesource", "Can test storage connection"),
+                    ("create_site", "Can create a site"),
+                    ("delete_site", "Can delete a site"),
+                    ("update_site", "Can update a site"),
+                    ("create_device", "Can create a device"),
+                    ("delete_device", "Can delete a device"),
+                    ("update_device", "Can update a device"),
+                    ("view_private_data", "Can view private data"),
+                    ("trigger_exports", "Can trigger data exports"),
+                ],
+            },
+        ),
+    ]

ami/main/migrations/0077_merge_20251014_1426.py

@@ -0,0 +1,12 @@
+# Generated by Django 4.2.10 on 2025-10-14 14:26
+
+from django.db import migrations
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("main", "0075_alter_project_options"),
+        ("main", "0076_add_occurrence_composite_indexes"),
+    ]
+
+    operations = []

ami/main/models.py

@@ -332,6 +332,7 @@ class Permissions:
         RUN_POPULATE_CAPTURES_COLLECTION_JOB = "run_populate_captures_collection_job"
         RUN_DATA_STORAGE_SYNC_JOB = "run_data_storage_sync_job"
         RUN_DATA_EXPORT_JOB = "run_data_export_job"
+        RUN_POST_PROCESSING_JOB = "run_post_processing_job"
         DELETE_JOB = "delete_job"
 
         # Deployment permissions
@@ -394,6 +395,7 @@ class Meta:
             ("run_data_storage_sync_job", "Can run/retry/cancel Data Storage Sync jobs"),
             ("run_data_export_job", "Can run/retry/cancel Data Export jobs"),
             ("run_single_image_ml_job", "Can process a single capture"),
+            ("run_post_processing_job", "Can run/retry/cancel Post-Processing jobs"),
             ("delete_job", "Can delete a job"),
             # Deployment permissions
             ("create_deployment", "Can create a deployment"),

ami/ml/migrations/0025_alter_algorithm_task_type.py

@@ -0,0 +1,42 @@
+# Generated by Django 4.2.10 on 2025-10-14 05:01
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("ml", "0024_fix_classifications_missing_category_maps"),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name="algorithm",
+            name="task_type",
+            field=models.CharField(
+                choices=[
+                    ("detection", "Detection"),
+                    ("localization", "Localization"),
+                    ("segmentation", "Segmentation"),
+                    ("classification", "Classification"),
+                    ("embedding", "Embedding"),
+                    ("tracking", "Tracking"),
+                    ("tagging", "Tagging"),
+                    ("regression", "Regression"),
+                    ("captioning", "Captioning"),
+                    ("generation", "Generation"),
+                    ("translation", "Translation"),
+                    ("summarization", "Summarization"),
+                    ("question_answering", "Question Answering"),
+                    ("depth_estimation", "Depth Estimation"),
+                    ("pose_estimation", "Pose Estimation"),
+                    ("size_estimation", "Size Estimation"),
+                    ("post_processing", "Post Processing"),
+                    ("other", "Other"),
+                    ("unknown", "Unknown"),
+                ],
+                default="unknown",
+                max_length=255,
+                null=True,
+            ),
+        ),
+    ]

ami/ml/models/algorithm.py

@@ -176,6 +176,7 @@ class AlgorithmTaskType(str, enum.Enum):
     DEPTH_ESTIMATION = "depth_estimation"
     POSE_ESTIMATION = "pose_estimation"
     SIZE_ESTIMATION = "size_estimation"
+    POST_PROCESSING = "post_processing"
     OTHER = "other"
     UNKNOWN = "unknown"
 

ami/ml/post_processing/__init__.py

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

ami/ml/post_processing/base.py

@@ -0,0 +1,92 @@
+# ami/ml/post_processing/base.py
+
+import abc
+import logging
+from typing import Any
+
+from ami.jobs.models import Job
+from ami.ml.models import Algorithm
+from ami.ml.models.algorithm import AlgorithmTaskType
+
+# Registry of available post-processing tasks
+POSTPROCESSING_TASKS: dict[str, type["BasePostProcessingTask"]] = {}
+
+
+def register_postprocessing_task(task_cls: type["BasePostProcessingTask"]):
+    """
+    Decorator to register a post-processing task in the global registry.
+    Each task must define a unique `key`.
+    Ensures an Algorithm entry exists for this task.
+    """
+    if not hasattr(task_cls, "key") or not task_cls.key:
+        raise ValueError(f"Task {task_cls.__name__} missing required 'key' attribute")
+
+    # Register the task
+    POSTPROCESSING_TASKS[task_cls.key] = task_cls
+    return task_cls
+
+
+def get_postprocessing_task(name: str) -> type["BasePostProcessingTask"] | None:
+    """
+    Get a task class by its registry key.
+    Returns None if not found.
+    """
+    return POSTPROCESSING_TASKS.get(name)
+
+
+class BasePostProcessingTask(abc.ABC):
+    """
+    Abstract base class for all post-processing tasks.
+    """
+
+    # Each task must override these
+    key: str = ""
+    name: str = ""
+
+    def __init__(
+        self,
+        job: Job | None = None,
+        logger: logging.Logger | None = None,
+        **config: Any,
+    ):
+        self.job = job
+        self.config = config
+        # Choose the right logger
+        if logger is not None:
+            self.logger = logger
+        elif job is not None:
+            self.logger = job.logger
+        else:
+            self.logger = logging.getLogger(f"ami.post_processing.{self.key}")
+
+        algorithm, _ = Algorithm.objects.get_or_create(
+            name=self.__class__.__name__,
+            defaults={
+                "description": f"Post-processing task: {self.key}",
+                "task_type": AlgorithmTaskType.POST_PROCESSING.value,
+            },
+        )
+        self.algorithm: Algorithm = algorithm
+
+        self.logger.info(f"Initialized {self.__class__.__name__} with config={self.config}, job={job}")
+
+    def update_progress(self, progress: float):
+        """
+        Update progress if job is present, otherwise just log.
+        """
+
+        if self.job:
+            self.job.progress.update_stage(self.job.job_type_key, progress=progress)
+            self.job.save(update_fields=["progress"])
+
+        else:
+            # No job object — fallback to plain logging
+            self.logger.info(f"[{self.name}] Progress {progress:.0%}")
+
+    @abc.abstractmethod
+    def run(self) -> None:
+        """
+        Run the task logic.
+        Must be implemented by subclasses.
+        """
+        raise NotImplementedError("BasePostProcessingTask subclasses must implement run()")