From b2d260176afc79fd08d951e639e55ea7f6f27408 Mon Sep 17 00:00:00 2001
From: bendichter <ben.dichter@gmail.com>
Date: Tue, 23 Jan 2024 11:13:04 -0500
Subject: [PATCH] add inscopix extractor for cell_set

---
 .../extractors/inscopix/__init__.py           |  1 +
 .../inscopix/inscopixsegmentationextractor.py | 71 +++++++++++++++++++
 2 files changed, 72 insertions(+)
 create mode 100644 src/roiextractors/extractors/inscopix/__init__.py
 create mode 100644 src/roiextractors/extractors/inscopix/inscopixsegmentationextractor.py

diff --git a/src/roiextractors/extractors/inscopix/__init__.py b/src/roiextractors/extractors/inscopix/__init__.py
new file mode 100644
index 00000000..765ecbc4
--- /dev/null
+++ b/src/roiextractors/extractors/inscopix/__init__.py
@@ -0,0 +1 @@
+from .inscopixsegmentationextractor import InscopixSegmentationExtractor
\ No newline at end of file
diff --git a/src/roiextractors/extractors/inscopix/inscopixsegmentationextractor.py b/src/roiextractors/extractors/inscopix/inscopixsegmentationextractor.py
new file mode 100644
index 00000000..2ad4a651
--- /dev/null
+++ b/src/roiextractors/extractors/inscopix/inscopixsegmentationextractor.py
@@ -0,0 +1,71 @@
+from typing import Optional
+import numpy as np
+
+from ...extraction_tools import PathType, ArrayType
+from ...segmentationextractor import SegmentationExtractor
+
+
+class InscopixSegmentationExtractor(SegmentationExtractor):
+    """A segmentation extractor for Inscopix."""
+
+    extractor_name = "InscopixSegmentationExtractor"
+    installed = True  # check at class level if installed or not
+    is_writable = False
+    mode = "file"
+    installation_mesg = ""  # error message when not installed
+
+    def __init__(self, file_path: PathType, sampling_frequency: Optional[float] = None):
+        """Initialize a InscopixSegmentationExtractor instance.
+
+        Parameters
+        ----------
+        file_path: str
+            The location of the folder containing Inscopix *.mat output file.
+        """
+        import isx
+
+        SegmentationExtractor.__init__(self)
+        self.file_path = file_path
+        self.cell_set = isx.CellSet.read("cellset.isxd")
+
+    def get_num_rois(self):
+        return self.cell_set.num_cells
+
+    def get_roi_image_masks(self, roi_ids=None) -> np.ndarray:
+        if roi_ids is None:
+            roi_idx_ = range(self.get_num_rois())
+        else:
+            all_ids = self.get_roi_ids()
+            roi_idx_ = [all_ids.index(i) for i in roi_ids]
+        return np.hstack([self.cell_set.get_cell_image_data(roi_id) for roi_id in roi_idx_])
+
+    def get_roi_ids(self) -> list:
+        return [self.cell_set.get_cell_name(x) for x in self.get_num_rois()]
+
+    def get_image_size(self) -> ArrayType:
+        num_pixels = self.cell_set.footer["spacingInfo"]["numPixels"]
+        return num_pixels["x"], num_pixels["y"]
+
+    def get_accepted_list(self) -> list:
+        return [x for x in self.get_num_rois() if self.cell_set.get_cell_status(x) == "accepted"]
+
+    def get_rejected_list(self) -> list:
+        return [x for x in self.get_num_rois() if self.cell_set.get_cell_status(x) == "rejected"]
+
+    def get_traces(self, roi_ids=None, start_frame=None, end_frame=None, name="raw") -> ArrayType:
+        if roi_ids is None:
+            roi_idx_ = range(self.get_num_rois())
+        else:
+            all_ids = self.get_roi_ids()
+            roi_idx_ = [all_ids.index(i) for i in roi_ids]
+        if start_frame is None:
+            start_frame = 0
+        if end_frame is None:
+            end_frame = self.cell_set.get_cell_num_frames(0)
+        return np.hstack([self.cell_set.get_cell_trace_data(roi_id)[start_frame, end_frame] for roi_id in roi_idx_])
+
+    def get_num_frames(self) -> int:
+        return self.cell_set.footer["timingInfo"]["numTimes"]
+
+    def get_sampling_frequency(self):
+        return 1 / self.cell_set.timing.period.secs_float
\ No newline at end of file