Synnunet

README.md

@@ -1,8 +1,7 @@
 English | [简体中文](README_CN.md)
 
 # MedicalSeg
-MedicalSeg is an easy-to-use 3D medical image segmentation toolkit that supports the whole segmentation process including data preprocessing, model training, and model deployment. Specially, We provide data preprocessing acceleration, high precision model on [COVID-19 CT scans](https://www.kaggle.com/andrewmvd/covid19-ct-scans) lung dataset and [MRISpineSeg](https://aistudio.baidu.com/aistudio/datasetdetail/81211) spine dataset, support for multiple datasets including [MSD](http://medicaldecathlon.com/), [Promise12](https://promise12.grand-challenge.org/), [Prostate_mri](https://liuquande.github.io/SAML/) and etc, and a [3D visualization demo](visualize.ipynb) based on [itkwidgets](https://github.com/InsightSoftwareConsortium/itkwidgets). The following image visualize the segmentation results on these two datasets:
-
+MedicalSeg is an easy-to-use 3D medical image segmentation toolkit that supports GPU acceleration from data preparation to deployment, one-click data preprocessing on five datasets. We also provided high precision model on [COVID-19 CT scans](https://www.kaggle.com/andrewmvd/covid19-ct-scans) dataset and [MRISpineSeg](https://aistudio.baidu.com/aistudio/datasetdetail/81211) dataset, and a [3D visualization demo]((visualize.ipynb) ) based on [itkwidgets](https://github.com/InsightSoftwareConsortium/itkwidgets). The following image visualize the segmentation results on these two datasets:
 
 <p align="center">
 <img src="https://github.com/shiyutang/files/raw/main/ezgif.com-gif-maker%20(1).gif" width="30.6%" height="20%"><img src="https://github.com/shiyutang/files/raw/main/ezgif.com-gif-maker.gif" width="40.6%" height="20%">
@@ -12,19 +11,13 @@ MedicalSeg is an easy-to-use 3D medical image segmentation toolkit that supports
 </p>
 
 
-**MedicalSeg is currently under development! If you find any problem using it or want to share any future develop suggestions, please open a github issue or join us by scanning the following wechat QR code.**
 
-<p align="center">
-<img src="https://user-images.githubusercontent.com/48433081/162115375-2dba8796-5184-4793-8efa-b142734fe734.png" width="20%" height="20%">
-</p>
 
 
 ## Contents
 1. [Performance](##Performance)
 2. [Quick Start](##QuickStart)
 3. [Structure](#Structure)
-4. [TODO](#TODO)
-5. [Acknowledgement](#Acknowledgement)
 
 ## Performance
 
@@ -74,7 +67,7 @@ This part introduce a easy to use the demo on COVID-19 CT scans dataset. This de
     ```
 - (Optional) Install CuPY if you want to accelerate the preprocess process. [CuPY installation guide](https://docs.cupy.dev/en/latest/install.html)
 
-- Get and preprocess the data. Remember to replace prepare_lung_coronavirus.py with different python script that you need [here](./tools):
+- Get and preprocess the data:
     - change the GPU setting [here](tools/preprocess_globals.yml) to True if you installed CuPY and want to use GPU to accelerate.
     ```
     python tools/prepare_lung_coronavirus.py
@@ -106,13 +99,5 @@ This part shows you the whole picture of our repository, which is easy to expand
 └── visualize.ipynb # You can try to visualize the result use this file.
 ```
 
-## TODO
-We have several thoughts in mind about what should our repo focus on. Your contribution will be very much welcomed.
-- [ ] Add PP-nnunet with acceleration in preprocess, automatic configuration for all dataset and better performance compared to nnunet.
-- [ ] Add top 1 liver segmentation algorithm on LITS challenge.
-- [ ] Add 3D Vertebral Measurement System.
-- [ ] Add pretrain model on various dataset.
 
-## Acknowledgement
-- Many thanks to [Lin Han](https://github.com/linhandev), [Lang Du](https://github.com/justld), [onecatcn](https://github.com/onecatcn) for their contribution in  our repository
-- Many thanks to [itkwidgets](https://github.com/InsightSoftwareConsortium/itkwidgets) for their powerful visualization toolkit that we used to present our visualizations.
+

README_CN.md

@@ -1,7 +1,7 @@
 [English](README.md) | 简体中文
 
 # MedicalSeg 介绍
-MedicalSeg 是一个简单易使用的全流程 3D 医学图像分割工具包，它支持从数据预处理、训练评估、再到模型部署的全套分割流程。特别的，我们还提供了数据预处理加速，在肺部数据 [COVID-19 CT scans](https://www.kaggle.com/andrewmvd/covid19-ct-scans) 和椎骨数据 [MRISpineSeg](https://aistudio.baidu.com/aistudio/datasetdetail/81211) 上的高精度模型， 对于[MSD](http://medicaldecathlon.com/)、[Promise12](https://promise12.grand-challenge.org/)、[Prostate_mri](https://liuquande.github.io/SAML/)等数据集的支持，以及基于[itkwidgets](https://github.com/InsightSoftwareConsortium/itkwidgets) 的 3D 可视化[Demo](visualize.ipynb)。如图所示是基于 MedicalSeg 在 Vnet 上训练之后的可视化结果：
+MedicalSeg 是一个简单易使用的 3D 医学图像分割工具包，支持从数据准备到部署的全流程 GPU 加速、五个数据集上的一键数据预处理，并提供了在 [COVID-19 CT scans](https://www.kaggle.com/andrewmvd/covid19-ct-scans) 数据集和 [MRISpineSeg](https://aistudio.baidu.com/aistudio/datasetdetail/81211) 数据集上的高精度模型，以及基于 [itkwidgets](https://github.com/InsightSoftwareConsortium/itkwidgets) 的 3D 可视化[Demo]((visualize.ipynb))。如图所示是基于 Vnet 训练之后的可视化结果：
 
 <p align="center">
 <img src="https://github.com/shiyutang/files/raw/main/ezgif.com-gif-maker%20(1).gif" width="30.6%" height="20%"><img src="https://github.com/shiyutang/files/raw/main/ezgif.com-gif-maker.gif" width="40.6%" height="20%">
@@ -10,18 +10,15 @@ MedicalSeg 是一个简单易使用的全流程 3D 医学图像分割工具包
 </p>
 </p>
 
-**MedicalSeg 目前正在开发中！如果您在使用中发现任何问题，或想分享任何开发建议，请提交 github issue 或扫描以下微信二维码加入我们。**
 
-<p align="center">
-<img src="https://user-images.githubusercontent.com/48433081/162115375-2dba8796-5184-4793-8efa-b142734fe734.png" width="20%" height="20%">
-</p>
+
+
 
 ## Contents
 1. [模型性能](##模型性能)
 2. [快速开始](##快速开始)
 3. [代码结构](#代码结构)
-4. [TODO](#TODO)
-5. [致谢](#致谢)
+
 
 ## 模型性能
 
@@ -31,16 +28,14 @@ MedicalSeg 是一个简单易使用的全流程 3D 医学图像分割工具包
 
 #### **COVID-19 CT scans 上的分割结果**
 
-
-| 骨干网络 | 分辨率 | 学习率 | 训练轮数 | mDice | 链接 |
+| 主干网络 | 分辨率 | 学习率 | 训练轮数 | mDice | 链接 |
 |:-:|:-:|:-:|:-:|:-:|:-:|
 |-|128x128x128|0.001|15000|97.04%|[model](https://bj.bcebos.com/paddleseg/paddleseg3d/lung_coronavirus/vnet_lung_coronavirus_128_128_128_15k_1e-3/model.pdparams) \| [log](https://bj.bcebos.com/paddleseg/paddleseg3d/lung_coronavirus/vnet_lung_coronavirus_128_128_128_15k_1e-3/train.log) \| [vdl](https://paddlepaddle.org.cn/paddle/visualdl/service/app?id=9db5c1e11ebc82f9a470f01a9114bd3c)|
 |-|128x128x128|0.0003|15000|92.70%|[model](https://bj.bcebos.com/paddleseg/paddleseg3d/lung_coronavirus/vnet_lung_coronavirus_128_128_128_15k_3e-4/model.pdparams) \| [log](https://bj.bcebos.com/paddleseg/paddleseg3d/lung_coronavirus/vnet_lung_coronavirus_128_128_128_15k_3e-4/train.log) \| [vdl](https://www.paddlepaddle.org.cn/paddle/visualdl/service/app/scalar?id=0fb90ee5a6ea8821c0d61a6857ba4614)|
 
 #### **MRISpineSeg 上的分割结果**
 
-
-| 骨干网络 | 分辨率 | 学习率 | 训练轮数 | mDice(20 classes) | Dice(16 classes) | 链接 |
+| 主干网络 | 分辨率 | 学习率 | 训练轮数 | mDice(20 classes) | Dice(16 classes) | 链接 |
 |:-:|:-:|:-:|:-:|:-:|:-:|:-:|
 |-|512x512x12|0.1|15000|74.41%| 88.17% |[model](https://bj.bcebos.com/paddleseg/paddleseg3d/mri_spine_seg/vnet_mri_spine_seg_512_512_12_15k_1e-1/model.pdparams) \| [log](https://bj.bcebos.com/paddleseg/paddleseg3d/mri_spine_seg/vnet_mri_spine_seg_512_512_12_15k_1e-1/train.log) \| [vdl](https://www.paddlepaddle.org.cn/paddle/visualdl/service/app/scalar?id=36504064c740e28506f991815bd21cc7)|
 |-|512x512x12|0.5|15000|74.69%| 89.14% |[model](https://bj.bcebos.com/paddleseg/paddleseg3d/mri_spine_seg/vnet_mri_spine_seg_512_512_12_15k_5e-1/model.pdparams) \| [log](https://bj.bcebos.com/paddleseg/paddleseg3d/mri_spine_seg/vnet_mri_spine_seg_512_512_12_15k_5e-1/train.log) \| [vdl](https://www.paddlepaddle.org.cn/paddle/visualdl/service/app/index?id=08b0f9f62ebb255cdfc93fd6bd8f2c06)|
@@ -73,7 +68,7 @@ MedicalSeg 是一个简单易使用的全流程 3D 医学图像分割工具包
     ```
 - (可选) 如果需要GPU加速，则可以参考[教程](https://docs.cupy.dev/en/latest/install.html) 安装 CuPY。
 
-- 一键数据预处理。如果不是准备肺部数据，可以在这个[目录](./tools)下，替换你需要的其他数据：
+- 一键数据预处理：
     - 如果你安装了CuPY并且想要 GPU 加速，修改[这里](tools/preprocess_globals.yml)的 use_gpu 配置为 True。
     ```
     python tools/prepare_lung_coronavirus.py
@@ -105,14 +100,3 @@ MedicalSeg 是一个简单易使用的全流程 3D 医学图像分割工具包
 └── visualize.ipynb # 用于进行 3D 可视化
 ```
 
-## TODO
-未来，我们想在这几个方面来发展 MedicalSeg，欢迎加入我们的开发者小组。
-- [ ] 增加带有预训练加速，自动化参数配置的高精度 PP-nnunet 模型。
-- [ ] 增加在 LITs 挑战中的 Top 1 肝脏分割算法。
-- [ ] 增加 3D 椎骨可视化测量系统。
-- [ ] 增加在多个数据上训练的预训练模型。
-
-
-## 致谢
-- 非常感谢 [Lin Han](https://github.com/linhandev), [Lang Du](https://github.com/justld), [onecatcn](https://github.com/onecatcn) 对我们仓库的贡献。
-- 非常感谢 [itkwidgets](https://github.com/InsightSoftwareConsortium/itkwidgets) 强大的3D可视化功能。

configs/lung_coronavirus/vnet_lung_coronavirus_128_128_128_15k.yml

@@ -5,4 +5,4 @@ model:
   elu: False
   in_channels: 1
   num_classes: 3
-  pretrained: https://bj.bcebos.com/paddleseg/dygraph/lung_coronavirus/vnet_lung_coronavirus_128_128_128_15k/pretrain/model.pdparams
+  pretrained: null

configs/mri_spine_seg/mri_spine_seg_1e-2_big_rmresizecrop.yml

@@ -0,0 +1,46 @@
+_base_: '../_base_/global_configs.yml'
+
+batch_size: 4
+iters: 15000
+
+train_dataset:
+  type: MRISpineSeg
+  dataset_root: MRSpineSeg/MRI_spine_seg_phase0_class3_big_12
+  result_dir: MRSpineSeg/MRI_spine_seg_phase1
+  transforms:
+    - type: RandomRotation3D
+      degrees: 30
+    - type: RandomFlip3D
+  mode: train
+  num_classes: 3
+
+val_dataset:
+  type: MRISpineSeg
+  dataset_root: MRSpineSeg/MRI_spine_seg_phase0_class3_big_12
+  result_dir: MRSpineSeg/MRI_spine_seg_phase1
+  num_classes: 3
+  transforms: []
+  mode: val
+  dataset_json_path: "data/MRSpineSeg/MRI_spine_seg_raw/dataset.json"
+
+optimizer:
+  type: sgd
+  momentum: 0.9
+  weight_decay: 1.0e-4
+
+lr_scheduler:
+  type: PolynomialDecay
+  decay_steps: 15000
+  learning_rate: 0.01
+  end_lr: 0
+  power: 0.9
+
+loss:
+  types:
+    - type: MixedLoss
+      losses:
+        - type: CrossEntropyLoss
+          weight: Null
+        - type: DiceLoss
+      coef: [1, 1]
+  coef: [1]

configs/mri_spine_seg/mri_spine_seg_1e-2_big_rmresizecrop_class20.yml

@@ -0,0 +1,46 @@
+_base_: '../_base_/global_configs.yml'
+
+batch_size: 3
+iters: 15000
+
+train_dataset:
+  type: MRISpineSeg
+  dataset_root: MRSpineSeg/MRI_spine_seg_phase0_class20_big_12
+  result_dir: MRSpineSeg/MRI_spine_seg_phase1
+  transforms:
+    - type: RandomRotation3D
+      degrees: 30
+    - type: RandomFlip3D
+  mode: train
+  num_classes: 20
+
+val_dataset:
+  type: MRISpineSeg
+  dataset_root: MRSpineSeg/MRI_spine_seg_phase0_class20_big_12
+  result_dir: MRSpineSeg/MRI_spine_seg_phase1
+  num_classes: 20
+  transforms: []
+  mode: val
+  dataset_json_path: "data/MRSpineSeg/MRI_spine_seg_raw/dataset.json"
+
+optimizer:
+  type: sgd
+  momentum: 0.9
+  weight_decay: 1.0e-4
+
+lr_scheduler:
+  type: PolynomialDecay
+  decay_steps: 15000
+  learning_rate: 0.1
+  end_lr: 0
+  power: 0.9
+
+loss:
+  types:
+    - type: MixedLoss
+      losses:
+        - type: CrossEntropyLoss
+          weight: Null
+        - type: DiceLoss
+      coef: [1, 1]
+  coef: [1]

configuration.py

@@ -0,0 +1,5 @@
+import os
+
+default_num_threads = 8 if 'nnUNet_def_n_proc' not in os.environ else int(os.environ['nnUNet_def_n_proc'])
+RESAMPLING_SEPARATE_Z_ANISO_THRESHOLD = 3  # determines what threshold to use for resampling the low resolution axis
+# separately (with NN)

default_configuration.py

@@ -0,0 +1,68 @@
+
+
+
+import medicalseg.core
+from paths import network_training_output_dir, preprocessing_output_dir, default_plans_identifier
+from tools.batchgenerators.utilities.file_and_folder_operations import *
+from tools.experiment_planning.summarize_plans import summarize_plans
+from medicalseg.core.training.model_restore import recursive_find_python_class
+
+
+def get_configuration_from_output_folder(folder):
+    # split off network_training_output_dir
+    folder = folder[len(network_training_output_dir):]
+    if folder.startswith("/"):
+        folder = folder[1:]
+
+    configuration, task, trainer_and_plans_identifier = folder.split("/")
+    trainer, plans_identifier = trainer_and_plans_identifier.split("__")
+    return configuration, task, trainer, plans_identifier
+
+
+def get_default_configuration(network, task, network_trainer, plans_identifier=default_plans_identifier,
+                              search_in=(medicalseg.core.__path__[0], "training", "network_training"),
+                              base_module='medicalseg.core.training.network_training'):
+    assert network in ['2d', '3d_lowres', '3d_fullres', '3d_cascade_fullres'], \
+        "network can only be one of the following: \'3d_lowres\', \'3d_fullres\', \'3d_cascade_fullres\'"
+
+    dataset_directory = join(preprocessing_output_dir, task)
+
+    if network == '2d':
+        plans_file = join(preprocessing_output_dir, task, plans_identifier + "_plans_2D.pkl")
+    else:
+        plans_file = join(preprocessing_output_dir, task, plans_identifier + "_plans_3D.pkl")
+
+    plans = load_pickle(plans_file)
+    possible_stages = list(plans['plans_per_stage'].keys())
+
+    if (network == '3d_cascade_fullres' or network == "3d_lowres") and len(possible_stages) == 1:
+        raise RuntimeError("3d_lowres/3d_cascade_fullres only applies if there is more than one stage. This task does "
+                           "not require the cascade. Run 3d_fullres instead")
+
+    if network == '2d' or network == "3d_lowres":
+        stage = 0
+    else:
+        stage = possible_stages[-1]
+
+    trainer_class = recursive_find_python_class([join(*search_in)], network_trainer,
+                                                current_module=base_module)
+
+    output_folder_name = join(network_training_output_dir, network, task, network_trainer + "__" + plans_identifier)
+
+    print("###############################################")
+    print("I am running the following nnUNet: %s" % network)
+    print("My trainer class is: ", trainer_class)
+    print("For that I will be using the following configuration:")
+    summarize_plans(plans_file)
+    print("I am using stage %d from these plans" % stage)
+
+    if (network == '2d' or len(possible_stages) > 1) and not network == '3d_lowres':
+        batch_dice = True
+        print("I am using batch dice + CE loss")
+    else:
+        batch_dice = False
+        print("I am using sample dice + CE loss")
+
+    print("\nI am using data from this folder: ", join(dataset_directory, plans['data_identifier']))
+    print("###############################################")
+    return plans_file, output_folder_name, dataset_directory, batch_dice, stage, trainer_class

deploy/python/infer.py

@@ -12,18 +12,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import argparse
+import codecs
 import os
 import sys
-import codecs
-import warnings
-import argparse
 
 LOCAL_PATH = os.path.dirname(os.path.abspath(__file__))
 sys.path.append(os.path.join(LOCAL_PATH, '..', '..'))
 
 import yaml
-import functools
 import numpy as np
+import functools
 
 from paddle.inference import create_predictor, PrecisionType
 from paddle.inference import Config as PredictConfig
@@ -384,32 +383,29 @@ def _preprocess(self, img):
         """
         if not "npy" in img:
             image_files = get_image_list(img, None, None)
-            warnings.warn("The image path is {}, please make sure this is the images you want to infer".format(image_files))
             savepath = os.path.dirname(img)
             pre = [
                 HUnorm,
                 functools.partial(
-                    resample,  # TODO: config preprocess in deply.yaml(export) to set params
+                    resample,  # TODO: config preprocess in deply.yaml to set params
                     new_shape=[128, 128, 128],
                     order=1)
             ]
 
-            for f in image_files:
-                f_nps = Prep.load_medical_data(f)
-                for f_np in f_nps:
-                    if pre is not None:
-                        for op in pre:
-                            f_np = op(f_np)
-
-                    # Set image to a uniform format before save. 
-                    if isinstance(f_np, tuple):
-                        f_np = f_np[0]             
-                    f_np = f_np.astype("float32")
-
-                    np.save(
-                        os.path.join(
-                            savepath, f.split("/")[-1].split(
-                                ".", maxsplit=1)[0]), f_np)
+            for f in tqdm(image_files, total=len(image_files)):
+                f_np = Prep.load_medical_data(f)
+                if pre is not None:
+                    for op in pre:
+                        f_np = op(f_np)
+
+                # Set image to a uniform format before save.             
+                f_np = f_np.astype("float32")
+
+                np.save(
+                    os.path.join(
+                        savepath, f.split("/")[-1].split(
+                            ".", maxsplit=1)[0]),
+                    f_np)
 
             img = img.split(".", maxsplit=1)[0] + ".npy"