fix(overlay-menu): Adding, changing and removing segmentation overlays should update the overlay menu instantly (#5181)

extensions/cornerstone/src/components/ViewportDataOverlaySettingMenu/ViewportDataOverlayMenu.tsx

@@ -36,6 +36,9 @@ function ViewportDataOverlayMenu({ viewportId }: withAppTypes<{ viewportId: stri
     foregroundDisplaySets,
   } = useViewportDisplaySets(viewportId);
 
+  const [optimisticOverlayDisplaySets, setOptimisticOverlayDisplaySets] =
+    useState(overlayDisplaySets);
+
   const [thresholdOpacityEnabled, setThresholdOpacityEnabled] = useState(false);
 
   /**
@@ -77,6 +80,18 @@ function ViewportDataOverlayMenu({ viewportId }: withAppTypes<{ viewportId: stri
    * Remove a display set layer
    */
   const handleRemoveDisplaySetLayer = (displaySetInstanceUID: string) => {
+    const optimisticOverlayDisplaySetsIndex = optimisticOverlayDisplaySets.findIndex(
+      displaySet => displaySet.displaySetInstanceUID === displaySetInstanceUID
+    );
+
+    if (optimisticOverlayDisplaySetsIndex !== -1) {
+      setOptimisticOverlayDisplaySets(prevOptimisticOverlayDisplaySets => {
+        return prevOptimisticOverlayDisplaySets.filter(
+          displaySet => displaySet.displaySetInstanceUID !== displaySetInstanceUID
+        );
+      });
+    }
+
     commandsManager.runCommand('removeDisplaySetLayer', {
       viewportId,
       displaySetInstanceUID,
@@ -110,6 +125,17 @@ function ViewportDataOverlayMenu({ viewportId }: withAppTypes<{ viewportId: stri
     );
 
     if (selectedDisplaySet) {
+      setOptimisticOverlayDisplaySets(prevOptimisticOverlayDisplaySets => {
+        const currentDisplaySetIndex = prevOptimisticOverlayDisplaySets.findIndex(
+          displaySet => displaySet.displaySetInstanceUID === currentDisplaySet.displaySetInstanceUID
+        );
+        return [
+          ...prevOptimisticOverlayDisplaySets.slice(0, currentDisplaySetIndex),
+          selectedDisplaySet,
+          ...prevOptimisticOverlayDisplaySets.slice(currentDisplaySetIndex + 1),
+        ];
+      });
+
       handleReplaceDisplaySetLayer(
         currentDisplaySet.displaySetInstanceUID,
         selectedDisplaySet.displaySetInstanceUID
@@ -150,6 +176,10 @@ function ViewportDataOverlayMenu({ viewportId }: withAppTypes<{ viewportId: stri
     );
 
     if (selectedDisplaySet) {
+      setOptimisticOverlayDisplaySets(prevOptimisticOverlayDisplaySets => [
+        ...prevOptimisticOverlayDisplaySets,
+        selectedDisplaySet,
+      ]);
       handleAddDisplaySetAsLayer(selectedDisplaySet.displaySetInstanceUID);
       // Remove this pending segmentation from the list
       setPendingSegmentations(pendingSegmentations.filter(id => id !== pendingId));
@@ -184,7 +214,10 @@ function ViewportDataOverlayMenu({ viewportId }: withAppTypes<{ viewportId: stri
   };
 
   return (
-    <div className="bg-popover flex h-full w-[275px] flex-col rounded rounded-md p-1.5">
+    <div
+      className="bg-popover flex h-full w-[275px] flex-col rounded rounded-md p-1.5"
+      data-cy={`viewport-data-overlay-menu-${viewportId}`}
+    >
       {/* Top buttons row */}
       <div className={`flex`}>
         <Button
@@ -216,7 +249,7 @@ function ViewportDataOverlayMenu({ viewportId }: withAppTypes<{ viewportId: stri
       <div className="">
         {/* Overlays Segmentation section */}
         <div className="my-2 ml-1">
-          {overlayDisplaySets.map((displaySet, index) => (
+          {optimisticOverlayDisplaySets.map(displaySet => (
             <div
               key={displaySet.displaySetInstanceUID}
               className="mb-1 flex items-center"
@@ -227,7 +260,11 @@ function ViewportDataOverlayMenu({ viewportId }: withAppTypes<{ viewportId: stri
                 onValueChange={value => handleOverlaySelectionChange(displaySet, value)}
               >
                 <SelectTrigger className="flex-1">
-                  <SelectValue>{displaySet.label?.toUpperCase()}</SelectValue>
+                  <SelectValue
+                    data-cy={`overlay-ds-select-value-${displaySet.label?.toUpperCase()}`}
+                  >
+                    {displaySet.label?.toUpperCase()}
+                  </SelectValue>
                 </SelectTrigger>
                 <SelectContent>
                   {/* Include both potential overlays and the current overlay */}
@@ -255,12 +292,14 @@ function ViewportDataOverlayMenu({ viewportId }: withAppTypes<{ viewportId: stri
                     variant="ghost"
                     size="icon"
                     className="ml-2 flex-shrink-0"
+                    dataCY={`overlay-ds-more-button-${displaySet.label?.toUpperCase()}`}
                   >
                     <Icons.More className="h-4 w-4" />
                   </Button>
                 </DropdownMenuTrigger>
                 <DropdownMenuContent align="start">
                   <DropdownMenuItem
+                    data-cy={`overlay-ds-remove-button-${displaySet.label?.toUpperCase()}`}
                     onClick={() => handleRemoveDisplaySetLayer(displaySet.displaySetInstanceUID)}
                   >
                     Remove

extensions/cornerstone/src/hooks/useViewportDisplaySets.ts

@@ -1,4 +1,4 @@
-import { useMemo } from 'react';
+import { useEffect, useMemo, useState } from 'react';
 import { useSystem, utils } from '@ohif/core';
 import { useViewportGrid } from '@ohif/ui-next';
 import {
@@ -122,12 +122,34 @@ export function useViewportDisplaySets(
 
   // Get all available segmentations (only if needed)
   const needsSegmentations = includeOverlay;
-  const segmentationRepresentations = useMemo(
-    () =>
-      needsSegmentations ? segmentationService.getSegmentationRepresentations(viewportIdToUse) : [],
-    [segmentationService, viewportIdToUse, needsSegmentations]
+
+  const [segmentationRepresentations, setSegmentationRepresentations] = useState(
+    needsSegmentations ? segmentationService.getSegmentationRepresentations(viewportIdToUse) : []
   );
 
+  useEffect(() => {
+    setSegmentationRepresentations(
+      needsSegmentations ? segmentationService.getSegmentationRepresentations(viewportIdToUse) : []
+    );
+
+    const unsubscribeArr = needsSegmentations
+      ? [
+          segmentationService.EVENTS.SEGMENTATION_REPRESENTATION_MODIFIED,
+          segmentationService.EVENTS.SEGMENTATION_REPRESENTATION_REMOVED,
+        ].map(event =>
+          segmentationService.subscribe(event, () => {
+            setSegmentationRepresentations(
+              segmentationService.getSegmentationRepresentations(viewportIdToUse)
+            );
+          })
+        )
+      : [];
+
+    return () => {
+      unsubscribeArr.forEach(item => item.unsubscribe());
+    };
+  }, [segmentationService, viewportIdToUse, needsSegmentations]);
+
   const overlayDisplaySets = useMemo(() => {
     if (!includeOverlay) {
       return [];

tests/DataOverlayMenu.spec.ts

@@ -0,0 +1,78 @@
+import { test } from 'playwright-test-coverage';
+import { visitStudy, checkForScreenshot, screenShotPaths } from './utils';
+import { press } from './utils/keyboardUtils';
+
+test('should display added, selected and removed segmentation promptly', async ({ page }) => {
+  const studyInstanceUID = '1.3.6.1.4.1.32722.99.99.239341353911714368772597187099978969331';
+  const mode = 'segmentation';
+  await visitStudy(page, studyInstanceUID, mode, 2000);
+
+  // Add a segmentation overlay and ensure the overlay menu reflects this change.
+  await page.getByTestId('dataOverlayMenu-default-btn').click();
+  await page.getByTestId('AddSegmentationDataOverlay-default').click();
+  await page.getByText('SELECT A SEGMENTATION').click();
+  await page.getByTestId('2d-tta_nnU-Net_Segmentation').click();
+
+  await checkForScreenshot({
+    page,
+    screenshotPath:
+      screenShotPaths.dataOverlayMenu.overlayMenuWith2d_tta_nnU_Net_SegmentationSelected,
+  });
+
+  // Hide the overlay menu.
+  await page.getByTestId('dataOverlayMenu-default-btn').click();
+
+  // navigate to the 51st image and ensure the correct overlay is displayed
+  await press({ page, key: 'ArrowDown', nTimes: 50 });
+
+  await checkForScreenshot({
+    page,
+    screenshotPath: screenShotPaths.dataOverlayMenu.overlay2d_tta_nnU_Net_Segmentation,
+  });
+
+  // Show the overlay menu.
+  await page.getByTestId('dataOverlayMenu-default-btn').click();
+
+  // Change the segmentation overlay to a different one and ensure the overlay menu reflects this change.
+  await page.getByTestId('overlay-ds-select-value-2D-TTA_NNU-NET_SEGMENTATION').click();
+  await page.getByTestId('Segmentation-SEG').click();
+
+  await checkForScreenshot({
+    page,
+    screenshotPath: screenShotPaths.dataOverlayMenu.overlayMenuWithSegmentationSelected,
+  });
+
+  // Hide the overlay menu.
+  await page.getByTestId('dataOverlayMenu-default-btn').click();
+
+  // navigate to the 51st image and ensure the correct overlay is displayed
+  await press({ page, key: 'ArrowDown', nTimes: 50 });
+
+  await checkForScreenshot({
+    page,
+    screenshotPath: screenShotPaths.dataOverlayMenu.overlaySegmentation,
+  });
+
+  // Show the overlay menu.
+  await page.getByTestId('dataOverlayMenu-default-btn').click();
+
+  // Remove the segmentation overlay and ensure the overlay menu reflects this change.
+  await page.getByTestId('overlay-ds-more-button-SEGMENTATION').click();
+  await page.getByTestId('overlay-ds-remove-button-SEGMENTATION').click();
+
+  await checkForScreenshot({
+    page,
+    screenshotPath: screenShotPaths.dataOverlayMenu.overlayMenuWithSegmentationOverlaysRemoved,
+  });
+
+  // Hide the overlay menu.
+  await page.getByTestId('dataOverlayMenu-default-btn').click();
+
+  // navigate to the 51st image and ensure no overlay is displayed
+  await press({ page, key: 'ArrowDown', nTimes: 50 });
+
+  await checkForScreenshot({
+    page,
+    screenshotPath: screenShotPaths.dataOverlayMenu.noOverlay,
+  });
+});

tests/MPRThenRTOverlayNoHydration.spec.ts

@@ -10,11 +10,12 @@ test.beforeEach(async ({ page }) => {
 test('should launch MPR with unhydrated RTSTRUCT chosen from the data overlay menu', async ({
   page,
 }) => {
-  await page.getByTestId('side-panel-header-right').click();
-
   await page.getByTestId('Layout').click();
   await page.getByTestId('MPR').click();
 
+  // Wait 5 seconds for MPR to load. This is necessary in particular when screen shots are added or replaced.
+  await page.waitForTimeout(5000);
+
   await checkForScreenshot(
     page,
     page,
@@ -31,6 +32,9 @@ test('should launch MPR with unhydrated RTSTRUCT chosen from the data overlay me
   // Hide the overlay menu.
   await page.getByTestId('dataOverlayMenu-mpr-sagittal-btn').click();
 
+  // Wait 5 seconds for RT to load. This is necessary in particular when screen shots are added or replaced.
+  await page.waitForTimeout(5000);
+
   await checkForScreenshot(
     page,
     page,

tests/utils/checkForScreenshot.ts

@@ -3,7 +3,7 @@ import { Locator, Page } from 'playwright';
 
 type CheckForScreenshotProps = {
   page: Page;
-  locator: Locator | Page;
+  locator?: Locator | Page;
   screenshotPath: string;
   attempts?: number;
   delay?: number;
@@ -14,7 +14,7 @@ type CheckForScreenshotProps = {
 const _checkForScreenshot = async (props: CheckForScreenshotProps) => {
   const {
     page,
-    locator,
+    locator = page,
     screenshotPath,
     attempts = 10,
     delay = 500,

tests/utils/screenShotPaths.ts

@@ -164,6 +164,15 @@ const screenShotPaths = {
     jumpToMeasurementAfterSeriesChange:
       'jumpToMeasurementMPR-jumpToMeasurementAfterSeriesChange.png',
   },
+  dataOverlayMenu: {
+    overlayMenuWithSegmentationSelected: 'overlayMenuWithSegmentationSelected.png',
+    overlayMenuWith2d_tta_nnU_Net_SegmentationSelected:
+      'overlayMenuWith2d_tta_nnU_Net_SegmentationSelected.png',
+    overlayMenuWithSegmentationOverlaysRemoved: 'overlayMenuWithSegmentationOverlaysRemoved.png',
+    overlay2d_tta_nnU_Net_Segmentation: 'overlay2d_tta_nnU_Net_Segmentation.png',
+    overlaySegmentation: 'overlaySegmentation.png',
+    noOverlay: 'noOverlay.png',
+  },
 };
 
 export { screenShotPaths };