--- a/dmri.py	Fri Nov 03 14:45:38 2017 -0400
+++ b/dmri.py	Fri Nov 03 14:45:44 2017 -0400
@@ -11,7 +11,7 @@
 
 parser = argparse.ArgumentParser()
 parser.add_argument('--input', dest='input', help='Input dataset')
-parser.add_argument('--output_gtab', dest='output_gtab', help='Output gradient table dataset')
+parser.add_argument('--output_nifti1', dest='output_nifti1', help='Output Nifti1 dataset')
 parser.add_argument('--output_png', dest='output_png', help='Output dataset')
 
 args = parser.parse_args()
@@ -40,5 +40,9 @@
 # Load the b-values and b-vectors.
 bvals, bvecs = read_bvals_bvecs(fbval, fbvec)
 gtab = gradient_table(bvals, bvecs)
-with open(args.output_gtab, 'w') as fh:
-    fh.write(gtab.info)
\ No newline at end of file
+# 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.
+nib.save(nib.Nifti1Image(S0s, img.affine), 'HARDI193_S0.nii')
+shutil.move('HARDI193_S0.nii', args.output_nifti1)