qim3d/io/convert.py

@@ -2,16 +2,17 @@ import difflib
 import os
 from itertools import product
 
+import nibabel as nib
 import numpy as np
 import tifffile as tiff
 import zarr
 from tqdm import tqdm
 
 from qim3d.utils.internal_tools import stringify_path
+from qim3d.io.saving import save
 
 
 class Convert:
-
     def __init__(self, **kwargs):
         """Utility class to convert files to other formats without loading the entire file into memory
 
@@ -21,16 +22,24 @@ class Convert:
         self.chunk_shape = kwargs.get("chunk_shape", (64, 64, 64))
 
     def convert(self, input_path, output_path):
+        def get_file_extension(file_path):
+            root, ext = os.path.splitext(file_path)
+            if ext in ['.gz', '.bz2', '.xz']:  # handle common compressed extensions
+                root, ext2 = os.path.splitext(root)
+                ext = ext2 + ext
+            return ext
         # Stringify path in case it is not already a string
         input_path = stringify_path(input_path)
-        input_ext = os.path.splitext(input_path)[1]
-        output_ext = os.path.splitext(output_path)[1]
+        input_ext = get_file_extension(input_path)
+        output_ext = get_file_extension(output_path)
         output_path = stringify_path(output_path)
 
         if os.path.isfile(input_path):
             match input_ext, output_ext:
                 case (".tif", ".zarr") | (".tiff", ".zarr"):
                     return self.convert_tif_to_zarr(input_path, output_path)
+                case (".nii", ".zarr") | (".nii.gz", ".zarr"):
+                    return self.convert_nifti_to_zarr(input_path, output_path)
                 case _:
                     raise ValueError("Unsupported file format")
         # Load a directory
@@ -38,6 +47,10 @@ class Convert:
             match input_ext, output_ext:
                 case (".zarr", ".tif") | (".zarr", ".tiff"):
                     return self.convert_zarr_to_tif(input_path, output_path)
+                case (".zarr", ".nii"):
+                    return self.convert_zarr_to_nifti(input_path, output_path)
+                case (".zarr", ".nii.gz"):
+                    return self.convert_zarr_to_nifti(input_path, output_path, compression=True)
                 case _:
                     raise ValueError("Unsupported file format")
         # Fail
@@ -54,7 +67,7 @@ class Convert:
             else:
                 raise ValueError("Invalid path")
 
-    def convert_tif_to_zarr(self, tif_path, zarr_path, chunks=(64, 64, 64)):
+    def convert_tif_to_zarr(self, tif_path, zarr_path):
         """Convert a tiff file to a zarr file
 
         Args:
@@ -66,12 +79,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.chunk_shape, 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
@@ -89,7 +108,49 @@ class Convert:
             None
         """
         z = zarr.open(zarr_path)
-        tiff.imwrite(tif_path, z)
+        save(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.chunk_shape, 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_zarr_to_nifti(self, zarr_path, nifti_path, compression=False):
+        """Convert a zarr file to a nifti file
+
+        Args:
+            zarr_path (str): path to the zarr file
+            nifti_path (str): path to the nifti file
+
+        Returns:
+            None
+        """
+        z = zarr.open(zarr_path)
+        save(nifti_path, z, compression=compression)
         
 
 def convert(input_path: str, output_path: str, chunk_shape: tuple = (64, 64, 64)):
@@ -100,6 +161,5 @@ def convert(input_path: str, output_path: str, chunk_shape: tuple = (64, 64, 64)
         output_path (str): path to the output file
         chunk_shape (tuple, optional): chunk size for the zarr file. Defaults to (64, 64, 64).
     """
-
     converter = Convert(chunk_shape=chunk_shape)
     converter.convert(input_path, output_path)

qim3d/utils/cli.py

 import argparse
 import webbrowser
+import outputformat as ouf
 
 from qim3d.gui import annotation_tool, data_explorer, iso3d, local_thickness
 from qim3d.io.loading import DataLoader
 from qim3d.utils import image_preview
 from qim3d import __version__ as version
-import outputformat as ouf
-import qim3d
+import qim3d.io
 
 QIM_TITLE = ouf.rainbow(
     f"\n         _          _____     __ \n  ____ _(_)___ ___ |__  /____/ / \n / __ `/ / __ `__ \ /_ </ __  /  \n/ /_/ / / / / / / /__/ / /_/ /   \n\__, /_/_/ /_/ /_/____/\__,_/    \n  /_/                 v{version}\n\n",
@@ -14,6 +14,9 @@ QIM_TITLE = ouf.rainbow(
     cmap="hot",
 )
 
+def parse_tuple(arg):
+    # Remove parentheses if they are included and split by comma
+    return tuple(map(int, arg.strip('()').split(',')))
 
 def main():
     parser = argparse.ArgumentParser(description="Qim3d command-line interface.")
@@ -87,6 +90,30 @@ def main():
         help="By default set the maximum value to be 255 so the contrast is strong. This turns it off.",
     )
 
+    # File Convert
+    convert_parser = subparsers.add_parser(
+        "convert",
+        help="Convert files to different formats without loading the entire file into memory",
+    )
+    convert_parser.add_argument(
+        "input_path",
+        type=str,
+        metavar="Input path",
+        help="Path to image that will be converted",
+    )
+    convert_parser.add_argument(
+        "output_path",
+        type=str,
+        metavar="Output path",
+        help="Path to save converted image",
+    )
+    convert_parser.add_argument(
+        "--chunks",
+        type=parse_tuple,
+        metavar="Chunk shape",
+        default=(64,64,64),
+        help="Chunk size for the zarr file. Defaults to (64, 64, 64).",
+    )
     args = parser.parse_args()
 
     if args.subcommand == "gui":
@@ -118,7 +145,6 @@ def main():
                 interface = local_thickness.Interface()
                 interface.launch(inbrowser=inbrowser, force_light_mode=False)
 
-
     elif args.subcommand == "viz":
         if not args.source:
             print("Please specify a source file using the argument --source")
@@ -139,6 +165,7 @@ def main():
 
     elif args.subcommand == "preview":
         image = DataLoader().load(args.filename)
+
         image_preview(
             image,
             image_width=args.resolution,
@@ -147,6 +174,9 @@ def main():
             relative_intensity=args.absolute_values,
         )
 
+    elif args.subcommand == "convert":
+        qim3d.io.convert(args.input_path, args.output_path, chunk_shape=args.chunks)
+
     elif args.subcommand is None:
         print(QIM_TITLE)
         welcome_text = (
@@ -159,6 +189,9 @@ def main():
         parser.print_help()
         print("\n")
 
+    elif args.subcommand == 'convert':
+        qim3d.io.convert(args.input_path, args.output_path)
+
 
 if __name__ == "__main__":
     main()