# HG changeset patch
# User greg
# Date 1512052771 18000
# Node ID 6c1fd6d2c07ed0156ec3e88efc9d16f43f5330e0
# Parent  da507620bd7487ac66467754da1cb10b8b3d1ade
Uploaded

diff -r da507620bd74 -r 6c1fd6d2c07e dmri.py
--- a/dmri.py	Thu Nov 30 09:39:24 2017 -0500
+++ b/dmri.py	Thu Nov 30 09:39:31 2017 -0500
@@ -12,6 +12,7 @@
 
 parser = argparse.ArgumentParser()
 parser.add_argument('--drmi_dataset', dest='drmi_dataset', help='Input dataset')
+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')
@@ -34,18 +35,15 @@
 
 # Get input data.
 input_dir = args.drmi_dataset
-if input_dir == 'stanford_hardi':
+#if input_dir == 'stanford_hardi':
+if args.drmi_dataset_type == "dataset":
     fetch_stanford_hardi()
     fdwi = os.path.join(input_dir, 'HARDI150.nii.gz')
     fbval = os.path.join(input_dir, 'HARDI150.bval')
     fbvec = os.path.join(input_dir, 'HARDI150.bvec')
-elif input_dir == 'sherbrooke_3shell':
-    fetch_sherbrooke_3shell()
-    fdwi = os.path.join(input_dir, 'HARDI193.nii.gz')
-    fbval = os.path.join(input_dir, 'HARDI193.bval')
-    fbvec = os.path.join(input_dir, 'HARDI193.bvec')
-# Load the dMRI datasets.
-img = nibabel.load(fdwi)
+    img = nibabel.load(fdwi)
+else:
+    img, gtab, labels = read_stanford_labels()
 data = img.get_data()
 # 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