feat(state): support remote URIs for segment groups

Author: PaulHax

src/io/state-file/schema.ts

@@ -310,11 +310,19 @@ export const SegmentGroupMetadata = z.object({
   }),
 });
 
-export const LabelMap = z.object({
-  id: z.string(),
-  path: z.string(),
-  metadata: SegmentGroupMetadata,
-});
+export const LabelMap = z
+  .object({
+    id: z.string(),
+    path: z.string().optional(),
+    dataSourceId: z.number().optional(),
+    metadata: SegmentGroupMetadata,
+  })
+  .refine(
+    (data) => data.path !== undefined || data.dataSourceId !== undefined,
+    {
+      message: 'Either path or dataSourceId is required',
+    }
+  );
 
 export type LabelMap = z.infer<typeof LabelMap>;
 

src/store/segmentGroups.ts

@@ -24,10 +24,12 @@ import {
   StateFile,
   Manifest,
   SegmentGroupMetadata,
+  LabelMap,
 } from '../io/state-file/schema';
 import { FileEntry } from '../io/types';
 import { ensureSameSpace } from '../io/resample/resample';
 import { untilLoaded } from '../composables/untilLoaded';
+import { useDatasetStore } from './datasets';
 
 const LabelmapArrayType = Uint8Array;
 export type LabelmapArrayType = Uint8Array;
@@ -498,6 +500,7 @@ export const useSegmentGroupStore = defineStore('segmentGroup', () => {
     dataIDMap: Record<string, string>
   ) {
     const { labelMaps } = manifest;
+    const datasetStore = useDatasetStore();
 
     const segmentGroupIDMap: Record<string, string> = {};
 
@@ -508,33 +511,48 @@ export const useSegmentGroupStore = defineStore('segmentGroup', () => {
     // First restore the data, then restore the store.
     // This preserves ordering from orderByParent.
 
-    const newLabelmapIDs = await Promise.all(
-      labelMaps.map(async (labelMap) => {
-        const [file] = stateFiles
-          .filter((entry) => entry.archivePath === normalize(labelMap.path))
-          .map((entry) => entry.file);
+    async function loadLabelMapImage(labelMap: LabelMap) {
+      if (labelMap.dataSourceId !== undefined) {
+        const storeId = dataIDMap[String(labelMap.dataSourceId)];
+        await untilLoaded(storeId);
+        const image = imageCacheStore.getVtkImageData(storeId);
+        if (!image) {
+          throw new Error(
+            `Could not get image data for dataSourceId ${labelMap.dataSourceId}`
+          );
+        }
+        return { image, storeIdToRemove: storeId };
+      }
+
+      const file = stateFiles.find(
+        (entry) => entry.archivePath === normalize(labelMap.path!)
+      )?.file;
+      return { image: await readImage(file!) };
+    }
 
-        const vtkImage = await readImage(file);
-        const labelmapImage = toLabelMap(vtkImage);
+    const labelmapResults = await Promise.all(
+      labelMaps.map(async (labelMap) => {
+        const { image, storeIdToRemove } = await loadLabelMapImage(labelMap);
+        const labelmapImage = toLabelMap(image);
 
         const id = useIdStore().nextId();
         dataIndex[id] = labelmapImage;
-        return id;
+        return { id, storeIdToRemove };
       })
     );
 
     labelMaps.forEach((labelMap, index) => {
-      const { metadata } = labelMap;
-      // map parent id to new id
-      const parentImage = dataIDMap[metadata.parentImage];
-      metadata.parentImage = parentImage;
-
-      const newID = newLabelmapIDs[index];
+      const { id: newID, storeIdToRemove } = labelmapResults[index];
       segmentGroupIDMap[labelMap.id] = newID;
 
-      metadataByID[newID] = metadata;
+      const parentImage = dataIDMap[labelMap.metadata.parentImage];
+      metadataByID[newID] = { ...labelMap.metadata, parentImage };
       orderByParent.value[parentImage] ??= [];
       orderByParent.value[parentImage].push(newID);
+
+      if (storeIdToRemove) {
+        datasetStore.remove(storeIdToRemove);
+      }
     });
 
     return segmentGroupIDMap;

tests/specs/configTestUtils.ts

@@ -21,6 +21,11 @@ export const PROSTATEX_DATASET = {
   name: 'prostate.zip',
 } as const;
 
+export const PROSTATE_SEGMENT_GROUP = {
+  url: 'https://data.kitware.com/api/v1/file/692f13ed80eaefe49a4abb72/download',
+  name: 'prostate-total.seg.nii.gz',
+} as const;
+
 export const MRA_HEAD_NECK_DATASET = {
   url: 'https://data.kitware.com/api/v1/item/6352a2b311dab8142820a33b/download',
   name: 'MRA-Head_and_Neck.zip',

tests/specs/sparse-manifest.e2e.ts

@@ -1,4 +1,8 @@
-import { MINIMAL_DICOM } from './configTestUtils';
+import {
+  MINIMAL_DICOM,
+  PROSTATEX_DATASET,
+  PROSTATE_SEGMENT_GROUP,
+} from './configTestUtils';
 import {
   downloadFile,
   openVolViewPage,
@@ -109,4 +113,71 @@ describe('Sparse manifest.json', () => {
     await writeManifestToFile(sparseManifest, fileName);
     await openVolViewPage(fileName);
   });
+
+  it('loads remote segment group from URI', async () => {
+    await downloadFile(PROSTATEX_DATASET.url, PROSTATEX_DATASET.name);
+    await downloadFile(PROSTATE_SEGMENT_GROUP.url, PROSTATE_SEGMENT_GROUP.name);
+
+    const sparseManifest = {
+      version: '6.1.0',
+      dataSources: [
+        {
+          id: 0,
+          type: 'uri',
+          uri: `/tmp/${PROSTATEX_DATASET.name}`,
+        },
+        {
+          id: 1,
+          type: 'uri',
+          uri: `/tmp/${PROSTATE_SEGMENT_GROUP.name}`,
+        },
+      ],
+      labelMaps: [
+        {
+          id: 'seg-1',
+          dataSourceId: 1,
+          metadata: {
+            name: 'Prostate Segmentation',
+            parentImage: '0',
+            segments: {
+              order: [1],
+              byValue: {
+                '1': {
+                  value: 1,
+                  name: 'Prostate',
+                  color: [255, 0, 0, 255],
+                  visible: true,
+                },
+              },
+            },
+          },
+        },
+      ],
+    };
+
+    const fileName = 'remote-segment-group.volview.json';
+    await writeManifestToFile(sparseManifest, fileName);
+    await openVolViewPage(fileName);
+
+    const annotationsTab = await $(
+      'button[data-testid="module-tab-Annotations"]'
+    );
+    await annotationsTab.click();
+
+    const segmentGroupsTab = await $('button.v-tab*=Segment Groups');
+    await segmentGroupsTab.waitForClickable();
+    await segmentGroupsTab.click();
+
+    await browser.waitUntil(
+      async () => {
+        const segmentGroups = await $$('.segment-group-list .v-list-item');
+        const count = await segmentGroups.length;
+        return count >= 1;
+      },
+      {
+        timeout: 15000,
+        timeoutMsg: 'Segment group not found in segment groups list',
+      }
+    );
+  });
 });