Add ZarrWrapper

lens-biophotonics · May 5, 2023 · 8bb3e95 · 8bb3e95
1 parent e5d3701
commit 8bb3e95
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 0 deletions.
diff --git a/zetastitcher/io/inputfile.py b/zetastitcher/io/inputfile.py
@@ -10,6 +10,7 @@
 from .ffmpeg_wrapper import FFMPEGWrapper
 from .tiffwrapper import TiffWrapper
 from tifffile import TiffFileError
+from .zarr_wrapper import ZarrWrapper
 from .zipwrapper import ZipWrapper
 from .mhdwrapper import MHDWrapper
 
@@ -124,6 +125,13 @@ def _open(self):
         if not self.path.exists():
             raise FileNotFoundError(self.path)
 
+        if '.zarr' in str(self.path):
+            try:
+                self.wrapper = ZarrWrapper(self.file_path)
+                return
+            except:
+                pass
+
         try:
             self.wrapper = TiffWrapper(self.file_path)
             return

diff --git a/zetastitcher/io/zarr_wrapper.py b/zetastitcher/io/zarr_wrapper.py
@@ -0,0 +1,26 @@
+import zarr
+
+from zetastitcher.io.inputfile_mixin import InputFileMixin
+
+
+class ZarrWrapper(InputFileMixin):
+    def __init__(self, file_path=None):
+        super().__init__()
+        self.file_path = file_path
+        self.zarr_file = None
+        self.open(file_path)
+
+    def open(self, file_path):
+        z = zarr.open(file_path, mode='r')
+        self.zarr_file = next(z.arrays())[1]
+
+        self.nfrms = self.zarr_file.shape[0]
+        self.ysize = self.zarr_file.shape[-2]
+        self.xsize = self.zarr_file.shape[-1]
+        self.dtype = self.zarr_file.dtype
+
+    def frame(self, index, dtype=None):
+        a = self.zarr_file[index]
+        if dtype is not None:
+            a = a.astype(dtype)
+        return a