PaddleCV-SIG / MedicalSeg

MedicalSeg is an easy-to-use 3D medical image segmentation toolkit that supports the whole segmentation process. Specially, We provide data preprocessing acceleration, high precision model on COVID-19 CT scans dataset and MRISpineSeg spine dataset, and a 3D visualization demo based on itkwidgets.
Apache License 2.0
68 stars 14 forks source link

[Feature] Add Experiment plan and Preocess #77

Closed justld closed 2 years ago

justld commented 2 years ago

添加plan生成,数据集preprocess类,如下代码可以得到plans和预处理后的数据集。数据集预处理和plans功能已完善。

from paddleseg3d.datasets.preprocess_utils import verify_dataset_integrity, convert_to_decathlon, crop, ExperimentPlanner2D_v21, ExperimentPlanner3D_v21, DatasetAnalyzer
import json
import pickle
import os

raw_data_dir = '/home/aistudio/data/Task06_Lung'
decathlon_dir = '/home/aistudio/data/Task006_Lung'
num_threads = 4
cropped_data_dir = '/home/aistudio/data/cropped/Task006_Lung'
preprocessed_dir = '/home/aistudio/data/preprocessed/Task006_Lung'

print('convert_to_decathlon')
convert_to_decathlon(input_folder=raw_data_dir, output_folder=decathlon_dir, num_processes=num_threads)

print('verify_dataset_integrity')
verify_dataset_integrity(decathlon_dir, default_num_threads=num_threads)

print('crop')
crop(decathlon_dir, cropped_data_dir, override=False, num_threads=num_threads)
with open(os.path.join(os.path.join(cropped_data_dir, 'dataset.json')), 'rb') as f:
    dataset_json = json.load(f)
modalities = list(dataset_json["modality"].values())
collect_intensityproperties = True if (("CT" in modalities) or ("ct" in modalities)) else False

print('DatasetAnalyzer')
dataset_analyzer = DatasetAnalyzer(cropped_data_dir, overwrite=True, num_processes=num_threads)
dataset_analyzer.analyze_dataset(collect_intensityproperties)

print('ExperimentPlanner2D_v21')
exp_planner_2d = ExperimentPlanner2D_v21(cropped_data_dir, preprocessed_dir)
if not os.path.exists(exp_planner_2d.plans_fname):
    print("{} not exists, generating plans, please wait a minute.".format(exp_planner_2d.plans_fname))
    exp_planner_2d.plan_experiment()
    exp_planner_2d.run_preprocessing(num_threads)

print('ExperimentPlanner3D_v21')

exp_planner_3d = ExperimentPlanner3D_v21(cropped_data_dir, preprocessed_dir)
if not os.path.exists(exp_planner_3d.plans_fname):
    print("{} already exists, generating plans, please wait a minute.".format(exp_planner_3d.plans_fname))
    exp_planner_3d.plan_experiment()
    exp_planner_3d.run_preprocessing(num_threads)