Merge pull request #194 from catalystneuro/fix_get_images_dict_Extrac…

…tSegmentationExtractor [Fix] get segmentation images in `NewExtractSegmentationExtractor`
catalystneuro · Aug 15, 2022 · 5f5dfb3 · 5f5dfb3
2 parents 5a547c5 + f676459
commit 5f5dfb3
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/src/roiextractors/extractors/schnitzerextractor/extractsegmentationextractor.py b/src/roiextractors/extractors/schnitzerextractor/extractsegmentationextractor.py
@@ -240,7 +240,8 @@ def get_images_dict(self):
         images_dict: dict
             dictionary with key, values representing different types of Images
         """
-        images_dict = dict(
+        images_dict = super().get_images_dict()
+        images_dict.update(
             summary_image=self._info_struct["summary_image"][:],
             f_per_pixel=self._info_struct["F_per_pixel"][:],
             max_image=self._info_struct["max_image"][:],

diff --git a/tests/test_extractsegmentationextractor.py b/tests/test_extractsegmentationextractor.py
@@ -199,7 +199,10 @@ def test_extractor_get_images_dict(self):
             )[:]
 
         images_dict = self.extractor.get_images_dict()
-        self.assertEqual(len(images_dict), 3)
+        self.assertEqual(len(images_dict), 5)
+
+        self.assertEqual(images_dict["correlation"], None)
+        self.assertEqual(images_dict["mean"], None)
 
         self.assertEqual(images_dict["summary_image"].shape, summary_image.shape)
         self.assertEqual(images_dict["max_image"].shape, max_image.shape)