feature: Draw motion vectors

src/Viewer.tsx

@@ -32,7 +32,7 @@ import {
 import { AnalyticsEvent, triggerAnalyticsEvent } from "./colorizer/utils/analytics";
 import { getColorMap, getInRangeLUT, thresholdMatchFinder, validateThresholds } from "./colorizer/utils/data_utils";
 import { numberToStringDecimal } from "./colorizer/utils/math_utils";
-import { useConstructor, useDebounce, useRecentCollections } from "./colorizer/utils/react_utils";
+import { useConstructor, useDebounce, useMotionDeltas, useRecentCollections } from "./colorizer/utils/react_utils";
 import * as urlUtils from "./colorizer/utils/url_utils";
 import { SCATTERPLOT_TIME_FEATURE } from "./components/Tabs/scatter_plot_data_utils";
 import { DEFAULT_PLAYBACK_FPS } from "./constants";
@@ -107,6 +107,8 @@ function Viewer(): ReactElement {
     getDefaultScatterPlotConfig()
   );
 
+  const motionDeltas = useMotionDeltas(dataset, workerPool, config.vectorConfig);
+
   const [isInitialDatasetLoaded, setIsInitialDatasetLoaded] = useState(false);
   const [isDatasetLoading, setIsDatasetLoading] = useState(false);
   const [datasetLoadProgress, setDatasetLoadProgress] = useState<number | null>(null);
@@ -1017,6 +1019,7 @@ function Viewer(): ReactElement {
                   loadingProgress={datasetLoadProgress}
                   canv={canv}
                   collection={collection || null}
+                  vectorData={motionDeltas}
                   dataset={dataset}
                   datasetKey={datasetKey}
                   featureKey={featureKey}

src/colorizer/ColorizeCanvas.ts

@@ -29,12 +29,14 @@ import {
   OUT_OF_RANGE_COLOR_DEFAULT,
   OUTLIER_COLOR_DEFAULT,
   OUTLINE_COLOR_DEFAULT,
+  VectorConfig,
 } from "./types";
 import { packDataTexture } from "./utils/texture_utils";
 
 import ColorRamp from "./ColorRamp";
 import Dataset from "./Dataset";
 import Track from "./Track";
+import VectorField from "./VectorField";
 
 import pickFragmentShader from "./shaders/cellId_RGBA8U.frag";
 import vertexShader from "./shaders/colorize.vert";
@@ -122,12 +124,14 @@ export default class ColorizeCanvas {
   private mesh: Mesh;
   private pickMesh: Mesh;
 
+  private vectorField: VectorField;
   // Rendered track line that shows the trajectory of a cell.
   private line: Line;
+  private points: Float32Array;
   private showTrackPath: boolean;
 
   protected frameSizeInCanvasCoordinates: Vector2;
-  private frameToCanvasCoordinates: Vector2;
+  protected frameToCanvasCoordinates: Vector2;
 
   /**
    * The zoom level of the frame in the canvas. At default zoom level 1, the frame will be
@@ -150,7 +154,6 @@ export default class ColorizeCanvas {
 
   protected dataset: Dataset | null;
   protected track: Track | null;
-  private points: Float32Array;
   protected canvasResolution: Vector2 | null;
 
   protected featureKey: string | null;
@@ -187,10 +190,14 @@ export default class ColorizeCanvas {
     this.camera = new OrthographicCamera(-1, 1, 1, -1, 0, 1);
     this.scene = new Scene();
     this.scene.add(this.mesh);
+
+    this.vectorField = new VectorField();
+    this.scene.add(this.vectorField.sceneObject);
+
     this.pickScene = new Scene();
     this.pickScene.add(this.pickMesh);
 
-    // Configure lines
+    // Configure track lines
     this.points = new Float32Array([0, 0, 0]);
 
     const lineGeometry = new BufferGeometry();
@@ -210,7 +217,7 @@ export default class ColorizeCanvas {
     this.pickRenderTarget = new WebGLRenderTarget(1, 1, {
       depthBuffer: false,
     });
-    this.renderer = new WebGLRenderer();
+    this.renderer = new WebGLRenderer({ antialias: true });
     this.checkPixelRatio();
 
     this.dataset = null;
@@ -286,6 +293,7 @@ export default class ColorizeCanvas {
       2 * this.panOffset.y * this.frameToCanvasCoordinates.y,
       0
     );
+    this.vectorField.setPosition(this.panOffset, this.frameToCanvasCoordinates);
     this.render();
   }
 
@@ -329,6 +337,8 @@ export default class ColorizeCanvas {
       2 * this.panOffset.y * this.frameToCanvasCoordinates.y,
       0
     );
+    this.vectorField.setPosition(this.panOffset, this.frameToCanvasCoordinates);
+    this.vectorField.setScale(this.frameToCanvasCoordinates, this.canvasResolution || new Vector2(1, 1));
   }
 
   public async setDataset(dataset: Dataset): Promise<void> {
@@ -346,6 +356,7 @@ export default class ColorizeCanvas {
     this.currentFrame = -1;
     await this.setFrame(frame);
     this.updateScaling(this.dataset.frameResolution, this.canvasResolution);
+    this.vectorField.setDataset(dataset);
     this.render();
   }
 
@@ -404,6 +415,14 @@ export default class ColorizeCanvas {
     }
   }
 
+  setVectorData(vectorData: Float32Array | null): void {
+    this.vectorField.setVectorData(vectorData);
+  }
+
+  setVectorFieldConfig(config: VectorConfig): void {
+    this.vectorField.setConfig(config);
+  }
+
   setSelectedTrack(track: Track | null): void {
     if (this.track && this.track?.trackId === track?.trackId) {
       return;
@@ -627,6 +646,7 @@ export default class ColorizeCanvas {
     this.onFrameChangeCallback(isMissingFile);
     // Force rescale in case frame dimensions changed
     this.updateScaling(this.dataset?.frameResolution || null, this.canvasResolution);
+    this.vectorField.setFrame(this.currentFrame);
   }
 
   /** Switches the coloring between the categorical and color ramps depending on the currently

src/colorizer/Dataset.ts

@@ -525,6 +525,19 @@ export default class Dataset {
     return this.trackIds?.[index] || 0;
   }
 
+  /**
+   * Returns the 2D centroid of a given object id.
+   */
+  public getCentroid(objectId: number): [number, number] | undefined {
+    const index = objectId * 2;
+    const x = this.centroids?.[index];
+    const y = this.centroids?.[index + 1];
+    if (x && y) {
+      return [x, y];
+    }
+    return undefined;
+  }
+
   private getIdsOfTrack(trackId: number): number[] {
     return this.trackIds?.reduce((arr: number[], elem: number, ind: number) => {
       if (elem === trackId) arr.push(ind);

src/colorizer/Track.ts

@@ -1,8 +1,12 @@
 export default class Track {
   public trackId: number;
+  /** Times, in ascending order. */
   public times: number[];
+  /** Object IDs in the track, in ascending order of time. */
   public ids: number[];
+  /** Centroids in the track, in ascending order of time. */
   public centroids: number[];
+  /** Bounds in the track, in ascending order of time. */
   public bounds: number[];
 
   constructor(trackId: number, times: number[], ids: number[], centroids: number[], bounds: number[]) {
@@ -19,17 +23,17 @@ export default class Track {
       indices.sort((a, b) => (times[a] < times[b] ? -1 : times[a] === times[b] ? 0 : 1));
       this.times = indices.map((i) => times[i]);
       this.ids = indices.map((i) => ids[i]);
+
       this.centroids = indices.reduce((result, i) => {
         result.push(centroids[i * 2], centroids[i * 2 + 1]);
         return result;
       }, [] as number[]);
+
+      this.bounds = indices.reduce((result, i) => {
+        result.push(bounds[i * 4], bounds[i * 4 + 1], bounds[i * 4 + 2], bounds[i * 4 + 3]);
+        return result;
+      }, [] as number[]);
     }
-    console.log(
-      `Track ${trackId} has ${this.times.length} objects over ${this.duration()} timepoints starting from ${
-        this.times[0]
-      } to ${this.times[this.times.length - 1]}`
-    );
-    console.log(this.ids);
   }
 
   getIdAtTime(t: number): number {