From c2d791d6b0928ff29c6e4ed1044c4236af6f8950 Mon Sep 17 00:00:00 2001
From: Paul Elliott <paul.elliott@kitware.com>
Date: Thu, 11 Apr 2024 18:17:49 -0400
Subject: [PATCH] fix(datasets-dicom): rebuild volume if image deleted

And only build volume if not already requested to build.
---
 src/store/datasets-dicom.ts | 37 ++++++++++++++++++++++++++++---------
 src/store/load-data.ts      |  2 +-
 src/store/segmentGroups.ts  |  2 +-
 3 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/src/store/datasets-dicom.ts b/src/store/datasets-dicom.ts
index 7f05daae5..6520654af 100644
--- 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
@@ -122,6 +126,7 @@ export const useDICOMStore = defineStore('dicom', {
     sliceData: {},
     volumeToImageID: {},
     imageIDToVolumeKey: {},
+    volumeImageData: {},
     patientInfo: {},
     patientStudies: {},
     studyInfo: {},
@@ -254,6 +259,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);
@@ -354,15 +363,17 @@ export const useDICOMStore = defineStore('dicom', {
       return this.getVolumeSlice(volumeKey, middleSlice, true);
     },
 
-    async buildVolume(volumeKey: string, forceRebuild: boolean = false) {
+    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 existingImageID = this.volumeToImageID[volumeKey];
+      const imageExists =
+        existingImageID && imageStore.dataIndex[existingImageID];
+      if (!rebuild && imageExists) {
+        return imageStore.dataIndex[existingImageID];
       }
 
       const fileStore = useFileStore();
@@ -372,21 +383,29 @@ export const useDICOMStore = defineStore('dicom', {
         await dicomIO.buildImage(files)
       );
 
-      const existingImageID = this.volumeToImageID[volumeKey];
-      if (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;
     },
+
+    buildVolume(volumeKey: string) {
+      if (volumeKey in this.volumeImageData) {
+        return this.volumeImageData[volumeKey];
+      }
+      this.volumeImageData[volumeKey] = this._buildVolume(volumeKey);
+      return this.volumeImageData[volumeKey];
+    },
   },
 });
diff --git a/src/store/load-data.ts b/src/store/load-data.ts
index f6d915383..794ef591c 100644
--- 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
index dc1f3c73b..59a37903d 100644
--- 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) {