add column in dataset df about subjects' split

sct-pipeline · Apr 16, 2024 · b774fd7 · b774fd7
1 parent 48d5bd2
commit b774fd7
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/monai/create_msd_data.py b/monai/create_msd_data.py
@@ -233,6 +233,12 @@ def main():
     # sort the subjects
     train_subjects, val_subjects, test_subjects = sorted(train_subjects), sorted(val_subjects), sorted(test_subjects)
 
+    # add a column specifying whether the subject is in train, val or test split
+    df['split'] = 'none'
+    df.loc[df['subjectID'].isin(train_subjects), 'split'] = 'train'
+    df.loc[df['subjectID'].isin(val_subjects), 'split'] = 'validation'
+    df.loc[df['subjectID'].isin(test_subjects), 'split'] = 'test'
+
     # get boilerplate json
     params = get_boilerplate_json(dataset_name, dataset_commits)