|
- import os
- import json
- import shutil
- import SimpleITK as sitk
-
-
- origin_imagesTr_path = "/home/data2/wan/picai/workdir/nnUNet_raw_data/Task2201_picai_baseline/imagesTr"
- origin_labelsTr_path = "/home/data2/wan/picai/workdir/nnUNet_raw_data/Task2201_picai_baseline/labelsTr"
-
- nnUNet_raw_data_path = "/home/data2/wan/picai_nnunetv2/nnUNet_raw"
- os.makedirs(nnUNet_raw_data_path, exist_ok=True)
- task_path = os.path.join(nnUNet_raw_data_path, "Dataset602_adc")
- os.makedirs(task_path, exist_ok=True)
- dst_imagesTr_path = os.path.join(task_path, "imagesTr")
- os.makedirs(dst_imagesTr_path, exist_ok=True)
- dst_labelsTr_path = os.path.join(task_path, "labelsTr")
- os.makedirs(dst_labelsTr_path, exist_ok=True)
- dst_imagesTs_path = os.path.join(task_path, "imagesTs")
- os.makedirs(dst_imagesTs_path, exist_ok=True)
- dst_labelsTs_path = os.path.join(task_path, "labelsTs")
- os.makedirs(dst_labelsTs_path, exist_ok=True)
-
- count = 0
- for image_name in sorted(os.listdir(origin_imagesTr_path)):
- case_id = image_name.split("_")[1]
- modality_id = image_name.split("_")[2]
- if modality_id != "0001.nii.gz":
- continue
-
- origin_image_path = os.path.join(origin_imagesTr_path, image_name)
- origin_label_path = os.path.join(origin_labelsTr_path, image_name.split("_")[0] + "_" + image_name.split("_")[1] + ".nii.gz")
-
- new_image_name = "Prostate_" + case_id + "_" + "0000.nii.gz"
- new_label_name = "Prostate_" + case_id + ".nii.gz"
-
- label_sitk = sitk.ReadImage(origin_label_path)
- label_array = sitk.GetArrayFromImage(label_sitk)
-
- if label_array.max() > 0 and count < 200:
- count += 1
- new_imageTr_path = os.path.join(dst_imagesTr_path, new_image_name)
- new_labelTr_path = os.path.join(dst_labelsTr_path, new_label_name)
-
- shutil.copy(origin_image_path, new_imageTr_path)
- if not os.path.exists(new_labelTr_path):
- shutil.copy(origin_label_path, new_labelTr_path)
-
- elif label_array.max() > 0:
- count += 1
- new_imageTs_path = os.path.join(dst_imagesTs_path, new_image_name)
- new_labelTs_path = os.path.join(dst_labelsTs_path, new_label_name)
-
- shutil.copy(origin_image_path, new_imageTs_path)
- if not os.path.exists(new_labelTs_path):
- shutil.copy(origin_label_path, new_labelTs_path)
-
- print(count)
|