Merge branch 'save_tiff_stack' into 'main'

Save tiff stack

See merge request !30

qim3d/io/save.py

@@ -4,17 +4,23 @@ import os
 import tifffile
 import numpy as np
 from qim3d.io.logger import log
+from qim3d.utils.internal_tools import sizeof, stringify_path
+
 
 class DataSaver:
     """Utility class for saving data to different file formats.
 
     Attributes:
         replace (bool): Specifies if an existing file with identical path is replaced.
-        compression (bool): Specifies if the file is with Deflate compression.
+        compression (bool): Specifies if the file is saved with Deflate compression (lossless).
+        basename (str): Specifies the basename for a TIFF stack saved as several files
+            (only relevant for TIFF stacks).
+        sliced_dim (int): Specifies the dimension that is sliced in case a TIFF stack is saved
+            as several files (only relevant for TIFF stacks)
 
     Methods:
         save_tiff(path,data): Save data to a TIFF file to the given path.
-        load(path,data): Save data to the given path. 
+        load(path,data): Save data to the given path.
 
     Example:
         image = qim3d.examples.blobs_256x256
@@ -26,33 +32,75 @@ class DataSaver:
         """Initializes a new instance of the DataSaver class.
 
         Args:
-            replace (bool, optional): Specifies if an existing file with identical path should be replaced. 
+            replace (bool, optional): Specifies if an existing file with identical path should be replaced.
                 Default is False.
             compression (bool, optional): Specifies if the file should be saved with Deflate compression.
                 Default is False.
+            basename (str, optional): Specifies the basename for a TIFF stack saved as several files
+                (only relevant for TIFF stacks). Default is None
+            sliced_dim (int, optional): Specifies the dimension that is sliced in case a TIFF stack is saved
+                as several files (only relevant for TIFF stacks). Default is 0, i.e., the first dimension.
         """
-        self.replace = kwargs.get("replace",False)
-        self.compression = kwargs.get("compression",False)
+        self.replace = kwargs.get("replace", False)
+        self.compression = kwargs.get("compression", False)
+        self.basename = kwargs.get("basename", None)
+        self.sliced_dim = kwargs.get("sliced_dim", 0)
 
-    def save_tiff(self,path,data):
+    def save_tiff(self, path, data):
         """Save data to a TIFF file to the given path.
 
-        Args: 
+        Args:
             path (str): The path to save file to
-            data (numpy.ndarray): The data to be saved    
+            data (numpy.ndarray): The data to be saved
+        """
+        tifffile.imwrite(path, data, compression=self.compression)
+
+    def save_tiff_stack(self, path, data):
+        """Save data as a TIFF stack containing slices in separate files to the given path.
+        The slices will be named according to the basename plus a suffix with a zero-filled
+        value corresponding to the slice number
+
+        Args:
+            path (str): The directory to save files to
+            data (numpy.ndarray): The data to be saved
         """
-        tifffile.imwrite(path,data,compression=self.compression)
+        extension = ".tif"
+
+        if data.ndim <= 2:
+            path = os.path.join(path, self.basename, ".tif")
+            self.save_tiff(path, data)
+        else:
+            # get number of total slices
+            no_slices = data.shape[self.sliced_dim]
+            # Get number of zero-fill values as the number of digits in the total number of slices
+            zfill_val = len(str(no_slices))
+
+            # Create index list
+            idx = [slice(None)] * data.ndim
+
+            # Iterate through each slice and save
+            for i in range(no_slices):
+                idx[self.sliced_dim] = i
+                sliced = data[tuple(idx)]
+                filename = self.basename + str(i).zfill(zfill_val) + extension
+                filepath = os.path.join(path, filename)
+                self.save_tiff(filepath, sliced)
+
+            pattern_string = filepath[:-(len(extension)+zfill_val)] + "-"*zfill_val + extension
+
+            log.info(f"Total of {no_slices} files saved following the pattern '{pattern_string}'")
 
     def save(self, path, data):
         """Save data to the given path.
 
-        Args: 
+        Args:
             path (str): The path to save file to
             data (numpy.ndarray): The data to be saved
-        
+
         Raises:
+            ValueError: If the provided path is an existing directory and self.basename is not provided
             ValueError: If the file format is not supported.
-            ValueError: If the specified folder does not exist.
+            ValueError: If the provided path does not exist and self.basename is not provided
             ValueError: If a file extension is not provided.
             ValueError: if a file with the specified path already exists and replace=False.
 
@@ -60,47 +108,94 @@ class DataSaver:
             image = qim3d.examples.blobs_256x256
             saver = qim3d.io.DataSaver(compression=True)
             saver.save("image.tif",image)
+
         """
 
-        folder = os.path.dirname(path) or '.'
-        # Check if directory exists
-        if os.path.isdir(folder):
-            _, ext = os.path.splitext(path)
-            # Check if provided path contains file extension
-            if ext:
-                # Check if a file with the given path already exists
-                if os.path.isfile(path) and not self.replace:
-                    raise ValueError("A file with the provided path already exists. To replace it set 'replace=True'")
-                
-                if path.endswith((".tif",".tiff")):
-                    return self.save_tiff(path,data)
-                else:
-                    raise ValueError("Unsupported file format")
+        path = stringify_path(path)
+        isdir = os.path.isdir(path)
+        _, ext = os.path.splitext(path)
+
+        # If path is an existing directory
+        if isdir:
+            # If basename is provided
+            if self.basename:
+                # Save as tiff stack
+                return self.save_tiff_stack(path, data)
+            # If basename is not provided
             else:
-                raise ValueError('Please provide a file extension')
+                raise ValueError(
+                    f"To save a stack as several TIFF files to the directory '{path}', please provide the keyword argument 'basename'. "
+                    + "Otherwise, to save a single file, please provide a full path with a filename and valid extension."
+                )
+
+        # If path is not an existing directory
         else:
-            raise ValueError(f'The directory {folder} does not exist. Please provide a valid directory')
-        
-def save(path,
-        data,
-        replace=False,
-        compression=False,
-        **kwargs
-        ):
+            # If there is no file extension in path and basename is provided
+            if not ext and self.basename:
+                # Make directory and save as tiff stack
+                os.mkdir(path)
+                log.info("Created directory '%s'!", path)
+                return self.save_tiff_stack(path, data)
+
+            # Check if a parent directory exists
+            parentdir = os.path.dirname(path) or "."
+            if os.path.isdir(parentdir):
+                # If there is a file extension in the path
+                if ext:
+                    # If there is a basename
+                    if self.basename:
+                        # It will be unused and the user is informed accordingly
+                        log.info("'basename' argument is unused")
+                    # Check if a file with the given path already exists
+                    if os.path.isfile(path) and not self.replace:
+                        raise ValueError(
+                            "A file with the provided path already exists. To replace it set 'replace=True'"
+                        )
+
+                    if path.endswith((".tif", ".tiff")):
+                        return self.save_tiff(path, data)
+                    else:
+                        raise ValueError("Unsupported file format")
+                # If there is no file extension in the path
+                else:
+                    raise ValueError(
+                        "Please provide a file extension if you want to save as a single file."
+                        + " Otherwise, please provide a basename to save as a TIFF stack"
+                    )
+            else:
+                raise ValueError(
+                    f"The directory '{parentdir}' does not exist.\n"
+                    + "Please provide a valid directory or specify a basename if you want to save a tiff stack as several files to a folder that does not yet exist"
+                )
+
+
+def save(
+    path, data, replace=False, compression=False, basename=None, sliced_dim=0, **kwargs
+):
     """Save data to a specified file path.
 
     Args:
         path (str): The path to save file to
-        data (numpy.ndarray): The data to be saved  
-        replace (bool, optional): Specifies if an existing file with identical path should be replaced. 
+        data (numpy.ndarray): The data to be saved
+        replace (bool, optional): Specifies if an existing file with identical path should be replaced.
             Default is False.
         compression (bool, optional): Specifies if the file should be saved with Deflate compression (lossless).
             Default is False.
+        basename (str, optional): Specifies the basename for a TIFF stack saved as several files
+            (only relevant for TIFF stacks). Default is None
+        sliced_dim (int, optional): Specifies the dimension that is sliced in case a TIFF stack is saved
+            as several files (only relevant for TIFF stacks). Default is 0, i.e., the first dimension.
         **kwargs: Additional keyword arguments to be passed to the DataSaver constructor
-    
+
     Example:
         image = qim3d.examples.blobs_256x256
         qim3d.io.save("image.tif",image,compression=True)
     """
-        
-    DataSaver(replace=replace, compression=compression, **kwargs).save(path, data)
+
+    DataSaver(
\ No newline at end of file
+        replace=replace,
+        compression=compression,
+        basename=basename,
+        sliced_dim=sliced_dim,
+        **kwargs,
+    ).save(path, data)

qim3d/tests/io/test_save.py

@@ -127,11 +127,13 @@ def test_no_file_ext():
     # Create filename without extension
     filename = 'test_image'
 
+    message = 'Please provide a file extension if you want to save as a single file. Otherwise, please provide a basename to save as a TIFF stack'
+
     # Create temporary directory
     with tempfile.TemporaryDirectory() as temp_dir:
         image_path = os.path.join(temp_dir,filename)
 
-        with pytest.raises(ValueError,match='Please provide a file extension'):
+        with pytest.raises(ValueError,match=message):
             # Try to save the test image to a path witout file extension
             qim3d.io.save(image_path,test_image)
 
@@ -142,8 +144,10 @@ def test_folder_doesnt_exist():
     # Create invalid path 
     invalid_path = os.path.join('this','path','doesnt','exist.tif')
 
-    message = f'The directory {re.escape(os.path.dirname(invalid_path))} does not exist. Please provide a valid directory'
-    
+    #message = f'The directory {re.escape(os.path.dirname(invalid_path))} does not exist. Please provide a valid directory'
+    message = f"""The directory {re.escape(os.path.dirname(invalid_path))} does not exist. Please provide a valid directory or specify a basename
+ if you want to save a tiff stack as several files to a folder that does not yet exist"""
+
     with pytest.raises(ValueError,match=message):
         # Try to save test image to an invalid path
         qim3d.io.save(invalid_path,test_image)
@@ -162,6 +166,62 @@ def test_unsupported_file_format():
             # Try to save test image with an unsupported file extension
             qim3d.io.save(image_path,test_image)
 
+def test_no_basename():
+    # Create random test image
+    test_image = np.random.randint(0,256,(100,100,100),'uint8')
+
+    # Create temporary directory
+    with tempfile.TemporaryDirectory() as temp_dir:
+        message = f"""Please provide a basename with the keyword argument 'basename' if you want to save
+ a TIFF stack as several files to '{re.escape(temp_dir)}'. Otherwise, please provide a path with a valid filename"""
+
+        with pytest.raises(ValueError,match=message):
+            # Try to save test image to an existing directory (indicating
+            # that you want to save as several files) without providing a basename
+            qim3d.io.save(temp_dir,test_image)
+
+def test_mkdir_tiff_stack():
+    # Create random test image
+    test_image = np.random.randint(0,256,(10,100,100),'uint8')
+
+    # create temporary directory
+    with tempfile.TemporaryDirectory() as temp_dir:
+        # Define a folder that does not yet exist
+        path2save= os.path.join(temp_dir,'tempfolder')
+       
+        # Save to this folder with a basename
+        qim3d.io.save(path2save,test_image,basename='test')
+
+        # Assert that folder is created
+        assert os.path.isdir(path2save)
+
+def test_tiff_stack_naming():
+    # Create random test image
+    test_image = np.random.randint(0,256,(10,100,100),'uint8')
+
+    # Define expected filenames
+    basename = 'test'
+    expected_filenames = [basename + str(i).zfill(2) + '.tif' for i,_ in enumerate(test_image)]
+    
+    # create temporary directory
+    with tempfile.TemporaryDirectory() as temp_dir:
+        qim3d.io.save(temp_dir,test_image,basename=basename)
+
+        assert expected_filenames==sorted(os.listdir(temp_dir))
+    
+
+def test_tiff_stack_slicing_dim():
+    # Create random test image where the three dimensions are not the same length
+    test_image = np.random.randint(0,256,(5,10,15),'uint8')
+    
+    with tempfile.TemporaryDirectory() as temp_dir:
+        # Iterate thorugh all three dims and save the image as slices in 
+        # each dimension in separate folder and assert the number of files
+        # match the shape of the image
+        for dim in range(3):
+            path2save = os.path.join(temp_dir,'dim'+str(dim))
+            qim3d.io.save(path2save,test_image,basename='test',sliced_dim=dim)
+            assert len(os.listdir(path2save))==test_image.shape[dim]
 
 def calculate_image_hash(image): 
     image_bytes = image.tobytes()