Repository: catalystneuro/roiextractors (language: python)

Issue:
EXTRACT-PUBLIC .mat file support
this PR is to support the file type for the output of the [EXTRACT-PUBLIC](https://github.com/schnitzer-lab/EXTRACT-public#frequently-asked-questions) method of Schnitzer lab. 


Discussion / hints:
@bendichter What should be done about this PR?

There is no example file for it: https://gin.g-node.org/CatalystNeuro/ophys_testing_data/src/main/segmentation_datasets/extract

and no PR history for adding it back when this was WIP.

I would be hesitant to review/merge this without tests.
It looks like the file is the output of the tutorials [here](https://github.com/schnitzer-lab/EXTRACT-public/tree/master/Learning%20materials/Lecture%20Tutorials/Tutorial%201%20-%20Starting%20code)
> It looks like the file is the output of the tutorials [here](https://github.com/schnitzer-lab/EXTRACT-public/tree/master/Learning%20materials/Lecture%20Tutorials/Tutorial%201%20-%20Starting%20code)

Nice, I'll get it added, then.
@bendichter PR up on GIN https://gin.g-node.org/CatalystNeuro/ophys_testing_data/pulls/6
Granted, this branch has so many conflicts I'll probably just port this over to a new one to get this done.
Actually, I'm re-examining the example data.

The 'example' .mat file from that tutorial is actually just a couple of MATLAB variables that are used in running the tutorial itself, I do not believe they actually contain the structured format of the output of the call to the EXTRACT library functions. Instead, these objects are transformations of the output (see https://github.com/schnitzer-lab/EXTRACT-public/blob/master/Learning%20materials/Lecture%20Tutorials/Tutorial%201%20-%20Starting%20code/example_tutorial.m#L38-L39 for one such object)

However, by running the tutorial myself, I can just save the `output` variable from the EXTACT method run on that example data and use that instead.

Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9a26092b..b824af1d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,6 +1,10 @@
 # Upcoming
 
 ### Features
+* Add support for newer versions of EXTRACT output files.
+The `ExtractSegmentationExtractor` class is now abstract and redirects to the newer or older
+extractor depending on the version of the file.
+The `ExtractSegmentatioExtractor` class `write_segmentation` method has now been deprecated. [PR #170](https://github.com/catalystneuro/roiextractors/pull/170)
 
 ### Improvements
 * Add `frame_to_time` to `SegmentationExtractor`, `get_roi_ids` is now a class method. [PR #187](https://github.com/catalystneuro/roiextractors/pull/187)
diff --git a/src/roiextractors/extractorlist.py b/src/roiextractors/extractorlist.py
index 20d468d4..b5faeea3 100644
--- a/src/roiextractors/extractorlist.py
+++ b/src/roiextractors/extractorlist.py
@@ -1,3 +1,4 @@
+"""Listing of available formats for extraction."""
 from .extractors.caiman import CaimanSegmentationExtractor
 from .extractors.hdf5imagingextractor import Hdf5ImagingExtractor
 from .extractors.numpyextractors import (
diff --git a/src/roiextractors/extractors/schnitzerextractor/__init__.py b/src/roiextractors/extractors/schnitzerextractor/__init__.py
index 6a735a89..7a10c775 100644
--- a/src/roiextractors/extractors/schnitzerextractor/__init__.py
+++ b/src/roiextractors/extractors/schnitzerextractor/__init__.py
@@ -1,2 +1,6 @@
 from .cnmfesegmentationextractor import CnmfeSegmentationExtractor
-from .extractsegmentationextractor import ExtractSegmentationExtractor
+from .extractsegmentationextractor import (
+    LegacyExtractSegmentationExtractor,
+    ExtractSegmentationExtractor,
+    NewExtractSegmentationExtractor,
+)
diff --git a/src/roiextractors/extractors/schnitzerextractor/extractsegmentationextractor.py b/src/roiextractors/extractors/schnitzerextractor/extractsegmentationextractor.py
index 99cbd3b9..c1f18e0b 100644
--- a/src/roiextractors/extractors/schnitzerextractor/extractsegmentationextractor.py
+++ b/src/roiextractors/extractors/schnitzerextractor/extractsegmentationextractor.py
@@ -1,4 +1,11 @@
+"""Extractor for reading the segmentation data that results from calls to EXTRACT."""
+from abc import ABC
 from pathlib import Path
+from typing import Optional
+
+import numpy as np
+from lazy_ops import DatasetView
+from packaging import version
 
 try:
     import h5py
@@ -7,22 +14,249 @@
 except ImportError:
     HAVE_H5PY = False
 
-import numpy as np
-from lazy_ops import DatasetView
 
-from ...extraction_tools import PathType
+from ...extraction_tools import PathType, ArrayType
 from ...multisegmentationextractor import MultiSegmentationExtractor
 from ...segmentationextractor import SegmentationExtractor
 
 
-class ExtractSegmentationExtractor(SegmentationExtractor):
+def _decode_h5py_array(unicode_int_array: np.ndarray) -> str:
+    """Auxiliary function to decode a numpy array of unicode ints to a string."""
+    return "".join(chr(unicode_int) for unicode_int in unicode_int_array)
+
+
+class ExtractSegmentationExtractor(ABC):
+    """Abstract class that defines which extractor class to use for a given file."""
+
+    extractor_name = "ExtractSegmentation"
+    installed = HAVE_H5PY  # check at class level if installed or not
+    installation_mesg = "To use ExtractSegmentationExtractor install h5py: \n\n pip install h5py \n\n"  # error message when not installed
+
+    def __new__(
+        cls,
+        file_path: PathType,
+        sampling_frequency: float,
+        output_struct_name: str = "output",
+    ):
+        """Abstract class that defines which extractor class to use for a given file.
+        For newer versions of the EXTRACT algorithm, the extractor class redirects to
+        NewExtractSegmentationExtractor. For older versions, the extractor class
+        redirects to LegacyExtractSegmentationExtractor.
+
+        Parameters
+        ----------
+        file_path: str
+            The location of the folder containing the .mat file.
+        output_struct_name: str
+            The name of output struct in the .mat file, default is "output".
+        sampling_frequency: float
+            The sampling frequency in units of Hz.
+        """
+        self = super().__new__(cls)
+        self.file_path = file_path
+        self.output_struct_name = output_struct_name
+        # Check if the file is a .mat file
+        cls._assert_file_is_mat(self)
+
+        # Check the version of the .mat file
+        if cls._check_extract_file_version(self):
+            # For newer versions of the .mat file, use the newer extractor
+            return NewExtractSegmentationExtractor(
+                file_path=file_path,
+                sampling_frequency=sampling_frequency,
+                output_struct_name=output_struct_name,
+            )
+
+        # For older versions of the .mat file, use the legacy extractor
+        return LegacyExtractSegmentationExtractor(file_path=file_path)
+
+    def _assert_file_is_mat(self):
+        """Check that the file exists and is a .mat file."""
+        file_path = Path(self.file_path)
+        assert file_path.exists(), f"File {file_path} does not exist."
+        assert file_path.suffix == ".mat", f"File {file_path} must be a .mat file."
+
+    def _check_extract_file_version(self) -> bool:
+        """Check the version of the extract file.
+        If the file was created with a newer version of the EXTRACT algorithm, the
+        function will return True, otherwise it will return False."""
+        with h5py.File(name=self.file_path, mode="r") as mat_file:
+            if self.output_struct_name not in mat_file:
+                return False
+            dataset_version = mat_file[self.output_struct_name]["info"]["version"][:]
+            dataset_version = np.ravel(dataset_version)
+            # dataset_version is an HDF5 dataset of encoded characters
+            version_name = _decode_h5py_array(dataset_version)
+
+            return version.Version(version_name) >= version.Version("1.1.0")
+
+
+class NewExtractSegmentationExtractor(SegmentationExtractor):
+    """
+    This class inherits from the SegmentationExtractor class, having all
+    its functionality specifically applied to the dataset output from
+    the \'EXTRACT\' ROI segmentation method.
+    """
+
+    extractor_name = "NewExtractSegmentation"
+    installed = HAVE_H5PY  # check at class level if installed or not
+    installation_mesg = (
+        "To use NewExtractSegmentation install h5py: \n\n pip install h5py \n\n"
+        # error message when not installed
+    )
+    is_writable = False
+    mode = "file"
+
+    def __init__(
+        self,
+        file_path: PathType,
+        sampling_frequency: float,
+        output_struct_name: str = "output",
+    ):
+        """
+        Load a SegmentationExtractor from a .mat file containing the output and config structs of the EXTRACT algorithm.
+        For regular timing, supply the sampling frequency. For irregular timing, supply the timestamps.
+
+        Parameters
+        ----------
+        file_path: PathType
+            Path to the .mat file containing the structs.
+        sampling_frequency: float
+            The sampling frequency in units of Hz. Supply if timing is regular.
+        output_struct_name: str, optional
+            The user has control over the names of the variables that return from `extraction(images, config)`.
+            The tutorials for EXTRACT follow the naming convention of 'output', which we assume as the default.
+        """
+        super().__init__()
+
+        self.output_struct_name = output_struct_name
+        self.file_path = file_path
+
+        if sampling_frequency is None:
+            raise AssertionError("The sampling_frequency must be provided.")
+
+        self._dataset_file = self._file_extractor_read()
+        assert output_struct_name in self._dataset_file, "Output struct not found in file."
+        self._output_struct = self._dataset_file[output_struct_name]
+
+        assert "config" in self._output_struct, "Config struct not found in file."
+        config_struct = self._output_struct["config"]
+        self.config = self._config_struct_to_dict(config_struct=config_struct)
+
+        traces = self._trace_extractor_read()
+        if self.config["preprocess"][0] == 1:
+            self._roi_response_dff = traces
+        else:
+            self._roi_response_raw = traces
+
+        self._sampling_frequency = sampling_frequency
+
+        self._image_masks = self._image_mask_extractor_read()
+
+        assert "info" in self._output_struct, "Info struct not found in file."
+        self._info_struct = self._output_struct["info"]
+
+    def close(self):
+        self._dataset_file.close()
+
+    def _file_extractor_read(self):
+        return h5py.File(self.file_path, "r")
+
+    def _config_struct_to_dict(self, config_struct: h5py.Group) -> dict:
+        """Flatten the config struct into a dictionary."""
+        config_dict = dict()
+        for property_name in config_struct:
+            if isinstance(config_struct[property_name], h5py.Dataset):
+                data = np.ravel(config_struct[property_name][:])
+                if property_name in ["trace_output_option", "cellfind_filter_type"]:
+                    data = _decode_h5py_array(data)
+                config_dict[property_name] = data
+            elif isinstance(config_struct[property_name], h5py.Group):
+                config_dict[property_name] = self._config_struct_to_dict(config_struct=config_struct[property_name])
+        return config_dict
+
+    def _image_mask_extractor_read(self) -> DatasetView:
+        """Returns the image masks with a shape of height, width, number of ROIs."""
+        return DatasetView(self._output_struct["spatial_weights"]).lazy_transpose()
+
+    def _trace_extractor_read(self) -> DatasetView:
+        """Returns the traces with a shape of number of ROIs and number of frames."""
+        return DatasetView(self._output_struct["temporal_weights"])
+
+    def get_accepted_list(self) -> list:
+        """
+        The ids of the ROIs which are accepted after manual verification of
+        ROIs.
+
+        Returns
+        -------
+        accepted_list: list
+            List of accepted ROIs
+        """
+        return [roi for roi in self.get_roi_ids() if np.any(self._image_masks[roi])]
+
+    def get_rejected_list(self) -> list:
+        """
+        The ids of the ROIs which are rejected after manual verification of
+        ROIs.
+
+        Returns
+        -------
+        rejected_list: list
+            List of rejected ROIs
+        """
+        accepted_list = self.get_accepted_list()
+        rejected_list = list(set(self.get_roi_ids()) - set(accepted_list))
+
+        return rejected_list
+
+    def get_roi_ids(self) -> list:
+        """Returns the list of ROI ids.
+
+        Returns
+        -------
+        roi_ids: list
+            ROI ids list.
+        """
+        return list(range(self.get_num_rois()))
+
+    def get_image_size(self) -> ArrayType:
+        """
+        Frame size of movie (height and width of image).
+
+        Returns
+        -------
+        image_size: array_like
+            2-D array: image height x image width
+        """
+        return self._image_masks.shape[:-1]
+
+    def get_images_dict(self):
+        """
+        Returns a dictionary with key, values representing different types of Images
+        used in segmentation.
+        Returns
+        -------
+        images_dict: dict
+            dictionary with key, values representing different types of Images
+        """
+        images_dict = dict(
+            summary_image=DatasetView(self._info_struct["summary_image"]),
+            f_per_pixel=DatasetView(self._info_struct["F_per_pixel"]),
+            max_image=DatasetView(self._info_struct["max_image"]),
+        )
+
+        return images_dict
+
+
+class LegacyExtractSegmentationExtractor(SegmentationExtractor):
     """
     This class inherits from the SegmentationExtractor class, having all
     its funtionality specifically applied to the dataset output from
     the \'EXTRACT\' ROI segmentation method.
     """
 
-    extractor_name = "ExtractSegmentation"
+    extractor_name = "LegacyExtractSegmentation"
     installed = HAVE_H5PY  # check at class level if installed or not
     is_writable = False
     mode = "file"
@@ -35,7 +269,7 @@ def __init__(self, file_path: PathType):
         file_path: str
             The location of the folder containing dataset.mat file.
         """
-        SegmentationExtractor.__init__(self)
+        super().__init__()
         self.file_path = file_path
         self._dataset_file, self._group0 = self._file_extractor_read()
         self._image_masks = self._image_mask_extractor_read()
@@ -54,7 +288,7 @@ def _file_extractor_read(self):
         return f, _group0
 
     def _image_mask_extractor_read(self):
-        return DatasetView(self._dataset_file[self._group0[0]]["filters"]).lazy_transpose([1, 2, 0]).dsetread()
+        return self._dataset_file[self._group0[0]]["filters"][:].transpose([1, 2, 0])
 
     def _trace_extractor_read(self):
         extracted_signals = DatasetView(self._dataset_file[self._group0[0]]["traces"])


Test patch:
diff --git a/tests/test_extractsegmentationextractor.py b/tests/test_extractsegmentationextractor.py
new file mode 100644
index 00000000..14bdfeb0
--- /dev/null
+++ b/tests/test_extractsegmentationextractor.py
@@ -0,0 +1,206 @@
+import h5py
+import numpy as np
+from hdmf.testing import TestCase
+from lazy_ops import DatasetView
+from numpy.testing import assert_array_equal
+from parameterized import parameterized, param
+
+from roiextractors import ExtractSegmentationExtractor
+from roiextractors.extractors.schnitzerextractor import (
+    NewExtractSegmentationExtractor,
+    LegacyExtractSegmentationExtractor,
+)
+
+from .setup_paths import OPHYS_DATA_PATH
+
+
+class TestExtractSegmentationExtractor(TestCase):
+    ophys_data_path = OPHYS_DATA_PATH / "segmentation_datasets" / "extract"
+
+    @classmethod
+    def setUpClass(cls):
+        cls.sampling_frequency = 30.0
+        cls.ophys_data_path = OPHYS_DATA_PATH / "segmentation_datasets" / "extract"
+
+    def test_extract_segmentation_extractor_file_path_does_not_exist(self):
+        """Test that the extractor raises an error if the file does not exist."""
+        not_a_mat_file_path = "not_a_mat_file.txt"
+        with self.assertRaisesWith(AssertionError, f"File {not_a_mat_file_path} does not exist."):
+            ExtractSegmentationExtractor(
+                file_path=not_a_mat_file_path,
+                sampling_frequency=self.sampling_frequency,
+            )
+
+    def test_extract_segmentation_extractor_file_path_is_not_a_mat_file(self):
+        """Test that the extractor raises an error if the file is not a .mat file."""
+        not_a_mat_file_path = OPHYS_DATA_PATH / "segmentation_datasets" / "nwb" / "nwb_test.nwb"
+        with self.assertRaisesWith(AssertionError, f"File {not_a_mat_file_path} must be a .mat file."):
+            ExtractSegmentationExtractor(
+                file_path=not_a_mat_file_path,
+                sampling_frequency=self.sampling_frequency,
+            )
+
+    def test_extract_segmentation_extractor_with_default_output_struct_name(self):
+        """Test that the extractor returns the NewExtractSegmentationExtractor
+        when the default "output" struct name is used."""
+        extractor = ExtractSegmentationExtractor(
+            file_path=self.ophys_data_path / "extract_public_output.mat",
+            sampling_frequency=self.sampling_frequency,
+        )
+
+        self.assertIsInstance(extractor, NewExtractSegmentationExtractor)
+
+    param_list = [
+        param(
+            file_path=ophys_data_path / "2014_04_01_p203_m19_check01_extractAnalysis.mat",
+            output_struct_name="extractAnalysisOutput",
+            extractor_class=LegacyExtractSegmentationExtractor,
+        ),
+        param(
+            file_path=ophys_data_path / "extract_public_output.mat",
+            output_struct_name="output",
+            extractor_class=NewExtractSegmentationExtractor,
+        ),
+    ]
+
+    @parameterized.expand(
+        param_list,
+    )
+    def test_extract_segmentation_extractor_redirects(self, file_path, output_struct_name, extractor_class):
+        """
+        Test that the extractor class redirects to the correct class
+        given the version of the .mat file.
+        """
+        extractor = ExtractSegmentationExtractor(
+            file_path=file_path,
+            output_struct_name=output_struct_name,
+            sampling_frequency=self.sampling_frequency,
+        )
+
+        self.assertIsInstance(extractor, extractor_class)
+
+
+class TestNewExtractSegmentationExtractor(TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.file_path = OPHYS_DATA_PATH / "segmentation_datasets" / "extract" / "extract_public_output.mat"
+        cls.output_struct_name = "output"
+        cls.sampling_frequency = 30.0
+
+    def setUp(self):
+        self.extractor = NewExtractSegmentationExtractor(
+            file_path=self.file_path,
+            output_struct_name=self.output_struct_name,
+            sampling_frequency=self.sampling_frequency,
+        )
+
+    def tearDown(self):
+        self.extractor.close()
+
+    def test_extractor_output_struct_assertion(self):
+        """Test that the extractor raises an error if the output struct name is not in the file."""
+        with self.assertRaisesWith(AssertionError, "Output struct not found in file."):
+            NewExtractSegmentationExtractor(
+                file_path=self.file_path,
+                output_struct_name="not_output",
+                sampling_frequency=self.sampling_frequency,
+            )
+
+    def test_extractor_no_timestamps_or_sampling_frequency(self):
+        """Test that the extractor raises an error if neither timestamps
+        nor sampling frequency are provided."""
+        with self.assertRaisesWith(exc_type=AssertionError, exc_msg=("The sampling_frequency must be provided.")):
+            NewExtractSegmentationExtractor(file_path=self.file_path, sampling_frequency=None)
+
+    def test_extractor_data_validity(self):
+        """Test that the extractor class returns the expected data."""
+
+        with h5py.File(self.file_path, "r") as segmentation_file:
+            spatial_weights = DatasetView(
+                segmentation_file[self.output_struct_name]["spatial_weights"]
+            ).lazy_transpose()
+            self.assertEqual(self.extractor._image_masks.shape, spatial_weights.shape)
+
+            temporal_weights = DatasetView(segmentation_file[self.output_struct_name]["temporal_weights"])
+            self.assertEqual(self.extractor._roi_response_dff.shape, temporal_weights.shape)
+
+            self.assertEqual(self.extractor._roi_response_raw, None)
+
+            self.assertEqual(self.extractor._sampling_frequency, self.sampling_frequency)
+            self.assertIsInstance(self.extractor.get_sampling_frequency(), float)
+
+            assert_array_equal(self.extractor.get_image_size(), [50, 50])
+
+            num_rois = temporal_weights.shape[0]
+            self.assertEqual(self.extractor.get_num_rois(), num_rois)
+
+            num_frames = temporal_weights.shape[1]
+            self.assertEqual(self.extractor.get_num_frames(), num_frames)
+
+            self.assertEqual(self.extractor.get_rejected_list(), [])
+            self.assertEqual(self.extractor.get_accepted_list(), list(range(num_rois)))
+
+    def test_extractor_config(self):
+        """Test that the extractor class returns the expected config."""
+
+        assert "preprocess" in self.extractor.config
+        self.assertEqual(self.extractor.config["preprocess"], [1])
+
+        assert "S_corr_thresh" in self.extractor.config
+        self.assertEqual(self.extractor.config["S_corr_thresh"], [0.1])
+
+        assert "thresholds" in self.extractor.config
+        self.assertEqual(self.extractor.config["thresholds"]["S_dup_corr_thresh"], [0.95])
+        self.assertEqual(self.extractor.config["thresholds"]["T_dup_corr_thresh"], [0.95])
+
+        assert "trace_output_option" in self.extractor.config
+        self.assertEqual(self.extractor.config["trace_output_option"], "raw")
+        assert "cellfind_filter_type" in self.extractor.config
+        self.assertEqual(self.extractor.config["cellfind_filter_type"], "none")
+
+    param_list = [
+        param(accepted_list=[4, 6, 12, 18]),
+        param(accepted_list=[]),
+        param(accepted_list=list(range(20))),
+    ]
+
+    @parameterized.expand(
+        param_list,
+    )
+    def test_extractor_accepted_list(self, accepted_list):
+        """Test that the extractor class returns the list of accepted and rejected ROIs
+        correctly given the list of non-zero ROIs."""
+        dummy_image_mask = np.zeros((20, 50, 50))
+        dummy_image_mask[accepted_list, ...] = 1
+
+        self.extractor._image_masks = dummy_image_mask
+
+        assert_array_equal(self.extractor.get_accepted_list(), accepted_list)
+        assert_array_equal(
+            self.extractor.get_rejected_list(),
+            list(set(range(20)) - set(accepted_list)),
+        )
+
+    def test_extractor_get_images_dict(self):
+        """Test that the extractor class returns the expected images dict."""
+        with h5py.File(self.file_path, "r") as segmentation_file:
+            summary_image = DatasetView(
+                segmentation_file[self.output_struct_name]["info"]["summary_image"],
+            )[:]
+            max_image = DatasetView(
+                segmentation_file[self.output_struct_name]["info"]["max_image"],
+            )[:]
+            f_per_pixel = DatasetView(
+                segmentation_file[self.output_struct_name]["info"]["F_per_pixel"],
+            )[:]
+
+        images_dict = self.extractor.get_images_dict()
+        self.assertEqual(len(images_dict), 3)
+
+        self.assertEqual(images_dict["summary_image"].shape, summary_image.shape)
+        self.assertEqual(images_dict["max_image"].shape, max_image.shape)
+        self.assertEqual(images_dict["f_per_pixel"].shape, f_per_pixel.shape)
+
+        assert_array_equal(images_dict["summary_image"], summary_image)
+        assert_array_equal(images_dict["max_image"], max_image)
+        assert_array_equal(images_dict["f_per_pixel"], f_per_pixel)
diff --git a/tests/test_io.py b/tests/test_io.py
index 1d914be7..06f98064 100644
--- a/tests/test_io.py
+++ b/tests/test_io.py
@@ -11,10 +11,10 @@
     Hdf5ImagingExtractor,
     SbxImagingExtractor,
     CaimanSegmentationExtractor,
-    ExtractSegmentationExtractor,
     Suite2pSegmentationExtractor,
     CnmfeSegmentationExtractor,
 )
+from roiextractors.extractors.schnitzerextractor import LegacyExtractSegmentationExtractor
 from roiextractors.testing import (
     check_imaging_equal,
     check_segmentations_equal,
@@ -126,7 +126,7 @@ def test_imaging_extractors_canonical_shape(self, extractor_class, extractor_kwa
             ),
         ),
         param(
-            extractor_class=ExtractSegmentationExtractor,
+            extractor_class=LegacyExtractSegmentationExtractor,
             extractor_kwargs=dict(
                 file_path=str(
                     OPHYS_DATA_PATH