Merge branch 'deployments/ood.antenna.insectai.org' of github.com:RolnickLab/antenna into feat/better-cluster-data

Author: mihow

ami/main/admin.py

@@ -27,6 +27,7 @@
     Site,
     SourceImage,
     SourceImageCollection,
+    Tag,
     TaxaList,
     Taxon,
 )
@@ -471,6 +472,7 @@ class TaxonAdmin(admin.ModelAdmin[Taxon]):
         "rank",
         "parent",
         "parent_names",
+        "tag_list",
         "list_names",
         "created_at",
         "updated_at",
@@ -491,10 +493,10 @@ def get_queryset(self, request):
 
         return qs.annotate(occurrence_count=models.Count("occurrences")).order_by("-occurrence_count")
 
-    @admin.display(
-        description="Occurrences",
-        ordering="occurrence_count",
-    )
+    @admin.display(description="Tags")
+    def tag_list(self, obj) -> str:
+        return ", ".join([tag.name for tag in obj.tags.all()])
+
     def occurrence_count(self, obj) -> int:
         return obj.occurrence_count
 
@@ -653,3 +655,10 @@ def cluster_detections(self, request: HttpRequest, queryset: QuerySet[SourceImag
 
     # Hide images many-to-many field from form. This would list all source images in the database.
     exclude = ("images",)
+
+
+@admin.register(Tag)
+class TagAdmin(admin.ModelAdmin):
+    list_display = ("id", "name", "project")
+    list_filter = ("project",)
+    search_fields = ("name",)

ami/main/api/serializers.py

@@ -9,7 +9,7 @@
 from ami.base.fields import DateStringField
 from ami.base.serializers import DefaultSerializer, MinimalNestedModelSerializer, get_current_user, reverse_with_params
 from ami.jobs.models import Job
-from ami.main.models import create_source_image_from_upload
+from ami.main.models import Tag, create_source_image_from_upload
 from ami.ml.models import Algorithm
 from ami.ml.serializers import AlgorithmSerializer
 from ami.users.models import User
@@ -500,12 +500,33 @@ class Meta:
         ]
 
 
+class TagSerializer(DefaultSerializer):
+    project = ProjectNestedSerializer(read_only=True)
+    project_id = serializers.PrimaryKeyRelatedField(queryset=Project.objects.all(), source="project", write_only=True)
+    taxa_ids = serializers.PrimaryKeyRelatedField(
+        queryset=Taxon.objects.all(), many=True, source="taxa", write_only=True, required=False
+    )
+    taxa = serializers.SerializerMethodField()
+
+    class Meta:
+        model = Tag
+        fields = ["id", "name", "project", "project_id", "taxa_ids", "taxa"]
+
+    def get_taxa(self, obj):
+        return [{"id": taxon.id, "name": taxon.name} for taxon in obj.taxa.all()]
+
+
 class TaxonListSerializer(DefaultSerializer):
     # latest_detection = DetectionNestedSerializer(read_only=True)
     occurrences = serializers.SerializerMethodField()
     parents = TaxonParentSerializer(many=True, read_only=True, source="parents_json")
     parent_id = serializers.PrimaryKeyRelatedField(queryset=Taxon.objects.all(), source="parent")
     cover_image_url = serializers.SerializerMethodField()
+    tags = serializers.SerializerMethodField()
+
+    def get_tags(self, obj):
+        tag_list = getattr(obj, "prefetched_tags", [])
+        return TagSerializer(tag_list, many=True, context=self.context).data
 
     class Meta:
         model = Taxon
@@ -518,6 +539,7 @@ class Meta:
             "details",
             "occurrences_count",
             "occurrences",
+            "tags",
             "last_detected",
             "best_determination_score",
             "cover_image_url",
@@ -737,6 +759,12 @@ class TaxonSerializer(DefaultSerializer):
     parent_id = serializers.PrimaryKeyRelatedField(queryset=Taxon.objects.all(), source="parent", write_only=True)
     parents = TaxonParentSerializer(many=True, read_only=True, source="parents_json")
     cover_image_url = serializers.SerializerMethodField()
+    tags = serializers.SerializerMethodField()
+
+    def get_tags(self, obj):
+        # Use prefetched tags
+        tag_list = getattr(obj, "prefetched_tags", [])
+        return TagSerializer(tag_list, many=True, context=self.context).data
 
     class Meta:
         model = Taxon
@@ -752,12 +780,12 @@ class Meta:
             "events_count",
             "occurrences",
             "gbif_taxon_key",
+            "tags",
             "last_detected",
             "fieldguide_id",
             "cover_image_url",
             "cover_image_credit",
             "unknown_species",
-            "last_detected",  # @TODO this has performance impact, review
         ]
 
     def get_cover_image_url(self, obj):

ami/main/api/views.py

@@ -43,7 +43,7 @@
 from ami.base.serializers import FilterParamsSerializer, SingleParamSerializer
 from ami.base.views import ProjectMixin
 from ami.jobs.models import DetectionClusteringJob, Job
-from ami.main.api.serializers import ClusterDetectionsSerializer
+from ami.main.api.serializers import ClusterDetectionsSerializer, TagSerializer
 from ami.utils.requests import get_active_classification_threshold, project_id_doc_param
 from ami.utils.storages import ConnectionTestResult
 
@@ -63,6 +63,7 @@
     SourceImage,
     SourceImageCollection,
     SourceImageUpload,
+    Tag,
     TaxaList,
     Taxon,
     User,
@@ -1187,6 +1188,29 @@ def filter_queryset(self, request, queryset, view):
         return queryset
 
 
+class TaxonTagFilter(filters.BaseFilterBackend):
+    """FilterBackend that allows OR-based filtering of taxa by tag ID."""
+
+    def filter_queryset(self, request, queryset, view):
+        tag_ids = request.query_params.getlist("tag_id")
+        if tag_ids:
+            queryset = queryset.filter(tags__id__in=tag_ids).distinct()
+        return queryset
+
+
+class TagInverseFilter(filters.BaseFilterBackend):
+    """
+    Exclude taxa that have any of the specified tag IDs using `not_tag_id`.
+    Example: /api/v2/taxa/?not_tag_id=1&not_tag_id=2
+    """
+
+    def filter_queryset(self, request, queryset, view):
+        not_tag_ids = request.query_params.getlist("not_tag_id")
+        if not_tag_ids:
+            queryset = queryset.exclude(tags__id__in=not_tag_ids)
+        return queryset.distinct()
+
+
 class TaxonViewSet(DefaultViewSet, ProjectMixin):
     """
     API endpoint that allows taxa to be viewed or edited.
@@ -1198,6 +1222,8 @@ class TaxonViewSet(DefaultViewSet, ProjectMixin):
         CustomTaxonFilter,
         TaxonCollectionFilter,
         TaxonTaxaListFilter,
+        TaxonTagFilter,
+        TagInverseFilter,
     ]
     filterset_fields = [
         "name",
@@ -1323,9 +1349,11 @@ def get_queryset(self) -> QuerySet:
         """
         qs = super().get_queryset()
         project = self.get_active_project()
+        qs = self.attach_tags_by_project(qs, project)
 
         if project:
             include_unobserved = True  # Show detail views for unobserved taxa instead of 404
+            # @TODO move to a QuerySet manager
             qs = qs.annotate(
                 best_detection_image_path=models.Subquery(
                     Occurrence.objects.filter(
@@ -1416,6 +1444,43 @@ def get_taxa_observed(self, qs: QuerySet, project: Project, include_unobserved=F
             )
         return qs
 
+    def attach_tags_by_project(self, qs: QuerySet, project: Project) -> QuerySet:
+        """
+        Prefetch and override the `.tags` attribute on each Taxon
+        with only the tags belonging to the given project.
+        """
+        # Include all tags if no project is passed
+        if project is None:
+            tag_qs = Tag.objects.all()
+        else:
+            # Prefetch only the tags that belong to the project or are global
+            tag_qs = Tag.objects.filter(models.Q(project=project) | models.Q(project__isnull=True))
+
+        tag_prefetch = Prefetch("tags", queryset=tag_qs, to_attr="prefetched_tags")
+
+        return qs.prefetch_related(tag_prefetch)
+
+    @action(detail=True, methods=["post"])
+    def assign_tags(self, request, pk=None):
+        """
+        Assign tags to a taxon
+        """
+        taxon = self.get_object()
+        tag_ids = request.data.get("tag_ids")
+        logger.info(f"Tag IDs: {tag_ids}")
+        if not isinstance(tag_ids, list):
+            return Response({"detail": "tag_ids must be a list of IDs."}, status=status.HTTP_400_BAD_REQUEST)
+
+        tags = Tag.objects.filter(id__in=tag_ids)
+        logger.info(f"Tags: {tags}, len: {len(tags)}")
+        taxon.tags.set(tags)  # replaces all tags for this taxon
+        taxon.save()
+        logger.info(f"Tags after assingment : {len(taxon.tags.all())}")
+        return Response(
+            {"taxon_id": taxon.id, "assigned_tag_ids": [tag.pk for tag in tags]},
+            status=status.HTTP_200_OK,
+        )
+
     @extend_schema(parameters=[project_id_doc_param])
     def list(self, request, *args, **kwargs):
         return super().list(request, *args, **kwargs)
@@ -1434,6 +1499,20 @@ def get_queryset(self):
     serializer_class = TaxaListSerializer
 
 
+class TagViewSet(DefaultViewSet, ProjectMixin):
+    queryset = Tag.objects.all()
+    serializer_class = TagSerializer
+    filterset_fields = ["taxa"]
+
+    def get_queryset(self):
+        qs = super().get_queryset()
+        project = self.get_active_project()
+        if project:
+            # Filter by project, but also include global tags
+            return qs.filter(models.Q(project=project) | models.Q(project__isnull=True))
+        return qs
+
+
 class ClassificationViewSet(DefaultViewSet, ProjectMixin):
     """
     API endpoint for viewing and adding classification results from a model.

ami/main/migrations/0067_tag_taxon_tags.py

@@ -0,0 +1,62 @@
+# Generated by Django 4.2.10 on 2025-05-15 21:23
+
+from django.db import migrations, models
+import django.db.models.deletion
+
+
+def add_inital_tags(apps, schema_editor):
+    """
+    Add initial tags to the database.
+    """
+    Tag = apps.get_model("main", "Tag")
+
+    # Make tags available for all projects
+    project = None
+
+    # Create initial tags
+    tags = [
+        "most wanted",
+        "also wanted",
+        "reviewed",
+        "collected",
+    ]
+
+    for tag in tags:
+        Tag.objects.get_or_create(name=tag, project=project)
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("main", "0066_populate_cached_occurence_fields"),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name="Tag",
+            fields=[
+                ("id", models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name="ID")),
+                ("created_at", models.DateTimeField(auto_now_add=True)),
+                ("updated_at", models.DateTimeField(auto_now=True)),
+                ("name", models.CharField(max_length=255)),
+                (
+                    "project",
+                    models.ForeignKey(
+                        blank=True,
+                        null=True,
+                        on_delete=django.db.models.deletion.CASCADE,
+                        related_name="tags",
+                        to="main.project",
+                    ),
+                ),
+            ],
+            options={
+                "unique_together": {("name", "project")},
+            },
+        ),
+        migrations.AddField(
+            model_name="taxon",
+            name="tags",
+            field=models.ManyToManyField(blank=True, related_name="taxa", to="main.tag"),
+        ),
+        migrations.RunPython(add_inital_tags, migrations.RunPython.noop),
+    ]

ami/main/models.py

@@ -142,6 +142,7 @@ class Project(BaseModel):
     devices: models.QuerySet["Device"]
     sites: models.QuerySet["Site"]
     jobs: models.QuerySet["Job"]
+    tags: models.QuerySet["Tag"]
 
     objects = ProjectManager()
 
@@ -2906,6 +2907,7 @@ class Taxon(BaseModel):
     authorship_date = models.DateField(null=True, blank=True, help_text="The date the taxon was described.")
     ordering = models.IntegerField(null=True, blank=True)
     sort_phylogeny = models.BigIntegerField(blank=True, null=True)
+    tags = models.ManyToManyField("Tag", related_name="taxa", blank=True)
     unknown_species = models.BooleanField(default=False, help_text="Is this a clustering-generated taxon")
     objects: TaxonManager = TaxonManager()
 
@@ -3166,6 +3168,19 @@ class Meta:
         verbose_name_plural = "Taxa Lists"
 
 
+@final
+class Tag(BaseModel):
+    """A tag for taxa"""
+
+    name = models.CharField(max_length=255)
+    project = models.ForeignKey(Project, on_delete=models.CASCADE, related_name="tags", null=True, blank=True)
+
+    taxa: models.QuerySet[Taxon]
+
+    class Meta:
+        unique_together = ("name", "project")
+
+
 @final
 class BlogPost(BaseModel):
     """

config/api_router.py

@@ -30,6 +30,7 @@
 router.register(r"occurrences", views.OccurrenceViewSet)
 router.register(r"taxa/lists", views.TaxaListViewSet)
 router.register(r"taxa", views.TaxonViewSet)
+router.register(r"tags", views.TagViewSet)
 router.register(r"ml/algorithms", ml_views.AlgorithmViewSet)
 router.register(r"ml/labels", ml_views.AlgorithmCategoryMapViewSet)
 router.register(r"ml/pipelines", ml_views.PipelineViewSet)

ui/src/components/filtering/filter-control.tsx

@@ -11,6 +11,7 @@ import { ScoreFilter } from './filters/score-filter'
 import { SessionFilter } from './filters/session-filter'
 import { StationFilter } from './filters/station-filter'
 import { StatusFilter } from './filters/status-filter'
+import { TagFilter } from './filters/tag-filter'
 import { TaxaListFilter } from './filters/taxa-list-filter'
 import { TaxonFilter } from './filters/taxon-filter'
 import { TypeFilter } from './filters/type-filter'
@@ -36,6 +37,8 @@ const ComponentMap: {
   source_image_collection: CollectionFilter,
   source_image_single: ImageFilter,
   status: StatusFilter,
+  tag_id: TagFilter,
+  not_tag_id: TagFilter,
   taxon: TaxonFilter,
   taxa_list_id: TaxaListFilter,
   verified_by_me: VerifiedByFilter,

ui/src/components/filtering/filters/tag-filter.tsx

@@ -0,0 +1,21 @@
+import { Select } from 'nova-ui-kit'
+import { FilterProps } from './types'
+
+export const TagFilter = ({ data = [], value, onAdd }: FilterProps) => {
+  const tags = data as { id: number; name: string }[]
+
+  return (
+    <Select.Root value={value ?? ''} onValueChange={onAdd}>
+      <Select.Trigger>
+        <Select.Value placeholder="Select a value" />
+      </Select.Trigger>
+      <Select.Content className="max-h-72">
+        {tags.map((option) => (
+          <Select.Item key={option.id} value={`${option.id}`}>
+            {option.name}
+          </Select.Item>
+        ))}
+      </Select.Content>
+    </Select.Root>
+  )
+}