Merge branch 'feature/annotation-visualization' into feature/annotati…

…on-export-csv

src/Viewer.tsx

@@ -808,7 +808,9 @@ function Viewer(): ReactElement {
       setFindTrackInput(track?.trackId.toString() || "");
       setSelectedTrack(track);
       if (track && annotationState.isAnnotationModeEnabled && annotationState.currentLabelIdx !== null) {
-        annotationState.toggleLabelOnId(annotationState.currentLabelIdx, track.getIdAtTime(currentFrame));
+        const id = track.getIdAtTime(currentFrame);
+        const isLabeled = annotationState.data.isLabelOnId(annotationState.currentLabelIdx, id);
+        annotationState.setLabelOnId(annotationState.currentLabelIdx, track.getIdAtTime(currentFrame), !isLabeled);
       }
     },
     [annotationState.isAnnotationModeEnabled, annotationState.currentLabelIdx, currentFrame]

src/colorizer/AnnotationData.ts

@@ -15,7 +15,7 @@ export type LabelData = {
   ids: Set<number>;
 };
 
-export interface IAnnotationData {
+export interface IAnnotationDataGetters {
   /**
    * Returns the array of label data objects.
    * @returns an array of `LabelData` objects, containing:
@@ -26,15 +26,6 @@ export interface IAnnotationData {
    */
   getLabels(): LabelData[];
 
-  /** Creates a new label and returns its index in the array of labels (as
-   * returned by `getLabels()`).
-   * @param name The name of the label. If no name is provided, a default name
-   * ("Label {number}") will be used.
-   * @param color The color to use for the label. If no color is provided, a
-   * color will be chosen from the default categorical palette.
-   * */
-  createNewLabel(name?: string, color?: Color): number;
-
   /**
    * Returns the indices of all labels that have been applied to the object id.
    * (Indices are ordered as returned by `getLabels()`.)
@@ -44,6 +35,15 @@ export interface IAnnotationData {
    */
   getLabelsAppliedToId(id: number): number[];
 
+  /**
+   * Returns whether the label has been applied to the object ID.
+   * @param labelIdx The index of the label to look up.
+   * @param id The object ID.
+   * @returns `true` if the label has been applied to the object ID, `false`
+   * otherwise.
+   */
+  isLabelOnId(labelIdx: number, id: number): boolean;
+
   /**
    * Returns all object IDs that the label has been applied to.
    * @param labelIdx The index of the label to look up object IDs for.
@@ -53,41 +53,58 @@ export interface IAnnotationData {
   getLabeledIds(labelIdx: number): number[];
 
   /**
-   * Returns a map from a frame number to the labeled object IDs present at that
-   * frame, represented as a record mapping label indices to IDs.
+   * Returns a time to label and ID map, in the format `{time: {labelId:
+   * [objectIds]}}`.
+   *
+   * Each time (by frame number) maps to a record of labeled object IDs present
+   * at that time. The record's keys are a label (by index), and the value is an
+   * array of the objects (by ID) present for that time frame that have that
+   * label applied. (If no objects have a label applied at a given time, it will
+   * not be present in the record.)
    *
-   * The record's keys are an index of a label (as returned by `getLabels()`),
-   * and the values are arrays of object IDs present for that time frame that
-   * have that label applied.
+   * A dataset with a single frame and an object with ID 0, labeled with label
+   * 0, would return `{0: {0: [0]}}`, in the format `{time: {labelId:
+   * [objectIds]}}`.
    *
    * @param dataset The dataset to use for time information.
    * @returns a map from time to a record of label indices to IDs.
    * @example
    * ```typescript
-   * // Let's say we have two labels (0 and 1). There are objects with IDs 1, 2,
-   * // and 3 at time 0, and IDs 4 and 5 at time 1.
-   * // Label 0 has been applied to objects 1, 2, 3, and 4.
-   * // Label 1 has been applied to objects 3 and 5.
+   * // Let's say we have two labels (0 and 1). There are objects with IDs 11, 12,
+   * // and 13 at time 234, and ID 14 and 15 at time 577.
+   * // Label 0 has been applied to objects 11, 12, and 13.
+   * // Label 1 has been applied to objects 13, 14, and 15.
    *
    * const timeToLabelMap = getTimeToLabelIdMap(some_dataset);
-   * timeToLabelMap.get(0); // { 0: [1, 2, 3], 1: [3] }
-   * timeToLabelMap.get(1); // { 0: [4], 1: [5]}
+   * timeToLabelMap.get(234); // { 0: [11, 12, 13], 1: [13] }
+   * timeToLabelMap.get(577); // { 1: [14, 15]}
    * ```
    * */
   getTimeToLabelIdMap(dataset: Dataset): Map<number, Record<number, number[]>>;
 
+  toCsv(dataset: Dataset, separator?: string): string;
+}
+
+export interface IAnnotationDataSetters {
+  /** Creates a new label and returns its index in the array of labels (as
+   * returned by `getLabels()`).
+   * @param name The name of the label. If no name is provided, a default name
+   * ("Label {number}") will be used.
+   * @param color The color to use for the label. If no color is provided, a
+   * color will be chosen from the default categorical palette.
+   * */
+  createNewLabel(name?: string, color?: Color): number;
+
   setLabelName(labelIdx: number, name: string): void;
   setLabelColor(labelIdx: number, color: Color): void;
   deleteLabel(labelIdx: number): void;
 
-  applyLabelToId(labelIdx: number, id: number): void;
-  removeLabelFromId(labelIdx: number, id: number): void;
-  toggleLabelOnId(labelIdx: number, id: number): void;
-
-  toCsv(dataset: Dataset, separator?: string): string;
+  setLabelOnId(labelIdx: number, id: number, value: boolean): void;
 }
 
-export class AnnotationData implements IAnnotationData {
+export type IAnnotationData = IAnnotationDataGetters & IAnnotationDataSetters;
+
+export class AnnotationData implements IAnnotationDataGetters, IAnnotationDataSetters {
   private labelData: LabelData[];
   private numLabelsCreated: number;
   /**
@@ -110,9 +127,8 @@ export class AnnotationData implements IAnnotationData {
     this.setLabelName = this.setLabelName.bind(this);
     this.setLabelColor = this.setLabelColor.bind(this);
     this.deleteLabel = this.deleteLabel.bind(this);
-    this.toggleLabelOnId = this.toggleLabelOnId.bind(this);
-    this.applyLabelToId = this.applyLabelToId.bind(this);
-    this.removeLabelFromId = this.removeLabelFromId.bind(this);
+    this.isLabelOnId = this.isLabelOnId.bind(this);
+    this.setLabelOnId = this.setLabelOnId.bind(this);
     this.toCsv = this.toCsv.bind(this);
   }
 
@@ -122,6 +138,11 @@ export class AnnotationData implements IAnnotationData {
     return [...this.labelData];
   }
 
+  isLabelOnId(labelIdx: number, id: number): boolean {
+    this.validateIndex(labelIdx);
+    return this.labelData[labelIdx].ids.has(id);
+  }
+
   getLabelsAppliedToId(id: number): number[] {
     const labelIdxs: number[] = [];
     for (let i = 0; i < this.labelData.length; i++) {
@@ -190,8 +211,8 @@ export class AnnotationData implements IAnnotationData {
     }
 
     this.labelData.push({
-      name: name,
-      color: color,
+      name,
+      color,
       ids: new Set(),
     });
 
@@ -221,24 +242,14 @@ export class AnnotationData implements IAnnotationData {
     this.timeToLabelIdMap = null;
   }
 
-  applyLabelToId(labelIdx: number, id: number): void {
-    this.validateIndex(labelIdx);
-    this.labelData[labelIdx].ids.add(id);
-    this.timeToLabelIdMap = null;
-  }
-
-  removeLabelFromId(labelIdx: number, id: number): void {
+  setLabelOnId(labelIdx: number, id: number, value: boolean): void {
     this.validateIndex(labelIdx);
-    this.labelData[labelIdx].ids.delete(id);
-    this.timeToLabelIdMap = null;
-  }
-
-  toggleLabelOnId(labelIdx: number, id: number): void {
-    if (!this.labelData[labelIdx].ids.has(id)) {
-      this.applyLabelToId(labelIdx, id);
+    if (value) {
+      this.labelData[labelIdx].ids.add(id);
     } else {
-      this.removeLabelFromId(labelIdx, id);
+      this.labelData[labelIdx].ids.delete(id);
     }
+    this.timeToLabelIdMap = null;
   }
 
   /**

src/colorizer/utils/react_utils.ts

@@ -1,10 +1,10 @@
-import React, { EventHandler, useEffect, useRef, useState } from "react";
+import React, { EventHandler, useEffect, useMemo, useRef, useState } from "react";
 import styled from "styled-components";
 import { useLocalStorage } from "usehooks-ts";
 
 import { VectorConfig } from "../types";
 
-import { AnnotationData, IAnnotationData } from "../AnnotationData";
+import { AnnotationData, IAnnotationDataGetters, IAnnotationDataSetters } from "../AnnotationData";
 import Dataset from "../Dataset";
 import SharedWorkerPool from "../workers/SharedWorkerPool";
 
@@ -308,22 +308,21 @@ export const useMotionDeltas = (
   return motionDeltas;
 };
 
-export type AnnotationState = IAnnotationData &  {
+export type AnnotationState =  {
   // Viewer state that lives outside the annotation data itself
   currentLabelIdx: number | null;
   setCurrentLabelIdx: (labelIdx: number) => void;
   isAnnotationModeEnabled: boolean;
   setIsAnnotationModeEnabled: (enabled: boolean) => void;
   visible: boolean;
   setVisibility: (visible: boolean) => void;
-  /*
-  * Increments every time the annotation data has changed (both label metadata
-  * and IDs the labels are applied to). Can be used as a dependency in React
-  * hooks like `useEffect` or `useMemo` to trigger updates when the annotation
-  * data changes.
-  */
-  annotationDataVersion: number;
-} ;
+  /**
+   * Contains annotation data getters. Use this object directly as a dependency
+   * in `useMemo` or `useCallback` to trigger updates when the underlying data
+   * changes.
+   */
+  data: IAnnotationDataGetters;
+} & IAnnotationDataSetters;
 
 export const useAnnotations = (): AnnotationState => {
   const annotationData = useConstructor(() => new AnnotationData());
@@ -357,13 +356,11 @@ export const useAnnotations = (): AnnotationState => {
   const wrapFunctionInUpdate = <F extends (...args: any[]) => void>(fn: F): F => {
     return <F>function (...args: any[]) {
       const result = fn(...args);
-      setDataUpdateCounter((value) => {return value + 1;});
+      setDataUpdateCounter((value) => value + 1);
       return result;
     };
   };
 
-
-
   const onDeleteLabel = (labelIdx: number): void => {
     if (currentLabelIdx === null) {
       return;
@@ -375,36 +372,38 @@ export const useAnnotations = (): AnnotationState => {
     } else if (currentLabelIdx === labelIdx) {
       setCurrentLabelIdx(null);
       setIsAnnotationEnabled(false);
-    }else if (currentLabelIdx > labelIdx) {
+    } else if (currentLabelIdx > labelIdx) {
       // Decrement because all indices will shift over
       setCurrentLabelIdx(currentLabelIdx - 1);
     }
     return annotationData.deleteLabel(labelIdx);
   };
 
+  const data = useMemo(() => ({
+      // Data getters
+      getLabels: annotationData.getLabels,
+      getLabelsAppliedToId: annotationData.getLabelsAppliedToId,
+      getLabeledIds: annotationData.getLabeledIds,
+      getTimeToLabelIdMap: annotationData.getTimeToLabelIdMap,
+      isLabelOnId: annotationData.isLabelOnId,
+      toCsv: annotationData.toCsv,
+    })
+  , [dataUpdateCounter]);
+
   return {
     // UI state
-    annotationDataVersion: dataUpdateCounter,
     currentLabelIdx,
     setCurrentLabelIdx,
     isAnnotationModeEnabled: isAnnotationEnabled,
     setIsAnnotationModeEnabled: setIsAnnotationEnabled,
     visible,
     setVisibility,
-    // Data getters
-    getLabels: annotationData.getLabels,
-    getLabelsAppliedToId: annotationData.getLabelsAppliedToId,
-    getLabeledIds: annotationData.getLabeledIds,
-    getTimeToLabelIdMap: annotationData.getTimeToLabelIdMap,
-    toCsv: annotationData.toCsv,
+    data,
     // Wrap state mutators
     createNewLabel: wrapFunctionInUpdate(annotationData.createNewLabel),
     setLabelName: wrapFunctionInUpdate(annotationData.setLabelName),
     setLabelColor: wrapFunctionInUpdate(annotationData.setLabelColor),
     deleteLabel: wrapFunctionInUpdate(onDeleteLabel),
-
-    applyLabelToId: wrapFunctionInUpdate(annotationData.applyLabelToId),
-    removeLabelFromId: wrapFunctionInUpdate(annotationData.removeLabelFromId),
-    toggleLabelOnId: wrapFunctionInUpdate(annotationData.toggleLabelOnId),
+    setLabelOnId: wrapFunctionInUpdate(annotationData.setLabelOnId),
   };
 };

src/components/CanvasWrapper.tsx

@@ -348,18 +348,13 @@ export default function CanvasWrapper(inputProps: CanvasWrapperProps): ReactElem
   }, [props.config.outlineColor]);
 
   useMemo(() => {
-    const annotationLabels = props.annotationState.getLabels();
+    const annotationLabels = props.annotationState.data.getLabels();
     const timeToAnnotationLabelIds = props.dataset
-      ? props.annotationState.getTimeToLabelIdMap(props.dataset)
+      ? props.annotationState.data.getTimeToLabelIdMap(props.dataset)
       : new Map();
     canv.setAnnotationData(annotationLabels, timeToAnnotationLabelIds, props.annotationState.currentLabelIdx);
     canv.isAnnotationVisible = props.annotationState.visible;
-  }, [
-    props.dataset,
-    props.annotationState.annotationDataVersion,
-    props.annotationState.currentLabelIdx,
-    props.annotationState.visible,
-  ]);
+  }, [props.dataset, props.annotationState.data, props.annotationState.currentLabelIdx, props.annotationState.visible]);
 
   // CANVAS RESIZING /////////////////////////////////////////////////
 
@@ -686,7 +681,7 @@ export default function CanvasWrapper(inputProps: CanvasWrapperProps): ReactElem
     )
   );
 
-  const labels = props.annotationState.getLabels();
+  const labels = props.annotationState.data.getLabels();
   const annotationTooltipContents: ReactNode[] = [];
   annotationTooltipContents.push(
     <span key="annotation-count">

src/components/Tabs/AnnotationTab.tsx

@@ -73,22 +73,20 @@ export default function AnnotationTab(props: AnnotationTabProps): ReactElement {
     setIsAnnotationModeEnabled,
     currentLabelIdx,
     setCurrentLabelIdx,
-    getLabels,
+    data: annotationData,
     createNewLabel,
     deleteLabel,
-    annotationDataVersion,
-    getLabeledIds,
     setLabelName,
     setLabelColor,
-    removeLabelFromId,
+    setLabelOnId,
   } = props.annotationState;
 
   const [showEditPopover, setShowEditPopover] = useState(false);
   const [editPopoverNameInput, setEditPopoverNameInput] = useState("");
   const editPopoverContainerRef = useRef<HTMLDivElement>(null);
   const editPopoverInputRef = useRef<InputRef>(null);
 
-  const labels = getLabels();
+  const labels = annotationData.getLabels();
   const selectedLabel: LabelData | undefined = labels[currentLabelIdx ?? -1];
 
   const onSelectLabel = (labelIdx: number): void => {
@@ -166,7 +164,7 @@ export default function AnnotationTab(props: AnnotationTabProps): ReactElement {
             // so we need to stop event propagation so that the row click event
             // doesn't fire.
             event.stopPropagation();
-            removeLabelFromId(currentLabelIdx!, record.id);
+            setLabelOnId(currentLabelIdx!, record.id, false);
           }}
         >
           <CloseOutlined />
@@ -181,15 +179,15 @@ export default function AnnotationTab(props: AnnotationTabProps): ReactElement {
   const tableData: TableDataType[] = useMemo(() => {
     const dataset = props.dataset;
     if (currentLabelIdx !== null && dataset) {
-      const ids = getLabeledIds(currentLabelIdx);
+      const ids = annotationData.getLabeledIds(currentLabelIdx);
       return ids.map((id) => {
         const track = dataset.getTrackId(id);
         const time = dataset.getTime(id);
         return { key: id.toString(), id, track, time };
       });
     }
     return [];
-  }, [annotationDataVersion, currentLabelIdx, props.dataset]);
+  }, [annotationData, currentLabelIdx, props.dataset]);
 
   // Options for the selection dropdown
   const selectLabelOptions: ItemType[] = labels.map((label, index) => {

src/components/Tooltips/CanvasHoverTooltip.tsx

@@ -115,8 +115,8 @@ export default function CanvasHoverTooltip(props: PropsWithChildren<CanvasHoverT
   }
 
   // Show all current labels applied to the hovered object
-  const labels = props.annotationState.getLabelsAppliedToId(lastHoveredId);
-  const labelData = props.annotationState.getLabels();
+  const labels = props.annotationState.data.getLabelsAppliedToId(lastHoveredId);
+  const labelData = props.annotationState.data.getLabels();
   if (labels.length > 0 && props.annotationState.visible) {
     objectInfoContent.push(
       <div style={{ lineHeight: "28px" }}>