Added integer handling to default get_frames method

catalystneuro · Jan 30, 2024 · ea67c66 · ea67c66
1 parent cfcc121
commit ea67c66
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/roiextractors/imagingextractor.py b/src/roiextractors/imagingextractor.py
@@ -142,6 +142,8 @@ def get_frames(self, frame_idxs: ArrayType, channel: Optional[int] = 0) -> np.nd
         frames: numpy.ndarray
             The video frames.
         """
+        if isinstance(frame_idxs, int):
+            frame_idxs = [frame_idxs]
         assert max(frame_idxs) <= self.get_num_frames(), "'frame_idxs' exceed number of frames"
         if np.all(np.diff(frame_idxs) == 0):
             return self.get_video(start_frame=frame_idxs[0], end_frame=frame_idxs[-1])