diff --git a/extensions/default/src/DicomWebDataSource/index.js b/extensions/default/src/DicomWebDataSource/index.js
index 053e1f841..ad0ca2199 100644
--- a/extensions/default/src/DicomWebDataSource/index.js
+++ b/extensions/default/src/DicomWebDataSource/index.js
@@ -125,11 +125,7 @@ function createDicomWebApi(dicomWebConfig) {
         },
       },
     },
-    getImageIdsForDisplaySet(displaySetInstanceUID, DisplaySetService) {
-      const displaySet = DisplaySetService.getDisplaySetByUID(
-        displaySetInstanceUID
-      );
-
+    getImageIdsForDisplaySet(displaySet) {
       const images = displaySet.images;
       const imageIds = [];
 
diff --git a/extensions/default/src/Panels/PanelStudyBrowser.jsx b/extensions/default/src/Panels/PanelStudyBrowser.jsx
index a2ae1a1bd..ad639780b 100644
--- a/extensions/default/src/Panels/PanelStudyBrowser.jsx
+++ b/extensions/default/src/Panels/PanelStudyBrowser.jsx
@@ -24,7 +24,7 @@ function PanelStudyBrowser({ servicesManager, getImageSrc, getStudiesByPatientId
         const dSets = displaySets[StudyInstanceUid];
         dSets.forEach(async (dset) => {
           const uid = dset.displaySetInstanceUID;
-          const imageIds = dataSource.getImageIdsForDisplaySet(uid, DisplaySetService);
+          const imageIds = dataSource.getImageIdsForDisplaySet(dset);
           const imageId = imageIds[Math.floor(imageIds.length / 2)];
           const imageSrc = await getImageSrc(imageId);
 
diff --git a/platform/core/src/utils/StackManager.js b/platform/core/src/utils/StackManager.js
index dfa4d1665..aed017a45 100644
--- a/platform/core/src/utils/StackManager.js
+++ b/platform/core/src/utils/StackManager.js
@@ -31,11 +31,7 @@ function createAndAddStack(
     return;
   }
 
-  // TODO: This is broken because it needs to give DisplaySetService to the dataSource
-  // DisplaySetService
-  // Should be something like:
-  // const imageIds = dataSource.getImageIdsForDisplaySet(displaySetInstanceUID, DisplaySetService);
-  const imageIds = dataSource.getImageIdsForDisplaySet(displaySetInstanceUID);
+  const imageIds = dataSource.getImageIdsForDisplaySet(displaySet);
 
   const stack = {
     StudyInstanceUID,