--- a/dmri.py	Thu Nov 30 11:09:00 2017 -0500
+++ b/dmri.py	Thu Nov 30 11:09:07 2017 -0500
@@ -16,7 +16,7 @@
 parser.add_argument('--drmi_dataset_type', dest='drmi_dataset_type', help='Input dataset type')
 parser.add_argument('--output_nifti1', dest='output_nifti1', help='Output Nifti1 dataset')
 parser.add_argument('--output_nifti1_extra_files', dest='output_nifti1_extra_files', help='Output Nifti1 extra files')
-parser.add_argument('--output_png', dest='output_png', help='Output dataset')
+parser.add_argument('--output_png', dest='output_png', default=None, help='Output dataset')
 
 args = parser.parse_args()
 
@@ -60,17 +60,16 @@
 # data is a 4D array where the first 3 dimensions are the i, j,
 # k voxel coordinates and the last dimension is the number of
 # non-weighted (S0s) and diffusion-weighted volumes.
-# Visualize the results using matplotlib.
-axial_middle = data.shape[2] // 2
-pyplot.subplot(1, 2, 1).set_axis_off()
+
 if args.drmi_dataset == 'stanford_hardi':
+    # Visualize the results using matplotlib.
+    axial_middle = data.shape[2] // 2
+    pyplot.subplot(1, 2, 1).set_axis_off()
     pyplot.imshow(data[:, :, axial_middle, 0].T, cmap='gray', origin='lower')
     pyplot.subplot(1, 2, 2).set_axis_off()
     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)
-
-if args.drmi_dataset == 'stanford_hardi':
+    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)
@@ -85,4 +84,3 @@
 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)
-