Lasso with point field

app/packages/embeddings/src/EmbeddingsPlot.tsx

@@ -92,10 +92,10 @@ export function EmbeddingsPlot({
               result[p.fullData.name].push(p.id);
               pointIds.push(p.id);
             }
-            handleSelected(pointIds);
+            handleSelected(pointIds, selected.lassoPoints);
           }}
           onDeselect={() => {
-            handleSelected(null);
+            handleSelected(null, null);
           }}
           config={{
             scrollZoom: true,

app/packages/embeddings/src/useBrainResult.tsx

@@ -6,6 +6,7 @@ import { useColorByField } from "./useLabelSelector";
 // a react hook that fetches a list of brain results
 // and has a loading state and an error state
 export const useBrainResult = () => usePanelStatePartial("brainResult", null);
+export const usePointsField = () => usePanelStatePartial("pointsField", null);
 export function useBrainResultsSelector() {
   const [selected, setSelected] = useBrainResult();
   const dataset = useRecoilValue(fos.dataset);

app/packages/embeddings/src/useExtendedStageEffect.tsx

@@ -3,6 +3,8 @@ import { useRecoilCallback, useRecoilValue, useSetRecoilState } from "recoil";
 import * as fos from "@fiftyone/state";
 import { usePanelStatePartial } from "@fiftyone/spaces";
 import { fetchExtendedStage } from "./fetch";
+import { atoms as selectionAtoms } from "./usePlotSelection";
+import { usePointsField } from "./useBrainResult";
 
 export default function useExtendedStageEffect() {
   const datasetName = useRecoilValue(fos.datasetName);
@@ -16,6 +18,8 @@ export default function useExtendedStageEffect() {
     return snapshot.getPromise(fos.datasetName);
   });
   const slices = useRecoilValue(fos.currentSlices(false));
+  const lassoPoints = useRecoilValue(selectionAtoms.lassoPoints);
+  const [pointsField] = usePointsField();
 
   useEffect(() => {
     if (loadedPlot && Array.isArray(selection)) {
@@ -25,6 +29,8 @@ export default function useExtendedStageEffect() {
         patchesField: loadedPlot.patches_field,
         selection,
         slices,
+        lassoPoints,
+        pointsField,
       }).then(async (res) => {
         const currentDataset = await getCurrentDataset();
         if (currentDataset !== datasetName) return;
@@ -33,5 +39,12 @@ export default function useExtendedStageEffect() {
         });
       });
     }
-  }, [datasetName, loadedPlot?.patches_field, view, selection]);
+  }, [
+    datasetName,
+    loadedPlot?.patches_field,
+    view,
+    selection,
+    pointsField,
+    lassoPoints,
+  ]);
 }

app/packages/embeddings/src/usePlotSelection.tsx

@@ -1,10 +1,22 @@
-import { useRecoilState, useRecoilValue, useSetRecoilState } from "recoil";
+import {
+  atom,
+  useRecoilState,
+  useRecoilValue,
+  useSetRecoilState,
+} from "recoil";
 import * as fos from "@fiftyone/state";
 import { usePanelStatePartial } from "@fiftyone/spaces";
 import { useBrainResultInfo } from "./useBrainResultInfo";
 import { SELECTION_SCOPE } from "./constants";
 import { useResetExtendedSelection } from "@fiftyone/state";
 
+export const atoms = {
+  lassoPoints: atom({
+    key: "lassoPoints",
+    default: { x: [], y: [] },
+  }),
+};
+
 export function usePlotSelection() {
   const brainResultInfo = useBrainResultInfo();
   const patchesField = brainResultInfo?.config?.patchesField;
@@ -20,14 +32,22 @@ export function usePlotSelection() {
     [],
     true
   );
+  const [lassoPoints, setLassoPoints] = useRecoilState(atoms.lassoPoints);
   const selectedPatchIds = useRecoilValue(fos.selectedPatchIds(patchesField));
   const selectedPatchSampleIds = useRecoilValue(fos.selectedPatchSamples);
-  function handleSelected(selectedResults) {
+  function handleSelected(
+    selectedResults,
+    lassoPoints: { x: number[]; y: number[] }
+  ) {
     setSelectedSamples(new Set());
     setExtendedSelection({
       selection: selectedResults,
       scope: SELECTION_SCOPE,
     });
+    if (lassoPoints) {
+      // TODO: this handleSelected is called without lassoPoints in some unkown cases
+      setLassoPoints(lassoPoints);
+    }
     if (selectedResults === null) {
       clearSelection();
     }

app/packages/embeddings/src/useViewChangeEffect.tsx

@@ -3,14 +3,15 @@ import * as fos from "@fiftyone/state";
 import { useEffect } from "react";
 import { useRecoilValue, useSetRecoilState } from "recoil";
 import { fetchPlot } from "./fetch";
-import { useBrainResult } from "./useBrainResult";
+import { useBrainResult, usePointsField } from "./useBrainResult";
 import { useColorByField } from "./useLabelSelector";
 import { useWarnings } from "./useWarnings";
 
 export function useViewChangeEffect() {
   const colorSeed = useRecoilValue(fos.colorSeed);
   const datasetName = useRecoilValue(fos.datasetName);
   const [brainKey, setBrainKey] = useBrainResult();
+  const [pointsField, setPointsField] = usePointsField();
   const [labelField] = useColorByField();
   const view = useRecoilValue(fos.view);
   const slices = useRecoilValue(fos.currentSlices(false));
@@ -72,6 +73,7 @@ export function useViewChangeEffect() {
 
         setLoadingPlotError(null);
         setLoadedPlot(res);
+        setPointsField(res.points_field);
       })
       .finally(() => setLoadingPlot(false));
   }, [datasetName, brainKey, labelField, view, colorSeed, slices, filters]);

fiftyone/core/stages.py

@@ -3295,7 +3295,8 @@ class GeoWithin(_GeoStage):
         strict (True): whether a sample's location data must strictly fall
             within boundary (True) in order to match, or whether any
             intersection suffices (False)
-
+        create_index (True): whether to create the required spherical index,
+            if necessary
     """
 
     def __init__(

fiftyone/server/routes/embeddings.py

@@ -240,6 +240,8 @@ def _post_sync(self, data):
         patches_field = data["patchesField"]  # patches field of plot, or None
         selected_ids = data["selection"]  # selected IDs in plot
         slices = data["slices"]
+        lasso_points = data.get("lassoPoints", None)
+        points_field = data.get("pointsField", None)
 
         view = fosv.get_view(
             dataset_name,
@@ -250,7 +252,39 @@ def _post_sync(self, data):
         is_patches_view = view._is_patches
         is_patches_plot = patches_field is not None
 
-        if not is_patches_view and not is_patches_plot:
+        if (
+            lasso_points is not None
+            and points_field is not None
+            and (patches_field is None or is_patches_view)
+        ):
+            # Lasso points via spatial index
+            # Unfortunately we can't use $geoWithin to filter nested arrays.
+            # The user must have switched to a patches view for this
+            if patches_field is not None:
+                _, points_field = view._get_label_field_path(
+                    patches_field, points_field
+                )
+
+            stage = fos.Mongo(
+                [
+                    {
+                        "$match": {
+                            points_field: {
+                                "$geoWithin": {
+                                    "$polygon": [
+                                        [x, y]
+                                        for x, y in zip(
+                                            lasso_points["x"],
+                                            lasso_points["y"],
+                                        )
+                                    ]
+                                }
+                            }
+                        }
+                    }
+                ]
+            )
+        elif not is_patches_view and not is_patches_plot:
             # Samples plot and samples view
             stage = fos.Select(selected_ids)
         elif is_patches_view and is_patches_plot: