--- a
+++ b/create_tf_dataset.py
@@ -0,0 +1,24 @@
+from Segmentation.utils.data_loader import create_OAI_dataset
+import os
+
+def create_tfrecords(folder="train", use_2d=False, crop_size=None, mid_folders=""):
+    train = (folder == 'train')
+    str_dim = "" if use_2d else "_3d"
+
+    if use_2d and crop_size is None:
+        crop_size = 144
+
+    if not os.path.exists(f'./Data{mid_folders}/tfrecords/'):
+        os.makedirs(f'./Data{mid_folders}/tfrecords/')
+    if not os.path.exists(f'./Data{mid_folders}/tfrecords/{folder}{str_dim}'):
+        os.makedirs(f'./Data{mid_folders}/tfrecords/{folder}{str_dim}')
+
+    create_OAI_dataset(data_folder=f"./Data{mid_folders}/" + folder,
+                       tfrecord_directory=f"./Data{mid_folders}/tfrecords/" + folder + str_dim,
+                       get_train=train,
+                       use_2d=use_2d,
+                       crop_size=crop_size)
+    
+if __name__ == "__main__":
+    # create_tfrecords("train", mid_folders='/mnt')
+    create_tfrecords("valid", mid_folders='/mnt')