tidy up tests

brainglobe_utils/image_io/load.py

@@ -43,7 +43,7 @@ def load_any(
 
     Parameters
     ----------
-    src_path : str or Path
+    src_path : str
         Can be the path of a nifty file, nrrd file, tiff file, tiff files
         folder, or text file containing a list of paths.
 
@@ -166,7 +166,7 @@ def load_img_stack(
 
     Parameters
     ----------
-    stack_path : str or Path
+    stack_path : str
         The path of the image to be loaded.
 
     x_scaling_factor : float

brainglobe_utils/image_io/save.py

@@ -20,7 +20,7 @@ def to_nii(img, dest_path, scale=None, affine_transform=None):
     img : nifty image object or np.ndarray
         A nifty image object or numpy array representing a brain.
 
-    dest_path : str or Path
+    dest_path : str
         The file path where to save the brain.
 
     scale : tuple of floats, optional
@@ -49,7 +49,7 @@ def to_tiff(img_volume, dest_path, photometric="minisblack"):
     img_volume : np.ndarray
         The image to be saved.
 
-    dest_path : str or Path
+    dest_path : str
         The file path where to save the tiff stack.
 
     photometric: str
@@ -99,7 +99,7 @@ def to_nrrd(img_volume, dest_path):
     img_volume : np.ndarray
         The image to be saved.
 
-    dest_path : str or Path
+    dest_path : str
         The file path where to save the nrrd image.
     """
     dest_path = str(dest_path)

tests/tests/test_image_io.py

@@ -35,6 +35,15 @@ def write_tiff_sequence_with_txt_file(txt_path, image_array, shuffle=False):
 
     The tiff sequence will be saved to a sub-folder inside the same folder
     as the text file.
+
+    Parameters
+    ----------
+    txt_path : pathlib.Path
+        Filepath of text file to create
+    image_array : np.ndarray
+        Image to write as sequence of tiffs
+    shuffle : bool
+        Whether to shuffle the order of filepaths in the text file
     """
     directory = txt_path.parent
 
@@ -56,6 +65,13 @@ def save_any(file_path, image_array):
     """
     Save image_array to given file path, using the save function matching
     its file extension
+
+    Parameters
+    ----------
+    file_path : pathlib.Path
+        File path of image to save
+    image_array : np.ndarray
+        Image to save
     """
     if file_path.is_dir():
         save.to_tiffs(image_array, str(file_path / "image_array"))
@@ -64,23 +80,23 @@ def save_any(file_path, image_array):
         write_tiff_sequence_with_txt_file(file_path, image_array)
 
     elif file_path.suffix == ".tif" or file_path.suffix == ".tiff":
-        save.to_tiff(image_array, file_path)
+        save.to_tiff(image_array, str(file_path))
 
     elif file_path.suffix == ".nrrd":
-        save.to_nrrd(image_array, file_path)
+        save.to_nrrd(image_array, str(file_path))
 
     elif file_path.suffix == ".nii":
-        save.to_nii(image_array, file_path)
+        save.to_nii(image_array, str(file_path), scale=(1, 1, 1))
 
 
 def test_tiff_io(tmp_path, image_array):
     """
     Test that a 2D/3D tiff can be written and read correctly
     """
     dest_path = tmp_path / "image_array.tiff"
-    save.to_tiff(image_array, dest_path)
+    save_any(dest_path, image_array)
 
-    reloaded = load.load_img_stack(dest_path, 1, 1, 1)
+    reloaded = load.load_img_stack(str(dest_path), 1, 1, 1)
     assert (reloaded == image_array).all()
 
 
@@ -95,10 +111,10 @@ def test_3d_tiff_scaling(
     Test that a 3D tiff is scaled correctly on loading
     """
     dest_path = tmp_path / "image_array.tiff"
-    save.to_tiff(array_3d, dest_path)
+    save_any(dest_path, array_3d)
 
     reloaded = load.load_img_stack(
-        dest_path, x_scaling_factor, y_scaling_factor, z_scaling_factor
+        str(dest_path), x_scaling_factor, y_scaling_factor, z_scaling_factor
     )
     assert reloaded.shape[0] == array_3d.shape[0] * z_scaling_factor
     assert reloaded.shape[1] == array_3d.shape[1] * y_scaling_factor
@@ -117,7 +133,7 @@ def test_tiff_sequence_io(tmp_path, array_3d, load_parallel):
     Test that a 3D image can be written and read correctly as a sequence
     of 2D tiffs (with or without parallel loading)
     """
-    save.to_tiffs(array_3d, str(tmp_path / "image_array"))
+    save_any(tmp_path, array_3d)
     reloaded_array = load.load_from_folder(
         str(tmp_path), 1, 1, 1, load_parallel=load_parallel
     )
@@ -134,7 +150,7 @@ def test_tiff_sequence_scaling(
     """
     Test that a tiff sequence is scaled correctly on loading
     """
-    save.to_tiffs(array_3d, str(tmp_path / "image_array"))
+    save_any(tmp_path, array_3d)
     reloaded_array = load.load_from_folder(
         str(tmp_path), x_scaling_factor, y_scaling_factor, z_scaling_factor
     )
@@ -150,7 +166,7 @@ def test_load_img_sequence_from_txt(tmp_path, array_3d):
     ordered list of the tiff file paths (one per line)
     """
     img_sequence_file = tmp_path / "imgs_file.txt"
-    write_tiff_sequence_with_txt_file(img_sequence_file, array_3d)
+    save_any(img_sequence_file, array_3d)
 
     # Load image from paths in text file
     reloaded_array = load.load_img_sequence(str(img_sequence_file), 1, 1, 1)
@@ -184,29 +200,29 @@ def test_nii_io(tmp_path, array_3d):
     """
     Test that a 3D image can be written and read correctly as nii
     """
-    nii_path = str(tmp_path / "test_array.nii")
-    save.to_nii(array_3d, nii_path, scale=(1, 1, 1))
-    assert (load.load_nii(nii_path).get_fdata() == array_3d).all()
+    nii_path = tmp_path / "test_array.nii"
+    save_any(nii_path, array_3d)
+    assert (load.load_nii(str(nii_path)).get_fdata() == array_3d).all()
 
 
 def test_nii_read_to_numpy(tmp_path, array_3d):
     """
     Test that conversion of loaded nii image to an in-memory numpy array works
     """
-    nii_path = str(tmp_path / "test_array.nii")
-    save.to_nii(array_3d, nii_path, scale=(1, 1, 1))
+    nii_path = tmp_path / "test_array.nii"
+    save_any(nii_path, array_3d)
 
-    reloaded_array = load.load_nii(nii_path, as_array=True, as_numpy=True)
+    reloaded_array = load.load_nii(str(nii_path), as_array=True, as_numpy=True)
     assert (reloaded_array == array_3d).all()
 
 
 def test_nrrd_io(tmp_path, array_3d):
     """
     Test that a 3D image can be written and read correctly as nrrd
     """
-    nrrd_path = str(tmp_path / "test_array.nrrd")
-    save.to_nrrd(array_3d, nrrd_path)
-    assert (load.load_nrrd(nrrd_path) == array_3d).all()
+    nrrd_path = tmp_path / "test_array.nrrd"
+    save_any(nrrd_path, array_3d)
+    assert (load.load_nrrd(str(nrrd_path)) == array_3d).all()
 
 
 @pytest.mark.parametrize(
@@ -227,15 +243,15 @@ def test_load_any(tmp_path, array_3d, file_name):
     src_path = tmp_path / file_name
     save_any(src_path, array_3d)
 
-    assert (load.load_any(src_path) == array_3d).all()
+    assert (load.load_any(str(src_path)) == array_3d).all()
 
 
 def test_load_any_error(tmp_path):
     """
     Test that load_any throws an error for an unknown file extension
     """
     with pytest.raises(NotImplementedError):
-        load.load_any(tmp_path / "test.unknown")
+        load.load_any(str(tmp_path / "test.unknown"))
 
 
 def test_scale_z(array_3d):
@@ -246,20 +262,22 @@ def test_scale_z(array_3d):
     assert utils.scale_z(array_3d, 2).shape[0] == array_3d.shape[0] * 2
 
 
-def test_image_size_dir(tmp_path, array_3d):
-    save.to_tiffs(array_3d, str(tmp_path / "image_array"))
-
-    image_shape = load.get_size_image_from_file_paths(str(tmp_path))
-    assert image_shape["x"] == array_3d.shape[2]
-    assert image_shape["y"] == array_3d.shape[1]
-    assert image_shape["z"] == array_3d.shape[0]
-
-
-def test_image_size_txt(tmp_path, array_3d):
-    txt_filepath = tmp_path / "image.txt"
-    write_tiff_sequence_with_txt_file(txt_filepath, array_3d)
+@pytest.mark.parametrize(
+    "file_name",
+    [
+        "test_array.txt",
+        pytest.param("", id="dir of tiffs"),
+    ],
+)
+def test_image_size(tmp_path, array_3d, file_name):
+    """
+    Test that image size can be detected from a directory of 2D tiffs, or
+    a text file containing the paths of a sequence of 2D tiffs
+    """
+    file_path = tmp_path / file_name
+    save_any(file_path, array_3d)
 
-    image_shape = load.get_size_image_from_file_paths(str(txt_filepath))
+    image_shape = load.get_size_image_from_file_paths(str(file_path))
     assert image_shape["x"] == array_3d.shape[2]
     assert image_shape["y"] == array_3d.shape[1]
     assert image_shape["z"] == array_3d.shape[0]