feat: Export annotations as a CSV

allen-cell-animated · Dec 23, 2024 · 05bae8b · 05bae8b
1 parent cf94f4f
commit 05bae8b
Show file tree

Hide file tree

Showing 4 changed files with 96 additions and 16 deletions.
diff --git 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<number, number[]> {
+    const idsToLabels = new Map<number, number[]>();
+    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
@@ -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
@@ -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 (
     <FlexColumnAlignCenter $gap={10}>
       {/* Top-level annotation edit toggle */}
-      <FlexRow style={{ width: "100%" }} $gap={6}>
-        <AnnotationModeButton
-          type="primary"
-          $color={isAnnotationModeEnabled ? "success" : "default"}
-          style={{ paddingLeft: "10px" }}
-          onClick={() => setIsAnnotationModeEnabled(!isAnnotationModeEnabled)}
+      <FlexRow style={{ width: "100%", justifyContent: "space-between" }}>
+        <FlexRow $gap={6}>
+          <AnnotationModeButton
+            type="primary"
+            $color={isAnnotationModeEnabled ? "success" : "default"}
+            style={{ paddingLeft: "10px" }}
+            onClick={() => setIsAnnotationModeEnabled(!isAnnotationModeEnabled)}
+          >
+            <FlexRowAlignCenter $gap={6}>
+              {isAnnotationModeEnabled ? <CheckOutlined /> : <TagIconSVG />}
+              {isAnnotationModeEnabled ? "Done editing" : "Create and edit"}
+            </FlexRowAlignCenter>
+          </AnnotationModeButton>
+          {isAnnotationModeEnabled && (
+            <p style={{ color: theme.color.text.hint }}>
+              <i>Editing in progress...</i>
+            </p>
+          )}
+        </FlexRow>
+        <Button
+          onClick={() => {
+            const csvData = props.annotationState.toCsv(props.dataset!);
+            download("annotations.csv", "data:text/csv;charset=utf-8," + encodeURIComponent(csvData));
+            console.log(csvData);
+          }}
         >
-          <FlexRowAlignCenter $gap={6}>
-            {isAnnotationModeEnabled ? <CheckOutlined /> : <TagIconSVG />}
-            {isAnnotationModeEnabled ? "Done editing" : "Create and edit"}
-          </FlexRowAlignCenter>
-        </AnnotationModeButton>
-        {isAnnotationModeEnabled && (
-          <p style={{ color: theme.color.text.hint }}>
-            <i>Editing in progress...</i>
-          </p>
-        )}
+          Export as CSV
+        </Button>
       </FlexRow>
 
       {/* Label selection and edit/create/delete buttons */}

diff --git a/src/utils/file_io.ts 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);
+}