diff --git a/extensions/cornerstone-dicom-seg/src/commandsModule.ts b/extensions/cornerstone-dicom-seg/src/commandsModule.ts
index aa104bc51..b9c9dec5d 100644
--- a/extensions/cornerstone-dicom-seg/src/commandsModule.ts
+++ b/extensions/cornerstone-dicom-seg/src/commandsModule.ts
@@ -94,6 +94,7 @@ const commandsModule = ({
       updateViewportsForSegmentationRendering({
         viewportId,
         servicesManager,
+        displaySet,
         loadFn: async () => {
           const currentSegmentations = segmentationService.getSegmentations();
           const segmentationId = await segmentationService.createSegmentationForDisplaySet(
diff --git a/extensions/cornerstone-dicom-seg/src/utils/hydrationUtils.ts b/extensions/cornerstone-dicom-seg/src/utils/hydrationUtils.ts
index 2fd6114db..d87c6e109 100644
--- a/extensions/cornerstone-dicom-seg/src/utils/hydrationUtils.ts
+++ b/extensions/cornerstone-dicom-seg/src/utils/hydrationUtils.ts
@@ -35,8 +35,8 @@ async function updateViewportsForSegmentationRendering({
     displaySet?.referencedDisplaySetInstanceUID || viewport?.displaySetInstanceUIDs[0];
 
   const updatedViewports = getUpdatedViewportsForSegmentation({
-    servicesManager,
     viewportId,
+    servicesManager,
     displaySet,
   });
 
@@ -175,7 +175,7 @@ function getUpdatedViewportsForSegmentation({
         viewportId,
         displaySetInstanceUIDs: viewport.displaySetInstanceUIDs,
         viewportOptions: {
-          viewportType: displaySet.Modality === 'RTSTRUCT' ? 'stack' : 'volume',
+          viewportType: displaySet?.Modality === 'RTSTRUCT' ? 'stack' : 'volume',
           needsRerendering: true,
         },
       });