LabelBOT

app/Annotation.php

@@ -39,6 +39,11 @@ abstract class Annotation extends Model implements AnnotationContract
         'points' => 'array',
     ];
 
+    /**
+     * The additional labels suggested by the LabelBOT.
+     */
+    public $labelBOTLabels = [];
+
     /**
      * Scope a query to only include annotations that are visible for a certain user.
      *
@@ -209,4 +214,14 @@ public function getFile(): VolumeFile
     {
         return $this->file;
     }
+
+    /**
+     * Get the LabelBOT suggested labels.
+     *
+     * @return array<int>
+     */
+    public function getLabelBOTLabelsAttribute(): array
+    {
+        return $this->labelBOTLabels;
+    }
 }

app/Http/Controllers/Api/ImageAnnotationController.php

@@ -7,12 +7,17 @@
 use Biigle\ImageAnnotation;
 use Biigle\ImageAnnotationLabel;
 use Biigle\Label;
+use Biigle\LabelTree;
+use Biigle\Modules\Largo\ImageAnnotationLabelFeatureVector;
+use Biigle\Project;
+use Biigle\Role;
 use Biigle\Shape;
 use DB;
 use Exception;
 use Generator;
 use Illuminate\Http\Request;
 use Illuminate\Validation\ValidationException;
+use Pgvector\Laravel\Vector;
 use Symfony\Component\HttpFoundation\StreamedJsonResponse;
 
 class ImageAnnotationController extends Controller
@@ -156,7 +161,8 @@
      * @apiParam {Number} id The image ID.
      *
      * @apiParam (Required arguments) {Number} shape_id ID of the shape of the new annotation.
-     * @apiParam (Required arguments) {Number} label_id ID of the initial category label of the new annotation.
+     * @apiParam (Required arguments) {Number} label_id ID of the initial category label of the new annotation. Required if 'feature_vector' is not provided.
+     * @apiParam (Required arguments) {Number[]} feature_vector A feature vector array of size 384 for label prediction. Required if 'label_id' is not provided.
      * @apiParam (Required arguments) {Number} confidence Confidence of the initial annotation label of the new annotation. Must be a value between 0 and 1.
      * @apiParam (Required arguments) {Number[]} points Array of the initial points of the annotation. Must contain at least one point. The points array is interpreted as alternating x and y coordinates like this `[x1, y1, x2, y2...]`. The interpretation of the points of the different shapes is as follows:
      * **Point:** The first point is the center of the annotation point.
@@ -212,16 +218,31 @@
 
         $annotation = new ImageAnnotation;
         $annotation->shape_id = $request->input('shape_id');
-        $annotation->image()->associate($request->image);
-
+        $image = $request->image;
+        $annotation->image()->associate($image);
         try {
             $annotation->validatePoints($points);
         } catch (Exception $e) {
             throw ValidationException::withMessages(['points' => [$e->getMessage()]]);
         }
 
         $annotation->points = $points;
-        $label = Label::findOrFail($request->input('label_id'));
+        $labelId = $request->input('label_id');
+        $topNLabels = [];
+        if (is_null($labelId) && $request->has('feature_vector')) {
+            // Get label tree id(s).
+            $trees = $this->getLabelTreeIds($request->user(), $image->volume_id);
+            // Perform ANN search.
+            $topNLabels = $this->performAnnSearch($request->input('feature_vector'), $trees);
+            // Perform KNN search as a fallback if ANN search returns no results.
+            if (empty($topNLabels)) {
+                $topNLabels = $this->performKnnSearch($request->input('feature_vector'), $trees);
+            }
+            // Set labelId to top 1 label.
+            $labelId = $topNLabels[0];
+        }
+
+        $label = Label::findOrFail($labelId);
 
         $this->authorize('attach-label', [$annotation, $label]);
 
@@ -236,6 +257,11 @@
 
         $annotation->load('labels.label', 'labels.user');
 
+        // Attach the other two labels if they exist.
+        for ($i = 1; $i < count($topNLabels); $i++) {
+            $annotation->labelBOTLabels[] = $topNLabels[$i];
+        }
+
         return $annotation;
     }
 
@@ -331,4 +357,108 @@
 
         return response('Deleted.', 200);
     }
+
+    /**
+     * Get all label trees that are used by all projects which are visible to the user.
+     *
+     * @param mixed $user
+     * @param int $volumeId
+     *
+     * @return array
+     */
+    protected function getLabelTreeIds($user, $volumeId)
+    {
+        if ($user->can('sudo')) {
+            // Global admins have no restrictions.
+            $projectIds = DB::table('project_volume')
+                ->where('volume_id', $volumeId)
+                ->pluck('project_id');
+        } else {
+            // Array of all project IDs that the user and the image have in common
+            // and where the user is editor, expert or admin.
+            $projectIds = Project::inCommon($user, $volumeId, [
+                Role::editorId(),
+                Role::expertId(),
+                Role::adminId(),
+            ])->pluck('id');
+        }
+        $trees = LabelTree::select('id', 'name', 'version_id')
+            ->with('labels', 'version')
+            ->whereIn('id', function ($query) use ($projectIds) {
+                $query->select('label_tree_id')
+                    ->from('label_tree_project')
+                    ->whereIn('project_id', $projectIds);
+            })
+            ->pluck('id')
+            ->toArray();
+
+        return $trees;
+    }
+
+    /**
+     * Perform ANN search (HNSW + Post-Subquery-Filtering).
+     *
+     * @param vector $featureVector
+     * @param int[] $trees
+     *
+     * @return array
+     */
+    protected function performAnnSearch($featureVector, $trees)
+    {
+        $featureVector = new Vector($featureVector);
+
+        $subquery = ImageAnnotationLabelFeatureVector::select('label_id', 'label_tree_id')
+            ->selectRaw('(vector <=> ?) AS distance', [$featureVector])
+            ->orderBy('distance')
+            // K = 100
+            ->limit(config('labelbot.K'));
+
+        return DB::table(DB::raw("({$subquery->toSql()}) as subquery"))
+            ->setBindings([$featureVector])
+            ->whereIn('label_tree_id', $trees)
+            ->select('label_id')
+            ->groupBy('label_id')
+            ->orderByRaw('MIN(distance)')
+            ->limit(config('labelbot.N')) // N = 3
+            ->pluck('label_id')
+            ->toArray();
+    }
+
+    /**
+     * Perform KNN search (B-Tree + Post-Filtering).
+     *
+     * @param Vector $featureVector
+     * @param int[] $trees
+     *
+     * @return array
+     */
+    protected function performKnnSearch($featureVector, $trees)
+    {
+        $featureVector = new Vector($featureVector);
+
+        $subquery = ImageAnnotationLabelFeatureVector::select('label_id', 'label_tree_id')
+            ->selectRaw('(vector <=> ?) AS distance', [$featureVector])
+            // filter by label tree id in subquery
+            // to use B-Tree index for filtering and speeding up the vector search
+            ->whereIn('label_tree_id', $trees)
+            ->orderBy('distance')
+            ->limit(config('labelbot.K')); // K = 100
+
+        // TODO: Drop HNSW index temporary
+        // DB::beginTransaction();
+
+        $topNLabels = DB::table(DB::raw("({$subquery->toSql()}) as subquery"))
+            ->setBindings(array_merge([$featureVector], $trees))
+            ->select('label_id')
+            ->groupBy('label_id')
+            ->orderByRaw('MIN(distance)')
+            ->limit(config('labelbot.N')) // N = 3
+            ->pluck('label_id')
+            ->toArray();
+
+        // TODO: Rollback the HNSW index drop
+        // DB::rollback();
+
+        return $topNLabels;
+    }
 }

app/Http/Requests/StoreImageAnnotation.php

@@ -5,7 +5,7 @@
 use Biigle\Image;
 use Biigle\Shape;
 
-class StoreImageAnnotation extends StoreImageAnnotationLabel
+class StoreImageAnnotation extends StoreImageAnnotationLabelFeatureVector
 {
     /**
      * The image on which the annotation should be created.

app/Http/Requests/StoreImageAnnotationLabel.php

@@ -43,7 +43,7 @@ public function authorize()
     public function rules()
     {
         return [
-            'label_id'    => 'required|integer|exists:labels,id',
+            'label_id' => 'required|integer|exists:labels,id',
             'confidence'  => 'required|numeric|between:0,1',
         ];
     }

app/Http/Requests/StoreImageAnnotationLabelFeatureVector.php

@@ -0,0 +1,51 @@
+<?php
+
+namespace Biigle\Http\Requests;
+
+use Biigle\ImageAnnotation;
+use Biigle\Label;
+use Illuminate\Foundation\Http\FormRequest;
+
+class StoreImageAnnotationLabelFeatureVector extends FormRequest
+{
+    /**
+     * The annotation to which the label should be attached.
+     *
+     * @var ImageAnnotation
+     */
+    public $annotation;
+
+    /**
+     * The label that should be attached.
+     *
+     * @var Label
+     */
+    public $label;
+
+    /**
+     * Determine if the user is authorized to make this request.
+     *
+     * @return bool
+     */
+    public function authorize()
+    {
+        $this->annotation = ImageAnnotation::findOrFail($this->route('id'));
+        $this->label = Label::findOrFail($this->input('label_id'));
+
+        return $this->user()->can('attach-label', [$this->annotation, $this->label]);
+    }
+
+    /**
+     * Get the validation rules that apply to the request.
+     *
+     * @return array
+     */
+    public function rules()
+    {
+        return [
+            'label_id'    => 'required_without:feature_vector|integer|exists:labels,id',
+            'feature_vector' => 'required_without:label_id|array|size:384',
+            'confidence'  => 'required|numeric|between:0,1',
+        ];
+    }
+}

app/ImageAnnotation.php

@@ -19,6 +19,13 @@ class ImageAnnotation extends Annotation
         'points' => 'array',
     ];
 
+    /**
+     * The attributes that should be included in the JSON response.
+     *
+     * @var array<int, string>
+     */
+    protected $appends = ['labelBOTLabels'];
+
     /**
      * The image, this annotation belongs to.
      *

config/labelbot.php

@@ -0,0 +1,14 @@
+<?php
+
+return [
+
+    /*
+    | K for KNN
+    */
+    'K' => 100,
+
+    /*
+    | N for top N labels
+    */
+    'N' => 3,
+];