feat(segmentGroups): support overlapping segments in SEG files

Each component in the image to convert is a non overlapping labelmap.
Create another segment group for each component of image to convert.

closes #672

src/io/dicom.ts

@@ -190,7 +190,7 @@ type ReadOverlappingSegmentationResultWithRealMeta =
     metaInfo: ReadOverlappingSegmentationMeta;
   };
 
-export async function buildLabelMap(file: File) {
+export async function buildSegmentGroups(file: File) {
   const inputImage = sanitizeFile(file);
   const result = (await readOverlappingSegmentation(inputImage, {
     webWorker: getWorker(),

src/store/datasets-dicom.ts

@@ -3,7 +3,6 @@ import { defineStore } from 'pinia';
 import { Image } from 'itk-wasm';
 import { DataSourceWithFile } from '@/src/io/import/dataSource';
 import * as DICOM from '@/src/io/dicom';
-import { pullComponent0 } from '@/src/utils/images';
 import { identity, pick, removeFromArray } from '../utils';
 import { useImageStore } from './datasets-images';
 import { useFileStore } from './datasets-files';
@@ -55,16 +54,10 @@ const buildImage = async (seriesFiles: File[], modality: string) => {
   const messages: string[] = [];
   if (modality === 'SEG') {
     const segFile = seriesFiles[0];
-    const results = await DICOM.buildLabelMap(segFile);
-    if (results.outputImage.imageType.components !== 1) {
-      messages.push(
-        `${segFile.name} SEG file has overlapping segments. Using first set.`
-      );
-      results.outputImage = pullComponent0(results.segImage);
-    }
+    const results = await DICOM.buildSegmentGroups(segFile);
     if (seriesFiles.length > 1)
       messages.push(
-        'SEG image has multiple components. Using only the first component.'
+        'Tried to make one volume from 2 SEG modality files. Using only the first file!'
       );
     return {
       modality: 'SEG',

src/store/segmentGroups.ts

@@ -17,6 +17,7 @@ import {
   getImage,
   isRegularImage,
 } from '@/src/utils/dataSelection';
+import vtkImageExtractComponents from '@/src/utils/imageExtractComponentsFilter';
 import vtkLabelMap from '../vtk/LabelMap';
 import {
   StateFile,
@@ -35,6 +36,7 @@ export const DEFAULT_SEGMENT_COLOR: RGBAColor = [255, 0, 0, 255];
 export const makeDefaultSegmentName = (value: number) => `Segment ${value}`;
 export const makeDefaultSegmentGroupName = (baseName: string, index: number) =>
   `Segment Group ${index} for ${baseName}`;
+const numberer = (index: number) => (index <= 1 ? '' : `${index}`); // start numbering at 2
 
 export interface SegmentGroupMetadata {
   name: string;
@@ -79,6 +81,20 @@ export function toLabelMap(imageData: vtkImageData) {
   return labelmap;
 }
 
+export function extractEachComponent(input: vtkImageData) {
+  const numComponents = input
+    .getPointData()
+    .getScalars()
+    .getNumberOfComponents();
+  const extractComponentsFilter = vtkImageExtractComponents.newInstance();
+  extractComponentsFilter.setInputData(input);
+  return Array.from({ length: numComponents }, (_, i) => {
+    extractComponentsFilter.setComponents([i]);
+    extractComponentsFilter.update();
+    return extractComponentsFilter.getOutputData() as vtkImageData;
+  });
+}
+
 export const useSegmentGroupStore = defineStore('segmentGroup', () => {
   type _This = ReturnType<typeof useSegmentGroupStore>;
 
@@ -156,6 +172,22 @@ export const useSegmentGroupStore = defineStore('segmentGroup', () => {
     });
   });
 
+  function pickUniqueName(
+    formatName: (index: number) => string,
+    parentID: string
+  ) {
+    const existingNames = new Set(
+      Object.values(metadataByID).map((meta) => meta.name)
+    );
+    let name = '';
+    do {
+      const nameIndex = nextDefaultIndex[parentID] ?? 1;
+      nextDefaultIndex[parentID] = nameIndex + 1;
+      name = formatName(nameIndex);
+    } while (existingNames.has(name));
+    return name;
+  }
+
   /**
    * Creates a new labelmap entry from a parent/source image.
    */
@@ -174,16 +206,10 @@ export const useSegmentGroupStore = defineStore('segmentGroup', () => {
       'value'
     );
 
-    // pick a unique name
-    let name = '';
-    const existingNames = new Set(
-      Object.values(metadataByID).map((meta) => meta.name)
+    const name = pickUniqueName(
+      (index: number) => makeDefaultSegmentGroupName(baseName, index),
+      parentID
     );
-    do {
-      const nameIndex = nextDefaultIndex[parentID] ?? 1;
-      nextDefaultIndex[parentID] = nameIndex + 1;
-      name = makeDefaultSegmentGroupName(baseName, nameIndex);
-    } while (existingNames.has(name));
 
     return addLabelmap.call(this, labelmap, {
       name,
@@ -272,27 +298,42 @@ export const useSegmentGroupStore = defineStore('segmentGroup', () => {
       );
     }
 
-    const name = imageStore.metadata[imageID].name;
-    // Don't remove image if DICOM as user may have selected segment group image as primary selection by now
+    // Don't remove image if DICOM. User may have selected segment group image as primary selection by now
     const deleteImage = isRegularImage(imageID);
     if (deleteImage) {
       imageStore.deleteData(imageID);
     }
 
-    const matchingParentSpace = await ensureSameSpace(
-      parentImage,
-      childImage,
-      true
-    );
-    const labelmapImage = toLabelMap(matchingParentSpace);
+    const componentCount = childImage
+      .getPointData()
+      .getScalars()
+      .getNumberOfComponents();
+    // for each component, create create new vtkImageData with just one component, pulled from each component of childImage
+    const images =
+      componentCount === 1 ? [childImage] : extractEachComponent(childImage);
+
+    const baseName = imageStore.metadata[imageID].name;
+    images.forEach(async (image) => {
+      const matchingParentSpace = await ensureSameSpace(
+        parentImage,
+        image,
+        true
+      );
+      const labelmapImage = toLabelMap(matchingParentSpace);
 
-    const segments = await decodeSegments(imageID, labelmapImage);
-    const { order, byKey } = normalizeForStore(segments, 'value');
-    const segmentGroupStore = useSegmentGroupStore();
-    segmentGroupStore.addLabelmap(labelmapImage, {
-      name,
-      parentImage: parentID,
-      segments: { order, byValue: byKey },
+      const segments = await decodeSegments(imageID, labelmapImage);
+      const { order, byKey } = normalizeForStore(segments, 'value');
+      const segmentGroupStore = useSegmentGroupStore();
+
+      const name = pickUniqueName(
+        (index: number) => `${baseName} ${numberer(index)}`,
+        parentID
+      );
+      segmentGroupStore.addLabelmap(labelmapImage, {
+        name,
+        parentImage: parentID,
+        segments: { order, byValue: byKey },
+      });
     });
   }
 

src/utils/__tests__/imageExtractComponentsFilter.ts

@@ -0,0 +1,56 @@
+import { describe, it, expect } from 'vitest';
+import vtkImageData from '@kitware/vtk.js/Common/DataModel/ImageData';
+import vtkDataArray from '@kitware/vtk.js/Common/Core/DataArray';
+import vtkImageExtractComponentsFilter from '../imageExtractComponentsFilter';
+
+describe('vtkImageExtractComponentsFilter', () => {
+  it('should extract specified components', () => {
+    // Create an image data with known scalar components
+    const imageData = vtkImageData.newInstance();
+    imageData.setDimensions([2, 2, 1]);
+
+    // Create scalar data with 3 components per voxel
+    const scalars = vtkDataArray.newInstance({
+      numberOfComponents: 3,
+      values: new Uint8Array([
+        // Voxel 0
+        10, 20, 30,
+        // Voxel 1
+        40, 50, 60,
+        // Voxel 2
+        70, 80, 90,
+        // Voxel 3
+        100, 110, 120,
+      ]),
+    });
+
+    imageData.getPointData().setScalars(scalars);
+
+    // Create the filter and set components to extract
+    const extractComponentsFilter =
+      vtkImageExtractComponentsFilter.newInstance();
+    extractComponentsFilter.setComponents([0, 2]); // Extract components 0 and 2
+    extractComponentsFilter.setInputData(imageData);
+    extractComponentsFilter.update();
+
+    const outputData = extractComponentsFilter.getOutputData();
+    const outputScalars = outputData.getPointData().getScalars();
+    const outputValues = outputScalars.getData();
+
+    // Expected output
+    const expectedValues = new Uint8Array([
+      // Voxel 0
+      10, 30,
+      // Voxel 1
+      40, 60,
+      // Voxel 2
+      70, 90,
+      // Voxel 3
+      100, 120,
+    ]);
+
+    // Check if output matches expected values
+    expect(outputScalars.getNumberOfComponents()).toBe(2);
+    expect(outputValues).toEqual(expectedValues);
+  });
+});

src/utils/imageExtractComponentsFilter.js

@@ -0,0 +1,86 @@
+/* eslint-disable no-param-reassign */
+import macro from '@kitware/vtk.js/macro';
+import vtkImageData from '@kitware/vtk.js/Common/DataModel/ImageData';
+import vtkDataArray from '@kitware/vtk.js/Common/Core/DataArray';
+
+function vtkImageExtractComponentsFilter(publicAPI, model) {
+  model.classHierarchy.push('vtkImageExtractComponentsFilter');
+
+  publicAPI.requestData = (inData, outData) => {
+    const inputData = inData[0];
+    const outputData = vtkImageData.newInstance();
+
+    const components = model.components;
+    if (!components || !Array.isArray(components) || components.length === 0) {
+      throw Error('No components specified for extraction.');
+    }
+
+    const inputScalars = inputData.getPointData().getScalars();
+    const numInputComponents = inputScalars.getNumberOfComponents();
+    components.forEach((c) => {
+      if (c < 0) {
+        throw Error('Component index must be greater than or equal to 0.');
+      }
+      if (c >= numInputComponents) {
+        throw Error(
+          'Component index must be less than the number of components in the input data.'
+        );
+      }
+    });
+
+    outputData.shallowCopy(inputData);
+
+    const inputArray = inputScalars.getData();
+    const numPixels = inputArray.length / numInputComponents;
+
+    const outputNumComponents = components.length;
+    const outputArray = new inputArray.constructor(
+      numPixels * outputNumComponents
+    );
+
+    for (let pixel = 0; pixel < numPixels; pixel++) {
+      for (let c = 0; c < components.length; c++) {
+        outputArray[pixel * outputNumComponents + c] =
+          inputArray[pixel * numInputComponents + components[c]];
+      }
+    }
+
+    outputData.getPointData().setScalars(
+      vtkDataArray.newInstance({
+        numberOfComponents: outputNumComponents,
+        values: outputArray,
+      })
+    );
+
+    outData[0] = outputData;
+  };
+}
+
+const DEFAULT_VALUES = {
+  components: [],
+};
+
+// ----------------------------------------------------------------------------
+
+export function extend(publicAPI, model, initialValues = {}) {
+  Object.assign(model, DEFAULT_VALUES, initialValues);
+
+  macro.obj(publicAPI, model);
+
+  macro.algo(publicAPI, model, 1, 1);
+
+  macro.setGet(publicAPI, model, ['components']);
+
+  vtkImageExtractComponentsFilter(publicAPI, model);
+}
+
+// ----------------------------------------------------------------------------
+
+export const newInstance = macro.newInstance(
+  extend,
+  'vtkImageExtractComponentsFilter'
+);
+
+// ----------------------------------------------------------------------------
+
+export default { newInstance, extend };