diff --git a/src/colorizer/AnnotationData.ts b/src/colorizer/AnnotationData.ts index 6156c52e..9761713f 100644 --- a/src/colorizer/AnnotationData.ts +++ b/src/colorizer/AnnotationData.ts @@ -4,6 +4,10 @@ import { DEFAULT_CATEGORICAL_PALETTE_KEY, KNOWN_CATEGORICAL_PALETTES } from "./c import Dataset from "./Dataset"; +const CSV_COL_ID = "ID"; +const CSV_COL_TIME = "Frame"; +const CSV_COL_TRACK = "Track"; + export type LabelData = { name: string; color: Color; @@ -78,6 +82,8 @@ export interface IAnnotationData { applyLabelToId(labelIdx: number, id: number): void; removeLabelFromId(labelIdx: number, id: number): void; toggleLabelOnId(labelIdx: number, id: number): void; + + toCsv(dataset: Dataset, separator?: string): string; } export class AnnotationData implements IAnnotationData { @@ -106,6 +112,7 @@ export class AnnotationData implements IAnnotationData { this.toggleLabelOnId = this.toggleLabelOnId.bind(this); this.applyLabelToId = this.applyLabelToId.bind(this); this.removeLabelFromId = this.removeLabelFromId.bind(this); + this.toCsv = this.toCsv.bind(this); } // Getters @@ -156,6 +163,19 @@ export class AnnotationData implements IAnnotationData { return timeToLabelIdMap; } + getIdsToLabels(): Map { + const idsToLabels = new Map(); + for (let labelIdx = 0; labelIdx < this.labelData.length; labelIdx++) { + for (const id of this.labelData[labelIdx].ids) { + if (!idsToLabels.has(id)) { + idsToLabels.set(id, []); + } + idsToLabels.get(id)!.push(labelIdx); + } + } + return idsToLabels; + } + // Setters /** Creates a new label and returns its index. */ @@ -219,4 +239,34 @@ export class AnnotationData implements IAnnotationData { this.removeLabelFromId(labelIdx, id); } } + + toCsv(dataset: Dataset, separator: string = ","): string { + const idsToLabels = this.getIdsToLabels(); + + // TODO: Sanitizing and parsing CSV data is pretty complex. + // Consider using `papaparse` or another library. + const headerRow = [CSV_COL_ID, CSV_COL_TRACK, CSV_COL_TIME]; + headerRow.push(...this.labelData.map((label) => `"${label.name}"`)); + + const csvRows: string[] = []; + for (const [id, labels] of idsToLabels) { + const track = dataset.getTrackId(id); + const time = dataset.getTime(id); + + const row = [id.toString(), track, time.toString()]; + for (let i = 0; i < this.labelData.length; i++) { + row.push(labels.includes(i) ? "1" : "0"); + } + csvRows.push(row.join(separator)); + } + + // TODO: Add additional metadata when we add support for importing + // annotation data. + // Add comment metadata for label name + colors + // const metadataRows = this.labelData.map((label) => { + // return `# ${label.name},#${label.color.getHexString()}`; + // }); + + return [headerRow.join(separator), ...csvRows].join("\n"); + } } diff --git a/src/colorizer/utils/react_utils.ts b/src/colorizer/utils/react_utils.ts index 7fb923d3..eb09af44 100644 --- a/src/colorizer/utils/react_utils.ts +++ b/src/colorizer/utils/react_utils.ts @@ -396,6 +396,7 @@ export const useAnnotations = (): AnnotationState => { getLabelsAppliedToId: annotationData.getLabelsAppliedToId, getLabeledIds: annotationData.getLabeledIds, getTimeToLabelIdMap: annotationData.getTimeToLabelIdMap, + toCsv: annotationData.toCsv, // Wrap state mutators createNewLabel: wrapFunctionInUpdate(annotationData.createNewLabel), setLabelName: wrapFunctionInUpdate(annotationData.setLabelName), diff --git a/src/components/Tabs/AnnotationTab.tsx b/src/components/Tabs/AnnotationTab.tsx index 6bb0b29b..1aab79af 100644 --- a/src/components/Tabs/AnnotationTab.tsx +++ b/src/components/Tabs/AnnotationTab.tsx @@ -10,6 +10,7 @@ import { TagIconSVG } from "../../assets"; import { Dataset } from "../../colorizer"; import { AnnotationState } from "../../colorizer/utils/react_utils"; import { FlexColumn, FlexColumnAlignCenter, FlexRow, FlexRowAlignCenter, VisuallyHidden } from "../../styles/utils"; +import { download } from "../../utils/file_io"; import { LabelData } from "../../colorizer/AnnotationData"; import { AppThemeContext } from "../AppStyle"; @@ -234,23 +235,34 @@ export default function AnnotationTab(props: AnnotationTabProps): ReactElement { return ( {/* Top-level annotation edit toggle */} - - setIsAnnotationModeEnabled(!isAnnotationModeEnabled)} + + + setIsAnnotationModeEnabled(!isAnnotationModeEnabled)} + > + + {isAnnotationModeEnabled ? : } + {isAnnotationModeEnabled ? "Done editing" : "Create and edit"} + + + {isAnnotationModeEnabled && ( +

+ Editing in progress... +

+ )} +
+
{/* Label selection and edit/create/delete buttons */} diff --git a/src/utils/file_io.ts b/src/utils/file_io.ts new file mode 100644 index 00000000..10b4c5ec --- /dev/null +++ b/src/utils/file_io.ts @@ -0,0 +1,17 @@ +/** + * Initiates a browser download of a file from a URL. + * @param filename The default filename to save the file as. + * @param url The `href` attribute of the download link. This can either be a data + * URL or + */ +export function download(filename: string, url: string): void { + // TODO: Add option to show save file picker? https://developer.mozilla.org/en-US/docs/Web/API/Window/showSaveFilePicker + const anchor = document.createElement("a"); + document.body.appendChild(anchor); + + anchor.href = url; + anchor.download = filename; + anchor.click(); + + document.body.removeChild(anchor); +}