CSV download

src/components/App.tsx

@@ -16,6 +16,7 @@
 import LeftSidebarWrapper from "./leftSidebar/LeftSidebarWrapper";
 import { TimestampOverlay } from "./overlays/TimestampOverlay";
 import { ColorMap } from "./overlays/ColorMap";
+import { TrackDownloadData } from "./DownloadButton";
 
 // Ideally we do this here so that we can use initial values as default values for React state.
 const initialViewerState = ViewerState.fromUrlHash(window.location.hash);
@@ -149,31 +150,31 @@
         // this fetches the entire lineage for each track
         const updateTracks = async () => {
             console.debug("updateTracks: ", canvas.selectedPointIds);
-            for (const pointId of canvas.selectedPointIds) {
-                if (canvas.fetchedPointIds.has(pointId)) continue;
+            canvas.selectedPointIds.forEach(async (pointId) => {
+                if (canvas.fetchedPointIds.has(pointId)) return;
                 setNumLoadingTracks((n) => n + 1);
                 canvas.fetchedPointIds.add(pointId);
                 const trackIds = await trackManager.fetchTrackIDsForPoint(pointId);
                 // TODO: points actually only belong to one track, so can get rid of the outer loop
-                for (const trackId of trackIds) {
-                    if (canvas.fetchedRootTrackIds.has(trackId)) continue;
+                trackIds.forEach(async (trackId) => {
+                    if (canvas.fetchedRootTrackIds.has(trackId)) return;
                     canvas.fetchedRootTrackIds.add(trackId);
-                    const lineage = await trackManager.fetchLineageForTrack(trackId);
-                    for (const relatedTrackId of lineage) {
-                        if (canvas.tracks.has(relatedTrackId)) continue;
+                    const [lineage, trackData] = await trackManager.fetchLineageForTrack(trackId);
+                    lineage.forEach(async (relatedTrackId: number, index) => {
+                        if (canvas.tracks.has(relatedTrackId)) return;
                         const [pos, ids] = await trackManager.fetchPointsForTrack(relatedTrackId);
                         // adding the track *in* the dispatcher creates issues with duplicate fetching
                         // but we refresh so the selected/loaded count is updated
-                        canvas.addTrack(relatedTrackId, pos, ids);
+                        canvas.addTrack(relatedTrackId, pos, ids, trackData[index]);
                         dispatchCanvas({ type: ActionType.REFRESH });
-                    }
-                }
+                    });
+                });
                 setNumLoadingTracks((n) => n - 1);
-            }
+            });
         };
         updateTracks();
         // TODO: add missing dependencies
     }, [trackManager, dispatchCanvas, canvas.selectedPointIds]);
 
     // playback time points
     // TODO: this is basic and may drop frames
@@ -194,6 +195,52 @@
         }
     }, [dispatchCanvas, numTimes, playing]);
 
+    const getTrackDownloadData = () => {
+        const trackData: TrackDownloadData[] = [];
+        canvas.tracks.forEach((track, trackID) => {
+            // Keep track of the timepoints we've seen in this track to avoid duplication
+            // This is necessary because if a track contains a single point, we set
+            // the start and end positions to be the same
+            const timepointsInTrack = new Set();
+
+            const startPositions = track.threeTrack.geometry.getAttribute("instanceStart");
+            const startTimes = track.threeTrack.geometry.getAttribute("instanceTimeStart");
+
+            for (let i = 0; i < startTimes.count; i++) {
+                timepointsInTrack.add(startTimes.getX(i));
+                trackData.push([
+                    // trackID is 1-indexed in input and output CSVs
+                    trackID + 1,
+                    startTimes.getX(i),
+                    startPositions.getX(i),
+                    startPositions.getY(i),
+                    startPositions.getZ(i),
+                    track.parentTrackID,
+                ]);
+            }
+            const endPositions = track.threeTrack.geometry.getAttribute("instanceEnd");
+            const endTimes = track.threeTrack.geometry.getAttribute("instanceTimeEnd");
+            const lastIndex = endPositions.count - 1;
+
+            // Only add the end position if it's not the same as the start position
+            if (!timepointsInTrack.has(endTimes.getX(lastIndex))) {
+                trackData.push([
+                    // trackID is 1-indexed in input and output CSVs
+                    trackID + 1,
+                    endTimes.getX(lastIndex),
+                    endPositions.getX(lastIndex),
+                    endPositions.getY(lastIndex),
+                    endPositions.getZ(lastIndex),
+                    track.parentTrackID,
+                ]);
+            }
+        });
+
+        // Round to 3 decimal places
+        const formatter = Intl.NumberFormat("en-US", { useGrouping: false });
+        return trackData.map((row) => row.map(formatter.format));
+    };
+
     return (
         <Box sx={{ display: "flex", width: "100%", height: "100%" }}>
             {/* TODO: components *could* go deeper still for organization */}
@@ -234,6 +281,7 @@
                             clearTracks={() => {
                                 dispatchCanvas({ type: ActionType.REMOVE_ALL_TRACKS });
                             }}
+                            getTrackDownloadData={getTrackDownloadData}
                             numSelectedCells={numTracksLoaded}
                             trackManager={trackManager}
                             pointBrightness={canvas.pointBrightness}

src/components/CellControls.tsx

@@ -4,9 +4,11 @@ import { FontS, SmallCapsButton, ControlLabel } from "@/components/Styled";
 
 import { PointSelectionMode } from "@/lib/PointSelector";
 import { TrackManager } from "@/lib/TrackManager";
+import { DownloadButton } from "./DownloadButton";
 
 interface CellControlsProps {
     clearTracks: () => void;
+    getTrackDownloadData: () => string[][];
     numSelectedCells?: number;
     trackManager: TrackManager | null;
     pointBrightness: number;
@@ -33,6 +35,7 @@ export default function CellControls(props: CellControlsProps) {
             <FontS>
                 <strong>{props.numSelectedCells ?? 0}</strong> cells selected
             </FontS>
+            {!!props.numSelectedCells && <DownloadButton getDownloadData={props.getTrackDownloadData} />}
             <label htmlFor="selection-mode-control">
                 <ControlLabel>Selection Mode</ControlLabel>
             </label>

src/components/DownloadButton.tsx

@@ -0,0 +1,34 @@
+import { Button } from "@czi-sds/components";
+
+// TrackDownloadData is a list for each point
+// It contains trackID, time, x, y, z, parentTrackID
+export type TrackDownloadData = [number, number, number, number, number, number];
+interface DownloadButtonProps {
+    getDownloadData: () => string[][];
+}
+
+const dataHeaders = ["track_id", "t", "x", "y", "z", "parent_track_id"];
+
+const convertToCSV = (nestedArray: string[][]) => {
+    return nestedArray.map((row) => row.join(",")).join("\r\n");
+};
+
+export const DownloadButton = (props: DownloadButtonProps) => {
+    const downloadCSV = () => {
+        const data = props.getDownloadData();
+        const csvData = new Blob([`${convertToCSV([dataHeaders])}\r\n${convertToCSV(data)}`], { type: "text/csv" });
+        const csvURL = URL.createObjectURL(csvData);
+        const link = document.createElement("a");
+        link.href = csvURL;
+        link.download = "points.csv";
+        document.body.appendChild(link);
+        link.click();
+        document.body.removeChild(link);
+    };
+
+    return (
+        <Button sdsStyle="square" sdsType="primary" onClick={downloadCSV}>
+            Export Cell Tracks
+        </Button>
+    );
+};

src/lib/PointCanvas.ts

@@ -24,7 +24,12 @@ import { Track } from "@/lib/three/Track";
 import { PointSelector, PointSelectionMode } from "@/lib/PointSelector";
 import { ViewerState } from "./ViewerState";
 
-type Tracks = Map<number, Track>;
+// TrackType is a place to store the visual information about a track and any track-specific attributes
+type TrackType = {
+    threeTrack: Track;
+    parentTrackID: number;
+};
+type Tracks = Map<number, TrackType>;
 
 export class PointCanvas {
     readonly scene: Scene;
@@ -222,30 +227,30 @@ export class PointCanvas {
         this.points.geometry.computeBoundingSphere();
     }
 
-    addTrack(trackID: number, positions: Float32Array, ids: Int32Array): Track | null {
+    addTrack(trackID: number, positions: Float32Array, ids: Int32Array, parentTrackID: number): Track | null {
         if (this.tracks.has(trackID)) {
             // this is a warning because it should alert us to duplicate fetching
             console.warn("Track with ID %d already exists", trackID);
             return null;
         }
-        const track = Track.new(positions, ids, this.maxPointsPerTimepoint);
-        track.updateAppearance(this.showTracks, this.showTrackHighlights, this.minTime, this.maxTime);
-        this.tracks.set(trackID, track);
-        this.scene.add(track);
-        return track;
+        const threeTrack = Track.new(positions, ids, this.maxPointsPerTimepoint);
+        threeTrack.updateAppearance(this.showTracks, this.showTrackHighlights, this.minTime, this.maxTime);
+        this.tracks.set(trackID, { threeTrack, parentTrackID });
+        this.scene.add(threeTrack);
+        return threeTrack;
     }
 
     updateAllTrackHighlights() {
-        for (const track of this.tracks.values()) {
-            track.updateAppearance(this.showTracks, this.showTrackHighlights, this.minTime, this.maxTime);
-        }
+        this.tracks.forEach((track) => {
+            track.threeTrack.updateAppearance(this.showTracks, this.showTrackHighlights, this.minTime, this.maxTime);
+        });
     }
 
     removeTrack(trackID: number) {
         const track = this.tracks.get(trackID);
         if (track) {
-            this.scene.remove(track);
-            track.dispose();
+            this.scene.remove(track.threeTrack);
+            track.threeTrack.dispose();
             this.tracks.delete(trackID);
         } else {
             console.warn("No track with ID %d to remove", trackID);

src/lib/TrackManager.ts

@@ -129,10 +129,15 @@ export class TrackManager {
         return [flatPoints, pointIDs];
     }
 
-    async fetchLineageForTrack(trackID: number): Promise<Int32Array> {
+    async fetchLineageForTrack(trackID: number): Promise<[Int32Array, Int32Array]> {
         const rowStartEnd = await this.tracksToTracks.getIndPtr(slice(trackID, trackID + 2));
-        const lineage = await this.tracksToTracks.indices.get([slice(rowStartEnd[0], rowStartEnd[1])]);
-        return lineage.data;
+        const lineage = await this.tracksToTracks.indices
+            .get([slice(rowStartEnd[0], rowStartEnd[1])])
+            .then((lineage: SparseZarrArray) => lineage.data);
+        const trackData = await this.tracksToTracks.data
+            .get([slice(rowStartEnd[0], rowStartEnd[1])])
+            .then((trackData: SparseZarrArray) => trackData.data);
+        return Promise.all([lineage, trackData]);
     }
 }
 
@@ -146,7 +151,7 @@ export async function loadTrackManager(url: string) {
         });
         const pointsToTracks = await openSparseZarrArray(url, "points_to_tracks", false);
         const tracksToPoints = await openSparseZarrArray(url, "tracks_to_points", true);
-        const tracksToTracks = await openSparseZarrArray(url, "tracks_to_tracks", false);
+        const tracksToTracks = await openSparseZarrArray(url, "tracks_to_tracks", true);
         trackManager = new TrackManager(url, points, pointsToTracks, tracksToPoints, tracksToTracks);
     } catch (err) {
         console.error("Error opening TrackManager: %s", err);

src/lib/three/Track.ts

@@ -40,6 +40,15 @@ export class Track extends Mesh {
             // TODO: use a LUT for the main track, too
             colors.push(((0.9 * (n - i)) / n) ** 3, ((0.9 * (n - i)) / n) ** 3, (0.9 * (n - i)) / n);
         }
+
+        // if this track has a single point, in order to keep this point in our data,
+        // we are using the point as both the start and end of the line segment
+        if (pos.length === 3) {
+            pos.push(pos[0], pos[1], pos[2]);
+            colors.push(colors[0], colors[1], colors[2]);
+            time.push(time[0]);
+        }
+
         track.geometry.setPositions(pos);
         track.geometry.setColors(colors);
         track.geometry.setTime(time);

tools/convert_tracks_csv_to_sparse_zarr.py

@@ -1,17 +1,34 @@
+import argparse
 import csv
 import time
 from collections import Counter
+from pathlib import Path
 
 import numpy as np
 import zarr
 from scipy.sparse import lil_matrix
 
-root_dir = "/Users/aandersoniii/Data/tracking/"
+parser = argparse.ArgumentParser(description="Convert a CSV of tracks to a sparse Zarr store")
+parser.add_argument("csv_file", type=str, help="Path to the CSV file")
+parser.add_argument(
+    "out_dir",
+    type=str,
+    help="Path to the output directory (optional, defaults to the parent dir of the CSV file)",
+    nargs="?",
+)
+args = parser.parse_args()
+
+csv_file = Path(args.csv_file)
+if args.out_dir is None:
+    out_dir = csv_file.parent
+else:
+    out_dir = Path(args.out_dir)
+zarr_path = out_dir / f"{csv_file.stem}_bundle.zarr"
 
 start = time.monotonic()
 points = []
 points_in_timepoint = Counter()
-with open(root_dir + "ZSNS001_tracks.csv", "r") as f:
+with open(csv_file, "r") as f:
     reader = csv.reader(f)
     next(reader)  # Skip the header
     # TrackID,t,z,y,x,parent_track_id
@@ -34,9 +51,19 @@
 points_to_tracks = lil_matrix((timepoints * max_points_in_timepoint, tracks), dtype=np.int32)
 tracks_to_children = lil_matrix((tracks, tracks), dtype=np.int32)
 tracks_to_parents = lil_matrix((tracks, tracks), dtype=np.int32)
+
+# create a map of the track_index to the parent_track_index
+# track_id and parent_track_id are 1-indexed, track_index and parent_track_index are 0-indexed
+direct_parent_index_map = {}  
+
 for point in points:
-    track_id, t, z, y, x, parent_track_id, n = point
-    point_id = t * max_points_in_timepoint + n
+    track_id, t, z, y, x, parent_track_id, n = point # n is the nth point in this timepoint
+    point_id = t * max_points_in_timepoint + n # creates a sequential ID for each point, but there is no guarantee that the points close together in space
+
+    track_index = track_id - 1
+    if track_index not in direct_parent_index_map:
+      # maps the track_index to the parent_track_index
+      direct_parent_index_map[track_id - 1] = parent_track_id - 1
 
     points_array[t, 3 * n:3 * (n + 1)] = [z, y, x]
 
@@ -45,7 +72,7 @@
     if parent_track_id > 0:
         tracks_to_parents[track_id - 1, parent_track_id - 1] = 1
         tracks_to_children[parent_track_id - 1, track_id - 1] = 1
-
+    
 print(f"Munged {len(points)} points in {time.monotonic() - start} seconds")
 
 tracks_to_parents.setdiag(1)
@@ -55,6 +82,8 @@
 
 start = time.monotonic()
 iter = 0
+# More info on sparse matrix: https://en.wikipedia.org/wiki/Sparse_matrix#Compressed_sparse_row_(CSR,_CRS_or_Yale_format)
+# Transitive closure: https://en.wikipedia.org/wiki/Transitive_closure
 while tracks_to_parents.nnz != (nxt := tracks_to_parents ** 2).nnz:
     tracks_to_parents = nxt
     iter += 1
@@ -71,6 +100,14 @@
 start = time.monotonic()
 
 tracks_to_tracks = tracks_to_parents + tracks_to_children
+tracks_to_tracks = tracks_to_tracks.tolil()
+non_zero = tracks_to_tracks.nonzero()
+
+for i in range(len(non_zero[0])):
+    # track_index = track_id - 1 since track_id is 1-indexed
+    track_index = non_zero[1][i]
+    parent_track_index = direct_parent_index_map[track_index]
+    tracks_to_tracks[non_zero[0][i], non_zero[1][i]] = parent_track_index + 1
 
 # Convert to CSR format for efficient row slicing
 tracks_to_points = points_to_tracks.T.tocsr()
@@ -82,7 +119,7 @@
 
 # save the points array (same format as ZSHS001_nodes.zarr)
 top_level_group = zarr.hierarchy.group(
-    zarr.storage.DirectoryStore(root_dir + "ZSNS001_tracks_bundle.zarr"),
+    zarr.storage.DirectoryStore(zarr_path.as_posix()),
     overwrite=True,
 )