diff --git a/src/colorizer/AnnotationData.ts b/src/colorizer/AnnotationData.ts new file mode 100644 index 000000000..02be5b61f --- /dev/null +++ b/src/colorizer/AnnotationData.ts @@ -0,0 +1,233 @@ +import { Color } from "three"; + +import { DEFAULT_CATEGORICAL_PALETTE_KEY, KNOWN_CATEGORICAL_PALETTES } from "./colors/categorical_palettes"; + +import Dataset from "./Dataset"; + +export type LabelData = { + name: string; + color: Color; + ids: Set; +}; + +export interface IAnnotationDataGetters { + /** + * Returns the array of label data objects. + * @returns an array of `LabelData` objects, containing: + * - `name`: The name of the label. + * - `color`: The color of the label. + * - `ids`: A set of object IDs that have the label applied. + * Object properties SHOULD NOT be modified directly. + */ + getLabels(): LabelData[]; + + /** + * Returns the indices of all labels that have been applied to the object id. + * (Indices are ordered as returned by `getLabels()`.) + * @param id The object ID to look up label indices for. + * @returns an array of label indices. Empty if no labels have been applied to + * the ID. + */ + 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. + * @returns an array of object IDs. Empty if the label has not been applied to + * any object IDs. + */ + getLabeledIds(labelIdx: number): number[]; + + /** + * 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.) + * + * 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 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(234); // { 0: [11, 12, 13], 1: [13] } + * timeToLabelMap.get(577); // { 1: [14, 15]} + * ``` + * */ + getTimeToLabelIdMap(dataset: Dataset): Map>; +} + +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; + + setLabelOnId(labelIdx: number, id: number, value: boolean): void; +} + +export type IAnnotationData = IAnnotationDataGetters & IAnnotationDataSetters; + +export class AnnotationData implements IAnnotationData { + private labelData: LabelData[]; + private numLabelsCreated: number; + /** + * Cached mapping from time to label indices to IDs. Must be invalidated when + * labels are removed, or if annotations are applied to or removed from + * objects. + */ + private timeToLabelIdMap: Map> | null; + + constructor() { + this.labelData = []; + this.numLabelsCreated = 0; + this.timeToLabelIdMap = null; + + this.getLabelsAppliedToId = this.getLabelsAppliedToId.bind(this); + this.getLabeledIds = this.getLabeledIds.bind(this); + this.getTimeToLabelIdMap = this.getTimeToLabelIdMap.bind(this); + this.getLabels = this.getLabels.bind(this); + this.createNewLabel = this.createNewLabel.bind(this); + this.setLabelName = this.setLabelName.bind(this); + this.setLabelColor = this.setLabelColor.bind(this); + this.deleteLabel = this.deleteLabel.bind(this); + this.isLabelOnId = this.isLabelOnId.bind(this); + this.setLabelOnId = this.setLabelOnId.bind(this); + } + + // Getters + + getLabels(): LabelData[] { + 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++) { + if (this.labelData[i].ids.has(id)) { + labelIdxs.push(i); + } + } + return labelIdxs; + } + + getLabeledIds(labelIdx: number): number[] { + this.validateIndex(labelIdx); + return Array.from(this.labelData[labelIdx].ids); + } + + getTimeToLabelIdMap(dataset: Dataset): Map> { + if (this.timeToLabelIdMap !== null) { + return this.timeToLabelIdMap; + } + + const timeToLabelIdMap = new Map>(); + + for (let labelIdx = 0; labelIdx < this.labelData.length; labelIdx++) { + const ids = this.labelData[labelIdx].ids; + for (const id of ids) { + const time = dataset.times?.[id]; + if (time === undefined) { + continue; + } + if (!timeToLabelIdMap.has(time)) { + timeToLabelIdMap.set(time, {}); + } + if (!timeToLabelIdMap.get(time)![labelIdx]) { + timeToLabelIdMap.get(time)![labelIdx] = []; + } + timeToLabelIdMap.get(time)![labelIdx].push(id); + } + } + this.timeToLabelIdMap = timeToLabelIdMap; + return timeToLabelIdMap; + } + + // Setters + + /** Creates a new label and returns its index. */ + createNewLabel(name?: string, color?: Color): number { + if (!color) { + const palette = KNOWN_CATEGORICAL_PALETTES.get(DEFAULT_CATEGORICAL_PALETTE_KEY)!; + color = new Color(palette.colorStops[this.numLabelsCreated % palette.colorStops.length]); + } + if (!name) { + name = `Label ${this.numLabelsCreated + 1}`; + } + + this.labelData.push({ + name, + color, + ids: new Set(), + }); + + this.numLabelsCreated++; + return this.labelData.length - 1; + } + + private validateIndex(idx: number): void { + if (idx < 0 || idx >= this.labelData.length) { + throw new Error(`Invalid label index: ${idx}`); + } + } + + setLabelName(labelIdx: number, name: string): void { + this.validateIndex(labelIdx); + this.labelData[labelIdx].name = name; + } + + setLabelColor(labelIdx: number, color: Color): void { + this.validateIndex(labelIdx); + this.labelData[labelIdx].color = color; + } + + deleteLabel(labelIdx: number): void { + this.validateIndex(labelIdx); + this.labelData.splice(labelIdx, 1); + this.timeToLabelIdMap = null; + } + + setLabelOnId(labelIdx: number, id: number, value: boolean): void { + this.validateIndex(labelIdx); + if (value) { + this.labelData[labelIdx].ids.add(id); + } else { + this.labelData[labelIdx].ids.delete(id); + } + this.timeToLabelIdMap = null; + } +} diff --git a/src/colorizer/utils/react_utils.ts b/src/colorizer/utils/react_utils.ts index 28c21872a..d50f7eab7 100644 --- a/src/colorizer/utils/react_utils.ts +++ b/src/colorizer/utils/react_utils.ts @@ -1,12 +1,16 @@ -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, IAnnotationDataGetters, IAnnotationDataSetters } from "../AnnotationData"; import Dataset from "../Dataset"; import SharedWorkerPool from "../workers/SharedWorkerPool"; +// TODO: Move this to a folder outside of `colorizer`. +// TODO: Split this up into multiple files. + /** * Delays changes to a value until no changes have occurred for the * set delay period, in milliseconds. This is useful for delaying changes to state @@ -250,14 +254,15 @@ export const useRecentCollections = (): [RecentCollection[], (collection: Recent }; /** - * Wrapper around the SharedWorkerPool method `getMotionDeltas`. Returns a debounced motion - * delta array for the given dataset and vector field configuration. + * Wrapper around the SharedWorkerPool method `getMotionDeltas`. Returns a + * debounced motion delta array for the given dataset and vector field + * configuration. * @param dataset The dataset to calculate motion deltas for. * @param workerPool The worker pool to use for asynchronous calculations. * @param config The vector field configuration to use. * @param debounceMs The debounce time in milliseconds. Defaults to 100ms. - * @returns The motion delta array or `null` if the dataset is invalid. Data will be - * asynchronously updated as calculations complete. + * @returns The motion delta array or `null` if the dataset is invalid. Data + * will be asynchronously updated as calculations complete. */ export const useMotionDeltas = ( dataset: Dataset | null, @@ -302,3 +307,103 @@ export const useMotionDeltas = ( return motionDeltas; }; + +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; + /** + * 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()); + + const [currentLabelIdx, setCurrentLabelIdx] = useState(null); + const [isAnnotationEnabled, _setIsAnnotationEnabled] = useState(false); + const [visible, _setVisibility] = useState(true); + + // Annotation mode can only be enabled if there is at least one label, so create + // one if necessary. + const setIsAnnotationEnabled = (enabled: boolean): void => { + if (enabled) { + _setVisibility(true); + if (annotationData.getLabels().length === 0) { + const newLabelIdx = annotationData.createNewLabel(); + setCurrentLabelIdx(newLabelIdx); + } + } + _setIsAnnotationEnabled(enabled); + }; + + const setVisibility = (visible: boolean): void => { + _setVisibility(visible); + if (!visible) { + setIsAnnotationEnabled(false); + } + }; + /** Increments every time a state update is required. */ + const [dataUpdateCounter, setDataUpdateCounter] = useState(0); + + const wrapFunctionInUpdate = void>(fn: F): F => { + return function (...args: any[]) { + const result = fn(...args); + setDataUpdateCounter((value) => value + 1); + return result; + }; + }; + + const onDeleteLabel = (labelIdx: number): void => { + if (currentLabelIdx === null) { + return; + } + // Update selected label index if necessary. + const labels = annotationData.getLabels(); + if (currentLabelIdx === labelIdx && labels.length > 1) { + setCurrentLabelIdx(Math.max(currentLabelIdx - 1, 0)); + } else if (currentLabelIdx === labelIdx) { + setCurrentLabelIdx(null); + setIsAnnotationEnabled(false); + } 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, + }) + , [dataUpdateCounter]); + + return { + // UI state + currentLabelIdx, + setCurrentLabelIdx, + isAnnotationModeEnabled: isAnnotationEnabled, + setIsAnnotationModeEnabled: setIsAnnotationEnabled, + visible, + setVisibility, + + data, + // Wrap state mutators + createNewLabel: wrapFunctionInUpdate(annotationData.createNewLabel), + setLabelName: wrapFunctionInUpdate(annotationData.setLabelName), + setLabelColor: wrapFunctionInUpdate(annotationData.setLabelColor), + deleteLabel: wrapFunctionInUpdate(onDeleteLabel), + setLabelOnId: wrapFunctionInUpdate(annotationData.setLabelOnId), + }; +}; diff --git a/tests/colorizer/AnnotationData.test.ts b/tests/colorizer/AnnotationData.test.ts new file mode 100644 index 000000000..716824047 --- /dev/null +++ b/tests/colorizer/AnnotationData.test.ts @@ -0,0 +1,144 @@ +import { Color } from "three"; +import { describe, expect, it } from "vitest"; + +import { Dataset, DEFAULT_CATEGORICAL_PALETTE_KEY, KNOWN_CATEGORICAL_PALETTES } from "../../src/colorizer"; + +import { AnnotationData } from "../../src/colorizer/AnnotationData"; + +describe("AnnotationData", () => { + const defaultPalette = KNOWN_CATEGORICAL_PALETTES.get(DEFAULT_CATEGORICAL_PALETTE_KEY)!; + + it("creates and returns index of new labels", () => { + const annotationData = new AnnotationData(); + const labelIdx1 = annotationData.createNewLabel(); + const labelIdx2 = annotationData.createNewLabel(); + const labelIdx3 = annotationData.createNewLabel(); + expect(labelIdx1).toBe(0); + expect(labelIdx2).toBe(1); + expect(labelIdx3).toBe(2); + }); + + it("gives default colors and names to labels", () => { + const annotationData = new AnnotationData(); + annotationData.createNewLabel(); + annotationData.createNewLabel(); + annotationData.createNewLabel(); + + expect(annotationData.getLabels()).to.deep.equal([ + { name: "Label 1", color: defaultPalette.colors[0], ids: new Set() }, + { name: "Label 2", color: defaultPalette.colors[1], ids: new Set() }, + { name: "Label 3", color: defaultPalette.colors[2], ids: new Set() }, + ]); + }); + + it("allows updating of label names and colors", () => { + const annotationData = new AnnotationData(); + annotationData.createNewLabel(); + annotationData.createNewLabel(); + annotationData.createNewLabel(); + + annotationData.setLabelName(0, "New Label Name"); + annotationData.setLabelColor(1, new Color("#FF0000")); + annotationData.setLabelName(2, "Another New Label Name"); + annotationData.setLabelColor(2, new Color("#00FF00")); + + expect(annotationData.getLabels()).to.deep.equal([ + { name: "New Label Name", color: defaultPalette.colors[0], ids: new Set() }, + { name: "Label 2", color: new Color("#FF0000"), ids: new Set() }, + { name: "Another New Label Name", color: new Color("#00FF00"), ids: new Set() }, + ]); + }); + + it("deletes labels", () => { + const annotationData = new AnnotationData(); + annotationData.createNewLabel(); + annotationData.createNewLabel(); + annotationData.createNewLabel(); + + annotationData.deleteLabel(1); + expect(annotationData.getLabels()).to.deep.equal([ + { name: "Label 1", color: defaultPalette.colors[0], ids: new Set() }, + { name: "Label 3", color: defaultPalette.colors[2], ids: new Set() }, + ]); + + // Creating new label should reuse deleted index and increment name by 1 + annotationData.createNewLabel(); + expect(annotationData.getLabels()).to.deep.equal([ + { name: "Label 1", color: defaultPalette.colors[0], ids: new Set() }, + { name: "Label 3", color: defaultPalette.colors[2], ids: new Set() }, + { name: "Label 4", color: defaultPalette.colors[3], ids: new Set() }, + ]); + }); + + it("can apply and remove labels from an ID", () => { + const annotationData = new AnnotationData(); + annotationData.createNewLabel(); + annotationData.createNewLabel(); + annotationData.setLabelOnId(0, 0, true); + annotationData.setLabelOnId(0, 35, true); + annotationData.setLabelOnId(0, 458, true); + annotationData.setLabelOnId(1, 35, true); + + expect(annotationData.getLabelsAppliedToId(0)).to.deep.equal([0]); + expect(annotationData.getLabelsAppliedToId(35)).to.deep.equal([0, 1]); + expect(annotationData.getLabelsAppliedToId(458)).to.deep.equal([0]); + + annotationData.setLabelOnId(0, 35, false); + expect(annotationData.getLabelsAppliedToId(0)).to.deep.equal([0]); + expect(annotationData.getLabelsAppliedToId(35)).to.deep.equal([1]); + expect(annotationData.getLabelsAppliedToId(458)).to.deep.equal([0]); + }); + + it("ignores duplicate calls to setLabelOnId", () => { + const annotationData = new AnnotationData(); + annotationData.createNewLabel(); + annotationData.createNewLabel(); + annotationData.setLabelOnId(0, 0, true); + annotationData.setLabelOnId(0, 0, true); + annotationData.setLabelOnId(0, 1, true); + + expect(annotationData.getLabelsAppliedToId(0)).to.deep.equal([0]); + expect(annotationData.isLabelOnId(0, 0)).toBe(true); + expect(annotationData.getLabelsAppliedToId(1)).to.deep.equal([0]); + expect(annotationData.isLabelOnId(0, 1)).toBe(true); + + annotationData.setLabelOnId(0, 0, false); + annotationData.setLabelOnId(0, 0, false); + annotationData.setLabelOnId(0, 1, false); + + expect(annotationData.getLabelsAppliedToId(0)).to.deep.equal([]); + expect(annotationData.isLabelOnId(0, 0)).toBe(false); + expect(annotationData.getLabelsAppliedToId(1)).to.deep.equal([]); + expect(annotationData.isLabelOnId(0, 1)).toBe(false); + }); + + it("can return mapping from time to labeled IDs", () => { + const mockDataset = { + times: [0, 1, 1, 2, 3, 4], + }; + const annotationData = new AnnotationData(); + annotationData.createNewLabel(); + annotationData.setLabelOnId(0, 0, true); + annotationData.setLabelOnId(0, 1, true); + annotationData.setLabelOnId(0, 2, true); + annotationData.setLabelOnId(0, 3, true); + annotationData.setLabelOnId(0, 4, true); + annotationData.setLabelOnId(0, 5, true); + + annotationData.createNewLabel(); + annotationData.setLabelOnId(1, 2, true); + + /* eslint-disable @typescript-eslint/naming-convention */ + // ESLint doesn't like "0" and "1" being property keys. + expect(annotationData.getTimeToLabelIdMap(mockDataset as unknown as Dataset)).to.deep.equal( + new Map([ + [0, { 0: [0] }], + [1, { 0: [1, 2], 1: [2] }], + [2, { 0: [3] }], + [3, { 0: [4] }], + [4, { 0: [5] }], + ]) + ); + /* eslint-enable @typescript-eslint/naming-convention */ + }); +});