| import argparse |
| import glob |
| import json |
| import os |
|
|
| import monai |
| from sklearn.model_selection import train_test_split |
|
|
|
|
| def produce_sample_dict(line: str): |
| names = os.listdir(line) |
| seg, t1ce, t1, t2, flair = [], [], [], [], [] |
| for name in names: |
| name = os.path.join(line, name) |
| if "_seg.nii" in name: |
| seg.append(name) |
| elif "_t1ce.nii" in name: |
| t1ce.append(name) |
| elif "_t1.nii" in name: |
| t1.append(name) |
| elif "_t2.nii" in name: |
| t2.append(name) |
| elif "_flair.nii" in name: |
| flair.append(name) |
|
|
| return {"label": seg[0], "image": t1ce + t1 + t2 + flair} |
|
|
|
|
| def produce_datalist(dataset_dir: str, train_size: int = 200): |
| """ |
| This function is used to split the dataset. |
| It will produce "train_size" number of samples for training, and the other samples |
| are divided equally into val and test sets. |
| """ |
|
|
| samples = sorted(glob.glob(os.path.join(dataset_dir, "*", "*"), recursive=True)) |
| datalist = [] |
| for line in samples: |
| datalist.append(produce_sample_dict(line)) |
| train_list, other_list = train_test_split(datalist, train_size=train_size) |
| val_list, test_list = train_test_split(other_list, train_size=0.5) |
|
|
| return {"training": train_list, "validation": val_list, "testing": test_list} |
|
|
|
|
| def main(args): |
| """ |
| split the dataset and output the data list into a json file. |
| """ |
| data_file_base_dir = os.path.join(os.path.abspath(args.path), "training") |
| |
| monai.utils.set_determinism(seed=123) |
| datalist = produce_datalist(dataset_dir=data_file_base_dir, train_size=args.train_size) |
| with open(args.output, "w") as f: |
| json.dump(datalist, f) |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser(description="") |
| parser.add_argument( |
| "--path", |
| type=str, |
| default="/workspace/data/medical/brats2018challenge", |
| help="root path of brats 2018 dataset.", |
| ) |
| parser.add_argument( |
| "--output", type=str, default="configs/datalist.json", help="relative path of output datalist json file." |
| ) |
| parser.add_argument("--train_size", type=int, default=200, help="number of training samples.") |
| args = parser.parse_args() |
|
|
| main(args) |
|
|
