From 1f02f96badf86536979b5c1195412c2181be4478 Mon Sep 17 00:00:00 2001
From: Forrest <forrest.li@kitware.com>
Date: Mon, 20 May 2024 13:19:23 -0400
Subject: [PATCH 1/2] feat: unify data IDs

Volume keys are now used as image data IDs. This simplifies going back
and forth between the image and dicom store, and simplifies the
structure of the selection object to be just an ID string.
---
 src/components/DataBrowser.vue                |  5 +-
 src/components/DicomQuickInfoButton.vue       |  5 +-
 src/components/ImageDataBrowser.vue           | 35 +++----
 src/components/LayerProperties.vue            |  5 +-
 src/components/PatientBrowser.vue             |  2 +-
 src/components/PatientStudyVolumeBrowser.vue  | 17 ++--
 src/components/SampleDataBrowser.vue          |  8 +-
 src/components/SliceViewer.vue                |  8 ++
 src/components/VolumeViewer.vue               |  8 ++
 .../tools/windowing/WindowLevelControls.vue   | 15 +--
 src/composables/useCurrentImage.ts            | 14 +--
 src/composables/useSliceConfigInitializer.ts  | 19 ++--
 .../useVolumeColoringInitializer.ts           |  6 +-
 .../useWindowingConfigInitializer.ts          |  5 +-
 src/io/import/importDataSources.ts            | 16 +---
 src/io/import/processors/restoreStateFile.ts  | 31 ++-----
 src/io/state-file/schema.ts                   | 23 ++---
 src/store/datasets-dicom.ts                   | 25 +----
 src/store/datasets-images.ts                  |  8 +-
 src/store/datasets-layers.ts                  | 88 ++++--------------
 src/store/datasets.ts                         | 56 ++++--------
 src/store/segmentGroups.ts                    | 30 +++---
 src/store/tools/useAnnotationTool.ts          |  6 +-
 src/store/view-configs/layers.ts              |  6 +-
 src/store/view-configs/volume-coloring.ts     |  5 +-
 src/utils/dataSelection.ts                    | 91 ++++---------------
 26 files changed, 178 insertions(+), 359 deletions(-)

diff --git a/src/components/DataBrowser.vue b/src/components/DataBrowser.vue
index 1de1cd274..954e84c47 100644
--- a/src/components/DataBrowser.vue
+++ b/src/components/DataBrowser.vue
@@ -1,5 +1,6 @@
 <script lang="ts">
 import { computed, defineComponent, ref, watch } from 'vue';
+import { isRegularImage } from '@/src/utils/dataSelection';
 import SampleDataBrowser from './SampleDataBrowser.vue';
 import { useDicomWebStore } from '../store/dicom-web/dicom-web-store';
 import ImageDataBrowser from './ImageDataBrowser.vue';
@@ -42,9 +43,7 @@ export default defineComponent({
     );
 
     const hasAnonymousImages = computed(
-      () =>
-        imageStore.idList.filter((id) => !(id in dicomStore.imageIDToVolumeKey))
-          .length > 0
+      () => imageStore.idList.filter((id) => isRegularImage(id)).length > 0
     );
 
     const panels = ref<string[]>([SAMPLE_DATA_KEY, DICOM_WEB_KEY]);
diff --git a/src/components/DicomQuickInfoButton.vue b/src/components/DicomQuickInfoButton.vue
index eb88e6f10..123e90cb9 100644
--- a/src/components/DicomQuickInfoButton.vue
+++ b/src/components/DicomQuickInfoButton.vue
@@ -1,6 +1,7 @@
 <script setup lang="ts">
 import { useDICOMStore } from '@/src/store/datasets-dicom';
 import { Maybe } from '@/src/types';
+import { isDicomImage } from '@/src/utils/dataSelection';
 import { computed, toRef } from 'vue';
 
 interface Props {
@@ -12,8 +13,8 @@ const imageId = toRef(props, 'imageId');
 
 const dicomStore = useDICOMStore();
 const dicomInfo = computed(() => {
-  if (imageId.value != null && imageId.value in dicomStore.imageIDToVolumeKey) {
-    const volumeKey = dicomStore.imageIDToVolumeKey[imageId.value];
+  const volumeKey = imageId.value;
+  if (volumeKey && isDicomImage(volumeKey)) {
     const volumeInfo = dicomStore.volumeInfo[volumeKey];
     const studyKey = dicomStore.volumeStudy[volumeKey];
     const studyInfo = dicomStore.studyInfo[studyKey];
diff --git a/src/components/ImageDataBrowser.vue b/src/components/ImageDataBrowser.vue
index 50df1e746..14d5ce142 100644
--- a/src/components/ImageDataBrowser.vue
+++ b/src/components/ImageDataBrowser.vue
@@ -6,12 +6,11 @@ import ImageListCard from '@/src/components/ImageListCard.vue';
 import { createVTKImageThumbnailer } from '@/src/core/thumbnailers/vtk-image';
 import { useSegmentGroupStore } from '@/src/store/segmentGroups';
 import {
-  DataSelection,
-  ImageSelection,
+  isRegularImage,
+  type DataSelection,
   selectionEquals,
 } from '@/src/utils/dataSelection';
 import { useImageStore } from '../store/datasets-images';
-import { useDICOMStore } from '../store/datasets-dicom';
 import { useDatasetStore } from '../store/datasets';
 
 import { useMultiSelection } from '../composables/useMultiSelection';
@@ -29,7 +28,6 @@ export default defineComponent({
   },
   setup() {
     const imageStore = useImageStore();
-    const dicomStore = useDICOMStore();
     const dataStore = useDatasetStore();
     const layersStore = useLayersStore();
     const segmentGroupStore = useSegmentGroupStore();
@@ -37,7 +35,7 @@ export default defineComponent({
     const primarySelection = computed(() => dataStore.primarySelection);
 
     const nonDICOMImages = computed(() =>
-      imageStore.idList.filter((id) => !(id in dicomStore.imageIDToVolumeKey))
+      imageStore.idList.filter((id) => isRegularImage(id))
     );
 
     const images = computed(() => {
@@ -45,23 +43,17 @@ export default defineComponent({
 
       const layerImages = layersStore
         .getLayers(primarySelection.value)
-        .filter(({ selection }) => selection.type === 'image');
-      const layerImageIDs = layerImages.map(
-        ({ selection }) => (selection as ImageSelection).dataID
-      );
+        .filter(({ selection }) => isRegularImage(selection));
+      const layerImageIDs = layerImages.map(({ selection }) => selection);
       const loadedLayerImageIDs = layerImages
         .filter(({ id }) => id in layersStore.layerImages)
-        .map(({ selection }) => (selection as ImageSelection).dataID);
+        .map(({ selection }) => selection);
 
       const selectedImageID =
-        primarySelection.value?.type === 'image' &&
-        primarySelection.value?.dataID;
+        isRegularImage(primarySelection.value) && primarySelection.value;
 
       return nonDICOMImages.value.map((id) => {
-        const selectionKey = {
-          type: 'image',
-          dataID: id,
-        } as DataSelection;
+        const selectionKey = id as DataSelection;
         const isLayer = layerImageIDs.includes(id);
         const layerLoaded = loadedLayerImageIDs.includes(id);
         const layerLoading = isLayer && !layerLoaded;
@@ -136,10 +128,7 @@ export default defineComponent({
 
     function convertToLabelMap(key: string) {
       if (primarySelection.value) {
-        segmentGroupStore.convertImageToLabelmap(
-          { type: 'image', dataID: key },
-          primarySelection.value
-        );
+        segmentGroupStore.convertImageToLabelmap(key, primarySelection.value);
       }
     }
 
@@ -261,9 +250,9 @@ export default defineComponent({
                       image.layerable ? convertToLabelMap(image.id) : null
                     "
                   >
-                    <v-icon v-if="!image.layerable" class="mr-1"
-                      >mdi-alert</v-icon
-                    >
+                    <v-icon v-if="!image.layerable" class="mr-1">
+                      mdi-alert
+                    </v-icon>
                     Convert to Segment Group
                     <v-tooltip
                       activator="parent"
diff --git a/src/components/LayerProperties.vue b/src/components/LayerProperties.vue
index b574d6748..6e17b10b3 100644
--- a/src/components/LayerProperties.vue
+++ b/src/components/LayerProperties.vue
@@ -1,6 +1,5 @@
 <script lang="ts">
 import { computed, defineComponent, PropType, toRefs } from 'vue';
-import { getImageID } from '@/src/utils/dataSelection';
 import { InitViewSpecs } from '../config';
 import { useImageStore } from '../store/datasets-images';
 import { BlendConfig } from '../types/views';
@@ -25,9 +24,7 @@ export default defineComponent({
 
     const imageName = computed(() => {
       const { selection } = props.layer;
-      const imageID = getImageID(selection);
-      if (imageID === undefined) throw new Error('imageID is undefined');
-      return imageStore.metadata[imageID].name;
+      return imageStore.metadata[selection].name;
     });
 
     const layerColoringStore = useLayerColoringStore();
diff --git a/src/components/PatientBrowser.vue b/src/components/PatientBrowser.vue
index f788df589..ab830c0f9 100644
--- a/src/components/PatientBrowser.vue
+++ b/src/components/PatientBrowser.vue
@@ -1,7 +1,7 @@
 <script lang="ts">
 import { computed, defineComponent, ref, toRefs, watch } from 'vue';
 import ItemGroup from '@/src/components/ItemGroup.vue';
-import { DataSelection, selectionEquals } from '@/src/utils/dataSelection';
+import { type DataSelection, selectionEquals } from '@/src/utils/dataSelection';
 import { useDICOMStore } from '../store/datasets-dicom';
 import { useDatasetStore } from '../store/datasets';
 import { useMultiSelection } from '../composables/useMultiSelection';
diff --git a/src/components/PatientStudyVolumeBrowser.vue b/src/components/PatientStudyVolumeBrowser.vue
index 1fba59c19..ffed6b9bc 100644
--- a/src/components/PatientStudyVolumeBrowser.vue
+++ b/src/components/PatientStudyVolumeBrowser.vue
@@ -3,7 +3,7 @@ import { computed, defineComponent, reactive, toRefs, watch } from 'vue';
 import { Image } from 'itk-wasm';
 import type { PropType } from 'vue';
 import GroupableItem from '@/src/components/GroupableItem.vue';
-import { DataSelection, DICOMSelection } from '@/src/utils/dataSelection';
+import { DataSelection, isDicomImage } from '@/src/utils/dataSelection';
 import { getDisplayName, useDICOMStore } from '../store/datasets-dicom';
 import { useDatasetStore } from '../store/datasets';
 import { useMultiSelection } from '../composables/useMultiSelection';
@@ -78,21 +78,16 @@ export default defineComponent({
       const primarySelection = primarySelectionRef.value;
       const layerVolumes = layersStore
         .getLayers(primarySelection)
-        .filter(({ selection }) => selection.type === 'dicom');
-      const layerVolumeKeys = layerVolumes.map(
-        ({ selection }) => (selection as DICOMSelection).volumeKey
-      );
+        .filter(({ selection }) => isDicomImage(selection));
+      const layerVolumeKeys = layerVolumes.map(({ selection }) => selection);
       const loadedLayerVolumeKeys = layerVolumes
         .filter(({ id }) => id in layersStore.layerImages)
-        .map(({ selection }) => (selection as DICOMSelection).volumeKey);
+        .map(({ selection }) => selection);
       const selectedVolumeKey =
-        primarySelection?.type === 'dicom' && primarySelection.volumeKey;
+        isDicomImage(primarySelection) && primarySelection;
 
       return volumeKeys.value.map((volumeKey) => {
-        const selectionKey = {
-          type: 'dicom',
-          volumeKey,
-        } as DataSelection;
+        const selectionKey = volumeKey as DataSelection;
         const isLayer = layerVolumeKeys.includes(volumeKey);
         const layerLoaded = loadedLayerVolumeKeys.includes(volumeKey);
         const layerLoading = isLayer && !layerLoaded;
diff --git a/src/components/SampleDataBrowser.vue b/src/components/SampleDataBrowser.vue
index 47d19cba1..3579799a5 100644
--- a/src/components/SampleDataBrowser.vue
+++ b/src/components/SampleDataBrowser.vue
@@ -102,12 +102,10 @@ export default defineComponent({
 
         const selection = convertSuccessResultToDataSelection(loadResult);
         if (selection) {
-          const id =
-            selection.type === 'image' ? selection.dataID : selection.volumeKey;
-          loaded.idToURL[id] = sample.url;
-          loaded.urlToID[sample.url] = id;
+          loaded.idToURL[selection] = sample.url;
+          loaded.urlToID[sample.url] = selection;
 
-          useVolumeColoringStore().setDefaults(id, {
+          useVolumeColoringStore().setDefaults(selection, {
             transferFunction: {
               preset: sample.defaults?.colorPreset,
             },
diff --git a/src/components/SliceViewer.vue b/src/components/SliceViewer.vue
index 6a20efb38..19b937e44 100644
--- a/src/components/SliceViewer.vue
+++ b/src/components/SliceViewer.vue
@@ -177,6 +177,7 @@ import VtkMouseInteractionManipulator from '@/src/components/vtk/VtkMouseInterac
 import vtkMouseCameraTrackballPanManipulator from '@kitware/vtk.js/Interaction/Manipulators/MouseCameraTrackballPanManipulator';
 import vtkMouseCameraTrackballZoomToMouseManipulator from '@kitware/vtk.js/Interaction/Manipulators/MouseCameraTrackballZoomToMouseManipulator';
 import { useResetViewsEvents } from '@/src/components/tools/ResetViews.vue';
+import { whenever } from '@vueuse/core';
 
 interface Props extends LayoutViewProps {
   viewDirection: LPSAxisDir;
@@ -216,6 +217,13 @@ const { slice: currentSlice, range: sliceRange } = useSliceConfig(
   currentImageID
 );
 
+whenever(
+  computed(() => !isImageLoading.value),
+  () => {
+    resetCamera();
+  }
+);
+
 // segmentations
 const segmentations = computed(() => {
   if (!currentImageID.value) return [];
diff --git a/src/components/VolumeViewer.vue b/src/components/VolumeViewer.vue
index a3289eeca..7456fcec4 100644
--- a/src/components/VolumeViewer.vue
+++ b/src/components/VolumeViewer.vue
@@ -74,6 +74,7 @@ import VtkOrientationMarker from '@/src/components/vtk/VtkOrientationMarker.vue'
 import ViewOverlayGrid from '@/src/components/ViewOverlayGrid.vue';
 import useVolumeColoringStore from '@/src/store/view-configs/volume-coloring';
 import { useResetViewsEvents } from '@/src/components/tools/ResetViews.vue';
+import { whenever } from '@vueuse/core';
 
 interface Props extends LayoutViewProps {
   viewDirection: LPSAxisDir;
@@ -100,6 +101,13 @@ useViewAnimationListener(vtkView, viewId, viewType);
 // base image
 const { currentImageID, isImageLoading } = useCurrentImage();
 
+whenever(
+  computed(() => !isImageLoading.value),
+  () => {
+    resetCamera();
+  }
+);
+
 // color preset
 const coloringStore = useVolumeColoringStore();
 const coloringConfig = computed(() =>
diff --git a/src/components/tools/windowing/WindowLevelControls.vue b/src/components/tools/windowing/WindowLevelControls.vue
index 0f0c41317..57997e0e9 100644
--- a/src/components/tools/windowing/WindowLevelControls.vue
+++ b/src/components/tools/windowing/WindowLevelControls.vue
@@ -7,6 +7,7 @@ import useWindowingStore, {
 import { useViewStore } from '@/src/store/views';
 import { WLAutoRanges, WLPresetsCT, WL_AUTO_DEFAULT } from '@/src/constants';
 import { getWindowLevels, useDICOMStore } from '@/src/store/datasets-dicom';
+import { isDicomImage } from '@/src/utils/dataSelection';
 
 export default defineComponent({
   setup() {
@@ -31,11 +32,8 @@ export default defineComponent({
     // --- CT Preset Options --- //
 
     const modality = computed(() => {
-      if (
-        currentImageID.value &&
-        currentImageID.value in dicomStore.imageIDToVolumeKey
-      ) {
-        const volKey = dicomStore.imageIDToVolumeKey[currentImageID.value];
+      if (currentImageID.value && isDicomImage(currentImageID.value)) {
+        const volKey = currentImageID.value;
         const { Modality } = dicomStore.volumeInfo[volKey];
         return Modality;
       }
@@ -98,11 +96,8 @@ export default defineComponent({
 
     // --- Tag WL Options --- //
     const tags = computed(() => {
-      if (
-        currentImageID.value &&
-        currentImageID.value in dicomStore.imageIDToVolumeKey
-      ) {
-        const volKey = dicomStore.imageIDToVolumeKey[currentImageID.value];
+      if (currentImageID.value && isDicomImage(currentImageID.value)) {
+        const volKey = currentImageID.value;
         return getWindowLevels(dicomStore.volumeInfo[volKey]);
       }
       return [];
diff --git a/src/composables/useCurrentImage.ts b/src/composables/useCurrentImage.ts
index 78498d38a..394939cc0 100644
--- a/src/composables/useCurrentImage.ts
+++ b/src/composables/useCurrentImage.ts
@@ -15,7 +15,6 @@ import {
 import { useLayersStore } from '@/src/store/datasets-layers';
 import { createLPSBounds, getAxisBounds } from '@/src/utils/lps';
 import { useDatasetStore } from '@/src/store/datasets';
-import { getDataSelection, getImageID } from '@/src/utils/dataSelection';
 import { storeToRefs } from 'pinia';
 
 export interface CurrentImageContext {
@@ -57,21 +56,16 @@ export function getImageData(imageID: Maybe<string>) {
 }
 
 export function getIsImageLoading(imageID: Maybe<string>) {
-  const dataStore = useDatasetStore();
-  if (!dataStore.primarySelection) return false;
-
-  const selectedImageID = getImageID(dataStore.primarySelection);
-  if (selectedImageID !== unref(imageID)) return false;
-
-  return !!dataStore.primarySelection && !dataStore.primaryDataset;
+  if (!imageID) return false;
+  const imageStore = useImageStore();
+  return !imageStore.dataIndex[imageID];
 }
 
 export function getImageLayers(imageID: Maybe<string>) {
   if (!imageID) return [];
-  const selection = getDataSelection(imageID);
   const layersStore = useLayersStore();
   return layersStore
-    .getLayers(selection)
+    .getLayers(imageID)
     .filter(({ id }) => id in layersStore.layerImages);
 }
 
diff --git a/src/composables/useSliceConfigInitializer.ts b/src/composables/useSliceConfigInitializer.ts
index 590070b56..68b5557ce 100644
--- a/src/composables/useSliceConfigInitializer.ts
+++ b/src/composables/useSliceConfigInitializer.ts
@@ -16,7 +16,7 @@ export function useSliceConfigInitializer(
 ) {
   const store = useViewSliceStore();
   const { config: sliceConfig } = useSliceConfig(viewID, imageID);
-  const { metadata } = useImage(imageID);
+  const { metadata, isLoading } = useImage(imageID);
 
   const viewAxis = computed(() => getLPSAxisFromDir(unref(viewDirection)));
   const sliceDomain = computed(() => {
@@ -33,17 +33,22 @@ export function useSliceConfigInitializer(
   });
 
   watchImmediate(
-    [toRef(sliceDomain), toRef(viewDirection)] as const,
-    ([domain, axisDirection]) => {
+    [
+      toRef(sliceDomain),
+      toRef(viewDirection),
+      toRef(imageID),
+      isLoading,
+    ] as const,
+    ([domain, axisDirection, id, loading]) => {
+      if (loading || !id) return;
+
       const configExisted = !!sliceConfig.value;
-      const imageIdVal = unref(imageID);
-      if (!imageIdVal) return;
-      store.updateConfig(unref(viewID), imageIdVal, {
+      store.updateConfig(unref(viewID), id, {
         ...domain,
         axisDirection,
       });
       if (!configExisted) {
-        store.resetSlice(unref(viewID), imageIdVal);
+        store.resetSlice(unref(viewID), id);
       }
     }
   );
diff --git a/src/composables/useVolumeColoringInitializer.ts b/src/composables/useVolumeColoringInitializer.ts
index 4c176893f..b2fe1568a 100644
--- a/src/composables/useVolumeColoringInitializer.ts
+++ b/src/composables/useVolumeColoringInitializer.ts
@@ -13,10 +13,10 @@ export function useVolumeColoringInitializer(
     store.getConfig(unref(viewId), unref(imageId))
   );
 
-  const { imageData } = useImage(imageId);
+  const { imageData, isLoading } = useImage(imageId);
 
-  watchImmediate(coloringConfig, (config) => {
-    if (config) return;
+  watchImmediate([coloringConfig, viewId, imageId, isLoading], () => {
+    if (coloringConfig.value || isLoading.value) return;
 
     const viewIdVal = unref(viewId);
     const imageIdVal = unref(imageId);
diff --git a/src/composables/useWindowingConfigInitializer.ts b/src/composables/useWindowingConfigInitializer.ts
index c2ef802a2..0208fc0b6 100644
--- a/src/composables/useWindowingConfigInitializer.ts
+++ b/src/composables/useWindowingConfigInitializer.ts
@@ -9,6 +9,7 @@ import { getWindowLevels, useDICOMStore } from '@/src/store/datasets-dicom';
 import useWindowingStore from '@/src/store/view-configs/windowing';
 import { Maybe } from '@/src/types';
 import { useResetViewsEvents } from '@/src/components/tools/ResetViews.vue';
+import { isDicomImage } from '@/src/utils/dataSelection';
 
 function useAutoRangeValues(imageID: MaybeRef<Maybe<string>>) {
   const { imageData } = useImage(imageID);
@@ -80,8 +81,8 @@ export function useWindowingConfigInitializer(
 
   const firstTag = computed(() => {
     const id = unref(imageID);
-    if (id && id in dicomStore.imageIDToVolumeKey) {
-      const volKey = dicomStore.imageIDToVolumeKey[id];
+    if (id && isDicomImage(id)) {
+      const volKey = id;
       const windowLevels = getWindowLevels(dicomStore.volumeInfo[volKey]);
       if (windowLevels.length) {
         return windowLevels[0];
diff --git a/src/io/import/importDataSources.ts b/src/io/import/importDataSources.ts
index b271b62d6..a92598259 100644
--- a/src/io/import/importDataSources.ts
+++ b/src/io/import/importDataSources.ts
@@ -22,10 +22,6 @@ import restoreStateFile from '@/src/io/import/processors/restoreStateFile';
 import updateFileMimeType from '@/src/io/import/processors/updateFileMimeType';
 import handleConfig from '@/src/io/import/processors/handleConfig';
 import { useDICOMStore } from '@/src/store/datasets-dicom';
-import {
-  makeDICOMSelection,
-  makeImageSelection,
-} from '@/src/utils/dataSelection';
 import { applyConfig } from '@/src/io/import/configJson';
 
 /**
@@ -177,16 +173,8 @@ export type ImportDataSourcesResult = Awaited<
 >[number];
 
 export function toDataSelection(loadable: VolumeResult) {
-  const { dataID, dataType } = loadable;
-  if (dataType === 'dicom') {
-    return makeDICOMSelection(dataID);
-  }
-  if (dataType === 'image') {
-    return makeImageSelection(dataID);
-  }
-
-  const _exhaustiveCheck: never = dataType;
-  throw new Error(`invalid loadable type ${_exhaustiveCheck}`);
+  const { dataID } = loadable;
+  return dataID;
 }
 
 export function convertSuccessResultToDataSelection(
diff --git a/src/io/import/processors/restoreStateFile.ts b/src/io/import/processors/restoreStateFile.ts
index cc577b575..2704086e1 100644
--- a/src/io/import/processors/restoreStateFile.ts
+++ b/src/io/import/processors/restoreStateFile.ts
@@ -22,11 +22,6 @@ import doneWithDataSource from '@/src/io/import/processors/doneWithDataSource';
 import { useDICOMStore } from '@/src/store/datasets-dicom';
 import { useViewStore } from '@/src/store/views';
 import { useDatasetStore } from '@/src/store/datasets';
-import {
-  DataSelection,
-  makeDICOMSelection,
-  makeImageSelection,
-} from '@/src/utils/dataSelection';
 import {
   makeDefaultSegmentGroupName,
   useSegmentGroupStore,
@@ -344,10 +339,7 @@ async function restoreDatasets(
   return stateIDToStoreID;
 }
 
-const restoreStateFile: ImportHandler = async (
-  dataSource,
-  { done, extra, execute }
-) => {
+const restoreStateFile: ImportHandler = async (dataSource, pipelineContext) => {
   const { fileSrc } = dataSource;
   if (fileSrc && (await isStateFile(fileSrc.file))) {
     const stateFileContents = await extractFilesFromZip(fileSrc.file);
@@ -369,24 +361,17 @@ const restoreStateFile: ImportHandler = async (
     // in the views as the data is loaded
     useViewStore().setLayout(manifest.layout);
 
-    const stateIDToStoreID = await restoreDatasets(manifest, restOfStateFile, {
-      extra,
-      execute,
-      done,
-    });
+    const stateIDToStoreID = await restoreDatasets(
+      manifest,
+      restOfStateFile,
+      pipelineContext
+    );
 
     // Restore the primary selection
     if (manifest.primarySelection !== undefined) {
       const selectedID = stateIDToStoreID[manifest.primarySelection];
-      let dataSelection: DataSelection | undefined;
-
-      if (selectedID in useDICOMStore().volumeInfo) {
-        dataSelection = makeDICOMSelection(selectedID);
-      } else {
-        dataSelection = makeImageSelection(selectedID);
-      }
 
-      useDatasetStore().setPrimarySelection(dataSelection);
+      useDatasetStore().setPrimarySelection(selectedID);
     }
 
     // Restore the views
@@ -404,7 +389,7 @@ const restoreStateFile: ImportHandler = async (
 
     useLayersStore().deserialize(manifest, stateIDToStoreID);
 
-    return done();
+    return pipelineContext.done();
   }
   return dataSource;
 };
diff --git a/src/io/state-file/schema.ts b/src/io/state-file/schema.ts
index e3608e65c..65a8cbc90 100644
--- a/src/io/state-file/schema.ts
+++ b/src/io/state-file/schema.ts
@@ -1,30 +1,31 @@
 import JSZip from 'jszip';
 import { z } from 'zod';
 import type { Vector3 } from '@kitware/vtk.js/types';
-import vtkPiecewiseFunctionProxy, {
+import vtkPiecewiseFunctionProxy from '@kitware/vtk.js/Proxy/Core/PiecewiseFunctionProxy';
+import type {
   PiecewiseGaussian,
   PiecewiseNode,
 } from '@kitware/vtk.js/Proxy/Core/PiecewiseFunctionProxy';
 
 import type { AnnotationTool, ToolID } from '@/src/types/annotation-tool';
 import { Tools as ToolsEnum } from '@/src/store/tools/types';
-import { Ruler } from '@/src/types/ruler';
-import { Rectangle } from '@/src/types/rectangle';
-import { Polygon } from '@/src/types/polygon';
-import { LPSCroppingPlanes } from '@/src/types/crop';
-import { FrameOfReference } from '@/src/utils/frameOfReference';
-import { Optional } from '@/src/types';
-
-import {
+import type { Ruler } from '@/src/types/ruler';
+import type { Rectangle } from '@/src/types/rectangle';
+import type { Polygon } from '@/src/types/polygon';
+import type { LPSCroppingPlanes } from '@/src/types/crop';
+import type { FrameOfReference } from '@/src/utils/frameOfReference';
+import type { Optional } from '@/src/types';
+
+import type {
   CameraConfig,
   SliceConfig,
   WindowLevelConfig,
   LayersConfig,
   VolumeColorConfig,
 } from '../../store/view-configs/types';
-import { LPSAxisDir, LPSAxis } from '../../types/lps';
+import type { LPSAxisDir, LPSAxis } from '../../types/lps';
 import { LayoutDirection } from '../../types/layout';
-import {
+import type {
   ColorBy,
   ColorTransferFunction,
   OpacityFunction,
diff --git a/src/store/datasets-dicom.ts b/src/store/datasets-dicom.ts
index 3917c0317..4e3667c8a 100644
--- a/src/store/datasets-dicom.ts
+++ b/src/store/datasets-dicom.ts
@@ -54,11 +54,6 @@ interface State {
   // volumeKey -> imageCacheMultiKey -> ITKImage
   sliceData: Record<string, Record<string, Image>>;
 
-  // volumeKey -> imageID
-  volumeToImageID: Record<string, string>;
-  // imageID -> volumeKey
-  imageIDToVolumeKey: Record<string, string>;
-
   // volume invalidation information
   needsRebuild: Record<string, boolean>;
 
@@ -155,8 +150,6 @@ const constructImage = async (dicomIO: DICOMIO, volumeKey: string) => {
 export const useDICOMStore = defineStore('dicom', {
   state: (): State => ({
     sliceData: {},
-    volumeToImageID: {},
-    imageIDToVolumeKey: {},
     volumeImageData: {},
     patientInfo: {},
     patientStudies: {},
@@ -236,7 +229,7 @@ export const useDICOMStore = defineStore('dicom', {
           }
 
           // invalidate any existing volume
-          if (volumeKey in this.volumeToImageID) {
+          if (volumeKey in useImageStore().dataIndex) {
             // buildVolume requestor uses this as a rebuild hint
             this.needsRebuild[volumeKey] = true;
           }
@@ -284,12 +277,6 @@ export const useDICOMStore = defineStore('dicom', {
         delete this.sliceData[volumeKey];
         delete this.volumeStudy[volumeKey];
 
-        if (volumeKey in this.volumeToImageID) {
-          const imageID = this.volumeToImageID[volumeKey]!;
-          delete this.volumeToImageID[volumeKey];
-          delete this.imageIDToVolumeKey[imageID];
-        }
-
         if (volumeKey in this.volumeImageData) {
           delete this.volumeImageData[volumeKey];
         }
@@ -415,18 +402,14 @@ export const useDICOMStore = defineStore('dicom', {
       const [image] = await Promise.all([newImagePromise, ...oldImagePromise]);
 
       // update image store
-      const existingImageID = this.volumeToImageID[volumeKey];
-      const imageExists =
-        existingImageID && imageStore.dataIndex[existingImageID];
+      const imageExists = imageStore.dataIndex[volumeKey];
       if (imageExists) {
         // was a rebuild
-        imageStore.updateData(existingImageID, image);
+        imageStore.updateData(volumeKey, image);
       } else {
         const info = this.volumeInfo[volumeKey];
         const name = getDisplayName(info);
-        const newImageID = imageStore.addVTKImageData(name, image);
-        this.imageIDToVolumeKey[newImageID] = volumeKey;
-        this.volumeToImageID[volumeKey] = newImageID;
+        imageStore.addVTKImageData(name, image, volumeKey);
       }
 
       return image;
diff --git a/src/store/datasets-images.ts b/src/store/datasets-images.ts
index 5d18ea739..817a7c067 100644
--- a/src/store/datasets-images.ts
+++ b/src/store/datasets-images.ts
@@ -36,8 +36,12 @@ export const useImageStore = defineStore('images', {
     metadata: Object.create(null),
   }),
   actions: {
-    addVTKImageData(name: string, imageData: vtkImageData) {
-      const id = useIdStore().nextId();
+    addVTKImageData(name: string, imageData: vtkImageData, useId?: string) {
+      if (useId && useId in this.dataIndex) {
+        throw new Error('ID already exists');
+      }
+
+      const id = useId || useIdStore().nextId();
 
       this.idList.push(id);
       this.dataIndex[id] = imageData;
diff --git a/src/store/datasets-layers.ts b/src/store/datasets-layers.ts
index 6293cb541..75b4b33a3 100644
--- a/src/store/datasets-layers.ts
+++ b/src/store/datasets-layers.ts
@@ -2,19 +2,13 @@ import { ref } from 'vue';
 import vtkImageData from '@kitware/vtk.js/Common/DataModel/ImageData';
 import vtkBoundingBox from '@kitware/vtk.js/Common/DataModel/BoundingBox';
 import { defineStore } from 'pinia';
-import {
-  DataSelection,
-  getImage,
-  makeDICOMSelection,
-  makeImageSelection,
-  selectionEquals,
-} from '@/src/utils/dataSelection';
+import { type DataSelection, getImage } from '@/src/utils/dataSelection';
+import { Maybe } from '@/src/types';
 import { ensureSameSpace } from '@/src/io/resample/resample';
 import { useErrorMessage } from '../composables/useErrorMessage';
 import { Manifest, StateFile } from '../io/state-file/schema';
 
 // differ from Image/Volume IDs with a branded type
-type DataSelectionKey = string & { __type: 'DataSelectionKey' };
 export type LayerID = string & { __type: 'LayerID' };
 
 export type Layer = {
@@ -22,34 +16,10 @@ export type Layer = {
   id: LayerID;
 };
 
-const assertNeverDataSelection = (selection: never): never => {
-  throw new Error(`Unknown DataSelection: ${selection}`);
-};
-
-const getID = (selection: DataSelection) => {
-  if (selection.type === 'dicom') return selection.volumeKey;
-  if (selection.type === 'image') return selection.dataID;
-  return assertNeverDataSelection(selection);
-};
-
-const toDataSelectionKey = (selection: DataSelection) => {
-  const id = getID(selection);
-  return `${selection.type}::${id}` as DataSelectionKey;
-};
-
-const toDataSelectionFromKey = (key: DataSelectionKey) => {
-  const [type, id] = key.split('::') as [DataSelection['type'], string];
-
-  if (type === 'dicom') return makeDICOMSelection(id);
-  if (type === 'image') return makeImageSelection(id);
-
-  return assertNeverDataSelection(type);
-};
-
 export const useLayersStore = defineStore('layer', () => {
   type _This = ReturnType<typeof useLayersStore>;
 
-  const parentToLayers = ref<Record<DataSelectionKey, Layer[]>>({});
+  const parentToLayers = ref<Record<string, Layer[]>>({});
   const layerImages = ref<Record<LayerID, vtkImageData>>({});
 
   async function _addLayer(
@@ -61,10 +31,9 @@ export const useLayersStore = defineStore('layer', () => {
       throw new Error('Tried to addLayer without parent data selection');
     }
 
-    const parentKey = toDataSelectionKey(parent);
-    const id = `${parentKey}::${toDataSelectionKey(source)}` as LayerID;
-    this.parentToLayers[parentKey] = [
-      ...(this.parentToLayers[parentKey] ?? []),
+    const id = `${parent}::${source}` as LayerID;
+    this.parentToLayers[parent] = [
+      ...(this.parentToLayers[parent] ?? []),
       { selection: source, id } as Layer,
     ];
 
@@ -102,9 +71,8 @@ export const useLayersStore = defineStore('layer', () => {
         await this._addLayer(parent, source);
       } catch (error) {
         // remove failed layer from parent's layer list
-        const parentKey = toDataSelectionKey(parent);
-        this.parentToLayers[parentKey] = this.parentToLayers[parentKey]?.filter(
-          ({ selection }) => !selectionEquals(selection, source)
+        this.parentToLayers[parent] = this.parentToLayers[parent]?.filter(
+          ({ selection }) => selection !== source
         );
         throw error;
       }
@@ -120,25 +88,21 @@ export const useLayersStore = defineStore('layer', () => {
       throw new Error('Tried to deleteLayer without parent data selection');
     }
 
-    const parentKey = toDataSelectionKey(parent);
-    const layers = this.parentToLayers[parentKey] ?? [];
+    const layers = this.parentToLayers[parent] ?? [];
 
-    const layerToDelete = layers.find(({ selection }) =>
-      selectionEquals(selection, source)
-    );
+    const layerToDelete = layers.find(({ selection }) => selection === source);
     if (!layerToDelete) return;
 
-    this.parentToLayers[parentKey] = layers.filter(
+    this.parentToLayers[parent] = layers.filter(
       (layer) => layer !== layerToDelete
     );
 
     delete this.layerImages[layerToDelete.id];
   }
 
-  function getLayers(key: DataSelection | null | undefined) {
+  function getLayers(key: Maybe<DataSelection>) {
     if (!key) return [];
-    const dataKey = toDataSelectionKey(key);
-    return parentToLayers.value[dataKey] ?? [];
+    return parentToLayers.value[key] ?? [];
   }
 
   const getLayer = (layerID: LayerID) =>
@@ -154,18 +118,16 @@ export const useLayersStore = defineStore('layer', () => {
       .getLayers(selectionToRemove)
       .forEach(({ selection }) => layerStore.deleteLayer(selection, selection));
     // delete from layer lists
-    Object.keys(layerStore.parentToLayers)
-      .map((key) => toDataSelectionFromKey(key as DataSelectionKey))
-      .forEach((parent) => layerStore.deleteLayer(parent, selectionToRemove));
+    Object.keys(layerStore.parentToLayers).forEach((parent) =>
+      layerStore.deleteLayer(parent, selectionToRemove)
+    );
   };
 
   function serialize(this: _This, state: StateFile) {
     state.manifest.parentToLayers = Object.entries(this.parentToLayers).map(
       ([selectionKey, layers]) => ({
         selectionKey,
-        sourceSelectionKeys: layers.map(({ selection }) =>
-          toDataSelectionKey(selection)
-        ),
+        sourceSelectionKeys: layers.map(({ selection }) => selection),
       })
     );
   }
@@ -176,25 +138,15 @@ export const useLayersStore = defineStore('layer', () => {
     dataIDMap: Record<string, string>
   ) {
     const remapSelection = (selection: DataSelection) => {
-      if (selection.type === 'dicom')
-        return makeDICOMSelection(dataIDMap[selection.volumeKey]);
-      if (selection.type === 'image')
-        return makeImageSelection(dataIDMap[selection.dataID]);
-
-      const _exhaustiveCheck: never = selection;
-      throw new Error(`invalid selection type ${_exhaustiveCheck}`);
+      return dataIDMap[selection];
     };
 
     const { parentToLayers: parentToLayersSerialized } = manifest;
     parentToLayersSerialized.forEach(
       ({ selectionKey, sourceSelectionKeys }) => {
-        const parent = remapSelection(
-          toDataSelectionFromKey(selectionKey as DataSelectionKey)
-        );
+        const parent = remapSelection(selectionKey);
         sourceSelectionKeys.forEach((sourceKey) => {
-          const source = remapSelection(
-            toDataSelectionFromKey(sourceKey as DataSelectionKey)
-          );
+          const source = remapSelection(sourceKey);
           this.addLayer(parent, source);
         });
       }
diff --git a/src/store/datasets.ts b/src/store/datasets.ts
index ee70539e5..e8f7e387f 100644
--- a/src/store/datasets.ts
+++ b/src/store/datasets.ts
@@ -2,10 +2,9 @@ import vtkImageData from '@kitware/vtk.js/Common/DataModel/ImageData';
 import { defineStore } from 'pinia';
 import { computed, ref } from 'vue';
 import {
-  DataSelection,
-  getDataSelection,
-  getImageID,
-  selectionEquals,
+  isDicomImage,
+  isRegularImage,
+  type DataSelection,
 } from '@/src/utils/dataSelection';
 import { useDICOMStore } from './datasets-dicom';
 import { useImageStore } from './datasets-images';
@@ -31,10 +30,7 @@ export const useDatasetStore = defineStore('dataset', () => {
 
   // --- getters --- //
 
-  const primaryImageID = computed(() => {
-    if (primarySelection.value) return getImageID(primarySelection.value);
-    return undefined;
-  });
+  const primaryImageID = primarySelection;
 
   const primaryDataset = computed<vtkImageData | null>(() => {
     const { dataIndex } = imageStore;
@@ -43,25 +39,20 @@ export const useDatasetStore = defineStore('dataset', () => {
 
   const idsAsSelections = computed(() => {
     const volumeKeys = Object.keys(dicomStore.volumeInfo);
-    const images = imageStore.idList.filter(
-      (id) => !(id in dicomStore.imageIDToVolumeKey)
-    );
-    return [...volumeKeys, ...images].map(getDataSelection);
+    const images = imageStore.idList.filter((id) => isRegularImage(id));
+    return [...volumeKeys, ...images];
   });
 
   // --- actions --- //
 
   function setPrimarySelection(sel: DataSelection | null) {
     primarySelection.value = sel;
-
-    if (sel === null) {
-      return;
-    }
+    if (!sel) return;
 
     // if selection is dicom, call buildVolume
-    if (sel.type === 'dicom') {
+    if (isDicomImage(sel)) {
       useErrorMessage('Failed to build volume', () =>
-        dicomStore.buildVolume(sel.volumeKey)
+        dicomStore.buildVolume(sel)
       );
     }
   }
@@ -70,39 +61,24 @@ export const useDatasetStore = defineStore('dataset', () => {
     await dicomStore.serialize(stateFile);
     await imageStore.serialize(stateFile);
 
-    if (primarySelection.value !== null) {
-      let id = null;
-      if (primarySelection.value.type === 'dicom') {
-        id = primarySelection.value.volumeKey;
-      } else {
-        id = primarySelection.value.dataID;
-      }
-
+    if (primarySelection.value) {
       const { manifest } = stateFile;
-      manifest.primarySelection = id;
+      manifest.primarySelection = primarySelection.value;
     }
   }
 
   const remove = (id: string) => {
-    const sel = getDataSelection(id);
-    if (
-      primarySelection.value &&
-      selectionEquals(sel, primarySelection.value)
-    ) {
+    if (id === primarySelection.value) {
       primarySelection.value = null;
     }
 
-    if (sel.type === 'dicom') {
-      if (sel.volumeKey in dicomStore.volumeToImageID) {
-        imageStore.deleteData(dicomStore.volumeToImageID[sel.volumeKey]);
-      }
-      dicomStore.deleteVolume(sel.volumeKey);
-    } else if (sel.type === 'image') {
-      imageStore.deleteData(id);
+    if (isDicomImage(id)) {
+      dicomStore.deleteVolume(id);
     }
+    imageStore.deleteData(id);
 
     fileStore.remove(id);
-    layersStore.remove(sel);
+    layersStore.remove(id);
   };
 
   return {
diff --git a/src/store/segmentGroups.ts b/src/store/segmentGroups.ts
index 312d7c7c0..53e928bdd 100644
--- a/src/store/segmentGroups.ts
+++ b/src/store/segmentGroups.ts
@@ -13,12 +13,9 @@ import { SegmentMask } from '@/src/types/segment';
 import { DEFAULT_SEGMENT_MASKS } from '@/src/config';
 import { readImage, writeImage } from '@/src/io/readWriteImage';
 import {
-  DataSelection,
+  type DataSelection,
   getImage,
-  selectionEquals,
-  findImageID,
-  getImageID,
-  getDataID,
+  isRegularImage,
 } from '@/src/utils/dataSelection';
 import vtkLabelMap from '../vtk/LabelMap';
 import {
@@ -214,12 +211,12 @@ export const useSegmentGroupStore = defineStore('segmentGroup', () => {
   }
 
   function decodeSegments(image: DataSelection) {
-    if (image.type === 'image') {
+    if (isRegularImage(image)) {
       return structuredClone(DEFAULT_SEGMENT_MASKS);
     }
 
     const dicomStore = useDICOMStore();
-    const volumeInfo = dicomStore.volumeInfo[image.volumeKey];
+    const volumeInfo = dicomStore.volumeInfo[image];
     const segmentSequence = undefined; // volumeInfo.SegmentSequence;
     if (!segmentSequence) {
       return [
@@ -245,20 +242,17 @@ export const useSegmentGroupStore = defineStore('segmentGroup', () => {
    * Converts an image to a labelmap.
    */
   async function convertImageToLabelmap(
-    image: DataSelection,
-    parent: DataSelection
+    imageID: DataSelection,
+    parentID: DataSelection
   ) {
-    if (selectionEquals(image, parent))
+    if (imageID === parentID)
       throw new Error('Cannot convert an image to be a labelmap of itself');
 
     // Build vtkImageData for DICOMs
     const [childImage, parentImage] = await Promise.all(
-      [image, parent].map(getImage)
+      [imageID, parentID].map(getImage)
     );
 
-    const imageID = getImageID(image);
-    const parentID = getImageID(parent);
-
     if (!imageID || !parentID)
       throw new Error('Image and/or parent datasets do not exist');
 
@@ -276,7 +270,7 @@ export const useSegmentGroupStore = defineStore('segmentGroup', () => {
 
     const name = imageStore.metadata[imageID].name;
     // Don't remove image if DICOM as user may have selected child image as primary selection by now
-    const deleteImage = image.type !== 'dicom';
+    const deleteImage = isRegularImage(imageID);
     if (deleteImage) {
       imageStore.deleteData(imageID);
     }
@@ -288,7 +282,7 @@ export const useSegmentGroupStore = defineStore('segmentGroup', () => {
       : resampled;
     const labelmapImage = toLabelMap(ownedMemoryImage);
 
-    const segments = decodeSegments(image);
+    const segments = decodeSegments(imageID);
     const { order, byKey } = normalizeForStore(segments, 'value');
     const segmentGroupStore = useSegmentGroupStore();
     segmentGroupStore.addLabelmap(labelmapImage, {
@@ -409,7 +403,7 @@ export const useSegmentGroupStore = defineStore('segmentGroup', () => {
           path: `labels/${id}.${saveFormat.value}`,
           metadata: {
             ...metadata,
-            parentImage: getDataID(metadata.parentImage),
+            parentImage: metadata.parentImage,
           },
         };
       });
@@ -465,7 +459,7 @@ export const useSegmentGroupStore = defineStore('segmentGroup', () => {
     labelMaps.forEach((labelMap, index) => {
       const { metadata } = labelMap;
       // map parent id to new id
-      const parentImage = findImageID(dataIDMap[metadata.parentImage]);
+      const parentImage = dataIDMap[metadata.parentImage];
       metadata.parentImage = parentImage;
 
       const newID = newLabelmapIDs[index];
diff --git a/src/store/tools/useAnnotationTool.ts b/src/store/tools/useAnnotationTool.ts
index def14ea9e..1524d6dc0 100644
--- a/src/store/tools/useAnnotationTool.ts
+++ b/src/store/tools/useAnnotationTool.ts
@@ -13,7 +13,6 @@ import { useViewStore } from '@/src/store/views';
 import { getLPSAxisFromDir } from '@/src/utils/lps';
 import { LPSAxisDir } from '@/src/types/lps';
 import { AnnotationTool, ToolID } from '@/src/types/annotation-tool';
-import { findImageID, getDataID } from '@/src/utils/dataSelection';
 import { useIdStore } from '@/src/store/id';
 import { useToolSelectionStore } from '@/src/store/tools/toolSelection';
 import type { IToolStore } from '@/src/store/tools/types';
@@ -171,8 +170,7 @@ export const useAnnotationTool = <
       .map((toolID) => toolByID.value[toolID])
       .filter((tool) => !tool.placing)
       .map(({ imageID, ...rest }) => ({
-        // If parent image is DICOM, save VolumeKey
-        imageID: getDataID(imageID),
+        imageID,
         ...rest,
       }));
 
@@ -205,7 +203,7 @@ export const useAnnotationTool = <
         ({ imageID, label, ...rest }) =>
           ({
             ...rest,
-            imageID: findImageID(dataIDMap[imageID]),
+            imageID: dataIDMap[imageID],
             label: (label && labelIDMap[label]) || '',
           } as ToolPatch)
       )
diff --git a/src/store/view-configs/layers.ts b/src/store/view-configs/layers.ts
index d24bc4817..66507e2bd 100644
--- a/src/store/view-configs/layers.ts
+++ b/src/store/view-configs/layers.ts
@@ -15,6 +15,7 @@ import {
 import { Maybe } from '@/src/types';
 import { identity } from '@/src/utils';
 import { LAYER_PRESET_BY_MODALITY, LAYER_PRESET_DEFAULT } from '@/src/config';
+import { isDicomImage } from '@/src/utils/dataSelection';
 import { createViewConfigSerializer } from './common';
 import { ViewConfig } from '../../io/state-file/schema';
 import { LayersConfig } from './types';
@@ -28,10 +29,9 @@ function getPreset(id: LayerID) {
     throw new Error(`Layer ${id} not found`);
   }
 
-  if (layer.selection.type === 'dicom') {
+  if (isDicomImage(layer.selection)) {
     const dicomStore = useDICOMStore();
-    const { Modality = undefined } =
-      dicomStore.volumeInfo[layer.selection.volumeKey];
+    const { Modality = undefined } = dicomStore.volumeInfo[layer.selection];
     return (
       (Modality && LAYER_PRESET_BY_MODALITY[Modality]) || LAYER_PRESET_DEFAULT
     );
diff --git a/src/store/view-configs/volume-coloring.ts b/src/store/view-configs/volume-coloring.ts
index 24a8e749b..ed0922ddd 100644
--- a/src/store/view-configs/volume-coloring.ts
+++ b/src/store/view-configs/volume-coloring.ts
@@ -16,6 +16,7 @@ import {
 } from '@/src/utils/doubleKeyRecord';
 import { DeepPartial, Maybe } from '@/src/types';
 import { identity } from '@/src/utils';
+import { isDicomImage } from '@/src/utils/dataSelection';
 import { createViewConfigSerializer } from './common';
 import { ViewConfig } from '../../io/state-file/schema';
 import { VolumeColorConfig } from './types';
@@ -30,8 +31,8 @@ export const DEFAULT_SAMPLING_DISTANCE = 0.2;
 
 function getPresetFromImageModality(imageID: string) {
   const dicomStore = useDICOMStore();
-  if (imageID in dicomStore.imageIDToVolumeKey) {
-    const volKey = dicomStore.imageIDToVolumeKey[imageID];
+  if (isDicomImage(imageID)) {
+    const volKey = imageID;
     const { Modality } = dicomStore.volumeInfo[volKey];
     if (Modality in DEFAULT_PRESET_BY_MODALITY) {
       return DEFAULT_PRESET_BY_MODALITY[Modality];
diff --git a/src/utils/dataSelection.ts b/src/utils/dataSelection.ts
index 15a41b470..27b12d406 100644
--- a/src/utils/dataSelection.ts
+++ b/src/utils/dataSelection.ts
@@ -1,95 +1,42 @@
 import { getDisplayName, useDICOMStore } from '@/src/store/datasets-dicom';
 import { useImageStore } from '@/src/store/datasets-images';
 import { useErrorMessage } from '@/src/composables/useErrorMessage';
+import { Maybe } from '@/src/types';
 
-export const makeDICOMSelection = (volumeKey: string) =>
-  ({
-    type: 'dicom',
-    volumeKey,
-  } as const);
+export type DataSelection = string;
 
-export type DICOMSelection = ReturnType<typeof makeDICOMSelection>;
+export const selectionEquals = (a: DataSelection, b: DataSelection) => a === b;
 
-export const makeImageSelection = (imageID: string) =>
-  ({
-    type: 'image',
-    dataID: imageID,
-  } as const);
-
-export type ImageSelection = ReturnType<typeof makeImageSelection>;
-
-export type DataSelection = DICOMSelection | ImageSelection;
-
-export const getImageID = (selection: DataSelection) => {
-  if (selection.type === 'image') return selection.dataID;
-  if (selection.type === 'dicom') {
-    // possibly undefined because image may not have been made yet
-    return useDICOMStore().volumeToImageID[selection.volumeKey];
-  }
+export const isDicomImage = (imageID: Maybe<string>) => {
+  if (!imageID) return false;
+  const store = useDICOMStore();
+  return imageID in store.volumeInfo;
+};
 
-  const _exhaustiveCheck: never = selection;
-  throw new Error(`invalid selection type ${_exhaustiveCheck}`);
+export const isRegularImage = (imageID: Maybe<string>) => {
+  if (!imageID) return false;
+  return !isDicomImage(imageID);
 };
 
-export const getImage = async (selection: DataSelection) => {
+export const getImage = async (imageID: string) => {
   const images = useImageStore();
   const dicoms = useDICOMStore();
-  if (selection.type === 'dicom') {
+  if (isDicomImage(imageID)) {
     // ensure image data exists
     await useErrorMessage('Failed to build volume', () =>
-      dicoms.buildVolume(selection.volumeKey)
+      dicoms.buildVolume(imageID)
     );
-    return images.dataIndex[getImageID(selection)!];
   }
-  // just an image that should be in dataIndex
-  return images.dataIndex[selection.dataID];
-};
-
-// Converts imageID to VolumeKey if exists, else return input imageID
-export const getDataID = (imageID: string) => {
-  const dicomStore = useDICOMStore();
-  return dicomStore.imageIDToVolumeKey[imageID] ?? imageID;
+  return images.dataIndex[imageID];
 };
 
-// @param id - VolumeKey or ImageID
-export const getDataSelection = (id: string) => {
-  const dicomStore = useDICOMStore();
-  const volumeKey =
-    dicomStore.imageIDToVolumeKey[id] ??
-    (id in dicomStore.volumeInfo ? id : undefined);
-  if (volumeKey) {
-    return makeDICOMSelection(volumeKey);
-  }
-  return makeImageSelection(id);
-};
-
-// Pass VolumeKey or ImageID and get ImageID
-export const findImageID = (dataID: string) => {
-  const dicomStore = useDICOMStore();
-  return dicomStore.volumeToImageID[dataID] ?? dataID;
-};
-
-export function selectionEquals(s1: DataSelection, s2: DataSelection) {
-  if (s1.type === 'dicom' && s2.type === 'dicom') {
-    return s1.volumeKey === s2.volumeKey;
-  }
-  if (s1.type === 'image' && s2.type === 'image') {
-    return s1.dataID === s2.dataID;
-  }
-  return false;
-}
-
 const getImageName = (imageID: string) => {
   return useImageStore().metadata[imageID].name;
 };
 
-export const getSelectionName = (selection: DataSelection) => {
-  if (selection.type === 'image') return getImageName(selection.dataID);
-  if (selection.type === 'dicom') {
-    const imageID = getImageID(selection);
-    if (imageID) return getImageName(imageID);
-    return getDisplayName(useDICOMStore().volumeInfo[selection.volumeKey]);
+export const getSelectionName = (selection: string) => {
+  if (isDicomImage(selection)) {
+    return getDisplayName(useDICOMStore().volumeInfo[selection]);
   }
-  const _exhaustiveCheck: never = selection;
-  throw new Error(`invalid selection type ${_exhaustiveCheck}`);
+  return getImageName(selection);
 };

From e4da4865892d7a0c21554f5907ba510e4b552730 Mon Sep 17 00:00:00 2001
From: Forrest <forrest.li@kitware.com>
Date: Tue, 21 May 2024 11:02:54 -0400
Subject: [PATCH 2/2] feat(state): bump manifest version to 4.0.0

---
 src/io/state-file/index.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/io/state-file/index.ts b/src/io/state-file/index.ts
index eaab18c88..8a7bcd05d 100644
--- a/src/io/state-file/index.ts
+++ b/src/io/state-file/index.ts
@@ -12,7 +12,7 @@ import { retypeFile } from '../io';
 import { ARCHIVE_FILE_TYPES } from '../mimeTypes';
 
 export const MANIFEST = 'manifest.json';
-export const MANIFEST_VERSION = '3.0.0';
+export const MANIFEST_VERSION = '4.0.0';
 
 export async function serialize() {
   const datasetStore = useDatasetStore();