diff --git a/ami/base/filters.py b/ami/base/filters.py index a8aa07e59..f3f1e213a 100644 --- a/ami/base/filters.py +++ b/ami/base/filters.py @@ -21,6 +21,10 @@ class RelatedIdFilter(NumberFilter): class NullsLastOrderingFilter(OrderingFilter): def get_ordering(self, request, queryset, view): + # A view lists orderings it applies itself (they are not in ordering_fields). Without this + # the filter drops them and applies the view's default ordering over the view's own. + if request.query_params.get(self.ordering_param) in getattr(view, "self_applied_orderings", ()): + return None values = super().get_ordering(request, queryset, view) if not values: return values diff --git a/ami/jobs/models.py b/ami/jobs/models.py index ff65f31f2..ea0f03e71 100644 --- a/ami/jobs/models.py +++ b/ami/jobs/models.py @@ -500,6 +500,7 @@ def run(cls, job: "Job"): if not job.pipeline: raise ValueError("No pipeline specified to process images in ML job") + job.pipeline.raise_if_embedding_only() job.progress.update_stage( "collect", diff --git a/ami/main/admin.py b/ami/main/admin.py index 55c84701f..a96a3cc56 100644 --- a/ami/main/admin.py +++ b/ami/main/admin.py @@ -15,8 +15,10 @@ 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.feature_vectors_form import AddFeatureVectorsActionForm 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.feature_vectors import AddFeatureVectorsTask from ami.ml.post_processing.small_size_filter import SmallSizeFilterTask from ami.ml.tasks import remove_duplicate_classifications @@ -859,11 +861,20 @@ def populate_collection_async(self, request: HttpRequest, queryset: QuerySet[Sou f"Post-processing: {task_cls.name} on Capture Set {collection.pk}" ), ) + run_add_feature_vectors = make_post_processing_action( + AddFeatureVectorsTask, + AddFeatureVectorsActionForm, + 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_add_feature_vectors, ] # Hide images many-to-many field from form. This would list all source images in the database. diff --git a/ami/main/api/serializers.py b/ami/main/api/serializers.py index 1ca588980..d691e97cc 100644 --- a/ami/main/api/serializers.py +++ b/ami/main/api/serializers.py @@ -1,7 +1,7 @@ import collections import datetime -from django.db.models import QuerySet +from django.db.models import Exists, OuterRef, QuerySet from drf_spectacular.utils import extend_schema_field from guardian.shortcuts import get_perms from rest_framework import serializers @@ -12,8 +12,8 @@ from ami.base.serializers import DefaultSerializer, MinimalNestedModelSerializer, reverse_with_params from ami.base.views import get_active_project from ami.jobs.models import Job -from ami.main.models import Tag -from ami.ml.models import Algorithm, Pipeline +from ami.main.models import DEFAULT_EMBEDDING_KEY, Tag +from ami.ml.models import Algorithm, DetectionEmbedding, Pipeline from ami.ml.serializers import AlgorithmSerializer, PipelineNestedSerializer from ami.users.models import User from ami.users.roles import ProjectManager @@ -1630,14 +1630,23 @@ class OccurrenceSerializer(OccurrenceListSerializer): predictions = ClassificationNestedSerializer(many=True, read_only=True) deployment = DeploymentNestedSerializer(read_only=True) event = EventNestedSerializer(read_only=True) + embedding_algorithms = serializers.SerializerMethodField() # first_appearance = TaxonSourceImageNestedSerializer(read_only=True) + def get_embedding_algorithms(self, obj: Occurrence) -> list[dict]: + """The algorithms with a feature vector on one of this occurrence's detections, one query.""" + has_vector = DetectionEmbedding.objects.filter( + algorithm_id=OuterRef("pk"), key=DEFAULT_EMBEDDING_KEY, detection__occurrence_id=obj.pk + ) + return list(Algorithm.objects.filter(Exists(has_vector)).order_by("pk").values("id", "name")) + class Meta: model = Occurrence fields = OccurrenceListSerializer.Meta.fields + [ "determination_id", "detections", "predictions", + "embedding_algorithms", ] read_only_fields = [ "determination_score", diff --git a/ami/main/api/views.py b/ami/main/api/views.py index 99709eb29..b7092cddb 100644 --- a/ami/main/api/views.py +++ b/ami/main/api/views.py @@ -36,6 +36,7 @@ from ami.main.api.serializers import TagSerializer from ami.main.models_future.identifications import create_identifications_batch, resolve_occurrences from ami.main.models_future.occurrence import model_agreement_for_project, top_identifiers_for_project +from ami.ml.embeddings.reader import algorithm_with_most_vectors, representative_embeddings from ami.ml.models.algorithm import Algorithm from ami.ml.serializers import AlgorithmSerializer from ami.utils.requests import get_default_classification_threshold @@ -1501,6 +1502,9 @@ class Meta: fields = list(OCCURRENCE_FILTERSET_FIELDS) +VISUAL_SIMILARITY_ORDERINGS = ("visual_similarity", "-visual_similarity") + + class OccurrenceViewSet(DefaultViewSet, ProjectMixin): """ API endpoint that allows occurrences to be viewed or edited. @@ -1512,6 +1516,7 @@ class OccurrenceViewSet(DefaultViewSet, ProjectMixin): serializer_class = OccurrenceSerializer filter_backends = DefaultViewSetMixin.filter_backends + list(OCCURRENCE_FILTER_BACKENDS) filterset_class = OccurrenceFilterSet + self_applied_orderings = VISUAL_SIMILARITY_ORDERINGS ordering_fields = [ "created_at", "updated_at", @@ -1554,11 +1559,94 @@ def get_queryset(self) -> QuerySet["Occurrence"]: else: qs = qs.with_detail_prefetches() # type: ignore + ordering = self.request.query_params.get("ordering") + if ordering in VISUAL_SIMILARITY_ORDERINGS: + qs = self._order_by_visual_similarity(qs, project, descending=ordering.startswith("-")) + return qs + def _order_by_visual_similarity(self, qs: QuerySet["Occurrence"], project: Project | None, descending: bool): + """Sort by cosine distance from a seed occurrence's feature vector, most similar first. + + One algorithm's vectors only (``similarity_algorithm``, or the one with the most vectors + in the project): distances between algorithms are meaningless. The seed is + ``similar_to``, or the most recently updated occurrence with such a vector that the + default filters show. Occurrences without a vector sort last, in both directions. + Not an ``ordering_fields`` entry, so the ordering filter leaves this ordering alone. + """ + if project is None: + raise api_exceptions.ValidationError({"ordering": "Sorting by visual similarity requires a project_id."}) + if not Project.objects.visible_for_user(self.request.user).filter(pk=project.pk).exists(): + return qs # Already empty for this user, and the seed must not say more than that. + params = self.request.query_params + algorithm_id = SingleParamSerializer[int].clean( + "similarity_algorithm", serializers.IntegerField(required=False, min_value=1), params + ) + if algorithm_id is None: + algorithm_id = algorithm_with_most_vectors(project) + if algorithm_id is None: + raise api_exceptions.ValidationError( + {"ordering": "No feature vectors have been stored for this project yet."} + ) + visible = Occurrence.objects.visible_for_user(self.request.user).valid().filter(project=project) + seed_id = SingleParamSerializer[int].clean( + "similar_to", serializers.IntegerField(required=False, min_value=1), params + ) + if seed_id is None: + seed_id = ( + visible.apply_default_filters(project, self.request) + .with_vectors(algorithm_id) + .order_by("-updated_at") + .values_list("pk", flat=True) + .first() + ) + if seed_id is None: + raise api_exceptions.ValidationError( + {"ordering": f"No occurrence in this project has a feature vector from algorithm #{algorithm_id}."} + ) + elif not visible.filter(pk=seed_id).exists(): + raise api_exceptions.ValidationError({"similar_to": f"Occurrence #{seed_id} is not in this project."}) + seed_vector = representative_embeddings(seed_id, algorithm_id).values_list("vector", flat=True).first() + if seed_vector is None: + raise api_exceptions.ValidationError( + {"similar_to": f"Occurrence #{seed_id} has no feature vector from algorithm #{algorithm_id}."} + ) + distance = models.F("visual_similarity") + return qs.with_visual_similarity(seed_vector, algorithm_id).order_by( # type: ignore[attr-defined] + distance.desc(nulls_last=True) if descending else distance.asc(nulls_last=True), "-pk" + ) + @extend_schema( parameters=[ project_id_doc_param, + OpenApiParameter( + name="ordering", + description=( + "Besides the usual fields, `visual_similarity` (or `-visual_similarity`) sorts by cosine " + "distance from a seed occurrence's feature vector, most similar first; occurrences without " + "a vector come last." + ), + required=False, + type=OpenApiTypes.STR, + ), + OpenApiParameter( + name="similar_to", + description=( + "With `ordering=visual_similarity`: the id of the occurrence to compare against. Defaults to " + "the most recently updated occurrence that has a feature vector." + ), + required=False, + type=OpenApiTypes.INT, + ), + OpenApiParameter( + name="similarity_algorithm", + description=( + "With `ordering=visual_similarity`: the id of the algorithm whose feature vectors to compare. " + "Defaults to the algorithm with the most vectors in the project." + ), + required=False, + type=OpenApiTypes.INT, + ), OpenApiParameter( name="classification_threshold", description="Filter occurrences by minimum determination score.", diff --git a/ami/main/models.py b/ami/main/models.py index a4c0c82f1..41323cb3d 100644 --- a/ami/main/models.py +++ b/ami/main/models.py @@ -11,6 +11,8 @@ from io import BytesIO from typing import Final, final # noqa: F401 +import pgvector +import pgvector.django import PIL.Image import pydantic from django.apps import apps @@ -50,7 +52,7 @@ if typing.TYPE_CHECKING: from ami.jobs.models import Job - from ami.ml.models import Pipeline, ProcessingService + from ami.ml.models import DetectionEmbedding, Pipeline, ProcessingService logger = logging.getLogger(__name__) @@ -3220,6 +3222,7 @@ class Detection(BaseModel): # For type hints classifications: models.QuerySet["Classification"] + embeddings: models.QuerySet["DetectionEmbedding"] source_image_id: int detection_algorithm_id: int @@ -3367,6 +3370,9 @@ def __str__(self) -> str: return f"#{self.pk} from SourceImage #{self.source_image_id} with Algorithm #{self.detection_algorithm_id}" +DEFAULT_EMBEDDING_KEY = "embedding" + + class OccurrenceQuerySet(BaseQuerySet): def valid(self): """ @@ -3387,6 +3393,31 @@ def valid(self): def with_detections_count(self): return self.annotate(detections_count=models.Count("detections", distinct=True)) + def with_vectors(self, algorithm_id: int, key: str = DEFAULT_EMBEDDING_KEY): + """Occurrences with at least one detection that has a feature vector from the algorithm.""" + from ami.ml.embeddings.reader import representative_embeddings + + return self.filter(Exists(representative_embeddings(OuterRef("pk"), algorithm_id, key).order_by())) + + def with_visual_similarity(self, seed_vector, algorithm_id: int, key: str = DEFAULT_EMBEDDING_KEY): + """Annotate ``visual_similarity``: the cosine distance from ``seed_vector`` to each occurrence's vector. + + The occurrence's vector is its representative detection's (see + ``representative_embeddings``), from one algorithm only: distances between vectors of + different algorithms are meaningless. NULL for an occurrence without such a vector, so an + ascending ``nulls_last`` ordering puts the most similar first and vector-less ones last. + The distance is computed inside the subquery, so the aggregate annotations' GROUP BY + evaluates it once per occurrence rather than once per joined detection. + """ + from ami.ml.embeddings.reader import representative_embeddings + + nearest = ( + representative_embeddings(OuterRef("pk"), algorithm_id, key) + .annotate(distance=pgvector.django.CosineDistance("vector", seed_vector)) + .values("distance")[:1] + ) + return self.annotate(visual_similarity=models.Subquery(nearest)) + def _processed_by_algorithm_q(self, algorithm_ids) -> Exists: """Subquery matching occurrences with any result from the given algorithms — a detection made by one (detectors) or a classification from one (classifiers @@ -3417,15 +3448,19 @@ def not_processed_by_algorithm(self, algorithm_ids) -> "OccurrenceQuerySet": def created_or_updated_by_job(self, job_id: int) -> "OccurrenceQuerySet": """Occurrences created or updated by the given job. - An occurrence matches when the job created one of its detections, or a classification - on one of them, so several jobs can match the same occurrence. + An occurrence matches when the job created one of its detections, a classification on + one of them, or a feature vector for one of them, so several jobs can match the same + occurrence. - Identifications are not matched: people make them, not jobs. Two EXISTS subqueries + Identifications are not matched: people make them, not jobs. EXISTS subqueries return each occurrence once, where a join would return one row per matching result. """ + from ami.ml.models import DetectionEmbedding + return self.filter( Exists(Detection.objects.filter(occurrence=OuterRef("pk"), job_id=job_id)) | Exists(Classification.objects.filter(detection__occurrence=OuterRef("pk"), job_id=job_id)) + | Exists(DetectionEmbedding.objects.filter(detection__occurrence=OuterRef("pk"), job_id=job_id)) ) def with_timestamps(self): diff --git a/ami/main/test_visual_similarity.py b/ami/main/test_visual_similarity.py new file mode 100644 index 000000000..a3055d126 --- /dev/null +++ b/ami/main/test_visual_similarity.py @@ -0,0 +1,242 @@ +"""Sorting the occurrence list by visual similarity (``ordering=visual_similarity``). + +These pin the contract of the sort: the seed comes first and the nearest vectors next, +occurrences without a vector come last in either direction, one algorithm's vectors are +compared at a time, bad parameters are refused with 400, and the sort adds a fixed number +of queries however many rows a page has. See #1462. +""" + +from unittest import mock + +from cachalot.api import cachalot_disabled +from django.db import connection +from django.test.utils import CaptureQueriesContext +from rest_framework import status +from rest_framework.test import APITestCase + +from ami.main.api.views import OccurrenceViewSet +from ami.main.models import Occurrence, Project +from ami.ml.models import Algorithm, DetectionEmbedding +from ami.tests.fixtures.main import ( + create_captures, + create_occurrences, + create_taxa, + no_processing_service_http, + setup_test_project, +) +from ami.users.models import User + +SEED = [1.0, 0.0, 0.0, 0.0] +NEAR = [0.9, 0.1, 0.0, 0.0] +FAR = [0.0, 1.0, 0.0, 0.0] + + +class VisualSimilarityFixture(APITestCase): + """Five occurrences: a seed, a near and a far vector from the backbone, one with a vector + from another algorithm only, and one with no vector at all.""" + + @classmethod + def setUpTestData(cls) -> None: + with no_processing_service_http(): + cls.project, cls.deployment = setup_test_project(reuse=False) + cls.project.default_filters_score_threshold = 0.0 + cls.project.save() + create_taxa(project=cls.project) + create_captures(deployment=cls.deployment, num_nights=1, images_per_night=5) + create_occurrences(deployment=cls.deployment, num=5, determination_score=0.9) + cls.seed, cls.near, cls.far, cls.other_only, cls.no_vector = list( + Occurrence.objects.filter(project=cls.project).order_by("pk") + ) + cls.backbone = Algorithm.objects.create(name="Backbone", key="backbone", task_type="embedding") + cls.other = Algorithm.objects.create(name="Other backbone", key="other-backbone", task_type="embedding") + DetectionEmbedding.objects.store( + [ + cls._embedding(cls.seed, cls.backbone, SEED), + cls._embedding(cls.near, cls.backbone, NEAR), + cls._embedding(cls.far, cls.backbone, FAR), + cls._embedding(cls.other_only, cls.other, FAR), + cls._embedding(cls.seed, cls.other, SEED), + ] + ) + cls.url = f"/api/v2/occurrences/?project_id={cls.project.pk}" + + @staticmethod + def _embedding(occurrence: Occurrence, algorithm: Algorithm, vector: list[float]) -> DetectionEmbedding: + return DetectionEmbedding(detection=occurrence.detections.get(), algorithm=algorithm, vector=vector) + + def _ids(self, query: str, expected_status: int = status.HTTP_200_OK) -> list[int]: + response = self.client.get(f"{self.url}&{query}") + self.assertEqual(response.status_code, expected_status, response.content) + if expected_status != status.HTTP_200_OK: + return [] + return [int(row["id"]) for row in response.json()["results"]] + + +class TestVisualSimilarityOrdering(VisualSimilarityFixture): + def test_the_seed_comes_first_the_nearest_next_and_occurrences_without_a_vector_last(self): + ids = self._ids(f"ordering=visual_similarity&similar_to={self.seed.pk}") + self.assertEqual(ids[:3], [self.seed.pk, self.near.pk, self.far.pk]) + self.assertEqual(set(ids[3:]), {self.other_only.pk, self.no_vector.pk}) + + def test_descending_puts_the_most_different_first_and_still_the_vector_less_last(self): + ids = self._ids(f"ordering=-visual_similarity&similar_to={self.seed.pk}") + self.assertEqual(ids[:3], [self.far.pk, self.near.pk, self.seed.pk]) + self.assertEqual(set(ids[3:]), {self.other_only.pk, self.no_vector.pk}) + + def test_the_default_seed_is_the_most_recently_updated_occurrence_with_a_vector(self): + """The default list is ordered by -updated_at, so the seed is its top row that has a vector.""" + self.far.save() # Bumps updated_at. + ids = self._ids("ordering=visual_similarity") + self.assertEqual(ids[:3], [self.far.pk, self.near.pk, self.seed.pk]) + + self.no_vector.save() # The newest occurrence has no vector, so it is skipped as a seed. + self.assertEqual(self._ids("ordering=visual_similarity")[0], self.far.pk) + + def test_only_one_algorithms_vectors_are_compared(self): + """By default the algorithm with the most vectors; ``similarity_algorithm`` picks another. + An occurrence whose only vector comes from a different algorithm counts as having none.""" + ids = self._ids(f"ordering=visual_similarity&similar_to={self.seed.pk}") + self.assertIn(self.other_only.pk, ids[3:]) + + ids = self._ids(f"ordering=visual_similarity&similar_to={self.seed.pk}&similarity_algorithm={self.other.pk}") + self.assertEqual(ids[:2], [self.seed.pk, self.other_only.pk]) + self.assertEqual(set(ids[2:]), {self.near.pk, self.far.pk, self.no_vector.pk}) + + def test_the_sort_goes_through_the_default_filters(self): + self.project.default_filters_score_threshold = 0.95 + self.project.save() + self.assertEqual(self._ids(f"ordering=visual_similarity&similar_to={self.seed.pk}"), []) + self.assertEqual(len(self._ids("ordering=visual_similarity&apply_defaults=false")), 5) + + def test_bad_parameters_return_400(self): + other_project, other_deployment = setup_test_project(reuse=False) + create_taxa(project=other_project) + create_captures(deployment=other_deployment, num_nights=1, images_per_night=1) + create_occurrences(deployment=other_deployment, num=1) + foreign = Occurrence.objects.get(project=other_project) + + for query, key in [ + ("ordering=visual_similarity&similar_to=abc", "similar_to"), + ("ordering=visual_similarity&similar_to=0", "similar_to"), + (f"ordering=visual_similarity&similar_to={foreign.pk}", "similar_to"), + (f"ordering=visual_similarity&similar_to={self.no_vector.pk}", "similar_to"), + ("ordering=visual_similarity&similarity_algorithm=abc", "similarity_algorithm"), + (f"ordering=visual_similarity&similar_to={self.seed.pk}&similarity_algorithm=999999", "similar_to"), + ]: + with self.subTest(query=query): + response = self.client.get(f"{self.url}&{query}") + self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST, response.content) + self.assertIn(key, response.json()) + + def test_a_project_without_vectors_says_so(self): + DetectionEmbedding.objects.all().delete() + response = self.client.get(f"{self.url}&ordering=visual_similarity") + self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) + self.assertIn("No feature vectors", str(response.json()["ordering"])) + + def test_the_sort_adds_a_fixed_number_of_queries_however_many_rows(self): + """The seed and algorithm lookups are a handful of queries; nothing runs per row.""" + + def count(query: str) -> int: + with CaptureQueriesContext(connection) as ctx: + response = self.client.get(f"{self.url}&{query}") + self.assertEqual(response.status_code, status.HTTP_200_OK) + return len(ctx.captured_queries) + + similarity = "ordering=visual_similarity" + # The query cache would serve repeated lookups from the first call, so it is off while + # counting; the assertions stay outside the block so a failure cannot leave it off. + with cachalot_disabled(): + overhead_small = count(f"{similarity}&limit=2") - count("limit=2") + overhead_large = count(f"{similarity}&limit=5") - count("limit=5") + self.assertEqual(overhead_small, overhead_large) + # Project visibility, the algorithm with the most vectors, the default filters' two taxa + # lists for the default seed, the seed itself, and its vector. + self.assertLessEqual(overhead_large, 6) + + +class TestVisualSimilarityPermissions(VisualSimilarityFixture): + """The sort shows exactly what the plain list shows for each kind of user.""" + + @classmethod + def setUpTestData(cls) -> None: + super().setUpTestData() + cls.member = User.objects.create_user(email="similar-member@insectai.org") + cls.project.members.add(cls.member) + cls.non_member = User.objects.create_user(email="similar-outsider@insectai.org") + cls.superuser = User.objects.create_superuser(email="similar-admin@insectai.org", password="x") + + def _list(self, query: str) -> tuple[int, int]: + response = self.client.get(f"{self.url}&{query}") + return response.status_code, response.json().get("count", 0) if response.status_code == 200 else 0 + + def _assert_matches_plain_list(self): + for name, user in [ + ("anonymous", None), + ("non-member", self.non_member), + ("member", self.member), + ("superuser", self.superuser), + ]: + with self.subTest(user=name): + self.client.force_authenticate(user=user) + plain = self._list("") + sorted_ = self._list(f"ordering=visual_similarity&similar_to={self.seed.pk}") + self.assertEqual(sorted_[0], plain[0]) + self.assertEqual(sorted_[1], plain[1]) + + def test_public_project(self): + Project.objects.filter(pk=self.project.pk).update(draft=False) + self.client.force_authenticate(user=None) + self.assertEqual(self._list("")[1], 5) + self._assert_matches_plain_list() + + def test_draft_project(self): + Project.objects.filter(pk=self.project.pk).update(draft=True) + self.client.force_authenticate(user=self.member) + self.assertEqual(self._list("")[1], 5) + self._assert_matches_plain_list() + + +class TestSimilarityOrderSurvivesAViewDefault(VisualSimilarityFixture): + """A default ordering on the viewset must not replace the similarity order: the ordering filter + drops values that are not in ``ordering_fields`` and applies the view default instead.""" + + def test_the_similarity_order_wins_over_a_default_ordering(self): + with mock.patch.object(OccurrenceViewSet, "ordering", ["-created_at"], create=True): + ids = self._ids(f"ordering=visual_similarity&similar_to={self.seed.pk}") + reverse = self._ids(f"ordering=-visual_similarity&similar_to={self.seed.pk}") + self.assertEqual(ids[:3], [self.seed.pk, self.near.pk, self.far.pk]) + self.assertEqual(reverse[:3], [self.far.pk, self.near.pk, self.seed.pk]) + + +class TestOccurrenceDetailEmbeddingAlgorithms(VisualSimilarityFixture): + """The detail view lists the algorithms whose vector can seed a similarity sort, so the UI + can offer the link only when it would work.""" + + def _detail(self, occurrence: Occurrence) -> dict: + response = self.client.get(f"/api/v2/occurrences/{occurrence.pk}/?project_id={self.project.pk}") + self.assertEqual(response.status_code, status.HTTP_200_OK, response.content) + return response.json() + + def test_the_list_is_empty_without_vectors_and_names_each_algorithm_that_has_one(self): + self.assertEqual(self._detail(self.no_vector)["embedding_algorithms"], []) + self.assertEqual( + self._detail(self.seed)["embedding_algorithms"], + [{"id": self.backbone.pk, "name": "Backbone"}, {"id": self.other.pk, "name": "Other backbone"}], + ) + self.assertEqual( + self._detail(self.other_only)["embedding_algorithms"], [{"id": self.other.pk, "name": "Other backbone"}] + ) + + def test_the_field_costs_one_query_whatever_the_number_of_algorithms(self): + def vector_queries(occurrence: Occurrence) -> int: + with cachalot_disabled(), CaptureQueriesContext(connection) as queries: + self._detail(occurrence) + return sum("ml_detectionembedding" in query["sql"] for query in queries.captured_queries) + + self.assertEqual(vector_queries(self.no_vector), 1) + self.assertEqual(vector_queries(self.seed), 1) + + def test_the_list_view_does_not_carry_the_field(self): + response = self.client.get(f"{self.url}") + self.assertNotIn("embedding_algorithms", response.json()["results"][0]) diff --git a/ami/main/tests.py b/ami/main/tests.py index 27785fe7f..03e2f4f85 100644 --- a/ami/main/tests.py +++ b/ami/main/tests.py @@ -7245,7 +7245,7 @@ def test_exclude_is_the_complement_of_include(self): class TestOccurrenceJobFilter(APITestCase): """ Covers the ``?job=`` occurrence filter: occurrences with a detection or a - classification written by the given job. Each occurrence must appear once, however + classification written by the given job, or a feature vector it stored. Each occurrence must appear once, however many of its rows the job wrote, and a malformed id must be a 400, not a 500. """ @@ -7330,6 +7330,37 @@ def test_non_integer_job_is_a_bad_request(self): self.assertEqual(response.status_code, 400) self.assertIn("job", response.json()) + def test_vector_only_link_to_a_job_matches(self): + """A job that only stored feature vectors still finds its occurrences, and a vector from + another job does not match.""" + from ami.ml.models import DetectionEmbedding + + vector_job = Job.objects.create(project=self.project, name="Vector job") + # Detection and classification come from another job; only the vector links to vector_job. + only_vector = self._make_occurrence([(self.other_job, [self.other_job])]) + other_vector = self._make_occurrence([(self.other_job, [self.other_job])]) + DetectionEmbedding.objects.store( + [ + DetectionEmbedding( + detection=only_vector.detections.get(), algorithm=self.algorithm, vector=[0.5] * 4, job=vector_job + ), + DetectionEmbedding( + detection=other_vector.detections.get(), + algorithm=self.algorithm, + vector=[0.5] * 4, + job=self.other_job, + ), + ] + ) + ids = {row["id"] for row in self._list(vector_job.pk).json()["results"]} + self.assertEqual(ids, {only_vector.pk}) + queryset_ids = set( + Occurrence.objects.filter(project=self.project) + .created_or_updated_by_job(vector_job.pk) + .values_list("pk", flat=True) + ) + self.assertEqual(queryset_ids, {only_vector.pk}) + def test_filter_stays_inside_the_occurrence_query(self): """The filter adds no queries of its own: one statement for the rows and one for the count, with no ids read into Python first and no query per occurrence. Cachalot is off so every diff --git a/ami/ml/embeddings/__init__.py b/ami/ml/embeddings/__init__.py new file mode 100644 index 000000000..73b024f87 --- /dev/null +++ b/ami/ml/embeddings/__init__.py @@ -0,0 +1 @@ +"""Reading and writing the feature vectors stored for detections (see ``DetectionEmbedding``).""" diff --git a/ami/ml/embeddings/reader.py b/ami/ml/embeddings/reader.py new file mode 100644 index 000000000..d3e0e258e --- /dev/null +++ b/ami/ml/embeddings/reader.py @@ -0,0 +1,127 @@ +"""Feature vectors for detections, read one algorithm and key at a time. + +Vectors from different algorithms, or under different keys of one algorithm, are not +comparable, so every reader here is keyed by algorithm and key. ``DetectionEmbedding`` holds +one row per (detection, algorithm, key). Each function below serves one known query and says which +index it relies on. See #1462 and docs/claude/reference/feature-vectors.md. +""" + +from __future__ import annotations + +from collections.abc import Iterable, Iterator + +import numpy as np +from django.db.models import Count, Exists, OuterRef, QuerySet + +from ami.main.models import DEFAULT_EMBEDDING_KEY, Detection, Project +from ami.ml.models.embedding import DetectionEmbedding, as_half_precision + + +def vectors_for_detections( + detection_ids: Iterable[int], algorithm_id: int, key: str = DEFAULT_EMBEDDING_KEY +) -> dict[int, np.ndarray]: + """The vectors of one (algorithm, key) for a set of detections, as float32 arrays, in one query. + + Serves "the vectors of one model for these detections" (two adjacent captures for tracking, + verified detections for retraining) through the unique index on (detection, algorithm, key). + Call it once per (algorithm, key); detections without a vector are absent. All the vectors + have one length, because the writer holds each (algorithm, key) to one. + """ + rows = ( + DetectionEmbedding.objects.for_algorithm(algorithm_id, key) + .filter(detection_id__in=list(detection_ids)) + .order_by() + .values_list("detection_id", "vector") + ) + return {detection_id: as_half_precision(vector).astype(np.float32) for detection_id, vector in rows} + + +def project_vectors( + project_id: int, + algorithm_id: int, + key: str = DEFAULT_EMBEDDING_KEY, + *, + detection_ids: Iterable[int] | None = None, + chunk_size: int = 2000, +) -> Iterator[tuple[list[int], np.ndarray]]: + """Every vector of one (algorithm, key) in a project, in chunks ordered by detection id. + + Serves exports, clustering and retraining sets that need all of a model's vectors without + holding them all at once. Each chunk is (detection ids, float32 array of shape (n, length)). + The next chunk is fetched with ``detection_id > last`` (keyset pagination), which walks the + index on (project, algorithm, key, detection) in order, so memory stays at one chunk and no + page re-reads earlier rows. ``detection_ids`` narrows the scope. Call it once per + (algorithm, key): vectors of different models or lengths must not share an array. + """ + scope = DetectionEmbedding.objects.for_algorithm(algorithm_id, key).filter(project_id=project_id) + if detection_ids is not None: + scope = scope.filter(detection_id__in=list(detection_ids)) + last_id = 0 + while True: + rows = list( + scope.filter(detection_id__gt=last_id) + .order_by("detection_id") + .values_list("detection_id", "vector")[:chunk_size] + ) + if not rows: + return + last_id = rows[-1][0] + yield ( + [detection_id for detection_id, _ in rows], + np.stack([as_half_precision(vector).astype(np.float32) for _, vector in rows]), + ) + + +def vector_counts_by_algorithm(project_id: int, key: str | None = None) -> dict[tuple[int, str], int]: + """How many vectors each (algorithm id, key) has in a project, in one grouped query. + + Serves "which models have vectors here" (the default model for a similarity sort, the + algorithm filter). It reads only the index on (project, algorithm, key). ``key`` limits the + result to one output name; the result is keyed by pair so nothing is merged across models. + """ + rows = DetectionEmbedding.objects.filter(project_id=project_id) + if key is not None: + rows = rows.filter(key=key) + grouped = rows.order_by().values_list("algorithm_id", "key").annotate(n=Count("pk")) + return {(algorithm_id, row_key): n for algorithm_id, row_key, n in grouped} + + +def detections_missing_vectors( + detections: QuerySet[Detection], algorithm_id: int, key: str = DEFAULT_EMBEDDING_KEY +) -> QuerySet[Detection]: + """The given detections that have no vector yet from one (algorithm, key). + + Serves finding what a feature-extraction run still has to send. It adds a NOT EXISTS on the + unique index (detection, algorithm, key) to the caller's queryset, so the caller keeps its + own scope (project, capture set, station) and the filter adds no query of its own. + """ + has_vector = DetectionEmbedding.objects.for_algorithm(algorithm_id, key).filter(detection_id=OuterRef("pk")) + # exclude(), not filter(~Exists(...)): django-cachalot 2.6 does not see the tables inside a + # negated Exists, so writes to the vector table would not invalidate a cached result. + return detections.exclude(Exists(has_vector)) + + +def representative_embeddings(occurrence_id, algorithm_id: int, key: str = DEFAULT_EMBEDDING_KEY): + """One occurrence's vectors from one (algorithm, key), its representative detection's first. + + The representative detection is the earliest one that has such a vector, by frame number, + then time, then id: the detection whose crop the occurrence list shows, when that crop has + a vector. The seed of a similarity sort and every occurrence it ranks both go through this, + so they are compared the same way. ``occurrence_id`` may be an ``OuterRef``. + """ + return ( + DetectionEmbedding.objects.for_algorithm(algorithm_id, key) + .filter(detection__occurrence_id=occurrence_id) + .order_by("detection__frame_num", "detection__timestamp", "detection_id") + ) + + +def algorithm_with_most_vectors(project: Project, key: str = DEFAULT_EMBEDDING_KEY) -> int | None: + """The algorithm that stored the most vectors under ``key`` in the project, or None when there are none. + + Ties go to the lowest algorithm id. + """ + counts = vector_counts_by_algorithm(project.pk, key) + if not counts: + return None + return min(counts, key=lambda pair: (-counts[pair], pair[0]))[0] diff --git a/ami/ml/embeddings/writer.py b/ami/ml/embeddings/writer.py new file mode 100644 index 000000000..38321d22f --- /dev/null +++ b/ami/ml/embeddings/writer.py @@ -0,0 +1,212 @@ +"""Store the feature vectors a processing service returns with its detections.""" + +from __future__ import annotations + +import collections +import contextlib +import dataclasses +import logging +import typing +import zlib + +import numpy as np +from django.db import connection, transaction + +from ami.main.models import DEFAULT_EMBEDDING_KEY, Detection +from ami.ml.exceptions import PipelineNotConfigured +from ami.ml.models.algorithm import Algorithm +from ami.ml.models.embedding import DetectionEmbedding, as_half_precision +from ami.ml.schemas import DetectionResponse, PipelineResultsResponse + +if typing.TYPE_CHECKING: + from ami.jobs.models import Job + from ami.ml.models.pipeline import Pipeline + +logger = logging.getLogger(__name__) + + +def _box_key(source_image_id, coordinates) -> tuple: + """Exact coordinates, the identity ``get_or_create_detection`` reuses detections by.""" + return (str(source_image_id), tuple(float(value) for value in coordinates)) + + +class EmbeddingDimensionMismatch(PipelineNotConfigured): + """A vector's length differs from the length its algorithm has produced before.""" + + +def _lock_first_write(algorithm: Algorithm, key: str) -> None: + """Serialise first writers of one (algorithm, key) until the surrounding transaction ends.""" + with connection.cursor() as cursor: + cursor.execute("SELECT pg_advisory_xact_lock(%s, %s)", [algorithm.pk, zlib.crc32(key.encode()) & 0x7FFFFFFF]) + + +def _check_embedding_dimensions( + lengths_by_pair: dict[tuple[Algorithm, str], set[int]], stored_by_pair: dict[tuple[Algorithm, str], int | None] +) -> None: + """Refuse vectors whose length differs from the one already stored for their (algorithm, key). + + Vectors of different lengths can never be compared, so each (algorithm, key) keeps one length. + With none stored yet, the batch must agree on one. + """ + for (algorithm, key), lengths in lengths_by_pair.items(): + stored = stored_by_pair[(algorithm, key)] + if stored is None and len(lengths) > 1: + raise EmbeddingDimensionMismatch( + f"Algorithm {algorithm.key} sent vectors of several lengths under key '{key}' in one batch: " + f"{sorted(lengths)}" + ) + expected = stored if stored is not None else next(iter(lengths)) + wrong = lengths - {expected} + if wrong: + raise EmbeddingDimensionMismatch( + f"Algorithm {algorithm.key} produces {expected}-dimension vectors under key '{key}'; " + f"refusing vectors of length {sorted(wrong)}." + ) + + +@dataclasses.dataclass +class EmbeddingStoreResult: + """What one batch of returned vectors did: the rows sent to the store and how each vector fared.""" + + embeddings: list[DetectionEmbedding] = dataclasses.field(default_factory=list) + written: int = 0 + unchanged: int = 0 + unmatched: int = 0 + not_finite: int = 0 + + +def create_detection_embeddings( + detections: list[Detection], + detection_responses: list[DetectionResponse], + algorithms_known: dict[str, Algorithm], + logger: logging.Logger = logger, + job_id: int | None = None, +) -> list[DetectionEmbedding]: + """ + Store the feature vectors sent with each detection, one row per (detection, algorithm, key). + + Writes are insert-mostly (see ``DetectionEmbeddingQuerySet.store``): saving the same results + twice changes nothing, and a new vector for a pair replaces the old one. Only + ``DetectionEmbedding`` rows are written, never a classification, so no determination can + change. A vector with a value half precision cannot hold (NaN, infinity, beyond 65504) is + skipped with a warning. + + Responses are matched to ``detections`` by image and box (by exact coordinates), the + key ``get_or_create_detection`` reuses detections by, because detection creation returns + existing detections ahead of new ones and pairing by position would swap vectors. An + algorithm key the pipeline has not registered raises ``PipelineNotConfigured``, as it does + for classifications, and a vector whose length differs from its algorithm's raises + ``EmbeddingDimensionMismatch``. ``job_id`` records the job whose results stored each vector. + Returns the embeddings that were sent to the store (written or unchanged). + """ + return store_detection_embeddings( + detections, detection_responses, algorithms_known, DEFAULT_EMBEDDING_KEY, logger, job_id + ).embeddings + + +def store_detection_embeddings( + detections: list[Detection], + detection_responses: list[DetectionResponse], + algorithms_known: dict[str, Algorithm], + key: str, + logger: logging.Logger = logger, + job_id: int | None = None, +) -> EmbeddingStoreResult: + """Match returned vectors to ``detections`` by image and box and store them under ``key``. + + The matching, validation and storing rules are those of ``create_detection_embeddings``. + """ + by_box: dict[tuple, list[Detection]] = collections.defaultdict(list) + for detection in detections: + if detection.bbox is not None: + by_box[_box_key(detection.source_image_id, detection.bbox)].append(detection) + embeddings: dict[tuple[int, int, str], DetectionEmbedding] = {} + lengths_by_pair: dict[tuple[Algorithm, str], set[int]] = collections.defaultdict(set) + unmatched = not_finite = 0 + for detection_resp in detection_responses: + if detection_resp.bbox is None or not detection_resp.embeddings: + continue + box_key = _box_key(detection_resp.source_image_id, detection_resp.bbox.dict().values()) + candidates = by_box.get(box_key, []) + if len(candidates) > 1: + logger.warning( + f"Skipped the vectors of capture {box_key[0]} box {box_key[1]}: " + f"{len(candidates)} stored detections share that box." + ) + if len(candidates) != 1: + unmatched += 1 + continue + detection = candidates[0] + for embedding_resp in detection_resp.embeddings: + try: + algorithm = algorithms_known[embedding_resp.algorithm.key] + except KeyError as err: + raise PipelineNotConfigured( + f"Embedding algorithm {embedding_resp.algorithm.key} is not a known algorithm. " + "The processing service must declare it in the /info endpoint. " + f"Known algorithms: {list(algorithms_known.keys())}" + ) from err + if not np.isfinite(as_half_precision(embedding_resp.features)).all(): + not_finite += 1 + continue + lengths_by_pair[(algorithm, key)].add(len(embedding_resp.features)) + embeddings[(detection.pk, algorithm.pk, key)] = DetectionEmbedding( + detection=detection, + algorithm=algorithm, + key=key, + vector=embedding_resp.features, + job_id=job_id, + ) + + stored_by_pair = {pair: DetectionEmbedding.objects.stored_length(pair[0].pk, pair[1]) for pair in lengths_by_pair} + first_writes = [pair for pair, stored in stored_by_pair.items() if stored is None] + result = EmbeddingStoreResult(embeddings=list(embeddings.values()), unmatched=unmatched, not_finite=not_finite) + + with contextlib.ExitStack() as stack: + if first_writes: + # Hold the lock until the vectors are committed, so a concurrent first writer sees them. + stack.enter_context(transaction.atomic()) + for algorithm, pair_key in first_writes: + _lock_first_write(algorithm, pair_key) + stored_by_pair[(algorithm, pair_key)] = DetectionEmbedding.objects.stored_length( + algorithm.pk, pair_key + ) + _check_embedding_dimensions(lengths_by_pair, stored_by_pair) + + if unmatched: + logger.warning(f"Skipped the vectors of {unmatched} returned boxes that match no single stored detection.") + if not_finite: + logger.warning(f"Skipped {not_finite} vectors with values a half-precision vector cannot store.") + if embeddings: + result.written, result.unchanged = DetectionEmbedding.objects.store(embeddings.values()) + logger.info( + f"Stored {result.written} feature vectors ({result.unchanged} unchanged) " + f"for {len(detections)} detections." + ) + return result + + +def save_embedding_results( + response: PipelineResultsResponse, + job: Job | None, + pipeline: Pipeline, + key: str = DEFAULT_EMBEDDING_KEY, +) -> EmbeddingStoreResult: + """Store only the vectors of a feature-only response, on detections that already exist. + + A feature-only processing service returns the detections it was sent, unchanged and still + naming the detector that made them, with an embedding attached and no classifications. Each + returned detection is matched by capture and box to a stored valid detection; nothing else is + written: no detection, classification or occurrence is created, and the echoed detector + reference is ignored. A returned box with no stored detection is counted in ``unmatched`` and + logged, never created. An embedding that names an algorithm outside ``pipeline`` raises + ``PipelineNotConfigured``, and a vector of the wrong length for its (algorithm, key) raises + ``EmbeddingDimensionMismatch``. + """ + capture_ids = {int(detection.source_image_id) for detection in response.detections} + detections = list(Detection.objects.valid().filter(source_image_id__in=capture_ids)) + algorithms_known = {algorithm.key: algorithm for algorithm in pipeline.algorithms.all()} + job_logger = job.logger if job else logger + return store_detection_embeddings( + detections, response.detections, algorithms_known, key, job_logger, job.pk if job else None + ) diff --git a/ami/ml/migrations/0029_enable_pgvector.py b/ami/ml/migrations/0029_enable_pgvector.py new file mode 100644 index 000000000..5df64cde9 --- /dev/null +++ b/ami/ml/migrations/0029_enable_pgvector.py @@ -0,0 +1,56 @@ +# First install of the pgvector extension, which stores feature vectors (DetectionEmbedding). +# +# The extension package must already be installed on the PostgreSQL server (for example +# the postgresql-16-pgvector package); CREATE EXTENSION only registers it in this database. +# The check below runs before any SQL so a missing or outdated package stops the deploy +# with one clear message instead of a Postgres error about a control file. See #1462. + +from django.db import migrations + +MINIMUM_VERSION = (0, 8) + + +def check_pgvector_is_installed(apps, schema_editor): + """Stop before any SQL runs unless the server offers pgvector 0.8 or later. + + ``halfvec`` needs 0.7; 0.8 adds the iterative index scans a nearest-neighbour index + will rely on, so a first install starts there. + """ + with schema_editor.connection.cursor() as cursor: + cursor.execute("SELECT default_version FROM pg_available_extensions WHERE name = 'vector'") + row = cursor.fetchone() + minimum = ".".join(str(part) for part in MINIMUM_VERSION) + if row is None: + raise RuntimeError( + f"The pgvector extension is not installed on this PostgreSQL server. Install pgvector {minimum} " + "or later (for example the postgresql-16-pgvector package) on every database server, then " + "run this migration again." + ) + (version,) = row + try: + parsed = tuple(int(part) for part in version.split(".")[:2]) + except ValueError: + parsed = () + if parsed < MINIMUM_VERSION: + raise RuntimeError( + f"pgvector {minimum} or later must be installed on this PostgreSQL server; found {version}. " + "Upgrade the pgvector package on every database server, then run this migration again." + ) + + +class Migration(migrations.Migration): + dependencies = [ + ("ml", "0028_normalize_empty_endpoint_url_to_null"), + ] + + operations = [ + migrations.RunPython(check_pgvector_is_installed, migrations.RunPython.noop), + migrations.RunSQL( + # ALTER ... UPDATE is a no-op on a fresh install. It only matters for a development + # database where an experiment already created an older version of the extension. + sql="CREATE EXTENSION IF NOT EXISTS vector; ALTER EXTENSION vector UPDATE;", + # The extension may be shared with other databases on the server, and dropping it can + # be restricted in hosted environments, so the reverse leaves it in place. + reverse_sql=migrations.RunSQL.noop, + ), + ] diff --git a/ami/ml/migrations/0030_detection_embedding.py b/ami/ml/migrations/0030_detection_embedding.py new file mode 100644 index 000000000..77c50deb3 --- /dev/null +++ b/ami/ml/migrations/0030_detection_embedding.py @@ -0,0 +1,102 @@ +# Additive: an empty table. The vector column is stored uncompressed out of +# line (STORAGE EXTERNAL): vectors do not compress, and it keeps the table's rows small. See #1462. + +import django.db.models.deletion +import django.utils.timezone +import pgvector.django.halfvec +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("main", "0097_detection_and_classification_job_indexes"), + ("jobs", "0023_alter_job_job_type_key"), + ("ml", "0029_enable_pgvector"), + ] + + operations = [ + migrations.CreateModel( + name="DetectionEmbedding", + 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)), + ("key", models.CharField(default="embedding", max_length=255)), + ( + "vector", + pgvector.django.halfvec.HalfVectorField(help_text="The feature vector, in half precision."), + ), + ("timestamp", models.DateTimeField(default=django.utils.timezone.now)), + ( + "algorithm", + models.ForeignKey( + db_index=False, + on_delete=django.db.models.deletion.CASCADE, + related_name="detection_embeddings", + to="ml.algorithm", + ), + ), + ( + "detection", + models.ForeignKey( + db_index=False, + on_delete=django.db.models.deletion.CASCADE, + related_name="embeddings", + to="main.detection", + ), + ), + ( + "job", + models.ForeignKey( + blank=True, + db_index=False, + null=True, + on_delete=django.db.models.deletion.SET_NULL, + related_name="detection_embeddings", + to="jobs.job", + ), + ), + ( + "project", + models.ForeignKey( + db_index=False, + on_delete=django.db.models.deletion.CASCADE, + related_name="detection_embeddings", + to="main.project", + ), + ), + ], + ), + migrations.AddConstraint( + model_name="detectionembedding", + constraint=models.CheckConstraint( + check=models.Q(("key", ""), _negated=True), name="ml_detectionembedding_key_not_empty" + ), + ), + migrations.AddConstraint( + model_name="detectionembedding", + constraint=models.UniqueConstraint( + fields=("detection", "algorithm", "key"), name="ml_detectionembedding_unique_detection_algorithm_key" + ), + ), + migrations.AddIndex( + model_name="detectionembedding", + index=models.Index( + fields=["project", "algorithm", "key", "detection"], name="ml_detemb_proj_algo_key_det" + ), + ), + migrations.AddIndex( + model_name="detectionembedding", + index=models.Index(fields=["algorithm", "key", "detection"], name="ml_detemb_algo_key"), + ), + migrations.AddIndex( + model_name="detectionembedding", + index=models.Index( + condition=models.Q(("job__isnull", False)), fields=["job", "detection"], name="ml_detemb_job_det" + ), + ), + migrations.RunSQL( + sql="ALTER TABLE ml_detectionembedding ALTER COLUMN vector SET STORAGE EXTERNAL", + reverse_sql=migrations.RunSQL.noop, + ), + ] diff --git a/ami/ml/models/__init__.py b/ami/ml/models/__init__.py index 5000c7f53..80df44738 100644 --- a/ami/ml/models/__init__.py +++ b/ami/ml/models/__init__.py @@ -1,4 +1,5 @@ from ami.ml.models.algorithm import Algorithm, AlgorithmCategoryMap +from ami.ml.models.embedding import DetectionEmbedding from ami.ml.models.pipeline import Pipeline from ami.ml.models.processing_service import ProcessingService from ami.ml.models.project_pipeline_config import ProjectPipelineConfig @@ -6,6 +7,7 @@ __all__ = [ "Algorithm", "AlgorithmCategoryMap", + "DetectionEmbedding", "Pipeline", "ProcessingService", "ProjectPipelineConfig", diff --git a/ami/ml/models/embedding.py b/ami/ml/models/embedding.py new file mode 100644 index 000000000..d3ba4cd08 --- /dev/null +++ b/ami/ml/models/embedding.py @@ -0,0 +1,178 @@ +"""Feature vectors for detections, stored next to the other model outputs in the ml app.""" + +import typing + +import numpy as np +import pgvector +import pgvector.django +from django.db import models +from django.utils import timezone + +from ami.base.models import BaseModel, BaseQuerySet +from ami.main.models import DEFAULT_EMBEDDING_KEY, Detection + +# A 2048-d vector is roughly 20 KB of SQL text (estimate), so this keeps each INSERT to a few MB. +EMBEDDING_BATCH_SIZE = 200 + + +class VectorDims(models.Func): + """pgvector's ``vector_dims()``: the length of a vector or halfvec value.""" + + function = "vector_dims" + output_field = models.IntegerField() + + +def as_half_precision(vector) -> np.ndarray: + """A vector as the half-precision array the ``halfvec`` column stores.""" + if isinstance(vector, pgvector.HalfVector): + return vector.to_numpy() + # A value beyond half precision's range becomes infinity, which the writer then refuses. + with np.errstate(over="ignore"): + return np.asarray(vector, dtype=np.float16) + + +def same_vector(stored, new) -> bool: + return bool(np.array_equal(as_half_precision(stored), as_half_precision(new))) + + +def fill_embedding_project_ids(embeddings: list["DetectionEmbedding"]) -> None: + """Give each vector without a project its detection's capture's project, or the capture's station's. + + Raises ``ValueError`` when neither has one: a vector must always be reachable by project. + """ + missing = [embedding for embedding in embeddings if embedding.project_id is None] + if not missing: + return + projects = { + detection_id: capture_project_id or station_project_id + for detection_id, capture_project_id, station_project_id in Detection.objects.filter( + pk__in={embedding.detection_id for embedding in missing} + ).values_list("pk", "source_image__project_id", "source_image__deployment__project_id") + } + for embedding in missing: + embedding.project_id = projects.get(embedding.detection_id) + if embedding.project_id is None: + raise ValueError( + f"Detection #{embedding.detection_id} belongs to no project: neither its capture nor its " + "station has one, so its feature vector cannot be stored." + ) + + +class DetectionEmbeddingQuerySet(BaseQuerySet): + def for_algorithm(self, algorithm_id: int, key: str = DEFAULT_EMBEDDING_KEY): + """Vectors of one kind from one algorithm: the only set whose vectors may be compared.""" + return self.filter(algorithm_id=algorithm_id, key=key) + + def stored_length(self, algorithm_id: int, key: str = DEFAULT_EMBEDDING_KEY) -> int | None: + """The length of the vectors already stored for one (algorithm, key), or None when there are none. + + Reads the first row of the pair (an indexed lookup on (algorithm, key, detection) whether or not + the pair has rows), so a writer can hold each pair to one length without a stored field. The + ``order_by`` is what lets the planner walk that index instead of scanning the table until a row matches. + """ + rows = ( + self.for_algorithm(algorithm_id, key) + .order_by("detection_id") + .annotate(dims=VectorDims("vector")) + .values_list("dims", flat=True)[:1] + ) + return next(iter(rows), None) + + def store(self, embeddings) -> tuple[int, int]: + """Write vectors insert-mostly; returns (written, unchanged). + + A (detection, algorithm, key) that already holds the same vector is left alone, so saving + the same results twice changes nothing. One holding another vector is updated in place + (ON CONFLICT DO UPDATE), which also keeps two writers of the same pair from losing a row. + """ + embeddings = list(embeddings) + if not embeddings: + return 0, 0 + fill_embedding_project_ids(embeddings) + stored = { + (detection_id, algorithm_id, key): vector + for detection_id, algorithm_id, key, vector in self.filter( + detection_id__in={embedding.detection_id for embedding in embeddings}, + algorithm_id__in={embedding.algorithm_id for embedding in embeddings}, + key__in={embedding.key for embedding in embeddings}, + ) + .order_by() + .values_list("detection_id", "algorithm_id", "key", "vector") + } + to_write = [] + for embedding in embeddings: + existing = stored.get((embedding.detection_id, embedding.algorithm_id, embedding.key)) + if existing is None or not same_vector(existing, embedding.vector): + to_write.append(embedding) + if to_write: + self.bulk_create( + to_write, + batch_size=EMBEDDING_BATCH_SIZE, + update_conflicts=True, + update_fields=["vector", "job", "timestamp", "updated_at"], + unique_fields=["detection", "algorithm", "key"], + ) + return len(to_write), len(embeddings) - len(to_write) + + +@typing.final +class DetectionEmbedding(BaseModel): + """A feature vector for one detection, used to compare detections by appearance. + + Kept apart from classifications so every detection can have one, including those the + moth/non-moth filter rejected, without adding a prediction that could change a + determination. Vectors are comparable only within one (algorithm, key): ``key`` names + the output when one model yields several, and the column is an unsized ``halfvec`` + because extractors differ in length: each (algorithm, key) keeps one length, enforced by + the writer against an existing row. See #1462. + """ + + # None of the four foreign keys has a single-column index of its own: the unique constraint + # leads with detection, and the Meta indexes lead with algorithm, with project and with job + # (partial: a vector whose job was deleted has no job left to look up). Extra single-column + # indexes would only cost writes and mislead the planner. + detection = models.ForeignKey(Detection, on_delete=models.CASCADE, related_name="embeddings", db_index=False) + algorithm = models.ForeignKey( + "ml.Algorithm", on_delete=models.CASCADE, related_name="detection_embeddings", db_index=False + ) + # The job whose results stored the vector. Deleting the job keeps the vector. + job = models.ForeignKey( + "jobs.Job", + on_delete=models.SET_NULL, + null=True, + blank=True, + related_name="detection_embeddings", + db_index=False, + ) + # Copied from the detection's capture (see fill_embedding_project_ids) so project queries need no join. + project = models.ForeignKey( + "main.Project", on_delete=models.CASCADE, related_name="detection_embeddings", db_index=False + ) + key = models.CharField(max_length=255, default=DEFAULT_EMBEDDING_KEY) + vector = pgvector.django.HalfVectorField(help_text="The feature vector, in half precision.") + timestamp = models.DateTimeField(default=timezone.now) + + objects = DetectionEmbeddingQuerySet.as_manager() + + class Meta: + indexes = [ + # Project-scoped reads of one model: a project's vectors in detection order, and counts per model. + models.Index(fields=["project", "algorithm", "key", "detection"], name="ml_detemb_proj_algo_key_det"), + # The writer's per-model length lookup: the first row of a pair, by detection, with or without rows. + models.Index(fields=["algorithm", "key", "detection"], name="ml_detemb_algo_key"), + # "Occurrences this job touched" looks up a job's vectors by job, then detection. + models.Index(fields=["job", "detection"], name="ml_detemb_job_det", condition=models.Q(job__isnull=False)), + ] + constraints = [ + models.CheckConstraint(check=~models.Q(key=""), name="%(app_label)s_%(class)s_key_not_empty"), + models.UniqueConstraint( + fields=["detection", "algorithm", "key"], name="%(app_label)s_%(class)s_unique_detection_algorithm_key" + ), + ] + + def __str__(self) -> str: + return f"#{self.pk} {self.key} for Detection #{self.detection_id} from Algorithm #{self.algorithm_id}" + + def save(self, *args, **kwargs): + fill_embedding_project_ids([self]) + super().save(*args, **kwargs) diff --git a/ami/ml/models/pipeline.py b/ami/ml/models/pipeline.py index 3d5b77d58..838fae925 100644 --- a/ami/ml/models/pipeline.py +++ b/ami/ml/models/pipeline.py @@ -38,8 +38,9 @@ update_calculated_fields_for_events, update_occurrence_determination, ) +from ami.ml.embeddings.writer import create_detection_embeddings from ami.ml.exceptions import PipelineNotConfigured -from ami.ml.models.algorithm import Algorithm, AlgorithmCategoryMap +from ami.ml.models.algorithm import Algorithm, AlgorithmCategoryMap, AlgorithmTaskType from ami.ml.schemas import ( AlgorithmConfigResponse, AlgorithmReference, @@ -410,6 +411,22 @@ def process_images( return results +def detection_request(detection: Detection, source_image_request: SourceImageRequest) -> DetectionRequest | None: + """The request form of a stored detection, or None when it has no box or no detector to name.""" + bbox = detection.get_bbox() + if not bbox or not detection.detection_algorithm: + return None + return DetectionRequest( + source_image=source_image_request, + bbox=bbox, + crop_image_url=detection.url(), + algorithm=AlgorithmReference( + name=detection.detection_algorithm.name, + key=detection.detection_algorithm.key, + ), + ) + + def collect_detections( source_image: SourceImage, source_image_request: SourceImageRequest, @@ -417,24 +434,50 @@ def collect_detections( """ Collect existing detections for a source image and send them with pipeline request. """ - detection_requests: list[DetectionRequest] = [] # Re-process all existing detections if they exist - for detection in source_image.detections.all(): - bbox = detection.get_bbox() - if bbox and detection.detection_algorithm: - detection_requests.append( - DetectionRequest( - source_image=source_image_request, - bbox=bbox, - crop_image_url=detection.url(), - algorithm=AlgorithmReference( - name=detection.detection_algorithm.name, - key=detection.detection_algorithm.key, - ), - ) - ) + built = (detection_request(detection, source_image_request) for detection in source_image.detections.all()) + return [request for request in built if request is not None] + - return detection_requests +def process_detections( + pipeline: Pipeline, + endpoint_url: str, + detections: typing.Iterable[Detection], + project_id: int, +) -> PipelineResultsResponse: + """Send stored detections, and only those, to a processing service in one synchronous request. + + Serves feature-only pipelines, which embed the detections they are sent instead of detecting. + The detections' captures (``source_image`` with its deployment and data source) and detector + (``detection_algorithm``) should be loaded with the queryset, because each is read per detection. + A failed request raises ``requests.HTTPError`` rather than returning an empty response. + """ + source_image_requests: dict[int, SourceImageRequest] = {} + detection_requests: list[DetectionRequest] = [] + for detection in detections: + source_image = detection.source_image + url = source_image.public_url() + if not url: + continue + source_image_request = source_image_requests.setdefault( + source_image.pk, SourceImageRequest(id=str(source_image.pk), url=url) + ) + request = detection_request(detection, source_image_request) + if request is not None: + detection_requests.append(request) + + request_data = PipelineRequest( + pipeline=pipeline.slug, + source_images=list(source_image_requests.values()), + config=pipeline.get_config(project_id=project_id), + detections=detection_requests, + ) + resp = create_session().post(endpoint_url, json=request_data.dict()) + if not resp.ok: + raise requests.HTTPError( + f"Failed to process {request_data.summary()}: {extract_error_message_from_response(resp)}" + ) + return PipelineResultsResponse(**resp.json()) def get_or_create_algorithm_and_category_map( @@ -1087,6 +1130,16 @@ def save_results( job_id=job.pk if job else None, ) + # Before classifications, so an unregistered embedding algorithm stops the batch at the + # same point an unregistered classification algorithm does. + create_detection_embeddings( + detections=detections, + detection_responses=results.detections, + algorithms_known=algorithms_known, + logger=job_logger, + job_id=job.pk if job else None, + ) + classifications = create_classifications( detections=detections, detection_responses=results.detections, @@ -1282,6 +1335,26 @@ def collect_images( reprocess_all_images=reprocess_all_images, ) + def embedding_algorithms(self) -> models.QuerySet[Algorithm]: + """The algorithms of this pipeline that produce feature vectors.""" + return self.algorithms.filter(task_type=AlgorithmTaskType.EMBEDDING.value) + + def is_embedding_only(self) -> bool: + """True when every algorithm of the pipeline produces feature vectors, and there is at least one. + + Such a pipeline detects and classifies nothing, so it can only be run through the + "Add feature vectors" task, which stores vectors on detections that already exist. + """ + algorithms = self.algorithms.all() + return bool(algorithms) and all(a.task_type == AlgorithmTaskType.EMBEDDING.value for a in algorithms) + + def raise_if_embedding_only(self) -> None: + if self.is_embedding_only(): + raise PipelineNotConfigured( + f'Pipeline "{self.name}" only produces feature vectors, so it cannot run as an ML job. ' + 'Use the "Add feature vectors" action on a capture set in the admin instead.' + ) + def choose_processing_service_for_pipeline( self, job_id: int | None, pipeline_name: str, project_id: int ) -> ProcessingService: @@ -1337,6 +1410,7 @@ def process_images( job_id: int | None = None, reprocess_all_images: bool = False, ) -> PipelineResultsResponse: + self.raise_if_embedding_only() processing_service = self.choose_processing_service_for_pipeline(job_id, self.name, project_id) if not processing_service.endpoint_url: diff --git a/ami/ml/post_processing/admin/feature_vectors_form.py b/ami/ml/post_processing/admin/feature_vectors_form.py new file mode 100644 index 000000000..40cd11188 --- /dev/null +++ b/ami/ml/post_processing/admin/feature_vectors_form.py @@ -0,0 +1,51 @@ +from __future__ import annotations + +from django import forms + +from ami.main.models import DEFAULT_EMBEDDING_KEY +from ami.ml.models import Pipeline +from ami.ml.models.algorithm import AlgorithmTaskType +from ami.ml.post_processing.admin.forms import BasePostProcessingActionForm +from ami.ml.post_processing.feature_vectors import AddFeatureVectorsConfig + + +class AddFeatureVectorsActionForm(BasePostProcessingActionForm): + """Knobs surfaced when an admin triggers Add feature vectors. + + The pipeline choice lists only pipelines with an algorithm that produces feature vectors, + narrowed to those enabled for the projects of the selected capture sets. The valid range of + ``batch_size`` lives on ``AddFeatureVectorsConfig``; the admin action surfaces its errors inline. + """ + + pipeline_id = forms.ModelChoiceField( + queryset=Pipeline.objects.none(), + label="Pipeline", + help_text="A pipeline with a feature extractor. Only detections that lack its vectors are sent.", + ) + key = forms.CharField( + label="Vector name", + initial=DEFAULT_EMBEDDING_KEY, + help_text="The name the vectors are stored under. Keep the default unless the extractor has several outputs.", + ) + batch_size = forms.IntegerField( + label="Captures per request", + initial=AddFeatureVectorsConfig.__fields__["batch_size"].default, + help_text="How many captures, with their missing detections, go to the processing service in one request.", + ) + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + pipelines = Pipeline.objects.filter(algorithms__task_type=AlgorithmTaskType.EMBEDDING.value) + if self.scope_queryset is not None: + project_ids = set(self.scope_queryset.values_list("project_id", flat=True)) + pipelines = pipelines.filter( + project_pipeline_configs__enabled=True, project_pipeline_configs__project_id__in=project_ids + ) + self.fields["pipeline_id"].queryset = pipelines.distinct().order_by("name") + + def to_config(self) -> dict: + return { + "pipeline_id": self.cleaned_data["pipeline_id"].pk, + "key": self.cleaned_data["key"], + "batch_size": self.cleaned_data["batch_size"], + } diff --git a/ami/ml/post_processing/feature_vectors.py b/ami/ml/post_processing/feature_vectors.py new file mode 100644 index 000000000..ec3c82aa7 --- /dev/null +++ b/ami/ml/post_processing/feature_vectors.py @@ -0,0 +1,107 @@ +import functools +import operator +import typing +from urllib.parse import urljoin + +import pydantic + +from ami.main.models import DEFAULT_EMBEDDING_KEY, Detection, SourceImageCollection +from ami.ml.embeddings.reader import detections_missing_vectors +from ami.ml.embeddings.writer import save_embedding_results +from ami.ml.exceptions import PipelineNotConfigured +from ami.ml.models import Pipeline +from ami.ml.models.pipeline import process_detections +from ami.ml.post_processing.base import BasePostProcessingTask + + +class AddFeatureVectorsConfig(pydantic.BaseModel): + source_image_collection_id: int + pipeline_id: int + key: str = DEFAULT_EMBEDDING_KEY + # Captures per request. Each request carries every missing detection of its captures. + batch_size: int = pydantic.Field(default=10, ge=1, le=100) + + class Config: + extra = "forbid" + + +class AddFeatureVectorsTask(BasePostProcessingTask): + """Add feature vectors to detections that already exist, without detecting or classifying again. + + For the captures of one capture set, sends only the valid detections that lack a vector from + the pipeline's feature extractor to the processing service, and stores the returned vectors + on those detections. Running it twice sends nothing the second time. + """ + + key = "add_feature_vectors" + name = "Add feature vectors" + config_schema = AddFeatureVectorsConfig + + def _missing_detections(self, collection: SourceImageCollection, pipeline: Pipeline, key: str): + """Valid detections of the capture set that lack a vector from any of the pipeline's extractors.""" + detections = Detection.objects.valid().filter(source_image__collections=collection) + return functools.reduce( + operator.or_, + ( + detections_missing_vectors(detections, extractor.pk, key) + for extractor in pipeline.embedding_algorithms() + ), + ) + + def run(self) -> None: + config = typing.cast(AddFeatureVectorsConfig, self.config) + try: + collection = SourceImageCollection.objects.get(pk=config.source_image_collection_id) + pipeline = Pipeline.objects.get(pk=config.pipeline_id) + except (SourceImageCollection.DoesNotExist, Pipeline.DoesNotExist) as err: + self.logger.error(str(err)) + raise ValueError(str(err)) from err + if not pipeline.embedding_algorithms().exists(): + msg = f'Pipeline "{pipeline.name}" has no algorithm that produces feature vectors (task type "embedding").' + self.logger.error(msg) + raise ValueError(msg) + + missing = self._missing_detections(collection, pipeline, config.key) + capture_ids = sorted(missing.order_by().values_list("source_image_id", flat=True).distinct()) + self.logger.info( + f"=== Starting {self.name}: {len(capture_ids)} captures of capture set {collection.pk} " + f"have detections without a vector from pipeline {pipeline} ===" + ) + + if not capture_ids: + self.logger.info(f"=== Completed {self.name}: nothing to add ===") + return + processing_service = pipeline.choose_processing_service_for_pipeline( + self.job.pk if self.job else None, pipeline.name, collection.project_id + ) + if not processing_service.endpoint_url: + raise PipelineNotConfigured( + f"No endpoint URL configured for this pipeline's processing service ({processing_service})" + ) + endpoint_url = urljoin(processing_service.endpoint_url, "/process") + + totals = { + "Captures": 0, + "Detections sent": 0, + "Vectors stored": 0, + "Vectors unchanged": 0, + "Boxes unmatched": 0, + } + self.report_stage_metrics(totals) + for start in range(0, len(capture_ids), config.batch_size): + batch = list( + missing.filter(source_image_id__in=capture_ids[start : start + config.batch_size]) + .select_related("source_image__deployment__data_source", "detection_algorithm") + .order_by("source_image_id", "pk") + ) + response = process_detections(pipeline, endpoint_url, batch, collection.project_id) + result = save_embedding_results(response, self.job, pipeline, config.key) + totals["Captures"] += len({detection.source_image_id for detection in batch}) + totals["Detections sent"] += len(batch) + totals["Vectors stored"] += result.written + totals["Vectors unchanged"] += result.unchanged + totals["Boxes unmatched"] += result.unmatched + self.update_progress(min(start + config.batch_size, len(capture_ids)) / len(capture_ids)) + self.report_stage_metrics(totals) + + self.logger.info(f"=== Completed {self.name}: {totals} ===") diff --git a/ami/ml/post_processing/registry.py b/ami/ml/post_processing/registry.py index 308be18ae..c4a16db33 100644 --- a/ami/ml/post_processing/registry.py +++ b/ami/ml/post_processing/registry.py @@ -1,10 +1,12 @@ # Registry of available post-processing tasks from ami.ml.post_processing.class_masking import ClassMaskingTask +from ami.ml.post_processing.feature_vectors import AddFeatureVectorsTask from ami.ml.post_processing.small_size_filter import SmallSizeFilterTask POSTPROCESSING_TASKS = { SmallSizeFilterTask.key: SmallSizeFilterTask, ClassMaskingTask.key: ClassMaskingTask, + AddFeatureVectorsTask.key: AddFeatureVectorsTask, } diff --git a/ami/ml/post_processing/tests/test_feature_vectors.py b/ami/ml/post_processing/tests/test_feature_vectors.py new file mode 100644 index 000000000..150ccdcbc --- /dev/null +++ b/ami/ml/post_processing/tests/test_feature_vectors.py @@ -0,0 +1,373 @@ +"""Domain tests for the "Add feature vectors" post-processing task. + +The task sends detections that already exist, and only those that lack a vector, to a feature-only +processing service and stores the vectors that come back. The service is stubbed with a response in +the shape ADC's feature-only pipeline returns: each requested detection echoed unchanged, still +naming the detector, with an embedding attached and no classifications. +""" +import datetime +from unittest import mock + +from cachalot.api import cachalot_disabled +from django.contrib import admin as django_admin +from django.db import connection +from django.test import Client, TestCase +from django.test.utils import CaptureQueriesContext +from django.urls import reverse + +from ami.jobs.models import Job +from ami.main.models import Classification, Detection, Occurrence, SourceImage, SourceImageCollection +from ami.ml.exceptions import PipelineNotConfigured +from ami.ml.models import Algorithm, DetectionEmbedding, Pipeline, ProcessingService, ProjectPipelineConfig +from ami.ml.models.algorithm import AlgorithmTaskType +from ami.ml.models.pipeline import get_or_create_algorithm_and_category_map +from ami.ml.post_processing.feature_vectors import AddFeatureVectorsTask +from ami.tests.fixtures.main import no_processing_service_http, setup_test_project +from ami.tests.fixtures.ml import ALGORITHM_CHOICES +from ami.users.models import User + +DETECTOR = ALGORITHM_CHOICES["random-detector"] +LENGTH = 8 +VECTOR = [0.5] * LENGTH + + +class FakeFeaturePipelineService: + """Stands in for the processing service's synchronous ``/process`` endpoint. + + Mirrors ``run_feature_pipeline``: the requested detections come back with their box and detector + reference unchanged, an ``embeddings`` entry from the extractor, and no classifications. + """ + + def __init__(self, extractor_key: str): + self.extractor_key = extractor_key + self.requests: list[dict] = [] + self.extra_detections: list[dict] = [] + self.embedding_key_override: str | None = None + self.status_ok = True + + def post(self, url: str, json: dict): + self.requests.append(json) + key = self.embedding_key_override or self.extractor_key + detections = [ + { + "source_image_id": requested["source_image"]["id"], + "bbox": requested["bbox"], + "algorithm": requested["algorithm"], + "crop_image_url": requested["crop_image_url"], + "timestamp": datetime.datetime.now().isoformat(), + "classifications": [], + "embeddings": [{"algorithm": {"name": "Extractor", "key": key}, "features": VECTOR}], + } + for requested in json["detections"] + ] + self.extra_detections + body = { + "pipeline": json["pipeline"], + "total_time": 0.1, + "source_images": [{"id": image["id"], "url": image["url"]} for image in json["source_images"]], + "detections": detections, + } + return mock.Mock(ok=self.status_ok, json=lambda: body, status_code=200 if self.status_ok else 500) + + @property + def detections_sent(self) -> int: + return sum(len(request["detections"]) for request in self.requests) + + +class FeatureVectorsFixture: + """A project with a feature-only pipeline, its processing service stubbed, and captures with detections.""" + + @classmethod + def _set_up_data(cls) -> None: + # Creating a service checks its status over the network right away; there is no network here. + with no_processing_service_http(): + cls.project, cls.deployment = setup_test_project(reuse=False) + cls.detector = get_or_create_algorithm_and_category_map(DETECTOR) + cls.extractor = Algorithm.objects.create( + name="Extractor", key="extractor", task_type=AlgorithmTaskType.EMBEDDING.value + ) + cls.pipeline = Pipeline.objects.create(name="Feature pipeline") + cls.pipeline.algorithms.set([cls.extractor]) + cls.service = ProcessingService.objects.create( + name="Feature service", endpoint_url="http://features.test:2000", last_seen_live=True + ) + ProcessingService.objects.filter(pk=cls.service.pk).update(last_seen_live=True) + cls.service.projects.add(cls.project) + cls.service.pipelines.add(cls.pipeline) + cls.captures = 0 + + def _set_up_stubs(self) -> None: + """Per-test stand-ins for the processing service: a fake session and no status checks.""" + status_patcher = mock.patch.object(ProcessingService, "get_status") + status_patcher.start() + self.addCleanup(status_patcher.stop) # type: ignore[attr-defined] + self.fake = FakeFeaturePipelineService(self.extractor.key) + patcher = mock.patch("ami.ml.models.pipeline.create_session", return_value=self.fake) + patcher.start() + self.addCleanup(patcher.stop) # type: ignore[attr-defined] + + def _collection(self, captures: int, detections_per_capture: int = 2) -> SourceImageCollection: + images = [] + for _ in range(captures): + self.captures += 1 + image = SourceImage.objects.create( + path=f"fv-{self.captures}-20240101{self.captures:06d}.jpg", + deployment=self.deployment, + project=self.project, + public_base_url="http://images.test/", + ) + for i in range(detections_per_capture): + Detection.objects.create( + source_image=image, + bbox=[20.0 * i, 0.0, 20.0 * i + 10, 10.0], + detection_algorithm=self.detector, + ) + images.append(image) + collection = SourceImageCollection.objects.create( + name=f"fv collection {self.captures}", project=self.project, method="manual" + ) + collection.images.set(images) + return collection + + def _config(self, collection: SourceImageCollection, **extra) -> dict: + return {"source_image_collection_id": collection.pk, "pipeline_id": self.pipeline.pk, **extra} + + def _job(self, collection: SourceImageCollection, **extra) -> Job: + return Job.objects.create( + name="Add feature vectors", + project=self.project, + job_type_key="post_processing", + params={"task": AddFeatureVectorsTask.key, "config": self._config(collection, **extra)}, + ) + + def _run(self, collection: SourceImageCollection, **extra) -> Job: + job = self._job(collection, **extra) + job.run() + job.refresh_from_db() + return job + + def _metrics(self, job: Job) -> dict[str, int]: + return {param.name: param.value for param in job.progress.stages[0].params} + + +class TestAddFeatureVectorsTask(FeatureVectorsFixture, TestCase): + @classmethod + def setUpTestData(cls) -> None: + cls._set_up_data() + + def setUp(self) -> None: + self._set_up_stubs() + + def test_only_detections_missing_a_vector_are_sent(self): + collection = self._collection(captures=1, detections_per_capture=3) + stored, *missing = list(Detection.objects.filter(source_image__collections=collection).order_by("pk")) + DetectionEmbedding.objects.create(detection=stored, algorithm=self.extractor, key="embedding", vector=VECTOR) + + self._run(collection) + + sent = [d["bbox"] for request in self.fake.requests for d in request["detections"]] + self.assertCountEqual( + sent, [{"x1": m.bbox[0], "y1": m.bbox[1], "x2": m.bbox[2], "y2": m.bbox[3]} for m in missing] + ) + self.assertEqual(DetectionEmbedding.objects.filter(algorithm=self.extractor).count(), 3) + + def test_captures_with_nothing_missing_are_skipped(self): + collection = self._collection(captures=2) + done, todo = list(collection.images.order_by("pk")) + for detection in done.detections.all(): + DetectionEmbedding.objects.create( + detection=detection, algorithm=self.extractor, key="embedding", vector=VECTOR + ) + + job = self._run(collection) + + (request,) = self.fake.requests + self.assertEqual([image["id"] for image in request["source_images"]], [str(todo.pk)]) + self.assertEqual(self._metrics(job)["Captures"], 1) + + def test_vectors_are_stored_on_the_detections_they_were_sent_for_and_reported(self): + collection = self._collection(captures=3) + + job = self._run(collection) + + detections = Detection.objects.filter(source_image__collections=collection) + self.assertEqual( + set(DetectionEmbedding.objects.filter(algorithm=self.extractor).values_list("detection_id", flat=True)), + set(detections.values_list("pk", flat=True)), + ) + self.assertEqual( + self._metrics(job), + { + "Captures": 3, + "Detections sent": 6, + "Vectors stored": 6, + "Vectors unchanged": 0, + "Boxes unmatched": 0, + }, + ) + + def test_no_detections_classifications_or_occurrences_are_created(self): + collection = self._collection(captures=2) + before = (Detection.objects.count(), Classification.objects.count(), Occurrence.objects.count()) + + self._run(collection) + + self.assertEqual( + (Detection.objects.count(), Classification.objects.count(), Occurrence.objects.count()), before + ) + + def test_a_returned_box_with_no_stored_detection_is_counted_and_not_created(self): + collection = self._collection(captures=1) + image = collection.images.get() + self.fake.extra_detections = [ + { + "source_image_id": str(image.pk), + "bbox": {"x1": 500.0, "y1": 500.0, "x2": 510.0, "y2": 510.0}, + "algorithm": {"name": DETECTOR.name, "key": DETECTOR.key}, + "timestamp": datetime.datetime.now().isoformat(), + "embeddings": [{"algorithm": {"name": "Extractor", "key": self.extractor.key}, "features": VECTOR}], + } + ] + + job = self._run(collection) + + self.assertEqual(self._metrics(job)["Boxes unmatched"], 1) + self.assertEqual(Detection.objects.filter(source_image=image).count(), 2) + self.assertEqual(DetectionEmbedding.objects.count(), 2) + + def test_an_embedding_from_an_algorithm_outside_the_pipeline_raises_and_stores_nothing(self): + collection = self._collection(captures=1) + self.fake.embedding_key_override = "not-in-pipeline" + + with self.assertRaises(PipelineNotConfigured): + self._run(collection) + + self.assertEqual(DetectionEmbedding.objects.count(), 0) + + def test_a_second_run_sends_nothing(self): + collection = self._collection(captures=2) + self._run(collection) + requests_after_first = len(self.fake.requests) + + self._run(collection) + + self.assertEqual(len(self.fake.requests), requests_after_first) + self.assertEqual(DetectionEmbedding.objects.count(), 4) + + def test_a_pipeline_without_an_embedding_algorithm_is_refused_before_any_request(self): + collection = self._collection(captures=1) + self.pipeline.algorithms.set([self.detector]) + + with self.assertRaisesMessage(ValueError, "produces feature vectors"): + self._run(collection) + + self.assertEqual(self.fake.requests, []) + + def test_a_failed_request_raises_instead_of_reporting_success(self): + collection = self._collection(captures=1) + self.fake.status_ok = False + with mock.patch("ami.ml.models.pipeline.extract_error_message_from_response", return_value="boom"): + with self.assertRaisesMessage(Exception, "boom"): + self._run(collection) + self.assertEqual(DetectionEmbedding.objects.count(), 0) + + def test_requests_are_made_per_batch_of_captures(self): + collection = self._collection(captures=5) + self._run(collection, batch_size=2) + self.assertEqual([len(r["source_images"]) for r in self.fake.requests], [2, 2, 1]) + + def test_queries_do_not_grow_with_the_number_of_captures(self): + """A batch of 2 captures and a batch of 8 take the same number of queries.""" + self._run(self._collection(captures=1)) # first run creates the task's algorithm row and similar one-offs + small, large = self._collection(captures=2), self._collection(captures=8) + # Cold counts: a query the cache served for the first run would otherwise be missing from the second. + with cachalot_disabled(): + with CaptureQueriesContext(connection) as small_queries: + self._run(small, batch_size=100) + with CaptureQueriesContext(connection) as large_queries: + self._run(large, batch_size=100) + self.assertEqual(len(large_queries), len(small_queries)) + + +class TestEmbeddingOnlyPipelineGuard(FeatureVectorsFixture, TestCase): + @classmethod + def setUpTestData(cls) -> None: + cls._set_up_data() + + def setUp(self) -> None: + self._set_up_stubs() + + def test_pipeline_knows_when_it_only_produces_vectors(self): + self.assertTrue(self.pipeline.is_embedding_only()) + self.pipeline.algorithms.add(self.detector) + self.assertFalse(self.pipeline.is_embedding_only()) + self.assertFalse(Pipeline.objects.create(name="Empty").is_embedding_only()) + + def test_a_regular_ml_job_with_an_embedding_only_pipeline_fails_early_pointing_to_the_task(self): + collection = self._collection(captures=1) + job = Job.objects.create( + name="ML job", + project=self.project, + job_type_key="ml", + pipeline=self.pipeline, + source_image_collection=collection, + ) + with self.assertRaisesMessage(PipelineNotConfigured, "Add feature vectors"): + job.run() + self.assertEqual(self.fake.requests, []) + + +class TestAddFeatureVectorsAdmin(FeatureVectorsFixture, TestCase): + @classmethod + def setUpTestData(cls) -> None: + cls._set_up_data() + cls.superuser = User.objects.create_superuser(email="afv-admin@example.com", password="x") + ProjectPipelineConfig.objects.create(project=cls.project, pipeline=cls.pipeline, enabled=True) + cls.other_pipeline = Pipeline.objects.create(name="Classifier only pipeline") + cls.other_pipeline.algorithms.set([cls.detector]) + + def setUp(self) -> None: + self._set_up_stubs() + self.client = Client() + self.client.force_login(self.superuser) + + def _post(self, collections: list[SourceImageCollection], data: dict): + return self.client.post( + reverse("admin:main_sourceimagecollection_changelist"), + { + "action": "run_add_feature_vectors", + django_admin.helpers.ACTION_CHECKBOX_NAME: [str(c.pk) for c in collections], + **data, + }, + ) + + def test_the_form_lists_only_pipelines_with_an_embedding_algorithm(self): + response = self._post([self._collection(captures=1)], {}) + self.assertEqual(response.status_code, 200) + self.assertContains(response, "Run Add feature vectors") + self.assertContains(response, "Feature pipeline") + self.assertNotContains(response, "Classifier only pipeline") + + def test_it_enqueues_one_job_per_capture_set_with_the_config_on_the_job(self): + first, second = self._collection(captures=1), self._collection(captures=1) + + response = self._post( + [first, second], {"confirm": "1", "pipeline_id": self.pipeline.pk, "key": "embedding", "batch_size": 5} + ) + + self.assertEqual(response.status_code, 302) + jobs = Job.objects.filter(job_type_key="post_processing", params__task="add_feature_vectors") + self.assertEqual(jobs.count(), 2) + self.assertCountEqual( + [job.params["config"]["source_image_collection_id"] for job in jobs], [first.pk, second.pk] + ) + for job in jobs: + self.assertEqual(job.params["config"]["pipeline_id"], self.pipeline.pk) + self.assertEqual(job.params["config"]["batch_size"], 5) + + def test_an_out_of_range_batch_size_is_shown_on_the_form_and_creates_no_job(self): + response = self._post( + [self._collection(captures=1)], + {"confirm": "1", "pipeline_id": self.pipeline.pk, "key": "embedding", "batch_size": 0}, + ) + self.assertEqual(response.status_code, 200) + self.assertEqual(Job.objects.filter(params__task="add_feature_vectors").count(), 0) diff --git a/ami/ml/schemas.py b/ami/ml/schemas.py index 9322e4116..8df39267d 100644 --- a/ami/ml/schemas.py +++ b/ami/ml/schemas.py @@ -168,6 +168,23 @@ class DetectionRequest(pydantic.BaseModel): algorithm: AlgorithmReference +class EmbeddingResponse(pydantic.BaseModel): + """A feature vector for one detection, and the algorithm whose backbone produced it. + + It is carried on the detection rather than on a classification, so storing it can never + add a prediction. Only vectors from the same algorithm are comparable. + """ + + features: list[float] = pydantic.Field(description="The feature vector.") + algorithm: AlgorithmReference + + @pydantic.validator("features") + def _features_not_empty(cls, v): + if not v: + raise ValueError("features must contain at least one value") + return v + + class DetectionResponse(pydantic.BaseModel): source_image_id: str bbox: BoundingBox | None = None @@ -176,6 +193,13 @@ class DetectionResponse(pydantic.BaseModel): timestamp: datetime.datetime crop_image_url: str | None = None classifications: list[ClassificationResponse] = [] + embeddings: list[EmbeddingResponse] | None = pydantic.Field( + default=None, + description=( + "Feature vectors for this detection, at most one per algorithm, including detections " + "the moth/non-moth filter rejected. Only vectors from the same algorithm are comparable." + ), + ) class PipelineRequestConfigParameters(dict): diff --git a/ami/ml/test_detection_embeddings.py b/ami/ml/test_detection_embeddings.py new file mode 100644 index 000000000..e62cc7113 --- /dev/null +++ b/ami/ml/test_detection_embeddings.py @@ -0,0 +1,634 @@ +"""How detection feature vectors are stored and read back. + +These pin that a vector is stored once per detection, algorithm and key on the detection it +was sent with, never as a classification; that its project is copied from the detection's +capture; that writes are insert-mostly; and that readers return one (algorithm, key) at a time. +""" + +import contextlib +import datetime + +import numpy as np +import pydantic +from django.db import connection +from django.test import SimpleTestCase, TestCase +from django.test.utils import CaptureQueriesContext + +from ami.jobs.models import Job +from ami.main.models import Classification, Deployment, Detection, SourceImage +from ami.ml.embeddings.reader import ( + algorithm_with_most_vectors, + detections_missing_vectors, + project_vectors, + vector_counts_by_algorithm, + vectors_for_detections, +) +from ami.ml.embeddings.writer import EmbeddingDimensionMismatch, create_detection_embeddings +from ami.ml.exceptions import PipelineNotConfigured +from ami.ml.models import Algorithm, DetectionEmbedding, Pipeline +from ami.ml.models.pipeline import get_or_create_algorithm_and_category_map, save_results +from ami.ml.schemas import DetectionResponse, PipelineResultsResponse +from ami.tests.fixtures.main import no_processing_service_http, setup_test_project +from ami.tests.fixtures.ml import ALGORITHM_CHOICES + +DETECTOR = ALGORITHM_CHOICES["random-detector"] +BINARY = ALGORITHM_CHOICES["random-binary-classifier"] # labels: "Moth", "Not a moth" +SPECIES = ALGORITHM_CHOICES["random-species-classifier"] +LENGTH = 2048 + + +def _embedding_payload(vector: list[float], algorithm=SPECIES) -> list[dict]: + return [{"algorithm": {"name": algorithm.name, "key": algorithm.key}, "features": vector}] + + +class TestEmbeddingSchema(SimpleTestCase): + """The per-detection ``embeddings`` field of the processing-service results schema.""" + + def _detection(self, **extra) -> dict: + return { + "source_image_id": "1", + "bbox": {"x1": 0.0, "y1": 0.0, "x2": 10.0, "y2": 10.0}, + "algorithm": {"name": DETECTOR.name, "key": DETECTOR.key}, + "timestamp": datetime.datetime.now().isoformat(), + **extra, + } + + def test_a_detection_carries_each_vector_with_its_algorithm(self): + """The field is optional, so a service that sends no embeddings still parses.""" + parsed = DetectionResponse.parse_obj(self._detection(embeddings=_embedding_payload([0.5] * LENGTH))) + self.assertEqual( + [(e.algorithm.key, len(e.features)) for e in parsed.embeddings or []], [(SPECIES.key, LENGTH)] + ) + self.assertIsNone(DetectionResponse.parse_obj(self._detection()).embeddings) + + def test_an_empty_vector_is_refused(self): + """Any length parses (extractors differ; each algorithm's length is checked on save), but not none.""" + parsed = DetectionResponse.parse_obj(self._detection(embeddings=_embedding_payload([0.5] * 512))) + self.assertEqual(len(parsed.embeddings or []), 1) + with self.assertRaises(pydantic.ValidationError): + DetectionResponse.parse_obj(self._detection(embeddings=_embedding_payload([]))) + + +class ClassifierPipelineMixin: + """A classifier pipeline whose results carry a vector on each detection, as the processing service sends them.""" + + LOW = [0.25] * LENGTH + HIGH = [0.75] * LENGTH + + @classmethod + def _set_up_pipeline(cls) -> None: + with no_processing_service_http(): + cls.project, cls.deployment = setup_test_project(reuse=False) + cls.pipeline = Pipeline.objects.create(name="Embedding test pipeline") + cls.pipeline.algorithms.set( + [get_or_create_algorithm_and_category_map(algorithm) for algorithm in (DETECTOR, BINARY, SPECIES)] + ) + cls.species = Algorithm.objects.get(key=SPECIES.key) + cls.images = 0 + + def _image(self) -> SourceImage: + self.images += 1 + return SourceImage.objects.create( + path=f"emb-{self.images}-20240101000{self.images}00.jpg", deployment=self.deployment, project=self.project + ) + + @staticmethod + def _classification(algorithm, label: str, score: float, terminal: bool) -> dict: + return { + "classification": label, + "scores": [score], + "algorithm": {"name": algorithm.name, "key": algorithm.key}, + "terminal": terminal, + "timestamp": datetime.datetime.now().isoformat(), + } + + def _detection(self, image: SourceImage, classifications: list[dict], embeddings=None, box: float = 0.0) -> dict: + payload = { + "source_image_id": str(image.pk), + "bbox": {"x1": box, "y1": box, "x2": box + 10.0, "y2": box + 10.0}, + "algorithm": {"name": DETECTOR.name, "key": DETECTOR.key}, + "timestamp": datetime.datetime.now().isoformat(), + "classifications": classifications, + } + if embeddings is not None: + payload["embeddings"] = embeddings + return payload + + def _rejected(self, image: SourceImage, embeddings=None, box: float = 0.0) -> dict: + """A crop the moth/non-moth filter rejected, labelled as the service labels it: non-terminal.""" + label = self._classification(BINARY, "Not a moth", 0.8, terminal=False) + return self._detection(image, [label], embeddings, box) + + def _moth(self, image: SourceImage, embeddings=None, box: float = 0.0) -> dict: + labels = [ + self._classification(BINARY, "Moth", 0.95, terminal=False), + self._classification(SPECIES, "Vanessa cardui", 0.6, terminal=True), + ] + return self._detection(image, labels, embeddings, box) + + def _save(self, *detections: dict, job_id: int | None = None) -> None: + image_ids = list(dict.fromkeys(d["source_image_id"] for d in detections)) + payload = { + "pipeline": self.pipeline.slug, + "total_time": 0.01, + "source_images": [{"id": image_id, "url": f"test/{image_id}.jpg"} for image_id in image_ids], + "detections": list(detections), + } + save_results(PipelineResultsResponse.parse_obj(payload), job_id=job_id) + + @staticmethod + def _stored(image: SourceImage) -> dict[tuple[float, str], list[float]]: + """{(box corner, algorithm key): vector} for the image's stored embeddings.""" + rows = DetectionEmbedding.objects.filter(detection__source_image=image).select_related( + "detection", "algorithm" + ) + return {(row.detection.bbox[0], row.algorithm.key): list(row.vector) for row in rows} + + +class TestDetectionEmbeddings(ClassifierPipelineMixin, TestCase): + """Storing the feature vector a processing service sends with each detection. + + The vector is what similarity search and tracking compare, and it arrives for every + detection, including those the moth/non-moth filter rejected. What these pin is that it + is stored once per detection and algorithm, on the detection it was sent with, and that + storing it never adds a classification or moves a determination. + """ + + @classmethod + def setUpTestData(cls) -> None: + cls._set_up_pipeline() + + def test_every_detection_stores_one_vector_per_algorithm(self): + """Including the rejected crop, which has no species classification that could carry one.""" + image = self._image() + self._save( + self._moth(image, _embedding_payload(self.LOW)), + self._rejected(image, _embedding_payload(self.HIGH), box=100.0), + ) + self.assertEqual(self._stored(image), {(0.0, SPECIES.key): self.LOW, (100.0, SPECIES.key): self.HIGH}) + self.assertEqual(set(DetectionEmbedding.objects.values_list("project_id", flat=True)), {self.project.pk}) + + def test_a_vector_adds_no_classification_and_moves_no_determination(self): + """The rejected crop's binary label is non-terminal, so a species classification sent in + the vector's place would become its determination. An embedding must not.""" + control, treated = self._image(), self._image() + self._save(self._rejected(control)) + self._save(self._rejected(treated, _embedding_payload(self.HIGH))) + + for image in (control, treated): + detection = Detection.objects.select_related("occurrence__determination").get(source_image=image) + self.assertEqual(detection.occurrence.determination.name, "Not a moth") + self.assertEqual(detection.classifications.count(), 1) + self.assertEqual(DetectionEmbedding.objects.filter(detection__source_image=treated).count(), 1) + + def test_each_row_records_the_job_that_saved_it_and_outlives_the_job(self): + image = self._image() + first, second = ( + Job.objects.create(project=self.project, name=f"Embedding job {n}", pipeline=self.pipeline) for n in (1, 2) + ) + self._save(self._moth(image, _embedding_payload(self.LOW)), job_id=first.pk) + self._save(self._moth(image, _embedding_payload(self.HIGH)), job_id=second.pk) + + detection = Detection.objects.get(source_image=image) + embedding = DetectionEmbedding.objects.get(detection=detection) + self.assertEqual(embedding.job_id, second.pk, "A replaced vector records the job that replaced it") + + second.delete() + embedding.refresh_from_db() + self.assertIsNone(embedding.job_id) + self.assertEqual(list(embedding.vector), self.HIGH) + + def test_a_vector_lands_on_its_own_detection_when_some_detections_already_exist(self): + """Detection creation returns existing detections ahead of new ones, so pairing responses + with detections by position would swap these two vectors.""" + image = self._image() + self._save(self._rejected(image, box=0.0)) + self._save( + self._rejected(image, _embedding_payload(self.HIGH), box=100.0), + self._rejected(image, _embedding_payload(self.LOW), box=0.0), + ) + self.assertEqual(self._stored(image), {(0.0, SPECIES.key): self.LOW, (100.0, SPECIES.key): self.HIGH}) + + def test_a_vector_from_an_unregistered_algorithm_stops_the_batch_like_a_classification(self): + """It raises where an unregistered classification algorithm does: after detections are + saved and before any classification is.""" + image = self._image() + unregistered = {"algorithm": {"name": "Unregistered", "key": "unregistered-embedder"}, "features": self.LOW} + with self.assertRaises(PipelineNotConfigured): + self._save(self._rejected(image, [unregistered])) + self.assertFalse(DetectionEmbedding.objects.exists()) + self.assertFalse(Classification.objects.filter(detection__source_image=image).exists()) + + def test_saving_the_same_results_twice_changes_nothing_and_a_new_vector_replaces_the_row(self): + image = self._image() + self._save(self._moth(image, _embedding_payload(self.LOW))) + first = DetectionEmbedding.objects.get() + + self._save(self._moth(image, _embedding_payload(self.LOW))) + unchanged = DetectionEmbedding.objects.get() + self.assertEqual((unchanged.pk, unchanged.updated_at), (first.pk, first.updated_at)) + + self._save(self._moth(image, _embedding_payload(self.HIGH))) + replaced = DetectionEmbedding.objects.get() + self.assertEqual(replaced.pk, first.pk, "The row is updated in place, not deleted and re-inserted") + self.assertEqual(list(replaced.vector), self.HIGH) + self.assertGreater(replaced.updated_at, first.updated_at) + + def test_an_algorithm_and_key_keep_the_length_of_their_first_vector(self): + """Vectors of another length are refused, because they could never be compared.""" + image, other = self._image(), self._image() + self._save(self._moth(image, _embedding_payload([0.5] * 512))) + + with self.assertRaises(EmbeddingDimensionMismatch) as raised: + self._save(self._moth(other, _embedding_payload(self.LOW))) + self.assertIn(self.species.key, str(raised.exception)) + self.assertIn("512", str(raised.exception)) + self.assertIn(str(LENGTH), str(raised.exception)) + self.assertFalse(DetectionEmbedding.objects.filter(detection__source_image=other).exists()) + + def test_a_batch_with_no_stored_vector_must_agree_on_one_length(self): + image = self._image() + with self.assertRaises(EmbeddingDimensionMismatch): + self._save( + self._rejected(image, _embedding_payload([0.5] * 8), box=0.0), + self._rejected(image, _embedding_payload([0.5] * 4), box=100.0), + ) + self.assertFalse(DetectionEmbedding.objects.exists()) + + def test_a_vector_half_precision_cannot_hold_is_skipped_with_a_warning(self): + image = self._image() + too_large = [1e6] + [0.5] * (LENGTH - 1) + not_a_number = [float("nan")] + [0.5] * (LENGTH - 1) + with self.assertLogs("ami.ml.models.pipeline", level="WARNING") as logs: + self._save( + self._rejected(image, _embedding_payload(too_large), box=0.0), + self._rejected(image, _embedding_payload(not_a_number), box=100.0), + self._rejected(image, _embedding_payload(self.LOW), box=200.0), + ) + self.assertIn("Skipped 2 vectors", "\n".join(logs.output)) + self.assertEqual(self._stored(image), {(200.0, SPECIES.key): self.LOW}) + + def test_storing_vectors_takes_the_same_queries_however_many_detections(self): + image = self._image() + boxes = [float(20 * i) for i in range(5)] + self._save(*[self._rejected(image, box=box) for box in boxes]) + detections = list(Detection.objects.filter(source_image=image).order_by("bbox")) + parsed = [ + DetectionResponse.parse_obj(self._rejected(image, _embedding_payload(self.LOW), box=b)) for b in boxes + ] + algorithms_known = {algorithm.key: algorithm for algorithm in self.pipeline.algorithms.all()} + + # Measure after a first write, on other detections so no read is served from the query cache. + create_detection_embeddings(detections[:1], parsed[:1], algorithms_known) + DetectionEmbedding.objects.all().delete() + with CaptureQueriesContext(connection) as one: + create_detection_embeddings(detections[1:2], parsed[1:2], algorithms_known) + DetectionEmbedding.objects.all().delete() + with CaptureQueriesContext(connection) as five: + stored = create_detection_embeddings(detections, parsed, algorithms_known) + self.assertEqual(len(five), len(one)) + self.assertEqual(len(stored), 5) + + # A rerun with identical vectors reads the length and the stored rows and writes nothing. + with self.assertNumQueries(2): + create_detection_embeddings(detections, parsed, algorithms_known) + + def test_a_vector_matches_its_box_exactly_after_a_json_round_trip(self): + image = self._image() + self._save(self._rejected(image, box=0.1234567)) + sent = DetectionResponse.parse_obj(self._rejected(image, _embedding_payload(self.LOW), box=0.1234567)) + round_tripped = DetectionResponse.parse_raw(sent.json()) + detections = list(Detection.objects.filter(source_image=image)) + algorithms_known = {algorithm.key: algorithm for algorithm in self.pipeline.algorithms.all()} + + stored = create_detection_embeddings(detections, [round_tripped], algorithms_known) + self.assertEqual([embedding.detection_id for embedding in stored], [detections[0].pk]) + + def test_boxes_closer_than_three_decimals_each_get_their_own_vector(self): + image = self._image() + near, nearer = 10.0001, 10.0002 + self._save(self._rejected(image, box=near), self._rejected(image, box=nearer)) + self._save( + self._rejected(image, _embedding_payload(self.LOW), box=near), + self._rejected(image, _embedding_payload(self.HIGH), box=nearer), + ) + self.assertEqual(self._stored(image), {(near, SPECIES.key): self.LOW, (nearer, SPECIES.key): self.HIGH}) + + def test_a_box_stored_twice_is_skipped_and_logged_not_guessed(self): + image = self._image() + self._save(self._rejected(image, box=5.0)) + original = Detection.objects.get(source_image=image) + Detection.objects.create( + source_image=image, bbox=original.bbox, detection_algorithm=original.detection_algorithm + ) + sent = DetectionResponse.parse_obj(self._rejected(image, _embedding_payload(self.LOW), box=5.0)) + algorithms_known = {algorithm.key: algorithm for algorithm in self.pipeline.algorithms.all()} + + with self.assertLogs("ami.ml.embeddings.writer", level="WARNING") as logs: + stored = create_detection_embeddings( + list(Detection.objects.filter(source_image=image)), [sent], algorithms_known + ) + self.assertEqual(stored, []) + self.assertFalse(DetectionEmbedding.objects.exists()) + self.assertIn(f"capture {image.pk}", "\n".join(logs.output)) + + def test_only_the_first_write_of_an_algorithm_and_key_takes_the_advisory_lock(self): + image, other = self._image(), self._image() + self._save(self._rejected(image), self._rejected(other)) + + def locks(queries) -> int: + return sum("pg_advisory_xact_lock" in query["sql"] for query in queries) + + with CaptureQueriesContext(connection) as first: + self._save(self._rejected(image, _embedding_payload(self.LOW))) + with CaptureQueriesContext(connection) as later: + self._save(self._rejected(other, _embedding_payload(self.LOW))) + self.assertEqual(locks(first.captured_queries), 1) + self.assertEqual(locks(later.captured_queries), 0) + + def test_a_second_batch_of_another_length_is_refused_after_the_first_commits(self): + image, other = self._image(), self._image() + self._save(self._moth(image, _embedding_payload(self.LOW))) + with self.assertRaises(EmbeddingDimensionMismatch): + self._save(self._moth(other, _embedding_payload([0.5] * 8))) + + +class TestEmbeddingProject(TestCase): + """An embedding's project is its capture's, falling back to the capture's station's.""" + + @classmethod + def setUpTestData(cls) -> None: + with no_processing_service_http(): + cls.project, cls.deployment = setup_test_project(reuse=False) + cls.image = SourceImage.objects.create(path="project-fill.jpg", deployment=cls.deployment, project=cls.project) + cls.detection = Detection.objects.create(source_image=cls.image, bbox=[0.0, 0.0, 10.0, 10.0]) + cls.algorithm = Algorithm.objects.create(name="Backbone", key="backbone", task_type="embedding") + + def _embedding(self, **fields) -> DetectionEmbedding: + return DetectionEmbedding(detection=self.detection, algorithm=self.algorithm, vector=[0.5] * 4, **fields) + + def test_the_captures_project_is_copied_on_save_and_on_store(self): + self._embedding().save() + self.assertEqual(DetectionEmbedding.objects.get().project_id, self.project.pk) + + DetectionEmbedding.objects.all().delete() + DetectionEmbedding.objects.store([self._embedding()]) + self.assertEqual(DetectionEmbedding.objects.get().project_id, self.project.pk) + + def test_a_capture_without_a_project_falls_back_to_its_station(self): + SourceImage.objects.filter(pk=self.image.pk).update(project=None) + DetectionEmbedding.objects.store([self._embedding()]) + self.assertEqual(DetectionEmbedding.objects.get().project_id, self.project.pk) + + def test_no_project_anywhere_refuses_the_write(self): + SourceImage.objects.filter(pk=self.image.pk).update(project=None) + Deployment.objects.filter(pk=self.deployment.pk).update(project=None) + with self.assertRaises(ValueError): + DetectionEmbedding.objects.store([self._embedding()]) + with self.assertRaises(ValueError): + self._embedding().save() + self.assertFalse(DetectionEmbedding.objects.exists()) + + +class TestEmbeddingReaders(TestCase): + """Readers return vectors of one (algorithm, key) at a time, as float32 arrays.""" + + @classmethod + def setUpTestData(cls) -> None: + with no_processing_service_http(): + cls.project, cls.deployment = setup_test_project(reuse=False) + image = SourceImage.objects.create(path="readers.jpg", deployment=cls.deployment, project=cls.project) + cls.detections = [ + Detection.objects.create(source_image=image, bbox=[float(i), 0.0, float(i) + 10.0, 10.0]) for i in range(3) + ] + cls.backbone = Algorithm.objects.create(name="Backbone", key="backbone", task_type="embedding") + cls.other = Algorithm.objects.create(name="Other backbone", key="other-backbone", task_type="embedding") + DetectionEmbedding.objects.store( + [ + DetectionEmbedding(detection=cls.detections[0], algorithm=cls.backbone, vector=[0.25] * 4), + DetectionEmbedding(detection=cls.detections[1], algorithm=cls.backbone, vector=[0.5] * 4), + DetectionEmbedding(detection=cls.detections[0], algorithm=cls.other, vector=[0.75] * 8), + DetectionEmbedding( + detection=cls.detections[2], algorithm=cls.backbone, key="projection", vector=[1.0] * 2 + ), + ] + ) + + def test_vectors_come_back_for_one_algorithm_and_key_only(self): + ids = [detection.pk for detection in self.detections] + with self.assertNumQueries(1): + vectors = vectors_for_detections(ids, self.backbone.pk) + self.assertEqual(sorted(vectors), ids[:2]) + self.assertEqual({v.shape for v in vectors.values()}, {(4,)}) + self.assertEqual({v.dtype for v in vectors.values()}, {np.dtype(np.float32)}) + self.assertEqual(float(vectors[ids[1]][0]), 0.5) + + self.assertEqual(list(vectors_for_detections(ids, self.other.pk)), ids[:1]) + self.assertEqual(list(vectors_for_detections(ids, self.backbone.pk, key="projection")), ids[2:]) + + +@contextlib.contextmanager +def cache_off(): + """Run a block with cachalot off, restoring it even when the block raises, so query counts are real.""" + from cachalot.api import cachalot_disabled + + disabled = cachalot_disabled() + disabled.__enter__() + try: + yield + finally: + disabled.__exit__(None, None, None) + + +class TestQueryHelpers(TestCase): + """The functions in ``ami.ml.embeddings.reader``: each reads one (algorithm, key) in a fixed number of queries. + + The fixture mixes two models of different lengths and a second key of one model with a third + length, so a helper that merged pairs would return the wrong lengths or counts. + """ + + @classmethod + def setUpTestData(cls) -> None: + with no_processing_service_http(): + cls.project, cls.deployment = setup_test_project(reuse=False) + cls.image = SourceImage.objects.create(path="helpers.jpg", deployment=cls.deployment, project=cls.project) + cls.eight = Algorithm.objects.create(name="Eight", key="eight", task_type="embedding") + cls.four = Algorithm.objects.create(name="Four", key="four", task_type="embedding") + cls.detections: list[Detection] = [] + + def _detections(self, count: int) -> list[Detection]: + made = [ + Detection.objects.create( + source_image=self.image, bbox=[float(len(self.detections) + i), 0.0, 500.0 + i, 10.0] + ) + for i in range(count) + ] + self.detections += made + return made + + def _store(self, detections, algorithm, length, key="embedding"): + DetectionEmbedding.objects.store( + [ + DetectionEmbedding(detection=d, algorithm=algorithm, key=key, vector=[float(i % 7) / 4] * length) + for i, d in enumerate(detections) + ] + ) + + def _fill(self, count: int): + """``count`` detections with an 8-d vector from one model and a 4-d one from another, the + first of them also with a 2-d vector under a second key of the first model.""" + detections = self._detections(count) + self._store(detections, self.eight, 8) + self._store(detections[::2], self.four, 4) + self._store(detections[:1], self.eight, 2, key="projection") + return detections + + def test_an_algorithm_may_hold_a_different_length_under_each_key(self): + detections = self._fill(3) + ids = [d.pk for d in detections] + self.assertEqual({v.shape for v in vectors_for_detections(ids, self.eight.pk).values()}, {(8,)}) + projection = vectors_for_detections(ids, self.eight.pk, key="projection") + self.assertEqual({k: v.shape for k, v in projection.items()}, {ids[0]: (2,)}) + + def test_vectors_for_detections_takes_one_query_at_any_size(self): + counts = [] + for size in (3, 12): + ids = [d.pk for d in self._fill(size)] + with cache_off(), CaptureQueriesContext(connection) as queries: + vectors = vectors_for_detections(ids, self.eight.pk) + self.assertEqual(len(vectors), size) + counts.append(len(queries)) + self.assertEqual(counts, [1, 1]) + + def test_project_vectors_yields_every_row_once_in_id_order_in_bounded_chunks(self): + detections = self._fill(7) + chunks = list(project_vectors(self.project.pk, self.eight.pk, chunk_size=3)) + self.assertEqual([len(ids) for ids, _ in chunks], [3, 3, 1]) + all_ids = [i for ids, _ in chunks for i in ids] + self.assertEqual(all_ids, sorted(d.pk for d in detections)) + for ids, array in chunks: + self.assertEqual(array.shape, (len(ids), 8)) + self.assertEqual(array.dtype, np.dtype(np.float32)) + # A chunk size that divides the rows exactly ends without an empty chunk. + self.assertEqual([len(i) for i, _ in project_vectors(self.project.pk, self.eight.pk, chunk_size=7)], [7]) + self.assertEqual(list(project_vectors(self.project.pk, self.eight.pk, key="missing")), []) + + def test_project_vectors_keeps_models_apart_and_honours_the_id_scope(self): + detections = self._fill(6) + ((ids, array),) = project_vectors(self.project.pk, self.four.pk) + self.assertEqual(ids, [d.pk for d in detections[::2]]) + self.assertEqual(array.shape, (3, 4)) + wanted = [detections[0].pk, detections[1].pk, detections[2].pk] + ((ids, _),) = project_vectors(self.project.pk, self.four.pk, detection_ids=wanted) + self.assertEqual(ids, [detections[0].pk, detections[2].pk]) + + def test_project_vectors_takes_one_query_per_chunk_plus_the_empty_end(self): + counts = [] + for size in (6, 12): + self._fill(size) + with cache_off(), CaptureQueriesContext(connection) as queries: + chunks = list(project_vectors(self.project.pk, self.four.pk, chunk_size=1000)) + self.assertEqual(len(chunks), 1) + counts.append(len(queries)) + # One chunk, so a read for it and one that finds nothing after it; the size does not matter. + self.assertEqual(counts, [2, 2]) + + def test_counts_are_per_algorithm_and_key(self): + self._fill(5) + pk = self.project.pk + self.assertEqual( + vector_counts_by_algorithm(pk), + {(self.eight.pk, "embedding"): 5, (self.four.pk, "embedding"): 3, (self.eight.pk, "projection"): 1}, + ) + self.assertEqual( + vector_counts_by_algorithm(pk, key="embedding"), + {(self.eight.pk, "embedding"): 5, (self.four.pk, "embedding"): 3}, + ) + self.assertEqual(vector_counts_by_algorithm(pk, key="nothing"), {}) + self.assertEqual(algorithm_with_most_vectors(self.project), self.eight.pk) + + def test_the_default_algorithm_breaks_ties_by_lowest_id(self): + detections = self._detections(2) + self._store(detections, self.four, 4) + self._store(detections, self.eight, 8) + self.assertEqual(algorithm_with_most_vectors(self.project), min(self.four.pk, self.eight.pk)) + + def test_counts_take_one_query_at_any_size(self): + counts = [] + for size in (3, 12): + self._fill(size) + with cache_off(), CaptureQueriesContext(connection) as queries: + vector_counts_by_algorithm(self.project.pk) + algorithm_with_most_vectors(self.project) + counts.append(len(queries)) + self.assertEqual(counts, [2, 2]) + + def test_detections_missing_vectors_keeps_the_callers_scope(self): + detections = self._fill(4) + scope = Detection.objects.filter(source_image=self.image) + missing = detections_missing_vectors(scope, self.four.pk) + self.assertEqual({d.pk for d in missing}, {detections[1].pk, detections[3].pk}) + self.assertEqual(list(detections_missing_vectors(scope, self.eight.pk)), []) + self.assertEqual( + {d.pk for d in detections_missing_vectors(scope.filter(pk__lte=detections[1].pk), self.four.pk)}, + {detections[1].pk}, + ) + self.assertEqual(detections_missing_vectors(scope, self.eight.pk, key="projection").count(), 3) + + def test_detections_missing_vectors_is_one_query_at_any_size(self): + counts = [] + for size in (3, 12): + self._fill(size) + scope = Detection.objects.filter(source_image=self.image) + with cache_off(), CaptureQueriesContext(connection) as queries: + list(detections_missing_vectors(scope, self.four.pk)) + counts.append(len(queries)) + self.assertEqual(counts, [1, 1]) + + def test_detections_missing_vectors_is_invalidated_by_new_vectors(self): + """The query cache must know the filter reads the vector table. + + django-cachalot does not see tables inside a negated ``Exists`` in ``filter()``; if the + filter is written that way, a stored vector leaves the cached "missing" result unchanged + and the next feature-vector run sends the same detections again. + """ + from cachalot.utils import _get_tables + + query = detections_missing_vectors(Detection.objects.all(), self.four.pk).query + self.assertIn(DetectionEmbedding._meta.db_table, _get_tables(connection.alias, query)) + + +class TestPgvectorGuard(SimpleTestCase): + """The extension migration stops with one clear message unless the server offers pgvector 0.8+.""" + + @staticmethod + def _run(row): + import importlib + from unittest import mock + + migration = importlib.import_module("ami.ml.migrations.0029_enable_pgvector") + cursor = mock.MagicMock() + cursor.__enter__.return_value.fetchone.return_value = row + schema_editor = mock.Mock() + schema_editor.connection.cursor.return_value = cursor + migration.check_pgvector_is_installed(apps=None, schema_editor=schema_editor) + + def test_a_missing_or_old_package_is_refused_and_a_current_one_passes(self): + with self.assertRaisesRegex(RuntimeError, "not installed"): + self._run(None) + with self.assertRaisesRegex(RuntimeError, "found 0.7.4"): + self._run(("0.7.4",)) + self._run(("0.8.6",)) + self._run(("1.0.0",)) + + +class TestEmbeddingColumn(TestCase): + def test_the_vector_column_is_an_unsized_halfvec_stored_out_of_line_uncompressed(self): + with connection.cursor() as cursor: + cursor.execute( + """ + SELECT format_type(atttypid, atttypmod), attstorage FROM pg_attribute + WHERE attrelid = 'ml_detectionembedding'::regclass AND attname = 'vector' + """ + ) + self.assertEqual(cursor.fetchone(), ("halfvec", "e")) diff --git a/ami/tests/fixtures/main.py b/ami/tests/fixtures/main.py index 073f00099..939a33d35 100644 --- a/ami/tests/fixtures/main.py +++ b/ami/tests/fixtures/main.py @@ -1,9 +1,11 @@ +import contextlib import datetime import logging import os import pathlib import random import uuid +from unittest import mock from django.db import transaction from django.utils import timezone @@ -39,6 +41,13 @@ def update_site_settings(**kwargs): return site +@contextlib.contextmanager +def no_processing_service_http(): + """Skip the status check and pipeline registration that fixtures send to the test processing service.""" + with mock.patch.object(ProcessingService, "get_status"), mock.patch.object(ProcessingService, "create_pipelines"): + yield + + def create_processing_service(project: Project, name: str = "Test Processing Service") -> ProcessingService: processing_service_to_add = { "name": name, diff --git a/compose/local/postgres/Dockerfile b/compose/local/postgres/Dockerfile index 5f864a4a0..6f3a6881d 100644 --- a/compose/local/postgres/Dockerfile +++ b/compose/local/postgres/Dockerfile @@ -1,6 +1,12 @@ FROM postgres:16 # FROM esgn/pgtuned:latest +# pgvector stores feature vectors (DetectionEmbedding). Migration ml/0029 enforces the 0.8 floor, +# so the package is not pinned and a newer PGDG release does not break the build. +RUN apt-get update \ + && apt-get install -y --no-install-recommends postgresql-16-pgvector \ + && rm -rf /var/lib/apt/lists/* + COPY ./compose/local/postgres/maintenance /usr/local/bin/maintenance RUN chmod +x /usr/local/bin/maintenance/* RUN mv /usr/local/bin/maintenance/* /usr/local/bin \ diff --git a/docs/claude/INDEX.md b/docs/claude/INDEX.md index 61f0980d8..89cfa14cf 100644 --- a/docs/claude/INDEX.md +++ b/docs/claude/INDEX.md @@ -11,6 +11,7 @@ archived. |---|---| | `reference/canonical-patterns.md` | Existing helpers/patterns to reuse before writing new ones, with file:line refs (SingleParamSerializer, ProjectMixin, permissions, schemas, fixtures). Keywords: reuse, helpers, conventions, DRF | | `reference/query-patterns.md` | DB model relationship table, composite indexes, prefetch/select_related patterns, full custom QuerySet method catalog, query anti-patterns. Keywords: N+1, indexes, ORM, performance | +| `reference/feature-vectors.md` | How detection feature vectors are stored and queried: one row per (detection, algorithm, key), the one-model-per-query rule, query patterns Q1-Q6 with functions and indexes, anti-patterns, logits, HNSW, precision. Keywords: embeddings, pgvector, halfvec, similarity | | `reference/api-stats-pattern.md` | How to add aggregate/leaderboard/chart endpoints (`//stats//`): GenericViewSet + @action, pure querysets in models_future. Keywords: stats, charts, aggregation | | `reference/monitoring-async-jobs.md` | Monitoring/debugging async_api (NATS JetStream) jobs: ORM, REST, consumer state, Redis counters, worker logs. Keywords: NATS, async, jobs, monitoring | | `reference/worktree-testing.md` | Testing git-worktree changes against the main Docker stack: bind-mount route vs duplicate-stack route, caveats, cleanup. Keywords: worktree, docker compose, override | diff --git a/docs/claude/reference/canonical-patterns.md b/docs/claude/reference/canonical-patterns.md index 5cd217a78..3ae7cc2cd 100644 --- a/docs/claude/reference/canonical-patterns.md +++ b/docs/claude/reference/canonical-patterns.md @@ -22,6 +22,8 @@ canonical but was undocumented. | Query-count assertion | `ami/ml/tests.py:1006` | `assertNumQueries` example — pair with a multi-row fixture | | Permission matrix test | `ami/main/tests.py:1532-1590` | Template: owner / member / other-user tests asserting 200/403 per endpoint action | | `update_calculated_fields()` | `ami/base/models.py:165` (usage: `ami/tests/fixtures/main.py:166`) | Refresh cached aggregates after related-data changes / bulk operations | +| Feature vectors (`DetectionEmbedding`) | `ami/ml/models/embedding.py` (`DetectionEmbeddingQuerySet.store()`), reader in `ami/ml/embeddings/reader.py` | One row per (detection, algorithm, key), written from `DetectionResponse.embeddings` by `create_detection_embeddings()` in `ami/ml/embeddings/writer.py`. Read with `vectors_for_detections(ids, algorithm_id)`; never compare vectors across algorithms. Full guide: `docs/claude/reference/feature-vectors.md` | +| Similarity sort | `OccurrenceQuerySet.with_visual_similarity()` (usage: `OccurrenceViewSet._order_by_visual_similarity`, `ami/main/api/views.py`) | `?ordering=visual_similarity&similar_to=`: cosine distance to one algorithm's vectors, computed inside a correlated subquery so GROUP BY evaluates it once per occurrence; NULL (no vector) sorts last | ## Conventions diff --git a/docs/claude/reference/feature-vectors.md b/docs/claude/reference/feature-vectors.md new file mode 100644 index 000000000..661a40362 --- /dev/null +++ b/docs/claude/reference/feature-vectors.md @@ -0,0 +1,128 @@ +# Feature vectors: how they are stored and how to query them + +Feature vectors (embeddings) are stored once per detection, per model output, in one table. This +page explains the layout, the single rule every reader must follow, the queries the layout is built +for (with the function that serves each), what to avoid, and how to extend it. The code lives in +`ami/ml/models/embedding.py` (model and queryset), `ami/ml/embeddings/writer.py` (storing results) +and `ami/ml/embeddings/reader.py` (query functions). Introduced in #1462. + +## Storage model + +`DetectionEmbedding` has one row per (detection, algorithm, key): + +| Column | Meaning | +|---|---| +| `detection` | The detection the vector describes. Rows cascade with it. | +| `algorithm` | The model that produced the vector. Rows cascade with it. | +| `key` | The name of the output when one model returns several (default `embedding`). | +| `vector` | A `halfvec` column with no declared length, stored out of line and uncompressed (`STORAGE EXTERNAL`), because vectors do not compress and this keeps heap rows small. | +| `project` | Copied from the detection's capture (or its station), so project-scoped queries need no join. | +| `job` | The job whose results stored the vector; set to null if the job is deleted. | + +Indexes: the unique constraint on (detection, algorithm, key), an index on +(project, algorithm, key, detection), an index on (algorithm, key, detection) that serves the writer's +length lookup (the first row of a pair, ordered by detection), and a partial index on (job, detection) +where a job is set, for a job's list of occurrences. The foreign keys have no single-column indexes of +their own, because these already lead with them. + +Because the column has no declared length, models of different lengths share the table (for example +a 2048-dimension classifier backbone and a 1024-dimension image-text model). Each (algorithm, key) +keeps one length: the writer reads the length of one existing row of the pair and refuses vectors +of another length with `EmbeddingDimensionMismatch`. No length is stored on `Algorithm`. + +## Which algorithm a vector is stored under + +A vector is stored under the algorithm the processing service names for it in the response. A +classifier that returns its own backbone features therefore stores them under that classifier's +algorithm row, while a dedicated extractor (for example BioCLIP, task type `embedding`) has its own +row. To find which models have vectors in a project, use `vector_counts_by_algorithm(project_id)` +rather than assuming an algorithm. + +## The rule: one (algorithm, key) per query + +Vectors from different models, or from different keys of one model, are not comparable. Every +reader therefore takes an algorithm id and a key, and nothing merges pairs. A query that mixes +pairs returns noise, and it also mixes lengths, which the database rejects when distances are +computed. When a cross-model query is truly intended, say so in a comment at the call site. + +## Query patterns + +Each pattern names the function in `ami/ml/embeddings/reader.py` that serves it, the SQL shape, +and the index it uses. Check a plan with `EXPLAIN (ANALYZE, BUFFERS)` on realistic data before +claiming a query is cheap. + +| # | Need | Function | SQL shape | Index | +|---|---|---|---|---| +| Q1 | Vectors of one model for some detections (tracking over adjacent captures, retraining on verified detections) | `vectors_for_detections(ids, algorithm_id, key)` | `detection_id = ANY(..) AND algorithm_id = A AND key = K` | unique (detection, algorithm, key) for a few captures' worth; for thousands of ids the planner may prefer (algorithm, key, detection) or a sequential scan, measured at 17-19 ms for 5,000 ids on a 450k-row table (a sequential scan; 12 ms when forced onto (algorithm, key, detection)) | +| Q2 | Occurrences sorted by similarity to one occurrence | `OccurrenceQuerySet.with_visual_similarity()` using `representative_embeddings()` | per occurrence, the representative detection's vector, then cosine distance | (algorithm, key, detection), probed once per occurrence; an exact scan | +| Q3 | All of one model's vectors in a project (exports, clustering) | `project_vectors(project_id, algorithm_id, key, detection_ids=None, chunk_size=2000)` | `project_id = P AND algorithm_id = A AND key = K AND detection_id > last ORDER BY detection_id LIMIT n` | (project, algorithm, key, detection): rows come out in order, so no sort | +| Q4 | Which models have vectors in a project, and how many | `vector_counts_by_algorithm(project_id, key=None)`, `algorithm_with_most_vectors(project)` | `GROUP BY algorithm_id, key` within a project | the same index (index-only); when one project holds most of the table the planner may scan the table instead, measured at 36 ms for 359k rows | +| Q5 | Detections that still lack a vector from a model | `detections_missing_vectors(detections, algorithm_id, key)` | `NOT EXISTS` on (detection, algorithm, key) added to the caller's queryset | (algorithm, key, detection), index-only | +| Q6 | Nearest neighbours of one vector, when exact scans are too slow | not shipped | `ORDER BY vector::halfvec(D) <=> seed` within one (algorithm, key) | a partial HNSW index per (algorithm, key), see below | + +`project_vectors` pages with `detection_id > last` rather than `OFFSET`, so every page costs the +same and memory stays at one chunk. Callers keep their own scope by passing `detection_ids` or by +filtering the queryset given to `detections_missing_vectors`. + +## Anti-patterns + +- **Computing similarity in Python over unbounded rows.** Use the database's distance operators, + or the chunked reader for exports. Never `list()` a project's vectors. +- **Mixing models.** Always filter by algorithm and key; a missing filter is a silent correctness bug. +- **`.distinct()` or JSON round trips over vector rows.** De-duplicating sorts the vector column + and is very slow; the unique constraint already prevents duplicates. +- **Omitting the project filter** on project-scoped reads. The index leads with `project`, and the + filter also enforces visibility. +- **Reading a pair's rows by (algorithm, key) alone** at scale. Only the (algorithm, key, detection) + index leads with those columns; it is meant for the writer's one-row length lookup, and a scan of + all of a pair's rows still reads the table. + +## Logits are not vectors to search + +Per-class logits belong to a classification, not to this table. Measured on a copy of production +data: 308 thousand classifications carry logits, taking about 4.3 GB of out-of-line storage; one +classifier has 29,176 classes, which is above pgvector's 16,000-dimension storage limit; and +37,776 (detection, algorithm) groups hold several classifications that disagree, so logits must +stay tied to their own classification row. Keep them on `Classification` and defer them in list +queries. If they move, use a side table keyed by classification with a `real[]` or bytea column +stored out of line, or files in object storage for bulk training exports; do not put them here. +Store only the top-k when the full vector is not needed. For all but one small model the stored +scores equal the softmax of the logits, so they can be derived. + +## Reduced dimensions + +A reduction (PCA, UMAP, random projection) of a vector is a different model output. Give it its own +`Algorithm` row, recording the method, the parent algorithm and the fit parameters, and store its +vectors in this table under that algorithm id. Its length and any index are then its own. Two +dimensional layouts for display are derived data that can be recomputed; store them only if a page +needs them persistently. + +## Vectors for other things + +Capture-level or taxon-level vectors should get sibling tables of the same shape (for example +`SourceImageEmbedding`, `TaxonEmbedding`), not a polymorphic target column on this table. That +keeps foreign keys, cascades and indexes exact. + +## Nearest-neighbour indexes (HNSW) + +Exact scans were measured at about 16 to 19 ms per 1,000 candidate rows, which is fine for a +filtered list. When that stops being enough, add a partial HNSW index per (algorithm, key) on the +cast expression, for example +`CREATE INDEX ... USING hnsw ((vector::halfvec(1024)) halfvec_cosine_ops) WHERE algorithm_id = A AND key = 'embedding'`. +The cast is valid because each pair has one length, and halfvec HNSW supports up to 4,000 +dimensions. Measured costs at about 350 thousand rows: 0.9 to 2.7 GB and 3.5 to 15 minutes to +build per model, with recall between 0.68 and 0.98, so results are approximate and filtered +queries need care. Not created by default. + +## Precision + +`halfvec` stores 16-bit floats. On 600 real 2,048-dimension vectors the cosine similarity changed +by at most 1.5e-4 and the top-1 neighbour agreed in 98.2 percent of cases. Values outside the +half-precision range (beyond 65,504, NaN, infinity) cannot be stored; the writer skips them with a +warning. Keep full precision somewhere else only where that difference matters, and say so. + +## Deletion and retention + +Vectors are deleted with their detection or algorithm. Deleting a job keeps its vectors (the job +reference becomes null). Saving identical results twice changes nothing; a new vector for the same +(detection, algorithm, key) replaces the old one in place. diff --git a/requirements/base.txt b/requirements/base.txt index be5d24b67..04e128d23 100644 --- a/requirements/base.txt +++ b/requirements/base.txt @@ -19,6 +19,7 @@ django-pydantic-field==0.3.10 sentry-sdk==2.59.0 # https://github.com/getsentry/sentry-python django-cachalot==2.6.3 numpy==2.1 +pgvector==0.5.0 # https://github.com/pgvector/pgvector-python # Django # ------------------------------------------------------------------------------ diff --git a/ui/src/components/filtering/filter-control.tsx b/ui/src/components/filtering/filter-control.tsx index 742620c79..e9cc0eed5 100644 --- a/ui/src/components/filtering/filter-control.tsx +++ b/ui/src/components/filtering/filter-control.tsx @@ -9,6 +9,7 @@ import { CaptureSetFilter } from './filters/capture-set-filter' import { DateFilter } from './filters/date-filter' import { DeviceFilter } from './filters/device-filter' import { ImageFilter } from './filters/image-filter' +import { OccurrenceFilter } from './filters/occurrence-filter' import { PipelineFilter } from './filters/pipeline-filter' import { SessionFilter } from './filters/session-filter' import { SiteFilter } from './filters/site-filter' @@ -45,6 +46,7 @@ const ComponentMap: { not_tag_id: TagFilter, not_taxa_list_id: TaxaListFilter, pipeline: PipelineFilter, + similar_to: OccurrenceFilter, source_image_collection: CaptureSetFilter, source_image_single: ImageFilter, status: StatusFilter, diff --git a/ui/src/components/filtering/filters/occurrence-filter.tsx b/ui/src/components/filtering/filters/occurrence-filter.tsx new file mode 100644 index 000000000..a6167fc29 --- /dev/null +++ b/ui/src/components/filtering/filters/occurrence-filter.tsx @@ -0,0 +1,7 @@ +import { FilterProps } from './types' + +export const OccurrenceFilter = ({ value }: FilterProps) => ( +
+ {value ? `#${value}` : ''} +
+) diff --git a/ui/src/data-services/models/occurrence-details.ts b/ui/src/data-services/models/occurrence-details.ts index 39bcbfb39..92340ce8d 100644 --- a/ui/src/data-services/models/occurrence-details.ts +++ b/ui/src/data-services/models/occurrence-details.ts @@ -108,6 +108,11 @@ export class OccurrenceDetails extends Occurrence { return this._occurrence.details } + // Algorithms with a feature vector on this occurrence; a similarity sort needs one. + get embeddingAlgorithms(): { id: number; name: string }[] { + return this._occurrence.embedding_algorithms ?? [] + } + get detections(): string[] { return this._detections } diff --git a/ui/src/pages/occurrence-details/occurrence-details.tsx b/ui/src/pages/occurrence-details/occurrence-details.tsx index 5bc502d89..81fadf695 100644 --- a/ui/src/pages/occurrence-details/occurrence-details.tsx +++ b/ui/src/pages/occurrence-details/occurrence-details.tsx @@ -4,11 +4,12 @@ import { } from 'components/blueprint-collection/blueprint-collection' import { TaxonDetails } from 'components/taxon-details/taxon-details' import { OccurrenceDetails as Occurrence } from 'data-services/models/occurrence-details' -import { SearchIcon } from 'lucide-react' +import { ImagesIcon, SearchIcon } from 'lucide-react' import { BasicTooltip, Box, Button, + buttonVariants, CodeBlock, IdentificationScore, InfoBlockField, @@ -17,7 +18,7 @@ import { } from 'nova-ui-kit' import { useMemo, useRef, useState } from 'react' import { Helmet } from 'react-helmet-async' -import { useLocation, useNavigate, useParams } from 'react-router-dom' +import { Link, useLocation, useNavigate, useParams } from 'react-router-dom' import { APP_ROUTES } from 'utils/constants' import { getAppRoute } from 'utils/getAppRoute' import { STRING, translate } from 'utils/language' @@ -58,6 +59,14 @@ export const OccurrenceDetails = ({ const location = useLocation() const [suggestIdOpen, setSuggestIdOpen] = useState(false) const canUpdate = occurrence.userPermissions.includes(UserPermission.Update) + const similarityAlgorithm = occurrence.embeddingAlgorithms[0] + const similarOccurrencesRoute = `${APP_ROUTES.OCCURRENCES({ + projectId: projectId as string, + })}?${new URLSearchParams({ + ordering: 'visual_similarity', + similar_to: occurrence.id, + similarity_algorithm: `${similarityAlgorithm?.id}`, + })}` const blueprintItems = useMemo( () => @@ -167,6 +176,15 @@ export const OccurrenceDetails = ({ /> ) : null} + {similarityAlgorithm ? ( + + + {translate(STRING.SHOW_SIMILAR_OCCURRENCES)} + + ) : null} {canUpdate && ( <> [] => [ { id: 'snapshots', name: translate(STRING.FIELD_LABEL_SNAPSHOTS), + sortField: canSortBySimilarity ? 'visual_similarity' : undefined, styles: { textAlign: TextAlign.Center, }, diff --git a/ui/src/pages/occurrences/occurrences.tsx b/ui/src/pages/occurrences/occurrences.tsx index e0aed9b24..bbffc1e12 100644 --- a/ui/src/pages/occurrences/occurrences.tsx +++ b/ui/src/pages/occurrences/occurrences.tsx @@ -51,10 +51,10 @@ export const Occurrences = () => { ['updated-at']: true, } ) - const { sort, setSort } = useSort({ - field: 'updated_at', - order: 'desc', - }) + const { sort, setSort } = useSort( + { field: 'updated_at', order: 'desc' }, + { similar_to: 'visual_similarity' } + ) const { pagination, setPage } = usePagination() const { activeFilters, filters } = useFilters() const { occurrences, total, isLoading, isFetching, error } = useOccurrences({ @@ -72,6 +72,9 @@ export const Occurrences = () => { const tableColumns = columns({ projectId: projectId as string, showActions: selectedItems.length === 0, + canSortBySimilarity: !!filters.find( + (filter) => filter.field === 'similar_to' + )?.value?.length, }) useEffect(() => { @@ -95,6 +98,7 @@ export const Occurrences = () => { + {taxaLists.length > 0 && ( <> diff --git a/ui/src/utils/getAppRoute.ts b/ui/src/utils/getAppRoute.ts index 354b9ed36..71130cd56 100644 --- a/ui/src/utils/getAppRoute.ts +++ b/ui/src/utils/getAppRoute.ts @@ -11,6 +11,7 @@ type FilterType = | 'include_unobserved' | 'job' | 'occurrence' + | 'similar_to' | 'source_image_collection' | 'source_image_single' | 'taxa_list_id' diff --git a/ui/src/utils/language.ts b/ui/src/utils/language.ts index 38e20a5e4..729aef2fa 100644 --- a/ui/src/utils/language.ts +++ b/ui/src/utils/language.ts @@ -153,6 +153,7 @@ export enum STRING { FIELD_LABEL_SCORE, FIELD_LABEL_SESSION, FIELD_LABEL_SESSIONS, + FIELD_LABEL_SIMILAR_TO, FIELD_LABEL_SITE, FIELD_LABEL_SIZE, FIELD_LABEL_SLUG, @@ -288,6 +289,7 @@ export enum STRING { TOOLTIP_PROCESSING_SERVICE, TOOLTIP_SCORE, TOOLTIP_SESSION, + TOOLTIP_SIMILAR_TO, TOOLTIP_SITE, TOOLTIP_STATS, TOOLTIP_STATS_AGREEMENT_ANY_RANK, @@ -357,6 +359,7 @@ export enum STRING { SET_PASSWORD, SETTINGS, SHOW_DETECTIONS, + SHOW_SIMILAR_OCCURRENCES, SNAP_TO_DETECTIONS, SORT_BY, SORT_JOBS_ACTIVITY, @@ -517,6 +520,7 @@ const ENGLISH_STRINGS: { [key in STRING]: string } = { [STRING.FIELD_LABEL_SCORE]: 'Score', [STRING.FIELD_LABEL_SESSION]: 'Session', [STRING.FIELD_LABEL_SESSIONS]: 'Sessions', + [STRING.FIELD_LABEL_SIMILAR_TO]: 'Similar to occurrence', [STRING.FIELD_LABEL_SITE]: 'Site', [STRING.FIELD_LABEL_SIZE]: 'Size', [STRING.FIELD_LABEL_SLUG]: 'Slug', @@ -717,6 +721,8 @@ const ENGLISH_STRINGS: { [key in STRING]: string } = { 'This is a model derived prediction score, not a real-world probability. Think of it as a relative metric that will vary based on model calibration and available training data.', [STRING.TOOLTIP_SESSION]: 'A session is a fixed period of time of monitoring for one station. The period is typically one night.', + [STRING.TOOLTIP_SIMILAR_TO]: + 'Occurrences are sorted by how similar they look to this occurrence, most similar first. Occurrences without a feature vector come last.', [STRING.TOOLTIP_SITE]: 'A site is a physical location where monitoring is taking place. One or many stations can be connected to a site.', [STRING.TOOLTIP_STATS]: @@ -792,6 +798,7 @@ const ENGLISH_STRINGS: { [key in STRING]: string } = { [STRING.SET_PASSWORD]: 'Set password', [STRING.SETTINGS]: 'Settings', [STRING.SHOW_DETECTIONS]: 'Show detections', + [STRING.SHOW_SIMILAR_OCCURRENCES]: 'Show similar occurrences', [STRING.SNAP_TO_DETECTIONS]: 'Snap to captures with detections', [STRING.SORT_BY]: 'Sort by', [STRING.SORT_JOBS_ACTIVITY]: 'Jobs activity', diff --git a/ui/src/utils/useFilters.ts b/ui/src/utils/useFilters.ts index 48fa644ba..698ea37c1 100644 --- a/ui/src/utils/useFilters.ts +++ b/ui/src/utils/useFilters.ts @@ -194,6 +194,13 @@ export const AVAILABLE_FILTERS = (projectId: string): FilterConfig[] => [ label: 'Exclude tag', field: 'not_tag_id', }, + { + label: translate(STRING.FIELD_LABEL_SIMILAR_TO), + field: 'similar_to', + tooltip: { + text: translate(STRING.TOOLTIP_SIMILAR_TO), + }, + }, { label: 'Taxon', field: 'taxon', diff --git a/ui/src/utils/useSort.ts b/ui/src/utils/useSort.ts index 121b9ced9..9684653bb 100644 --- a/ui/src/utils/useSort.ts +++ b/ui/src/utils/useSort.ts @@ -3,7 +3,12 @@ import { useSearchParams } from 'react-router-dom' const SEARCH_PARAM_KEY_ORDERING = 'ordering' -export const useSort = (defaultSort?: TableSortSettings) => { +// `seedParams` maps a search param to the only sort field that uses it. Switching to +// any other sort removes the param, so the URL does not claim a seed that is ignored. +export const useSort = ( + defaultSort?: TableSortSettings, + seedParams?: { [param: string]: string } +) => { const [searchParams, setSearchParams] = useSearchParams() const ordering = searchParams.get(SEARCH_PARAM_KEY_ORDERING) @@ -21,6 +26,12 @@ export const useSort = (defaultSort?: TableSortSettings) => { const setSort = (sort: TableSortSettings | undefined) => { searchParams.delete(SEARCH_PARAM_KEY_ORDERING) + Object.entries(seedParams ?? {}).forEach(([param, field]) => { + if (sort?.field !== field) { + searchParams.delete(param) + } + }) + if (sort) { const newOrdering = `${sort.order === 'desc' ? '-' : ''}${sort.field}` searchParams.set(SEARCH_PARAM_KEY_ORDERING, newOrdering)