Merge pull request #1 from Australian-Imaging-Service/add-nifti

Add nifti module

conftest.py

@@ -1,5 +1,6 @@
 import os
 from pathlib import Path
+import tempfile
 from importlib import import_module
 import pytest
 
@@ -9,11 +10,15 @@
                  if p.stem != 'base']
 
 
+@pytest.fixture
+def work_dir():
+    return Path(tempfile.mkdtemp())
+
+
 @pytest.fixture(params=dicom_modules)
 def dicom_module(request):
     module_path = 'medimages4tests.dicom.' + '.'.join(request.param.split('__'))
     return import_module(module_path)
-
 
 
 # For debugging in IDE's don't catch raised exceptions and let the IDE

medimages4tests/__init__.py

@@ -0,0 +1,3 @@
+from pathlib import Path
+
+base_cache_dir = Path.home() / ".medimages4tests" / "cache"

medimages4tests/dicom/base.py

@@ -2,9 +2,9 @@
 import shutil
 from copy import copy
 import pydicom.dataset
+from .. import base_cache_dir
 
-
-cache_dir = Path.home() / ".medimages4tests" / "cache" / "dicom"
+cache_dir = base_cache_dir / "dicom"
 dicom_pkg_dir = Path(__file__).parent
 
 

medimages4tests/nifti.py

@@ -0,0 +1,43 @@
+from pathlib import Path
+import gzip
+import shutil
+import numpy as np
+import nibabel as nb
+
+
+def sample_image(
+    out_file: Path,
+    data: np.ndarray = None,
+    vox_sizes=(1.0, 1.0, 1.0),
+    qform=(1, 2, 3, 1),
+    compressed=False,
+) -> Path:
+    """Create a random Nifti file to satisfy BIDS parsers"""
+    if data is None:
+        data = np.random.randint(0, 1, size=[10, 10, 10])
+
+    uncompressed = out_file.with_suffix('.nii')
+
+    hdr = nb.Nifti1Header()
+    hdr.set_data_shape(data.shape)
+    hdr.set_zooms(vox_sizes)  # set voxel size
+    hdr.set_xyzt_units(2)  # millimeters
+    hdr.set_qform(np.diag(qform))
+    nb.save(
+        nb.Nifti1Image(
+            data,
+            hdr.get_best_affine(),
+            header=hdr,
+        ),
+        uncompressed,
+    )
+
+    if compressed:
+        out_file = out_file.with_suffix('.nii.gz')
+        with open(uncompressed, 'rb') as f_in:
+            with gzip.open(out_file, 'wb') as f_out:
+                shutil.copyfileobj(f_in, f_out)
+    else:
+        out_file = uncompressed
+
+    return out_file

setup.py

@@ -14,7 +14,8 @@
                  'to be used in image handling tests'),
     long_description=open('README.rst').read(),
     install_requires=[
-        'pydicom>=1.0.2'],
+        'pydicom>=1.0.2',
+        "nibabel>=4.0.1"],
     extras_require={
         'test': [
             'pytest>=5.4.3',

tests/test_nifti.py

@@ -0,0 +1,27 @@
+import gzip
+import shutil
+import nibabel as nb
+import numpy as np
+from medimages4tests.nifti import sample_image
+
+
+def test_nifti(work_dir):
+
+    nifti_fpath = sample_image(work_dir / "sample.nii")
+
+    nifti = nb.load(nifti_fpath)
+
+    assert np.array_equal(nifti.header["dim"][:4], [3, 10, 10, 10])
+
+
+def test_nifti_compressed(work_dir):
+
+    gz_fpath = sample_image(work_dir / "sample.nii.gz", compressed=True)
+    uncompressed_fpath = work_dir / "nifti.nii"
+
+    with gzip.open(gz_fpath, 'rb') as f_in:
+        with open(uncompressed_fpath, 'wb') as f_out:
+            shutil.copyfileobj(f_in, f_out)
+    nifti = nb.load(uncompressed_fpath)
+
+    assert np.array_equal(nifti.header["dim"][:4], [3, 10, 10, 10])