RF: Begin refactoring load into image classes

Author: effigies

nibabel/freesurfer/mghformat.py

@@ -13,6 +13,7 @@
 from os.path import splitext
 import numpy as np
 
+from ..imageglobals import valid_exts
 from ..volumeutils import (array_to_file, array_from_file, Recoder)
 from ..spatialimages import HeaderDataError, SpatialImage
 from ..fileholders import FileHolder,  copy_file_map
@@ -454,6 +455,7 @@ def writeftr_to(self, fileobj):
         fileobj.write(ftr_nd.tostring())
 
 
+@valid_exts('.mgh', '.mgz')
 class MGHImage(SpatialImage):
     """ Class for MGH format image
     """

nibabel/imageglobals.py

@@ -58,3 +58,13 @@ def __enter__(self):
     def __exit__(self, exc, value, tb):
         for handler in self.orig_handlers:
             logger.addHandler(handler)
+
+IMAGE_MAP = {}
+
+
+def valid_exts(*exts):
+    def decorate(klass):
+        for ext in exts:
+            IMAGE_MAP.setdefault(ext, []).append(klass)
+        return klass
+    return decorate

nibabel/loadsave.py

@@ -11,18 +11,13 @@
 
 import numpy as np
 
-from .filename_parser import types_filenames, splitext_addext
-from .volumeutils import BinOpener, Opener
-from .analyze import AnalyzeImage
-from .spm2analyze import Spm2AnalyzeImage
-from .nifti1 import Nifti1Image, Nifti1Pair, header_dtype as ni1_hdr_dtype
+from .filename_parser import splitext_addext
+from .volumeutils import BinOpener
+from .nifti1 import Nifti1Image, Nifti1Pair
 from .nifti2 import Nifti2Image, Nifti2Pair
-from .minc1 import Minc1Image
-from .minc2 import Minc2Image
-from .freesurfer import MGHImage
-from .fileholders import FileHolderError
 from .spatialimages import ImageFileError
 from .imageclasses import class_map, ext_map
+from .imageglobals import IMAGE_MAP
 from .arrayproxy import is_proxy
 
 
@@ -41,55 +36,20 @@ def load(filename, **kwargs):
     img : ``SpatialImage``
        Image of guessed type
     '''
-    return guessed_image_type(filename).from_filename(filename, **kwargs)
 
-
-def guessed_image_type(filename):
-    """ Guess image type from file `filename`
-
-    Parameters
-    ----------
-    filename : str
-        File name containing an image
-
-    Returns
-    -------
-    image_class : class
-        Class corresponding to guessed image type
-    """
     froot, ext, trailing = splitext_addext(filename, ('.gz', '.bz2'))
     lext = ext.lower()
-    try:
-        img_type = ext_map[lext]
-    except KeyError:
-        raise ImageFileError('Cannot work out file type of "%s"' %
-                             filename)
-    if lext in ('.mgh', '.mgz', '.par'):
-        klass = class_map[img_type]['class']
-    elif lext == '.mnc':
-        # Look for HDF5 signature for MINC2
-        # https://www.hdfgroup.org/HDF5/doc/H5.format.html
-        with Opener(filename) as fobj:
-            signature = fobj.read(4)
-            klass = Minc2Image if signature == b'\211HDF' else Minc1Image
-    elif lext == '.nii':
-        with BinOpener(filename) as fobj:
-            binaryblock = fobj.read(348)
-        ft = which_analyze_type(binaryblock)
-        klass = Nifti2Image if ft == 'nifti2' else Nifti1Image
-    else: # might be nifti 1 or 2 pair or analyze of some sort
-        files_types = (('image','.img'), ('header','.hdr'))
-        filenames = types_filenames(filename, files_types)
-        with BinOpener(filenames['header']) as fobj:
-            binaryblock = fobj.read(348)
-        ft = which_analyze_type(binaryblock)
-        if ft == 'nifti2':
-            klass = Nifti2Pair
-        elif ft == 'nifti1':
-            klass = Nifti1Pair
-        else:
-            klass = Spm2AnalyzeImage
-    return klass
+    potential_classes = IMAGE_MAP[lext]
+
+    if len(potential_classes) == 1:
+        return potential_classes[0].from_filename(filename, **kwargs)
+
+    # Allow image tests to cache data
+    sniff = None
+    for img_type in IMAGE_MAP[lext]:
+        is_valid, sniff = img_type.is_image(filename, sniff)
+        if is_valid:
+            return img_type.from_filename(filename, **kwargs)
 
 
 def save(img, filename):
@@ -213,43 +173,3 @@ def read_img_data(img, prefer='scaled'):
         if prefer == 'scaled':
             return hdr.data_from_fileobj(fileobj)
         return hdr.raw_data_from_fileobj(fileobj)
-
-
-def which_analyze_type(binaryblock):
-    """ Is `binaryblock` from NIfTI1, NIfTI2 or Analyze header?
-
-    Parameters
-    ----------
-    binaryblock : bytes
-        The `binaryblock` is 348 bytes that might be NIfTI1, NIfTI2, Analyze, or
-        None of the the above.
-
-    Returns
-    -------
-    hdr_type : str
-        * a nifti1 header (pair or single) -> return 'nifti1'
-        * a nifti2 header (pair or single) -> return 'nifti2'
-        * an Analyze header -> return 'analyze'
-        * None of the above -> return None
-
-    Notes
-    -----
-    Algorithm:
-
-    * read in the first 4 bytes from the file as 32-bit int ``sizeof_hdr``
-    * if ``sizeof_hdr`` is 540 or byteswapped 540 -> assume nifti2
-    * Check for 'ni1', 'n+1' magic -> assume nifti1
-    * if ``sizeof_hdr`` is 348 or byteswapped 348 assume Analyze
-    * Return None
-    """
-    hdr = np.ndarray(shape=(), dtype=ni1_hdr_dtype, buffer=binaryblock)
-    bs_hdr = hdr.byteswap()
-    sizeof_hdr = hdr['sizeof_hdr']
-    bs_sizeof_hdr = bs_hdr['sizeof_hdr']
-    if 540 in (sizeof_hdr, bs_sizeof_hdr):
-        return 'nifti2'
-    if hdr['magic'] in (b'ni1', b'n+1'):
-        return 'nifti1'
-    if 348 in (sizeof_hdr, bs_sizeof_hdr):
-        return 'analyze'
-    return None

nibabel/minc1.py

@@ -14,8 +14,11 @@
 
 from .externals.netcdf import netcdf_file
 
+from .filename_parser import splitext_addext
+from .imageglobals import valid_exts
 from .spatialimages import Header, SpatialImage
 from .fileslice import canonical_slicers
+from .volumeutils import BinOpener
 
 from .deprecated import FutureWarningMixin
 
@@ -280,6 +283,7 @@ def data_from_fileobj(self, fileobj):
         raise NotImplementedError
 
 
+@valid_exts('.mnc')
 class Minc1Image(SpatialImage):
     ''' Class for MINC1 format images
 
@@ -307,6 +311,26 @@ def from_file_map(klass, file_map):
             data = klass.ImageArrayProxy(minc_file)
         return klass(data, affine, header, extra=None, file_map=file_map)
 
+    @classmethod
+    def is_image(klass, filename, sniff=None):
+        ftypes = dict(klass.files_types)
+        froot, ext, trailing = splitext_addext(filename, klass._compressed_exts)
+        lext = ext.lower()
+
+        if lext not in ftypes.values():
+            return False, sniff
+
+        fname = froot + ftypes['header'] if 'header' in ftypes else filename
+        if not sniff:
+            with BinOpener(fname, 'rb') as fobj:
+                sniff = fobj.read(4)
+
+        return klass._minctest(sniff), sniff
+
+    @classmethod
+    def _minctest(klass, binaryblock):
+        return binaryblock != b'\211HDF'
+
 
 load = Minc1Image.load
 

nibabel/minc2.py

@@ -30,6 +30,7 @@
 from .optpkg import optional_package
 h5py, have_h5py, setup_module = optional_package('h5py')
 
+from .imageglobals import valid_exts
 from .minc1 import Minc1File, Minc1Image, MincError
 
 
@@ -134,6 +135,7 @@ def get_scaled_data(self, sliceobj=()):
         return self._normalize(raw_data, sliceobj)
 
 
+@valid_exts('.mnc')
 class Minc2Image(Minc1Image):
     ''' Class for MINC2 images
 
@@ -160,5 +162,9 @@ def from_file_map(klass, file_map):
         data = klass.ImageArrayProxy(minc_file)
         return klass(data, affine, header, extra=None, file_map=file_map)
 
+    @classmethod
+    def _minctest(klass, binaryblock):
+        return binaryblock[:4] == b'\211HDF'
+
 
 load = Minc2Image.load

nibabel/nifti1.py

@@ -18,6 +18,7 @@
 
 from .py3k import asstr
 from .volumeutils import Recoder, make_dt_codes, endian_codes
+from .imageglobals import valid_exts
 from .spatialimages import HeaderDataError, ImageFileError
 from .batteryrunners import Report
 from .quaternions import fillpositive, quat2mat, mat2quat
@@ -1592,13 +1593,20 @@ def _chk_xform_code(klass, code_type, hdr, fix):
             rep.fix_msg = 'setting to 0'
         return hdr, rep
 
+    @classmethod
+    def is_header(klass, binaryblock):
+        hdr = np.ndarray(shape=(), dtype=header_dtype,
+                         buffer=binaryblock[:348])
+        return hdr['magic'] in (b'ni1', b'n+1')
+
 
 class Nifti1PairHeader(Nifti1Header):
     ''' Class for NIfTI1 pair header '''
     # Signal whether this is single (header + data) file
     is_single = False
 
 
+@valid_exts('.img', '.hdr')
 class Nifti1Pair(analyze.AnalyzeImage):
     """ Class for NIfTI1 format image, header pair
     """
@@ -1822,6 +1830,7 @@ def set_sform(self, affine, code=None, **kwargs):
                 self._affine[:] = self._header.get_best_affine()
 
 
+@valid_exts('.nii')
 class Nifti1Image(Nifti1Pair):
     """ Class for single file NIfTI1 format image
     """

nibabel/nifti2.py

@@ -20,6 +20,7 @@
 import numpy as np
 
 from .analyze import AnalyzeHeader
+from .imageglobals import valid_exts
 from .batteryrunners import Report
 from .spatialimages import HeaderDataError, ImageFileError
 from .nifti1 import Nifti1Header, Nifti1Pair, Nifti1Image
@@ -221,19 +222,32 @@ def _chk_eol_check(hdr, fix=False):
             rep.fix_msg = 'setting EOL check to 13, 10, 26, 10'
         return hdr, rep
 
+    @classmethod
+    def is_header(klass, binaryblock):
+        if len(binaryblock) < 540:
+            return False
+
+        hdr = np.ndarray(shape=(), dtype=header_dtype,
+                         buffer=binaryblock[:540])
+        bs_hdr = hdr.byteswap()
+        return 540 in (hdr['sizeof_hdr'], bs_hdr['sizeof_hdr'])
+
+
 
 class Nifti2PairHeader(Nifti2Header):
     ''' Class for NIfTI2 pair header '''
     # Signal whether this is single (header + data) file
     is_single = False
 
 
+@valid_exts('.img', '.hdr')
 class Nifti2Pair(Nifti1Pair):
     """ Class for NIfTI2 format image, header pair
     """
     header_class = Nifti2PairHeader
 
 
+@valid_exts('.nii')
 class Nifti2Image(Nifti1Image):
     """ Class for single file NIfTI2 format image
     """

nibabel/parrec.py

@@ -100,6 +100,7 @@
 from .spatialimages import SpatialImage, Header
 from .eulerangles import euler2mat
 from .volumeutils import Recoder, array_from_file, BinOpener
+from .imageglobals import valid_exts
 from .affines import from_matvec, dot_reduce, apply_affine
 from .nifti1 import unit_codes
 from .fileslice import fileslice, strided_scalar
@@ -975,6 +976,7 @@ def get_sorted_slice_indices(self):
         return np.lexsort(keys)[:n_used]
 
 
+@valid_exts('.par', '.rec')
 class PARRECImage(SpatialImage):
     """PAR/REC image"""
     header_class = PARRECHeader

nibabel/spatialimages.py

@@ -141,9 +141,10 @@
 
 import numpy as np
 
-from .filename_parser import types_filenames, TypesFilenamesError
+from .filename_parser import types_filenames, TypesFilenamesError, \
+    splitext_addext
 from .fileholders import FileHolder
-from .volumeutils import shape_zoom_affine
+from .volumeutils import shape_zoom_affine, BinOpener
 
 
 class HeaderDataError(Exception):
@@ -866,6 +867,22 @@ def from_image(klass, img):
                      klass.header_class.from_header(img.header),
                      extra=img.extra.copy())
 
+    @classmethod
+    def is_image(klass, filename, sniff=None):
+        ftypes = dict(klass.files_types)
+        froot, ext, trailing = splitext_addext(filename, ('.gz', '.bz2'))
+        lext = ext.lower()
+
+        if lext not in ftypes.values():
+            return False, sniff
+
+        fname = froot + ftypes['header'] if 'header' in ftypes else filename
+        if not sniff:
+            with BinOpener(fname, 'rb') as fobj:
+                sniff = fobj.read(1024)
+
+        return klass.header_class.is_header(sniff), sniff
+
     def __getitem__(self):
         ''' No slicing or dictionary interface for images
         '''

nibabel/spm2analyze.py

@@ -10,6 +10,7 @@
 import numpy as np
 
 from .spatialimages import HeaderDataError
+from .imageglobals import valid_exts
 from .batteryrunners import Report
 from . import spm99analyze as spm99 # module import
 
@@ -114,7 +115,16 @@ def get_slope_inter(self):
             return slope, inter
         return None, None
 
+    @classmethod
+    def is_header(klass, binaryblock):
+        hdr = np.ndarray(shape=(), dtype=header_dtype,
+                         buffer=binaryblock[:348])
+        bs_hdr = hdr.byteswap()
+        return (binaryblock[344:348] not in (b'ni1\x00', b'n+1\x00') and
+                348 in (hdr['sizeof_hdr'], bs_hdr['sizeof_hdr']))
 
+
+@valid_exts('.img', '.hdr')
 class Spm2AnalyzeImage(spm99.Spm99AnalyzeImage):
     """ Class for SPM2 variant of basic Analyze image
     """