--- a/dmri.py	Thu Nov 30 10:54:46 2017 -0500
+++ b/dmri.py	Thu Nov 30 11:02:43 2017 -0500
@@ -6,6 +6,7 @@
 from dipy.core.gradients import gradient_table
 from dipy.data import fetch_stanford_hardi, fetch_stanford_t1, read_stanford_labels
 from dipy.io import read_bvals_bvecs
+from dipy.io.image import save_nifti
 from matplotlib import pyplot
 
 import nibabel
@@ -68,14 +69,20 @@
     pyplot.imshow(data[:, :, axial_middle, 10].T, cmap='gray', origin='lower')
 pyplot.savefig('data.png', bbox_inches='tight')
 shutil.move('data.png', args.output_png)
-# Load the b-values and b-vectors.
-bvals, bvecs = read_bvals_bvecs(fbval, fbvec)
-gtab = gradient_table(bvals, bvecs)
-# gtab can be used to tell what part of the data is the S0
-# volumes (volumes which correspond to b-values of 0).
-S0s = data[:, :, :, gtab.b0s_mask]
-# Save this in a new Nifti file.
-nibabel.save(nibabel.Nifti1Image(S0s, img.affine), 'output.nii')
+
+if args.drmi_dataset == 'stanford_hardi':
+    # Load the b-values and b-vectors.
+    bvals, bvecs = read_bvals_bvecs(fbval, fbvec)
+    gtab = gradient_table(bvals, bvecs)
+    # gtab can be used to tell what part of the data is the S0
+    # volumes (volumes which correspond to b-values of 0).
+    S0s = data[:, :, :, gtab.b0s_mask]
+    # Save this in a new Nifti file.
+    nibabel.save(nibabel.Nifti1Image(S0s, img.affine), 'output.nii')
+else:
+    save_nifti('output.nii', data, img.affine)
+
 shutil.move('output.nii', args.output_nifti1)
 # Move the entire contents of input_dir to output_nifti1_extra_files.
 move_directory_files(input_dir, args.output_nifti1_extra_files)
+