Implemented save for h5, nifti, vol, png, jpeg, jpg

qim3d/io/save.py

@@ -2,8 +2,10 @@
 
 import os
 
+import h5py
 import nibabel as nib
 import numpy as np
+import PIL
 import tifffile
 
 from qim3d.io.logger import log
@@ -126,21 +128,55 @@ class DataSaver:
             path (str): The path to save file to
             data (numpy.ndarray): The data to be saved
         """
+        # No support for compression yet
+        if self.compression:
+            raise NotImplementedError("Saving compressed .vol files is not yet supported")
 
         # Create custom .vgi metadata file
         metadata = ""
-        metadata += "{volume1}\n"
-        metadata += "[file1]"
-        metadata += "Size = 1000 1000 766"
-        metadata += "datatype = float\n"
-        metadata += "Name = WFW_200kV_6W_1mmSn_6micro_1s.vol"
+        metadata += "{volume1}\n" # .vgi organization 
+        metadata += "[file1]\n" # .vgi organization 
+        metadata += "Size = {} {} {}\n".format(data.shape[1], data.shape[2], data.shape[0]) # Swap axes to match .vol format
+        metadata += "Datatype = {}\n".format(str(data.dtype)) # Get datatype as string
+        metadata += "Name = {}.vol\n".format(path.rsplit('/', 1)[-1][:-4]) # Get filename without extension
 
         # Save metadata
-        with open(path[:-3] + ".vgi", "w") as f:
+        with open(path[:-4] + ".vgi", "w") as f:
             f.write(metadata)
 
         # Save data using numpy in binary format
-        np.save(path[:-3] + ".vol", data)
+        data.tofile(path[:-4] + ".vol")
+
+    def save_h5(self, path, data):
+        """ Save data to a HDF5 file to the given path.
+
+        Args:
+            path (str): The path to save file to
+            data (numpy.ndarray): The data to be saved
+        """
+
+        with h5py.File(path, "w") as f:
+            f.create_dataset("dataset", data=data, compression="gzip" if self.compression else None)
+        
+    def save_PIL(self, path, data):
+        """ Save data to a PIL file to the given path.
+
+        Args:
+            path (str): The path to save file to
+            data (numpy.ndarray): The data to be saved
+        """
+        # No support for compression yet
+        if self.compression and path.endswith(".png"):
+            raise NotImplementedError("png does not support compression")
+        elif not self.compression and path.endswith((".jpeg",".jpg")):
+            raise NotImplementedError("jpeg does not support no compression")
+
+        # Convert to PIL image
+        img = PIL.Image.fromarray(data)
+
+        # Save image
+        img.save(path)
+
         
     def save(self, path, data):
         """Save data to the given path.
@@ -210,10 +246,12 @@ class DataSaver:
                         return self.save_nifti(path, data)
                     elif path.endswith(("TXRM","XRM","TXM")):
                         raise NotImplementedError("Saving TXRM files is not yet supported")
-                    elif path.endswith(("h5")):
-                        raise NotImplementedError("Saving HDF5 files is not yet supported")
+                    elif path.endswith((".h5")):
+                        return self.save_h5(path, data)
                     elif path.endswith((".vol",".vgi")):
-                        raise NotImplementedError("Saving VOL files is not yet supported")
+                        return self.save_vol(path, data)
+                    elif path.endswith((".jpeg",".jpg", ".png")):
+                        return self.save_PIL(path, data)
                     else:
                         raise ValueError("Unsupported file format")
                 # If there is no file extension in the path

qim3d/tests/io/test_save.py

-import qim3d
+import hashlib
+import os
+import re
 import tempfile
+
 import numpy as np
-import os
-import hashlib
 import pytest
-import re
+
+import qim3d
+
 
 def test_image_exist():
     # Create random test image
@@ -221,6 +224,143 @@ def test_tiff_stack_slicing_dim():
             qim3d.io.save(path2save,test_image,basename='test',sliced_dim=dim)
             assert len(os.listdir(path2save))==test_image.shape[dim]
 
+def test_tiff_save_load():
+    # Create random test image
+    original_image = qim3d.examples.blobs_256x256
+
+    # Create temporary directory
+    with tempfile.TemporaryDirectory() as temp_dir:
+        image_path = os.path.join(temp_dir,"test_image.tif")
+
+        # Save to temporary directory
+        qim3d.io.save(image_path,original_image)
+
+        # Load from temporary directory
+        saved_image = qim3d.io.load(image_path)
+
+        # Get hashes
+        original_hash = calculate_image_hash(original_image)
+        saved_hash = calculate_image_hash(saved_image)
+
+        # Assert that original image is identical to saved_image
+        assert original_hash == saved_hash
+
+def test_vol_save_load():
+    # Create random test image
+    original_image = qim3d.examples.blobs_256x256x256
+
+    # Create temporary directory
+    with tempfile.TemporaryDirectory() as temp_dir:
+        image_path = os.path.join(temp_dir,"test_image.vol")
+
+        # Save to temporary directory
+        qim3d.io.save(image_path,original_image)
+
+        # Load from temporary directory
+        saved_image = qim3d.io.load(image_path)
+
+        # Get hashes
+        original_hash = calculate_image_hash(original_image)
+        saved_hash = calculate_image_hash(saved_image)
+
+        # Assert that original image is identical to saved_image
+        assert original_hash == saved_hash
+
+def test_pil_save_load():
+    # Create random test image
+    original_image = qim3d.examples.blobs_256x256
+
+    # Create temporary directory
+    with tempfile.TemporaryDirectory() as temp_dir:
+        image_path_png = os.path.join(temp_dir,"test_image.png")
+        image_path_jpg = os.path.join(temp_dir,"test_image.jpg")
+
+        # Save to temporary directory
+        qim3d.io.save(image_path_png,original_image)
+        qim3d.io.save(image_path_jpg,original_image, compression=True)
+
+        # Load from temporary directory
+        saved_image_png = qim3d.io.load(image_path_png)
+        saved_image_jpg = qim3d.io.load(image_path_jpg)
+
+        # Get hashes
+        original_hash = calculate_image_hash(original_image)
+        saved_png_hash = calculate_image_hash(saved_image_png)
+
+        # Assert that original image is identical to saved_image
+        assert original_hash == saved_png_hash
+        
+        # jpg is lossy so the hashes will not match, checks that the image is the same size and similar values
+        assert original_image.shape == saved_image_jpg.shape
+
+def test_nifti_save_load():
+    # Create random test image
+    original_image = qim3d.examples.blobs_256x256
+
+    # Create temporary directory
+    with tempfile.TemporaryDirectory() as temp_dir:
+        image_path = os.path.join(temp_dir,"test_image.nii")
+        image_path_compressed = os.path.join(temp_dir,"test_image_compressed.nii.gz")
+
+        # Save to temporary directory
+        qim3d.io.save(image_path,original_image)
+        qim3d.io.save(image_path_compressed,original_image, compression=True)
+
+        # Load from temporary directory
+        saved_image = qim3d.io.load(image_path)
+        saved_image_compressed = qim3d.io.load(image_path_compressed)
+
+        # Get hashes
+        original_hash = calculate_image_hash(original_image)
+        saved_hash = calculate_image_hash(saved_image)
+        saved_compressed_hash = calculate_image_hash(saved_image_compressed)
+
+        # Assert that original image is identical to saved_image
+        assert original_hash == saved_hash
+        assert original_hash == saved_compressed_hash
+
+        # Compute file sizes
+        file_size = os.path.getsize(image_path)
+        compressed_file_size = os.path.getsize(image_path_compressed)
+
+        # Assert that compressed file size is smaller than non-compressed file size
+        assert compressed_file_size < file_size
+
+
+def test_h5_save_load():
+    # Create random test image
+    original_image = qim3d.examples.blobs_256x256x256
+
+    # Create temporary directory
+    with tempfile.TemporaryDirectory() as temp_dir:
+        image_path = os.path.join(temp_dir,"test_image.h5")
+        image_path_compressed = os.path.join(temp_dir,"test_image_compressed.nii.gz")
+
+        # Save to temporary directory
+        qim3d.io.save(image_path,original_image)
+        qim3d.io.save(image_path_compressed,original_image, compression=True)
+
+        # Load from temporary directory
+        saved_image = qim3d.io.load(image_path)
+        saved_image_compressed = qim3d.io.load(image_path_compressed)
+
+        # Get hashes
+        original_hash = calculate_image_hash(original_image)
+        saved_hash = calculate_image_hash(saved_image)
+        saved_compressed_hash = calculate_image_hash(saved_image_compressed)
+
+        # Assert that original image is identical to saved_image
+        # Assert that original image is identical to saved_image
+        assert original_hash == saved_hash
+        assert original_hash == saved_compressed_hash
+
+        # Compute file sizes
+        file_size = os.path.getsize(image_path)
+        compressed_file_size = os.path.getsize(image_path_compressed)
+
+        # Assert that compressed file size is smaller than non-compressed file size
+        assert compressed_file_size < file_size
+
 def calculate_image_hash(image): 
     image_bytes = image.tobytes()
     hash_object = hashlib.md5(image_bytes)