handle nii.gz correctly

brainglobe_utils/IO/image/load.py

@@ -135,7 +135,7 @@ def load_any(
             z_scaling_factor,
             anti_aliasing=anti_aliasing,
         )
-    elif src_path.suffix in [".nii", ".nii.gz"]:
+    elif src_path.suffix == ".nii" or str(src_path).endswith(".nii.gz"):
         logging.debug("Data type is: NifTI")
         img = load_nii(src_path, as_array=True, as_numpy=as_numpy)
     else:

brainglobe_utils/IO/image/save.py

@@ -121,7 +121,7 @@ def save_any(img_volume, dest_path):
     elif dest_path.suffix in [".tif", ".tiff"]:
         to_tiff(img_volume, dest_path)
 
-    elif dest_path.suffix == ".nii":
+    elif dest_path.suffix == ".nii" or str(dest_path).endswith(".nii.gz"):
         to_nii(img_volume, dest_path)
 
     else:

tests/tests/test_IO/test_image_io.py

@@ -245,7 +245,10 @@ def test_sort_img_sequence_from_txt(shuffled_txt_path, array_3d, sort):
 
 
 @pytest.mark.parametrize("use_path", [True, False], ids=["Path", "String"])
-def test_nii_io(tmp_path, array_3d, use_path):
+@pytest.mark.parametrize(
+    "nifti_suffix", [".nii.gz", ".nii"], ids=["compressed", "uncompressed"]
+)
+def test_nii_io(tmp_path, array_3d, use_path, nifti_suffix):
     """
     Test that a 3D image can be written and read correctly as nii with scale
     (keeping it as a nifty object with no numpy conversion on loading).
@@ -265,11 +268,14 @@ def test_nii_io(tmp_path, array_3d, use_path):
     assert reloaded.header.get_zooms() == scale
 
 
-def test_nii_read_to_numpy(tmp_path, array_3d):
+@pytest.mark.parametrize(
+    "nifti_suffix", [".nii.gz", ".nii"], ids=["compressed", "uncompressed"]
+)
+def test_nii_read_to_numpy(tmp_path, array_3d, nifti_suffix):
     """
     Test that conversion of loaded nii image to an in-memory numpy array works
     """
-    nii_path = tmp_path / "test_array.nii"
+    nii_path = tmp_path / f"test_array{nifti_suffix}"
     save.save_any(array_3d, nii_path)
     reloaded_array = load.load_any(nii_path, as_numpy=True)