Refactored ScanImage to support multi-channel and multi-plane directl…

…y rather than using a separate class
catalystneuro · Sep 14, 2023 · 882cf4c · 882cf4c
1 parent a2e0caf
commit 882cf4c
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 115 deletions.
diff --git a/src/roiextractors/extractors/tiffimagingextractors/scanimagetiffimagingextractor.py b/src/roiextractors/extractors/tiffimagingextractors/scanimagetiffimagingextractor.py
@@ -38,14 +38,31 @@ def extract_extra_metadata(
     return extra_metadata
 
 
-class ScanImageTiffMultiPlaneImagingExtractor(
-    ImagingExtractor
-):  # TODO: Refactor to avoid copy-paste from ScanImageTiffImagingExtractor
-    extractor_name = "ScanImageTiffMultiPlaneImaging"
+class ScanImageTiffImagingExtractor(ImagingExtractor):
+    """Specialized extractor for reading TIFF files produced via ScanImage."""
+
+    extractor_name = "ScanImageTiffImaging"
     is_writable = True
     mode = "file"
 
     def __init__(self, file_path: PathType) -> None:
+        """Create a ScanImageTiffImagingExtractor instance from a TIFF file produced by ScanImage.
+
+        The underlying data is stored in a round-robin format collapsed into 3 dimensions (frames, rows, columns).
+        I.e. the first frame of each channel and each plane is stored, and then the second frame of each channel and
+        each plane, etc.
+        Ex. for 2 channels and 2 planes:
+        [channel_1_plane_1_frame_1, channel_2_plane_1_frame_1, channel_1_plane_2_frame_1, channel_2_plane_2_frame_1,
+        channel_1_plane_1_frame_2, channel_2_plane_1_frame_2, channel_1_plane_2_frame_2, channel_2_plane_2_frame_2, ...
+        channel_1_plane_1_frame_N, channel_2_plane_1_frame_N, channel_1_plane_2_frame_N, channel_2_plane_2_frame_N]
+        This file structure is sliced lazily using ScanImageTiffReader with the appropriate logic for specified
+        channels/frames.
+
+        Parameters
+        ----------
+        file_path : PathType
+            Path to the TIFF file.
+        """
         super().__init__()
         self.file_path = Path(file_path)
         ScanImageTiffReader = _get_scanimage_reader()
@@ -60,7 +77,8 @@ def __init__(self, file_path: PathType) -> None:
                 "https://github.com/catalystneuro/roiextractors/issues "
             )
         self._sampling_frequency = float(extra_metadata["SI.hRoiManager.scanVolumeRate"])
-        # SI.hChannels.channelName = "{'channel_name_1' 'channel_name_2' ... 'channel_name_N'}" where N is the number of channels (active or not)
+        # SI.hChannels.channelName = "{'channel_name_1' 'channel_name_2' ... 'channel_name_N'}"
+        #   where N is the number of channels (active or not)
         self._channel_names = extra_metadata["SI.hChannels.channelName"].split("'")[1::2][: self._num_channels]
 
         valid_suffixes = [".tiff", ".tif", ".TIFF", ".TIF"]
@@ -193,111 +211,3 @@ def check_frame_inputs(self, frame, channel, plane) -> None:
             raise ValueError(f"Channel index ({channel}) exceeds number of channels ({self._num_channels}).")
         if plane >= self._num_planes:
             raise ValueError(f"Plane index ({plane}) exceeds number of planes ({self._num_planes}).")
-
-
-class ScanImageTiffImagingExtractor(ImagingExtractor):
-    """Specialized extractor for reading TIFF files produced via ScanImage."""
-
-    extractor_name = "ScanImageTiffImaging"
-    is_writable = True
-    mode = "file"
-
-    def __init__(
-        self,
-        file_path: PathType,
-        sampling_frequency: FloatType,
-    ):
-        """Create a ScanImageTiffImagingExtractor instance from a TIFF file produced by ScanImage.
-
-        This extractor allows for lazy accessing of slices, unlike
-        :py:class:`~roiextractors.extractors.tiffimagingextractors.TiffImagingExtractor`.
-        However, direct slicing of the underlying data structure is not equivalent to a numpy memory map.
-
-        Parameters
-        ----------
-        file_path : PathType
-            Path to the TIFF file.
-        sampling_frequency : float
-            The frequency at which the frames were sampled, in Hz.
-        """
-        ScanImageTiffReader = _get_scanimage_reader()
-
-        super().__init__()
-        self.file_path = Path(file_path)
-        self._sampling_frequency = sampling_frequency
-        valid_suffixes = [".tiff", ".tif", ".TIFF", ".TIF"]
-        if self.file_path.suffix not in valid_suffixes:
-            suffix_string = ", ".join(valid_suffixes[:-1]) + f", or {valid_suffixes[-1]}"
-            warn(
-                f"Suffix ({self.file_path.suffix}) is not of type {suffix_string}! "
-                f"The {self.extractor_name}Extractor may not be appropriate for the file."
-            )
-
-        with ScanImageTiffReader(str(self.file_path)) as io:
-            shape = io.shape()  # [frames, rows, columns]
-        if len(shape) == 3:
-            self._num_frames, self._num_rows, self._num_columns = shape
-            self._num_channels = 1
-        else:  # no example file for multiple color channels or depths
-            raise NotImplementedError(
-                "Extractor cannot handle 4D TIFF data. Please raise an issue to request this feature: "
-                "https://github.com/catalystneuro/roiextractors/issues "
-            )
-
-    def get_frames(self, frame_idxs: ArrayType, channel: int = 0) -> np.ndarray:
-        ScanImageTiffReader = _get_scanimage_reader()
-
-        squeeze_data = False
-        if isinstance(frame_idxs, int):
-            squeeze_data = True
-            frame_idxs = [frame_idxs]
-
-        if not all(np.diff(frame_idxs) == 1):
-            return np.concatenate([self._get_single_frame(idx=idx) for idx in frame_idxs])
-        else:
-            with ScanImageTiffReader(filename=str(self.file_path)) as io:
-                frames = io.data(beg=frame_idxs[0], end=frame_idxs[-1] + 1)
-                if squeeze_data:
-                    frames = frames.squeeze()
-            return frames
-
-    # Data accessed through an open ScanImageTiffReader io gets scrambled if there are multiple calls.
-    # Thus, open fresh io in context each time something is needed.
-    def _get_single_frame(self, idx: int) -> np.ndarray:
-        """Get a single frame of data from the TIFF file.
-
-        Parameters
-        ----------
-        idx : int
-            The index of the frame to retrieve.
-
-        Returns
-        -------
-        frame: numpy.ndarray
-            The frame of data.
-        """
-        ScanImageTiffReader = _get_scanimage_reader()
-
-        with ScanImageTiffReader(str(self.file_path)) as io:
-            return io.data(beg=idx, end=idx + 1)
-
-    def get_video(self, start_frame=None, end_frame=None, channel: Optional[int] = 0) -> np.ndarray:
-        ScanImageTiffReader = _get_scanimage_reader()
-
-        with ScanImageTiffReader(filename=str(self.file_path)) as io:
-            return io.data(beg=start_frame, end=end_frame)
-
-    def get_image_size(self) -> Tuple[int, int]:
-        return (self._num_rows, self._num_columns)
-
-    def get_num_frames(self) -> int:
-        return self._num_frames
-
-    def get_sampling_frequency(self) -> float:
-        return self._sampling_frequency
-
-    def get_num_channels(self) -> int:
-        return self._num_channels
-
-    def get_channel_names(self) -> list:
-        pass
diff --git a/tests/test_multiplane_scanimage.py b/tests/test_multiplane_scanimage.py
@@ -1,12 +1,12 @@
 from roiextractors.extractors.tiffimagingextractors.scanimagetiffimagingextractor import (
-    ScanImageTiffMultiPlaneImagingExtractor,
+    ScanImageTiffImagingExtractor,
 )
 import matplotlib.pyplot as plt
 
 
 def main():
     file_path = "/Volumes/T7/CatalystNeuro/NWB/MouseV1/raw-tiffs/2ret/20230119_w57_1_2ret_00001.tif"
-    extractor = ScanImageTiffMultiPlaneImagingExtractor(file_path=file_path)
+    extractor = ScanImageTiffImagingExtractor(file_path=file_path)
     print(f"num_frames: {extractor.get_num_frames()}")
     print(f"num_planes: {extractor.get_num_planes()}")
     print(f"num_channels: {extractor.get_num_channels()}")