improve representative taxon photo lookup

inaturalist · Jan 2, 2025 · f709467 · f709467
1 parent e1df4d1
commit f709467
Showing 1 changed file with 46 additions and 42 deletions.
diff --git a/lib/controllers/v1/computervision_controller.js b/lib/controllers/v1/computervision_controller.js
@@ -228,6 +228,11 @@ const ComputervisionController = class ComputervisionController {
       formData.append( "lat", req.body.lat );
       formData.append( "lng", req.body.lng );
     }
+    if ( req.userSession?.isAdmin && (
+      req.body.include_representative_photos || req.query.include_representative_photos
+    ) ) {
+      formData.append( "return_embedding", "true" );
+    }
 
     const requestAbortController = new AbortController( );
     const requestTimeout = setTimeout( ( ) => {
@@ -267,57 +272,56 @@ const ComputervisionController = class ComputervisionController {
     if ( _.isEmpty( embedding ) ) {
       return;
     }
-    const taxonIDsToLookup = _.map( _.filter(
+    const resultsToLookup = _.filter(
       results, result => result?.taxon?.rank_level < 30
-    ), result => result?.taxon?.id );
-    if ( _.isEmpty( taxonIDsToLookup ) ) {
+    );
+    if ( _.isEmpty( resultsToLookup ) ) {
       return;
     }
-    const embeddingsResponse = await esClient.search( "taxon_photos", {
-      body: {
-        query: {
-          bool: {
-            filter: [
-              { terms: { taxon_ids: taxonIDsToLookup } }
+
+    const representativeTaxonPhotos = [];
+    const representativeTaxonPhotosIDs = [];
+    const promiseProducer = ( ) => {
+      if ( _.isEmpty( resultsToLookup ) ) {
+        return null;
+      }
+      const resultToLookup = resultsToLookup.shift( );
+      return ( async ( ) => {
+        const embeddingsResponse = await esClient.search( "taxon_photos", {
+          body: {
+            knn: {
+              field: "embedding",
+              query_vector: embedding,
+              k: 10,
+              num_candidates: 10,
+              filter: {
+                term: {
+                  ancestor_ids: resultToLookup.taxon.id
+                }
+              }
+            },
+            size: 1,
+            _source: [
+              "id",
+              "taxon_id",
+              "photo_id",
+              "ancestor_ids"
             ]
           }
-        },
-        knn: {
-          field: "embedding",
-          query_vector: embedding,
-          k: 500,
-          num_candidates: 5000
-        },
-        size: 500,
-        _source: [
-          "id",
-          "taxon_id",
-          "photo_id",
-          "ancestor_ids"
-        ]
-      }
-    } );
-    const embeddingsHits = _.map( embeddingsResponse.hits.hits, "_source" );
-    // TODO: remove this once the ancestor_ids data in the index has been fixed
-    _.each( embeddingsHits, hit => {
-      hit.ancestor_ids = _.filter( hit.ancestor_ids, ancestorID => ancestorID !== hit.id );
-      hit.ancestor_ids.push( hit.taxon_id );
-    } );
-    const representativeTaxonPhotos = [];
-    _.each( results, result => {
-      if ( result && result.taxon && result.taxon.default_photo ) {
-        const firstMatch = _.find( embeddingsHits, h => (
-          h?.photo_id && _.includes( h.ancestor_ids, result.taxon.id )
-        ) );
-        if ( !firstMatch ) {
+        } );
+        const firstHitSource = embeddingsResponse.hits.hits[0]?._source;
+        if ( !firstHitSource ) {
           return;
         }
         representativeTaxonPhotos.push( {
-          taxon: result.taxon,
-          photo_id: firstMatch.photo_id
+          taxon: resultToLookup.taxon,
+          photo_id: firstHitSource.photo_id
         } );
-      }
-    } );
+      } )( );
+    };
+    const pool = new PromisePool( promiseProducer, 2 );
+    await pool.start( );
+
     await ObservationPreload.assignObservationPhotoPhotos( representativeTaxonPhotos );
     _.each( representativeTaxonPhotos, taxonPhoto => {
       if ( taxonPhoto?.photo?.url ) {