--- a
+++ b/data/image_folder.py
@@ -0,0 +1,34 @@
+# Manuel A. Morales (moralesq@mit.edu)
+# Harvard-MIT Department of Health Sciences & Technology  
+# Athinoula A. Martinos Center for Biomedical Imaging
+
+import os
+import pydicom
+
+EXTENSIONS = {}
+EXTENSIONS['NIFTI'] = ['.nii.gz', '.nii']
+EXTENSIONS['DICOM'] = ['SAX']
+EXTENSIONS['H5PY']  = ['.h5']
+
+def is_data_file(filename, dformat="NIFTI"):
+    return any(filename.endswith(extension) for extension in EXTENSIONS[dformat])
+
+def make_dataset(dir, max_dataset_size=float("inf"), dformat="NIFTI"):
+    
+    filenames = []
+    assert os.path.isdir(dir), '%s is not a valid directory' % dir
+    
+    for root, _, fnames in sorted(os.walk(dir)):
+        for fname in fnames:
+            if dformat == 'DICOM':
+                try:
+                    protocol_name = pydicom.read_file(os.path.join(root, fname)).ProtocolName
+                except:
+                    continue
+                if any(protocol in protocol_name for protocol in EXTENSIONS[dformat]):
+                    path = os.path.join(root, fname)
+                    filenames.append(path)
+            elif is_data_file(fname, dformat=dformat):
+                path = os.path.join(root, fname)
+                filenames.append(path)
+    return filenames[:min(max_dataset_size, len(filenames))]
\ No newline at end of file