|
import argparse |
|
import glob |
|
import json |
|
import os |
|
|
|
import monai |
|
from sklearn.model_selection import train_test_split |
|
|
|
|
|
def produce_sample_dict(line: str): |
|
return {"label": line, "image": line.replace("labelsTr", "imagesTr")} |
|
|
|
|
|
def produce_datalist(dataset_dir: str): |
|
""" |
|
This function is used to split the dataset. |
|
It will produce 200 samples for training, and the other samples are divided equally |
|
into val and test sets. |
|
""" |
|
|
|
samples = sorted(glob.glob(os.path.join(dataset_dir, "labelsTr", "*"), recursive=True)) |
|
samples = [_item.replace(os.path.join(dataset_dir, "labelsTr"), "labelsTr") for _item in samples] |
|
datalist = [] |
|
for line in samples: |
|
datalist.append(produce_sample_dict(line)) |
|
train_list, other_list = train_test_split(datalist, train_size=196) |
|
val_list, test_list = train_test_split(other_list, train_size=0.66) |
|
|
|
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 = args.path |
|
output_json = args.output |
|
|
|
monai.utils.set_determinism(seed=123) |
|
datalist = produce_datalist(dataset_dir=data_file_base_dir) |
|
with open(output_json, "w") as f: |
|
json.dump(datalist, f, ensure_ascii=True, indent=4) |
|
|
|
|
|
if __name__ == "__main__": |
|
|
|
parser = argparse.ArgumentParser(description="") |
|
parser.add_argument( |
|
"--path", |
|
type=str, |
|
default="/workspace/data/msd/Task07_Pancreas", |
|
help="root path of MSD Task07_Pancreas dataset.", |
|
) |
|
parser.add_argument( |
|
"--output", type=str, default="dataset_0.json", help="relative path of output datalist json file." |
|
) |
|
args = parser.parse_args() |
|
|
|
main(args) |
|
|
