added function to convert nifti to zarr

qim3d/io/convert.py

@@ -2,6 +2,7 @@ import difflib
 import os
 from itertools import product
 
+import nibabel as nib
 import numpy as np
 import tifffile as tiff
 import zarr
@@ -11,9 +12,8 @@ from qim3d.utils.internal_tools import stringify_path
 
 
 class Convert:
-
-    def __init__(self,**kwargs):
-        """ Utility class to convert files to other formats without loading the entire file into memory
+    def __init__(self, **kwargs):
+        """Utility class to convert files to other formats without loading the entire file into memory
 
         Args:
             chunk_shape (tuple, optional): chunk size for the zarr file. Defaults to (64, 64, 64).
@@ -27,9 +27,9 @@ class Convert:
         output_ext = os.path.splitext(output_path)[1]
         output_path = stringify_path(output_path)
 
-        if os.path.isfile(input_path):  
+        if os.path.isfile(input_path):
             match input_ext, output_ext:
-                case (".tif", ".zarr")  | (".tiff", ".zarr"):
+                case (".tif", ".zarr") | (".tiff", ".zarr"):
                     return self.convert_tif_to_zarr(input_path, output_path)
                 case _:
                     raise ValueError("Unsupported file format")
@@ -54,8 +54,8 @@ class Convert:
             else:
                 raise ValueError("Invalid path")
 
-    def convert_tif_to_zarr(self, tif_path, zarr_path, chunks=(64, 64, 64)):
-        """ Convert a tiff file to a zarr file
+    def convert_tif_to_zarr(self, tif_path, zarr_path):
+        """Convert a tiff file to a zarr file
 
         Args:
             tif_path (str): path to the tiff file
@@ -66,12 +66,18 @@ class Convert:
             zarr.core.Array: zarr array containing the data from the tiff file
         """
         vol = tiff.memmap(tif_path)
-        z = zarr.open(zarr_path, mode='w', shape=vol.shape, chunks=chunks, dtype=vol.dtype)
+        z = zarr.open(
+            zarr_path, mode="w", shape=vol.shape, chunks=self.chunks, dtype=vol.dtype
+        )
         chunk_shape = tuple((s + c - 1) // c for s, c in zip(z.shape, z.chunks))
         # ! Fastest way is z[:] = vol[:], but does not have a progress bar
-        for chunk_indices in tqdm(product(*[range(n) for n in chunk_shape]), total=np.prod(chunk_shape)):
-            slices = tuple(slice(c * i, min(c * (i + 1), s))
-                        for s, c, i in zip(z.shape, z.chunks, chunk_indices))
+        for chunk_indices in tqdm(
+            product(*[range(n) for n in chunk_shape]), total=np.prod(chunk_shape)
+        ):
+            slices = tuple(
+                slice(c * i, min(c * (i + 1), s))
+                for s, c, i in zip(z.shape, z.chunks, chunk_indices)
+            )
             temp_data = vol[slices]
             # The assignment takes 99% of the cpu-time
             z.blocks[chunk_indices] = temp_data
@@ -79,7 +85,7 @@ class Convert:
         return z
 
     def convert_zarr_to_tif(self, zarr_path, tif_path):
-        """ Convert a zarr file to a tiff file
+        """Convert a zarr file to a tiff file
 
         Args:
             zarr_path (str): path to the zarr file
@@ -91,9 +97,38 @@ class Convert:
         z = zarr.open(zarr_path)
         tiff.imwrite(tif_path, z)
 
+    def convert_nifti_to_zarr(self, nifti_path, zarr_path):
+        """Convert a nifti file to a zarr file
+
+        Args:
+            nifti_path (str): path to the nifti file
+            zarr_path (str): path to the zarr file
+
+        Returns:
+            zarr.core.Array: zarr array containing the data from the nifti file
+        """
+        vol = nib.load(nifti_path).dataobj
+        z = zarr.open(
+            zarr_path, mode="w", shape=vol.shape, chunks=self.chunks, dtype=vol.dtype
+        )
+        chunk_shape = tuple((s + c - 1) // c for s, c in zip(z.shape, z.chunks))
+        # ! Fastest way is z[:] = vol[:], but does not have a progress bar
+        for chunk_indices in tqdm(
+            product(*[range(n) for n in chunk_shape]), total=np.prod(chunk_shape)
+        ):
+            slices = tuple(
+                slice(c * i, min(c * (i + 1), s))
+                for s, c, i in zip(z.shape, z.chunks, chunk_indices)
+            )
+            temp_data = vol[slices]
+            # The assignment takes 99% of the cpu-time
+            z.blocks[chunk_indices] = temp_data
+
+        return z
+
 
 def convert(input_path: str, output_path: str, chunk_shape: tuple = (64, 64, 64)):
-    """ Convert a file to another format without loading the entire file into memory
+    """Convert a file to another format without loading the entire file into memory
 
     Args:
         input_path (str): path to the input file