RF+TEST: fix / test saving with output dtype

matthew-brett · matthew-brett · commit 3f8056546aa2 · 2012-06-15T23:47:48.000-07:00
Test output dtype can come from data, header, or specified dtype.
diff --git a/nipy/io/files.py b/nipy/io/files.py
@@ -13,7 +13,10 @@
 
 import os
 
+import numpy as np
+
 import nibabel as nib
+from nibabel.spatialimages import HeaderDataError
 
 from ..core.image.image import is_image
 
@@ -60,14 +63,17 @@ def save(img, filename, dtype_from='data'):
     filename : string
         Should be a valid filename.
     dtype_from : {'data', 'header'} or dtype specifier, optional
-        dtype to save data to disk.  Set into header before save as hint for
-        output format to chose dtype to save to. Not every format supports every
-        dtype, so some values of this parameter will raise errors.  The default
-        (`io_dtype` == None) is to use the dtype of the image data.
+        Method of setting dtype to save data to disk. Value of 'data' (default),
+        means use data dtype to save.  'header' means use data dtype specified
+        in header, if available, otherwise use data dtype.  Can also be any
+        valid specifier for a numpy dtype, e.g. 'i4', ``np.float32``.  Not every
+        format supports every dtype, so some values of this parameter or data
+        dtypes will raise errors.
 
     Returns
     -------
     image : An `Image` object
+        Possibly modified by saving.
 
     See Also
     --------
@@ -102,7 +108,7 @@ def save(img, filename, dtype_from='data'):
     >>> saved_image3 = save_image(img, fname)
     Traceback (most recent call last):
        ...
-    ValueError: Cannot save file type "minc"
+    ValueError: Sorry, we cannot yet save as format "minc"
 
     Finally, we clear up our temporary files:
 
@@ -119,27 +125,28 @@ def save(img, filename, dtype_from='data'):
     * SPM Analyze : ['.img', '.img.gz']
     """
     # Try and get nifti
-    ni_img = nipy2nifti(img)
+    dt_from_is_str = isinstance(dtype_from, basestring)
+    if dt_from_is_str and dtype_from == 'header':
+        # All done
+        io_dtype = None
+    elif dt_from_is_str and dtype_from == 'data':
+        io_dtype = img.get_data().dtype
+    else:
+        io_dtype = np.dtype(dtype_from)
+    # make new image
+    ni_img = nipy2nifti(img, data_dtype = io_dtype)
     ftype = _type_from_filename(filename)
     if ftype.startswith('nifti1'):
-        saver = nib.nifti1.save
+        ni_img.to_filename(filename)
     elif ftype == 'analyze':
-        saver = nib.spm2analyze.save
+        try:
+            ana_img = nib.Spm2AnalyzeImage.from_image(ni_img)
+        except HeaderDataError:
+            raise HeaderDataError('SPM analyze does not support datatype %s' %
+                                  ni_img.get_header().get_data_dtype())
+        ana_img.to_filename(filename)
     else:
-        raise ValueError('Cannot save file type "%s"' % ftype)
-    dtype_from_str = isinstance(dtype_from, basestring)
-    if dtype_from_str and dtype_from == 'header':
-        saver(ni_img, filename)
-        return img
-    # Set output dtype if possible
-    hdr = ni_img.get_header()
-    if dtype_from_str and dtype_from == 'data':
-        io_dtype = img.get_data().dtype
-    else:
-        io_dtype = dtype_from
-    hdr.set_data_dtype(io_dtype)
-    # make new image
-    saver(ni_img, filename)
+        raise ValueError('Sorry, we cannot yet save as format "%s"' % ftype)
     return img
 
 
diff --git a/nipy/io/tests/test_image_io.py b/nipy/io/tests/test_image_io.py
@@ -3,7 +3,8 @@
 
 import numpy as np
 
-from nibabel.spatialimages import ImageFileError
+from nibabel.spatialimages import ImageFileError, HeaderDataError
+from nibabel import Nifti1Header
 
 from ..api import load_image, save_image, as_image
 from nipy.core.api import AffineTransform as AfT, Image, vox2mni
@@ -125,9 +126,65 @@ def test_scaling_io_dtype():
         del img
 
 
-def test_from_data():
-    # Default data dtype comes from data
-    pass
+def assert_dt_no_end_equal(a, b):
+    """ Assert two numpy dtype specifiers are equal apart from byte order
+
+    Avoids failed comparison between int32 / int64 and intp
+    """
+    a = np.dtype(a).newbyteorder('=')
+    b = np.dtype(b).newbyteorder('=')
+    assert_equal(a.str, b.str)
+
+
+def test_output_dtypes():
+    shape = (4, 2, 3)
+    rng = np.random.RandomState(19441217) # IN-S BD
+    data = rng.normal(4, 20, size=shape)
+    aff = np.diag([2.2, 3.3, 4.1, 1])
+    cmap = vox2mni(aff)
+    img = Image(data, cmap)
+    fname_root = 'my_file'
+    with InTemporaryDirectory():
+        for ext in 'img', 'nii':
+            out_fname = fname_root + '.' + ext
+            # Default is for data to come from data dtype
+            save_image(img, out_fname)
+            img_back = load_image(out_fname)
+            hdr = img_back.metadata['header']
+            assert_dt_no_end_equal(hdr.get_data_dtype(), np.float)
+            # All these types are OK for both output formats
+            for out_dt in 'i2', 'i4', np.int16, '<f4', '>f8':
+                # Specified output dtype
+                save_image(img, out_fname, out_dt)
+                img_back = load_image(out_fname)
+                hdr = img_back.metadata['header']
+                assert_dt_no_end_equal(hdr.get_data_dtype(), out_dt)
+                # Output comes from data by default
+                data_typed = data.astype(out_dt)
+                img_again = Image(data_typed, cmap)
+                save_image(img_again, out_fname)
+                img_back = load_image(out_fname)
+                hdr = img_back.metadata['header']
+                assert_dt_no_end_equal(hdr.get_data_dtype(), out_dt)
+                # Even if header specifies otherwise
+                in_hdr = Nifti1Header()
+                in_hdr.set_data_dtype(np.dtype('c8'))
+                img_more = Image(data_typed, cmap, metadata={'header': in_hdr})
+                save_image(img_more, out_fname)
+                img_back = load_image(out_fname)
+                hdr = img_back.metadata['header']
+                assert_dt_no_end_equal(hdr.get_data_dtype(), out_dt)
+                # But can come from header if specified
+                save_image(img_more, out_fname, dtype_from='header')
+                img_back = load_image(out_fname)
+                hdr = img_back.metadata['header']
+                assert_dt_no_end_equal(hdr.get_data_dtype(), 'c8')
+        # u2 only OK for nifti
+        save_image(img, 'my_file.nii', 'u2')
+        img_back = load_image('my_file.nii')
+        hdr = img_back.metadata['header']
+        assert_dt_no_end_equal(hdr.get_data_dtype(), 'u2')
+        assert_raises(HeaderDataError, save_image, img, 'my_file.img', 'u2')
 
 
 def test_header_roundtrip():
