OHIF with Cornerstone3D 2.0 (#4347)

extensions/cornerstone-dicom-seg/src/commandsModule.ts

@@ -14,18 +14,22 @@ import vtkImageMarchingSquares from '@kitware/vtk.js/Filters/General/ImageMarchi
 import vtkDataArray from '@kitware/vtk.js/Common/Core/DataArray';
 import vtkImageData from '@kitware/vtk.js/Common/DataModel/ImageData';
 
-import {
-  updateViewportsForSegmentationRendering,
-  getUpdatedViewportsForSegmentation,
-  getTargetViewport,
-} from './utils/hydrationUtils';
 const { segmentation: segmentationUtils } = utilities;
 
 const { datasetToBlob } = dcmjs.data;
 
+const getTargetViewport = ({ viewportId, viewportGridService }) => {
+  const { viewports, activeViewportId } = viewportGridService.getState();
+  const targetViewportId = viewportId || activeViewportId;
+
+  const viewport = viewports.get(targetViewportId);
+
+  return viewport;
+};
+
 const {
   Cornerstone3D: {
-    Segmentation: { generateLabelMaps2DFrom3D, generateSegmentation },
+    Segmentation: { generateSegmentation },
   },
 } = adaptersSEG;
 
@@ -42,87 +46,14 @@ const commandsModule = ({
   extensionManager,
 }: Types.Extensions.ExtensionParams): Types.Extensions.CommandsModule => {
   const {
-    uiNotificationService,
     segmentationService,
     uiDialogService,
     displaySetService,
     viewportGridService,
     toolGroupService,
-    cornerstoneViewportService,
-  } = servicesManager.services;
+  } = servicesManager.services as AppTypes.Services;
 
   const actions = {
-    /**
-     * Retrieves a list of viewports that require updates in preparation for segmentation rendering.
-     * This function evaluates viewports based on their compatibility with the provided segmentation's
-     * frame of reference UID and appends them to the updated list if they should render the segmentation.
-     *
-     * @param {Object} params - Parameters for the function.
-     * @param params.viewportId - the ID of the viewport to be updated.
-     * @param params.servicesManager - The services manager
-     * @param params.referencedDisplaySetInstanceUID - Optional UID for the referenced display set instance.
-     *
-     * @returns {Array} Returns an array of viewports that require updates for segmentation rendering.
-     */
-    getUpdatedViewportsForSegmentation,
-    /**
-     * Creates an empty segmentation for a specified viewport.
-     * It first checks if the display set associated with the viewport is reconstructable.
-     * If not, it raises a notification error. Otherwise, it creates a new segmentation
-     * for the display set after handling the necessary steps for making the viewport
-     * a volume viewport first
-     *
-     * @param {Object} params - Parameters for the function.
-     * @param params.viewportId - the target viewport ID.
-     *
-     */
-    createEmptySegmentationForViewport: async ({ viewportId }) => {
-      const viewport = getTargetViewport({ viewportId, viewportGridService });
-      // Todo: add support for multiple display sets
-      const displaySetInstanceUID = viewport.displaySetInstanceUIDs[0];
-
-      const displaySet = displaySetService.getDisplaySetByUID(displaySetInstanceUID);
-
-      if (!displaySet.isReconstructable) {
-        uiNotificationService.show({
-          title: 'Segmentation',
-          message: 'Segmentation is not supported for non-reconstructible displaysets yet',
-          type: 'error',
-        });
-        return;
-      }
-
-      updateViewportsForSegmentationRendering({
-        viewportId,
-        servicesManager,
-        displaySet,
-        loadFn: async () => {
-          const currentSegmentations = segmentationService.getSegmentations();
-          const segmentationId = await segmentationService.createSegmentationForDisplaySet(
-            displaySetInstanceUID,
-            { label: `Segmentation ${currentSegmentations.length + 1}` }
-          );
-
-          const toolGroupId = viewport.viewportOptions.toolGroupId;
-
-          await segmentationService.addSegmentationRepresentationToToolGroup(
-            toolGroupId,
-            segmentationId
-          );
-
-          // Add only one segment for now
-          segmentationService.addSegment(segmentationId, {
-            toolGroupId,
-            segmentIndex: 1,
-            properties: {
-              label: 'Segment 1',
-            },
-          });
-
-          return segmentationId;
-        },
-      });
-    },
     /**
      * Loads segmentations for a specified viewport.
      * The function prepares the viewport for rendering, then loads the segmentation details.
@@ -134,106 +65,30 @@ const commandsModule = ({
      *
      */
     loadSegmentationsForViewport: async ({ segmentations, viewportId }) => {
-      updateViewportsForSegmentationRendering({
-        viewportId,
-        servicesManager,
-        loadFn: async () => {
-          // Todo: handle adding more than one segmentation
-          const viewport = getTargetViewport({ viewportId, viewportGridService });
-          const displaySetInstanceUID = viewport.displaySetInstanceUIDs[0];
-
-          const segmentation = segmentations[0];
-          const segmentationId = segmentation.id;
-          const label = segmentation.label;
-          const segments = segmentation.segments;
-
-          delete segmentation.segments;
-
-          await segmentationService.createSegmentationForDisplaySet(displaySetInstanceUID, {
-            segmentationId,
-            label,
-          });
-
-          if (segmentation.scalarData) {
-            const labelmapVolume = segmentationService.getLabelmapVolume(segmentationId);
-            labelmapVolume.scalarData.set(segmentation.scalarData);
-          }
-
-          segmentationService.addOrUpdateSegmentation(segmentation);
-
-          const toolGroupId = viewport.viewportOptions.toolGroupId;
-          await segmentationService.addSegmentationRepresentationToToolGroup(
-            toolGroupId,
-            segmentationId
-          );
-
-          segments.forEach(segment => {
-            if (segment === null) {
-              return;
-            }
-            segmentationService.addSegment(segmentationId, {
-              segmentIndex: segment.segmentIndex,
-              toolGroupId,
-              properties: {
-                color: segment.color,
-                label: segment.label,
-                opacity: segment.opacity,
-                isLocked: segment.isLocked,
-                visibility: segment.isVisible,
-                active: segmentation.activeSegmentIndex === segment.segmentIndex,
-              },
-            });
-          });
-
-          if (segmentation.centroidsIJK) {
-            segmentationService.setCentroids(segmentation.id, segmentation.centroidsIJK);
-          }
-
-          return segmentationId;
-        },
-      });
-    },
-    /**
-     * Loads segmentation display sets for a specified viewport.
-     * Depending on the modality of the display set (SEG or RTSTRUCT),
-     * it chooses the appropriate service function to create
-     * the segmentation for the display set.
-     * The function then prepares the viewport for rendering segmentation.
-     *
-     * @param {Object} params - Parameters for the function.
-     * @param params.viewportId - ID of the viewport where the segmentation display sets should be loaded.
-     * @param params.displaySets - Array of display sets to be loaded for segmentation.
-     *
-     */
-    loadSegmentationDisplaySetsForViewport: async ({ viewportId, displaySets }) => {
       // Todo: handle adding more than one segmentation
-      const displaySet = displaySets[0];
-      const referencedDisplaySet = displaySetService.getDisplaySetByUID(
-        displaySet.referencedDisplaySetInstanceUID
-      );
-      const viewport = cornerstoneViewportService.getCornerstoneViewport(viewportId);
-      const initialSliceIndex = viewport.getSliceIndex();
+      const viewport = getTargetViewport({ viewportId, viewportGridService });
+      const displaySetInstanceUID = viewport.displaySetInstanceUIDs[0];
 
-      updateViewportsForSegmentationRendering({
-        viewportId,
-        servicesManager,
-        displaySet,
-        loadFn: async () => {
-          const segDisplaySet = displaySet;
-          const suppressEvents = false;
-          const serviceFunction =
-            segDisplaySet.Modality === 'SEG'
-              ? 'createSegmentationForSEGDisplaySet'
-              : 'createSegmentationForRTDisplaySet';
+      const segmentation = segmentations[0];
+      const segmentationId = segmentation.segmentationId;
+      const label = segmentation.config.label;
+      const segments = segmentation.config.segments;
 
-          const boundFn = segmentationService[serviceFunction].bind(segmentationService);
-          const segmentationId = await boundFn(segDisplaySet, null, suppressEvents);
-          const segmentation = segmentationService.getSegmentation(segmentationId);
-          segmentation.description = `S${referencedDisplaySet.SeriesNumber}: ${referencedDisplaySet.SeriesDescription}`;
-          return segmentationId;
-        },
-        initialSliceIndex,
+      const displaySet = displaySetService.getDisplaySetByUID(displaySetInstanceUID);
+
+      await segmentationService.createLabelmapForDisplaySet(displaySet, {
+        segmentationId,
+        segments,
+        label,
       });
+
+      segmentationService.addOrUpdateSegmentation(segmentation);
+
+      await segmentationService.addSegmentationRepresentation(viewport.viewportId, {
+        segmentationId,
+      });
+
+      return segmentationId;
     },
     /**
      * Generates a segmentation from a given segmentation ID.
@@ -251,25 +106,61 @@ const commandsModule = ({
     generateSegmentation: ({ segmentationId, options = {} }) => {
       const segmentation = cornerstoneToolsSegmentation.state.getSegmentation(segmentationId);
 
-      const { referencedVolumeId } = segmentation.representationData.LABELMAP;
+      const { imageIds } = segmentation.representationData.Labelmap;
 
-      const segmentationVolume = cache.getVolume(segmentationId);
-      const referencedVolume = cache.getVolume(referencedVolumeId);
-      const referencedImages = referencedVolume.getCornerstoneImages();
+      const segImages = imageIds.map(imageId => cache.getImage(imageId));
+      const referencedImages = segImages.map(image => cache.getImage(image.referencedImageId));
 
-      const labelmapObj = generateLabelMaps2DFrom3D(segmentationVolume);
+      const labelmaps2D = [];
 
-      // Generate fake metadata as an example
-      labelmapObj.metadata = [];
+      let z = 0;
+
+      for (const segImage of segImages) {
+        const segmentsOnLabelmap = new Set();
+        const pixelData = segImage.getPixelData();
+        const { rows, columns } = segImage;
+
+        // Use a single pass through the pixel data
+        for (let i = 0; i < pixelData.length; i++) {
+          const segment = pixelData[i];
+          if (segment !== 0) {
+            segmentsOnLabelmap.add(segment);
+          }
+        }
+
+        labelmaps2D[z++] = {
+          segmentsOnLabelmap: Array.from(segmentsOnLabelmap),
+          pixelData,
+          rows,
+          columns,
+        };
+      }
+
+      const allSegmentsOnLabelmap = labelmaps2D.map(labelmap => labelmap.segmentsOnLabelmap);
+
+      const labelmap3D = {
+        segmentsOnLabelmap: Array.from(new Set(allSegmentsOnLabelmap.flat())),
+        metadata: [],
+        labelmaps2D,
+      };
 
       const segmentationInOHIF = segmentationService.getSegmentation(segmentationId);
-      segmentationInOHIF.segments.forEach(segment => {
+      const representations = segmentationService.getRepresentationsForSegmentation(segmentationId);
+
+      Object.entries(segmentationInOHIF.segments).forEach(([segmentIndex, segment]) => {
         // segmentation service already has a color for each segment
         if (!segment) {
           return;
         }
-        const segmentIndex = segment.segmentIndex;
-        const { label, color } = segment;
+
+        const { label } = segment;
+
+        const firstRepresentation = representations[0];
+        const color = segmentationService.getSegmentColor(
+          firstRepresentation.viewportId,
+          segmentationId,
+          segment.segmentIndex
+        );
 
         const RecommendedDisplayCIELabValue = dcmjs.data.Colors.rgb2DICOMLAB(
           color.slice(0, 3).map(value => value / 255)
@@ -292,12 +183,12 @@ const commandsModule = ({
             CodeMeaning: 'Tissue',
           },
         };
-        labelmapObj.metadata[segmentIndex] = segmentMetadata;
+        labelmap3D.metadata[segmentIndex] = segmentMetadata;
       });
 
       const generatedSegmentation = generateSegmentation(
         referencedImages,
-        labelmapObj,
+        labelmap3D,
         metaData,
         options
       );
@@ -445,18 +336,19 @@ const commandsModule = ({
   };
 
   const definitions = {
-    getUpdatedViewportsForSegmentation: {
-      commandFn: actions.getUpdatedViewportsForSegmentation,
-    },
+    /**
+     * Obsolete?
+     */
     loadSegmentationDisplaySetsForViewport: {
       commandFn: actions.loadSegmentationDisplaySetsForViewport,
     },
+    /**
+     * Obsolete?
+     */
     loadSegmentationsForViewport: {
       commandFn: actions.loadSegmentationsForViewport,
     },
-    createEmptySegmentationForViewport: {
-      commandFn: actions.createEmptySegmentationForViewport,
-    },
+
     generateSegmentation: {
       commandFn: actions.generateSegmentation,
     },