--- a/dmri.py	Thu Nov 30 10:51:01 2017 -0500
+++ b/dmri.py	Thu Nov 30 10:54:46 2017 -0500
@@ -62,9 +62,10 @@
 # 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')
+if args.drmi_dataset == 'stanford_hardi':
+    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)
 # Load the b-values and b-vectors.