fix(datasets-dicom): rebuild volume if image deleted

And only build volume if not already requested to build.
Kitware · Apr 12, 2024 · 5b86b11 · 5b86b11
1 parent b089616
commit 5b86b11
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 22 deletions.
diff --git a/src/store/datasets-dicom.ts b/src/store/datasets-dicom.ts
@@ -1,4 +1,5 @@
 import vtkITKHelper from '@kitware/vtk.js/Common/DataModel/ITKHelper';
+import vtkImageData from '@kitware/vtk.js/Common/DataModel/ImageData';
 import { defineStore } from 'pinia';
 import { Image } from 'itk-wasm';
 import { DataSourceWithFile } from '@/src/io/import/dataSource';
@@ -61,6 +62,9 @@ interface State {
   // volume invalidation information
   needsRebuild: Record<string, boolean>;
 
+  // Promise resolving to vtkImageData
+  volumeImageData: Record<string, Promise<vtkImageData>>;
+
   // patientKey -> patient info
   patientInfo: Record<string, PatientInfo>;
   // patientKey -> array of studyKeys
@@ -117,11 +121,22 @@ export const getDisplayName = (info: VolumeInfo) => {
   );
 };
 
+const constructImage = async (dicomIO: DICOMIO, volumeKey: string) => {
+  const fileStore = useFileStore();
+  const files = fileStore.getFiles(volumeKey);
+  if (!files) throw new Error('No files for volume key');
+  const image = vtkITKHelper.convertItkToVtkImage(
+    await dicomIO.buildImage(files)
+  );
+  return image;
+};
+
 export const useDICOMStore = defineStore('dicom', {
   state: (): State => ({
     sliceData: {},
     volumeToImageID: {},
     imageIDToVolumeKey: {},
+    volumeImageData: {},
     patientInfo: {},
     patientStudies: {},
     studyInfo: {},
@@ -254,6 +269,10 @@ export const useDICOMStore = defineStore('dicom', {
           delete this.imageIDToVolumeKey[imageID];
         }
 
+        if (volumeKey in this.volumeImageData) {
+          delete this.volumeImageData[volumeKey];
+        }
+
         removeFromArray(this.studyVolumes[studyKey], volumeKey);
         if (this.studyVolumes[studyKey].length === 0) {
           this.deleteStudy(studyKey);
@@ -356,36 +375,40 @@ export const useDICOMStore = defineStore('dicom', {
 
     async buildVolume(volumeKey: string, forceRebuild: boolean = false) {
       const imageStore = useImageStore();
-      const dicomIO = this.$dicomIO;
-
-      const rebuild = forceRebuild || this.needsRebuild[volumeKey];
 
-      if (!rebuild && this.volumeToImageID[volumeKey]) {
-        const imageID = this.volumeToImageID[volumeKey]!;
-        return imageStore.dataIndex[imageID];
-      }
-
-      const fileStore = useFileStore();
-      const files = fileStore.getFiles(volumeKey);
-      if (!files) throw new Error('No files for volume key');
-      const image = vtkITKHelper.convertItkToVtkImage(
-        await dicomIO.buildImage(files)
-      );
+      const alreadyBuilt = volumeKey in this.volumeImageData;
+      const buildNeeded =
+        forceRebuild || this.needsRebuild[volumeKey] || !alreadyBuilt;
+      delete this.needsRebuild[volumeKey];
 
+      // wait for old buildVolume call so we can run imageStore update side effects after
+      const oldImagePromise = alreadyBuilt
+        ? [this.volumeImageData[volumeKey]]
+        : [];
+      // actually build volume or wait for existing build?
+      const newImagePromise = buildNeeded
+        ? constructImage(this.$dicomIO, volumeKey)
+        : this.volumeImageData[volumeKey];
+      // let other calls to buildVolume reuse this constructImage work
+      this.volumeImageData[volumeKey] = newImagePromise;
+      const [image] = await Promise.all([newImagePromise, ...oldImagePromise]);
+
+      // update image store
       const existingImageID = this.volumeToImageID[volumeKey];
-      if (existingImageID) {
+      const imageExists =
+        existingImageID && imageStore.dataIndex[existingImageID];
+      if (imageExists) {
+        // was a rebuild
         imageStore.updateData(existingImageID, image);
       } else {
         const info = this.volumeInfo[volumeKey];
         const name =
           cleanupName(info.SeriesDescription) || info.SeriesInstanceUID;
-        const imageID = imageStore.addVTKImageData(name, image);
-        this.imageIDToVolumeKey[imageID] = volumeKey;
-        this.volumeToImageID[volumeKey] = imageID;
+        const newImageID = imageStore.addVTKImageData(name, image);
+        this.imageIDToVolumeKey[newImageID] = volumeKey;
+        this.volumeToImageID[volumeKey] = newImageID;
       }
 
-      delete this.needsRebuild[volumeKey];
-
       return image;
     },
   },

diff --git a/src/store/load-data.ts b/src/store/load-data.ts
@@ -193,7 +193,7 @@ function loadLayers(
   layersStore.addLayer(primarySelection, layerSelection);
 }
 
-// Converts DICOM SEG modalities to segmentations if found
+// Loads DICOM SEG modalities to segmentations if found
 function loadSegmentations(
   primaryDataSource: VolumeResult,
   succeeded: Array<PipelineResultSuccess<ImportResult>>

diff --git a/src/store/segmentGroups.ts b/src/store/segmentGroups.ts
@@ -210,7 +210,7 @@ export const useSegmentGroupStore = defineStore('segmentGroup', () => {
   function getNextColor() {
     const color = DEFAULT_SEGMENT_MASKS[lastColorIndex].color;
     lastColorIndex = (lastColorIndex + 1) % DEFAULT_SEGMENT_MASKS.length;
-    return color;
+    return [...color];
   }
 
   function decodeSegments(image: DataSelection) {