ENH: set defaults for common suffixes

nipreps · Apr 2, 2020 · f57e2a1 · f57e2a1
1 parent d266e67
commit f57e2a1
Showing 1 changed file with 17 additions and 3 deletions.
diff --git a/niworkflows/interfaces/bids.py b/niworkflows/interfaces/bids.py
@@ -2,6 +2,7 @@
 # vi: set ft=python sts=4 ts=4 sw=4 et:
 """Interfaces for handling BIDS-like neuroimaging structures."""
 
+from collections import defaultdict
 from json import dumps
 from pathlib import Path
 from shutil import copytree, rmtree
@@ -27,6 +28,18 @@
 LOGGER = logging.getLogger('nipype.interface')
 
 
+def _none():
+    return None
+
+
+# Automatically coerce certain suffixes (DerivativesDataSink)
+DEFAULT_DTYPES = defaultdict(_none, (
+    ("mask", "uint8"),
+    ("dseg", "int16"),
+    ("probseg", "float32"))
+)
+
+
 class _BIDSBaseInputSpec(BaseInterfaceInputSpec):
     bids_dir = traits.Either(
         (None, Directory(exists=True)), usedefault=True,
@@ -467,7 +480,8 @@ def _run_interface(self, runtime):
             self._results['compression'].append(_copy_any(fname, out_file))
 
             is_nii = out_file.endswith(('.nii', '.nii.gz'))
-            if is_nii and any((self.inputs.check_hdr, self.inputs.data_dtype)):
+            data_dtype = self.inputs.data_dtype or DEFAULT_DTYPES[self.inputs.suffix]
+            if is_nii and any((self.inputs.check_hdr, data_dtype)):
                 # Do not use mmap; if we need to access the data at all, it will be to
                 # rewrite, risking a BusError
                 nii = nb.load(out_file, mmap=False)
@@ -497,11 +511,11 @@ def _run_interface(self, runtime):
                         # Rewrite file with new header
                         overwrite_header(nii, out_file)
 
-                if self.inputs.data_dtype:
+                if data_dtype:
                     if self.inputs.check_hdr:
                         # load updated NIfTI
                         nii = nb.load(out_file, mmap=False)
-                    data_dtype = np.dtype(self.inputs.data_dtype)
+                    data_dtype = np.dtype(data_dtype)
                     if nii.get_data_dtype() != data_dtype:
                         nii.set_data_dtype(data_dtype)
                         nii.to_filename(out_file)