In drop_outliers_fn, check for Nifti1 vs Nifti2 format

Author: Adam Richie-Halford

dmriprep/run.py

@@ -489,7 +489,16 @@ def drop_outliers_fn(in_file, in_bval, in_bvec, drop_scans):
         img = nib.load(op.abspath(in_file))
         img_data = img.get_fdata()
         img_data_thinned = np.delete(img_data, drop_scans, axis=3)
-        img_thinned = nib.Nifti1Image(img_data_thinned.astype(np.float64), img.affine, header=img.header)
+        if isinstance(img, nib.nifti1.Nifti1Image):
+            img_thinned = nib.Nifti1Image(img_data_thinned.astype(np.float64),
+                                          img.affine,
+                                          header=img.header)
+        elif isinstance(img, nib.nifti2.Nifti2Image):
+            img_thinned = nib.Nifti2Image(img_data_thinned.astype(np.float64),
+                                          img.affine,
+                                          header=img.header)
+        else:
+            raise TypeError("in_file does not contain Nifti image datatype.")
 
         out_file = fname_presuffix(in_file, suffix="_thinned", newpath=op.abspath('.'))
         nib.save(img_thinned, op.abspath(out_file))