Refactor MAISI bundle to monai 1.4 (#634)

Fixes # . ### Description refactor to monai1.4, add quality check, correct typo add download script ### Status **Ready/Work in progress/Hold** ### Please ensure all the checkboxes:  - [x] Codeformat tests passed locally by running `./runtests.sh --codeformat`. - [ ] In-line docstrings updated. - [ ] Update `version` and `changelog` in `metadata.json` if changing an existing bundle. - [ ] Please ensure the naming rules in config files meet our requirements (please refer to: `CONTRIBUTING.md`). - [ ] Ensure versions of packages such as `monai`, `pytorch` and `numpy` are correct in `metadata.json`. - [ ] Descriptions should be consistent with the content, such as `eval_metrics` of the provided weights and TorchScript modules. - [ ] Files larger than 25MB are excluded and replaced by providing download links in `large_file.yml`. - [ ] Avoid using path that contains personal information within config files (such as use `/home/your_name/` for `"bundle_root"`). --------- Signed-off-by: Can-Zhao <[email protected]> Signed-off-by: Yiheng Wang <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Yiheng Wang <[email protected]> Co-authored-by: Yiheng Wang <[email protected]>
Project-MONAI · Sep 4, 2024 · 2c5d564 · 2c5d564
1 parent 19f45ee
commit 2c5d564
Show file tree

Hide file tree

Showing 17 changed files with 1,242 additions and 3,606 deletions.
diff --git a/ci/unit_tests/test_maisi_ct_generative.py b/ci/unit_tests/test_maisi_ct_generative.py
@@ -93,7 +93,7 @@
         "body_region": ["head"],
         "anatomy_list": ["colon cancer primaries"],
     },
-    "Cannot find body region with given organ list.",
+    "Cannot find body region with given anatomy list.",
 ]
 
 TEST_CASE_INFER_ERROR_2 = [
@@ -156,7 +156,7 @@
         "body_region": ["chest"],
         "anatomy_list": ["colon", "spleen", "trachea", "left humerus", "sacrum", "heart"],
     },
-    "Cannot find body region with given organ list.",
+    "Cannot find body region with given anatomy list.",
 ]
 
 TEST_CASE_TRAIN = [

diff --git a/models/maisi_ct_generative/configs/image_median_statistics.json b/models/maisi_ct_generative/configs/image_median_statistics.json
@@ -0,0 +1,72 @@
+{
+    "liver": {
+        "min_median": -14.0,
+        "max_median": 1000.0,
+        "percentile_0_5": 9.530000000000001,
+        "percentile_99_5": 162.0,
+        "sigma_6_low": -21.596463547885904,
+        "sigma_6_high": 156.27881534763367,
+        "sigma_12_low": -110.53410299564568,
+        "sigma_12_high": 245.21645479539342
+    },
+    "spleen": {
+        "min_median": -69.0,
+        "max_median": 1000.0,
+        "percentile_0_5": 16.925000000000004,
+        "percentile_99_5": 184.07500000000073,
+        "sigma_6_low": -43.133891656525165,
+        "sigma_6_high": 177.40494997185993,
+        "sigma_12_low": -153.4033124707177,
+        "sigma_12_high": 287.6743707860525
+    },
+    "pancreas": {
+        "min_median": -124.0,
+        "max_median": 1000.0,
+        "percentile_0_5": -29.0,
+        "percentile_99_5": 145.92000000000007,
+        "sigma_6_low": -56.59382515620725,
+        "sigma_6_high": 149.50627399318438,
+        "sigma_12_low": -159.64387473090306,
+        "sigma_12_high": 252.5563235678802
+    },
+    "kidney": {
+        "min_median": -165.5,
+        "max_median": 819.0,
+        "percentile_0_5": -40.0,
+        "percentile_99_5": 254.61999999999898,
+        "sigma_6_low": -130.56375604853028,
+        "sigma_6_high": 267.28163511081016,
+        "sigma_12_low": -329.4864516282005,
+        "sigma_12_high": 466.20433069048045
+    },
+    "lung": {
+        "min_median": -1000.0,
+        "max_median": 65.0,
+        "percentile_0_5": -937.0,
+        "percentile_99_5": -366.9500000000007,
+        "sigma_6_low": -1088.5583843889117,
+        "sigma_6_high": -551.8503346949108,
+        "sigma_12_low": -1356.912409235912,
+        "sigma_12_high": -283.4963098479103
+    },
+    "bone": {
+        "min_median": 77.5,
+        "max_median": 1000.0,
+        "percentile_0_5": 136.45499999999998,
+        "percentile_99_5": 551.6350000000002,
+        "sigma_6_low": 71.39901958080469,
+        "sigma_6_high": 471.9957615639765,
+        "sigma_12_low": -128.8993514107812,
+        "sigma_12_high": 672.2941325555623
+    },
+    "brain": {
+        "min_median": -1000.0,
+        "max_median": 238.0,
+        "percentile_0_5": -951.0,
+        "percentile_99_5": 126.25,
+        "sigma_6_low": -304.8208236135867,
+        "sigma_6_high": 369.5118535139189,
+        "sigma_12_low": -641.9871621773394,
+        "sigma_12_high": 706.6781920776717
+    }
+}
diff --git a/models/maisi_ct_generative/configs/inference.json b/models/maisi_ct_generative/configs/inference.json
@@ -19,7 +19,7 @@
     "all_anatomy_size_condtions_json": "$@bundle_root + '/configs/all_anatomy_size_condtions.json'",
     "label_dict_json": "$@bundle_root + '/configs/label_dict.json'",
     "label_dict_remap_json": "$@bundle_root + '/configs/label_dict_124_to_132.json'",
-    "quality_check_args": null,
+    "real_img_median_statistics_file": "$@bundle_root + '/configs/image_median_statistics.json'",
     "num_output_samples": 1,
     "body_region": [
         "abdomen"
@@ -62,8 +62,14 @@
         64,
         64
     ],
+    "autoencoder_sliding_window_infer_size": [
+        96,
+        96,
+        96
+    ],
+    "autoencoder_sliding_window_infer_overlap": 0.6667,
     "autoencoder_def": {
-        "_target_": "scripts.custom_network_tp.AutoencoderKlckModifiedTp",
+        "_target_": "monai.apps.generation.maisi.networks.autoencoderkl_maisi.AutoencoderKlMaisi",
         "spatial_dims": "@spatial_dims",
         "in_channels": "@image_channels",
         "out_channels": "@image_channels",
@@ -88,10 +94,13 @@
         "with_encoder_nonlocal_attn": false,
         "with_decoder_nonlocal_attn": false,
         "use_checkpointing": false,
-        "use_convtranspose": false
+        "use_convtranspose": false,
+        "norm_float16": true,
+        "num_splits": 8,
+        "dim_split": 1
     },
-    "difusion_unet_def": {
-        "_target_": "scripts.custom_network_diffusion.CustomDiffusionModelUNet",
+    "diffusion_unet_def": {
+        "_target_": "monai.apps.generation.maisi.networks.diffusion_model_unet_maisi.DiffusionModelUNetMaisi",
         "spatial_dims": "@spatial_dims",
         "in_channels": "@latent_channels",
         "out_channels": "@latent_channels",
@@ -115,12 +124,12 @@
         ],
         "num_res_blocks": 2,
         "use_flash_attention": true,
-        "input_top_region_index": true,
-        "input_bottom_region_index": true,
-        "input_spacing": true
+        "include_top_region_index_input": true,
+        "include_bottom_region_index_input": true,
+        "include_spacing_input": true
     },
     "controlnet_def": {
-        "_target_": "scripts.custom_network_controlnet.CustomControlNet",
+        "_target_": "monai.apps.generation.maisi.networks.controlnet_maisi.ControlNetMaisi",
         "spatial_dims": "@spatial_dims",
         "in_channels": "@latent_channels",
         "num_channels": [
@@ -151,11 +160,11 @@
         ]
     },
     "mask_generation_autoencoder_def": {
-        "_target_": "generative.networks.nets.AutoencoderKL",
-        "spatial_dims": 3,
+        "_target_": "monai.apps.generation.maisi.networks.autoencoderkl_maisi.AutoencoderKlMaisi",
+        "spatial_dims": "@spatial_dims",
         "in_channels": 8,
         "out_channels": 125,
-        "latent_channels": 4,
+        "latent_channels": "@latent_channels",
         "num_channels": [
             32,
             64,
@@ -177,14 +186,17 @@
         "with_decoder_nonlocal_attn": false,
         "use_flash_attention": false,
         "use_checkpointing": true,
-        "use_convtranspose": true
+        "use_convtranspose": true,
+        "norm_float16": true,
+        "num_splits": 8,
+        "dim_split": 1
     },
     "mask_generation_diffusion_def": {
-        "_target_": "generative.networks.nets.DiffusionModelUNet",
-        "spatial_dims": 3,
-        "in_channels": 4,
-        "out_channels": 4,
-        "num_channels": [
+        "_target_": "monai.networks.nets.diffusion_model_unet.DiffusionModelUNet",
+        "spatial_dims": "@spatial_dims",
+        "in_channels": "@latent_channels",
+        "out_channels": "@latent_channels",
+        "channels": [
             64,
             128,
             256,
@@ -209,33 +221,33 @@
         "cross_attention_dim": 10
     },
     "autoencoder": "$@autoencoder_def.to(@device)",
-    "checkpoint_autoencoder": "$scripts.utils.load_autoencoder_ckpt(@trained_autoencoder_path)",
+    "checkpoint_autoencoder": "$torch.load(@trained_autoencoder_path)",
     "load_autoencoder": "[email protected]_state_dict(@checkpoint_autoencoder)",
-    "difusion_unet": "$@difusion_unet_def.to(@device)",
-    "checkpoint_difusion_unet": "$torch.load(@trained_diffusion_path)",
-    "load_diffusion": "$@difusion_unet.load_state_dict(@checkpoint_difusion_unet['unet_state_dict'])",
+    "diffusion_unet": "$@diffusion_unet_def.to(@device)",
+    "checkpoint_diffusion_unet": "$torch.load(@trained_diffusion_path)",
+    "load_diffusion": "$@diffusion_unet.load_state_dict(@checkpoint_diffusion_unet['unet_state_dict'])",
     "controlnet": "$@controlnet_def.to(@device)",
-    "copy_controlnet_state": "$monai.networks.utils.copy_model_state(@controlnet, @difusion_unet.state_dict())",
+    "copy_controlnet_state": "$monai.networks.utils.copy_model_state(@controlnet, @diffusion_unet.state_dict())",
     "checkpoint_controlnet": "$torch.load(@trained_controlnet_path)",
     "load_controlnet": "[email protected]_state_dict(@checkpoint_controlnet['controlnet_state_dict'], strict=True)",
-    "scale_factor": "$@checkpoint_difusion_unet['scale_factor'].to(@device)",
+    "scale_factor": "$@checkpoint_diffusion_unet['scale_factor'].to(@device)",
     "mask_generation_autoencoder": "$@mask_generation_autoencoder_def.to(@device)",
     "checkpoint_mask_generation_autoencoder": "$torch.load(@trained_mask_generation_autoencoder_path)",
     "load_mask_generation_autoencoder": "$@mask_generation_autoencoder.load_state_dict(@checkpoint_mask_generation_autoencoder, strict=True)",
-    "mask_generation_difusion_unet": "$@mask_generation_diffusion_def.to(@device)",
-    "checkpoint_mask_generation_difusion_unet": "$torch.load(@trained_mask_generation_diffusion_path)",
-    "load_mask_generation_diffusion": "$@mask_generation_difusion_unet.load_state_dict(@checkpoint_mask_generation_difusion_unet, strict=True)",
-    "mask_generation_scale_factor": 1.0055984258651733,
+    "mask_generation_diffusion_unet": "$@mask_generation_diffusion_def.to(@device)",
+    "checkpoint_mask_generation_diffusion_unet": "$torch.load(@trained_mask_generation_diffusion_path)",
+    "load_mask_generation_diffusion": "$@mask_generation_diffusion_unet.load_state_dict(@checkpoint_mask_generation_diffusion_unet['unet_state_dict'], strict=True)",
+    "mask_generation_scale_factor": "$@checkpoint_mask_generation_diffusion_unet['scale_factor']",
     "noise_scheduler": {
-        "_target_": "generative.networks.schedulers.DDPMScheduler",
+        "_target_": "monai.networks.schedulers.ddpm.DDPMScheduler",
         "num_train_timesteps": 1000,
         "beta_start": 0.0015,
         "beta_end": 0.0195,
         "schedule": "scaled_linear_beta",
         "clip_sample": false
     },
     "mask_generation_noise_scheduler": {
-        "_target_": "generative.networks.schedulers.DDPMScheduler",
+        "_target_": "monai.networks.schedulers.ddpm.DDPMScheduler",
         "num_train_timesteps": 1000,
         "beta_start": 0.0015,
         "beta_end": 0.0195,
@@ -263,27 +275,29 @@
         "label_dict_json": "@label_dict_json",
         "label_dict_remap_json": "@label_dict_remap_json",
         "autoencoder": "@autoencoder",
-        "difusion_unet": "@difusion_unet",
+        "diffusion_unet": "@diffusion_unet",
         "controlnet": "@controlnet",
         "scale_factor": "@scale_factor",
         "noise_scheduler": "@noise_scheduler",
         "mask_generation_autoencoder": "@mask_generation_autoencoder",
-        "mask_generation_difusion_unet": "@mask_generation_difusion_unet",
+        "mask_generation_diffusion_unet": "@mask_generation_diffusion_unet",
         "mask_generation_scale_factor": "@mask_generation_scale_factor",
         "mask_generation_noise_scheduler": "@mask_generation_noise_scheduler",
         "controllable_anatomy_size": "@controllable_anatomy_size",
         "image_output_ext": "@image_output_ext",
         "label_output_ext": "@label_output_ext",
+        "real_img_median_statistics": "@real_img_median_statistics_file",
         "device": "@device",
         "latent_shape": "@latent_shape",
         "mask_generation_latent_shape": "@mask_generation_latent_shape",
         "output_size": "@output_size",
-        "quality_check_args": "@quality_check_args",
         "spacing": "@spacing",
         "output_dir": "@output_dir",
         "num_inference_steps": "@num_inference_steps",
         "mask_generation_num_inference_steps": "@mask_generation_num_inference_steps",
-        "random_seed": "@random_seed"
+        "random_seed": "@random_seed",
+        "autoencoder_sliding_window_infer_size": "@autoencoder_sliding_window_infer_size",
+        "autoencoder_sliding_window_infer_overlap": "@autoencoder_sliding_window_infer_overlap"
     },
     "run": [
         "$@ldm_sampler.sample_multiple_images(@num_output_samples)"

diff --git a/models/maisi_ct_generative/configs/metadata.json b/models/maisi_ct_generative/configs/metadata.json
@@ -1,18 +1,17 @@
 {
     "schema": "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/meta_schema_generator_ldm_20240318.json",
-    "version": "0.3.6",
+    "version": "0.4.0",
     "changelog": {
+        "0.4.0": "update to use monai 1.4, model ckpt updated, rm GenerativeAI repo, add quality check",
         "0.3.6": "first oss version"
     },
-    "monai_version": "1.3.1",
-    "pytorch_version": "2.2.2",
+    "monai_version": "1.4.0",
+    "pytorch_version": "2.4.0",
     "numpy_version": "1.24.4",
     "optional_packages_version": {
         "fire": "0.6.0",
         "nibabel": "5.2.1",
-        "monai-generative": "0.2.3",
-        "tqdm": "4.66.2",
-        "xformers": "0.0.26"
+        "tqdm": "4.66.2"
     },
     "supported_apps": {
         "maisi-nim": ""

diff --git a/models/maisi_ct_generative/configs/train.json b/models/maisi_ct_generative/configs/train.json
@@ -28,8 +28,8 @@
     "spatial_dims": 3,
     "image_channels": 1,
     "latent_channels": 4,
-    "difusion_unet_def": {
-        "_target_": "scripts.custom_network_diffusion.CustomDiffusionModelUNet",
+    "diffusion_unet_def": {
+        "_target_": "monai.apps.generation.maisi.networks.diffusion_model_unet_maisi.DiffusionModelUNetMaisi",
         "spatial_dims": "@spatial_dims",
         "in_channels": "@latent_channels",
         "out_channels": "@latent_channels",
@@ -53,12 +53,12 @@
         ],
         "num_res_blocks": 2,
         "use_flash_attention": true,
-        "input_top_region_index": true,
-        "input_bottom_region_index": true,
-        "input_spacing": true
+        "include_top_region_index_input": true,
+        "include_bottom_region_index_input": true,
+        "include_spacing_input": true
     },
     "controlnet_def": {
-        "_target_": "scripts.custom_network_controlnet.CustomControlNet",
+        "_target_": "monai.apps.generation.maisi.networks.controlnet_maisi.ControlNetMaisi",
         "spatial_dims": "@spatial_dims",
         "in_channels": "@latent_channels",
         "num_channels": [
@@ -89,19 +89,19 @@
         ]
     },
     "noise_scheduler": {
-        "_target_": "generative.networks.schedulers.DDPMScheduler",
+        "_target_": "monai.networks.schedulers.ddpm.DDPMScheduler",
         "num_train_timesteps": 1000,
         "beta_start": 0.0015,
         "beta_end": 0.0195,
         "schedule": "scaled_linear_beta",
         "clip_sample": false
     },
-    "unzip_dataset": "scripts.utils.unzip_dataset(@dataset_dir)",
-    "difusion_unet": "$@difusion_unet_def.to(@device)",
-    "checkpoint_difusion_unet": "$torch.load(@trained_diffusion_path)",
-    "load_diffusion": "$@difusion_unet.load_state_dict(@checkpoint_difusion_unet['unet_state_dict'])",
+    "unzip_dataset": "$scripts.utils.unzip_dataset(@dataset_dir)",
+    "diffusion_unet": "$@diffusion_unet_def.to(@device)",
+    "checkpoint_diffusion_unet": "$torch.load(@trained_diffusion_path)",
+    "load_diffusion": "$@diffusion_unet.load_state_dict(@checkpoint_diffusion_unet['unet_state_dict'])",
     "controlnet": "$@controlnet_def.to(@device)",
-    "copy_controlnet_state": "$monai.networks.utils.copy_model_state(@controlnet, @difusion_unet.state_dict())",
+    "copy_controlnet_state": "$monai.networks.utils.copy_model_state(@controlnet, @diffusion_unet.state_dict())",
     "checkpoint_controlnet": "$torch.load(@trained_controlnet_path)",
     "load_controlnet": "[email protected]_state_dict(@checkpoint_controlnet['controlnet_state_dict'], strict=True)",
     "scale_factor": "$@checkpoint_controlnet['scale_factor'].to(@device)",
@@ -244,7 +244,7 @@
             "max_epochs": "@epochs",
             "device": "@device",
             "train_data_loader": "@train#dataloader",
-            "difusion_unet": "@difusion_unet",
+            "diffusion_unet": "@diffusion_unet",
             "controlnet": "@controlnet",
             "noise_scheduler": "@noise_scheduler",
             "loss_function": "@loss",

diff --git a/models/maisi_ct_generative/docs/README.md b/models/maisi_ct_generative/docs/README.md
@@ -23,6 +23,10 @@ The inference requires:
 - Disk Memory: at least 21GB disk memory
 
 ### Execute inference
+The model weights can be downloaded with
+```
+python -m scripts.download_files
+```
 The following code generates a synthetic image from a random sampled noise.
 ```
 python -m monai.bundle run --config_file configs/inference.json

diff --git a/models/maisi_ct_generative/large_files.yml b/models/maisi_ct_generative/large_files.yml
@@ -1,14 +1,14 @@
 large_files:
   - path: "models/autoencoder_epoch273.pt"
-    url: "https://drive.google.com/file/d/1jQefG0yJPzSvTG5rIJVHNqDReBTvVmZ0/view?usp=drive_link"
+    url: "https://drive.google.com/file/d/1Ojw25lFO8QbHkxazdK4CgZTyp3GFNZGz/view?usp=sharing"
   - path: "models/input_unet3d_data-all_steps1000size512ddpm_random_current_inputx_v1.pt"
-    url: "https://drive.google.com/file/d/1FtOHBGUF5dLZNHtiuhf5EH448EQGGs-_/view?usp=sharing"
+    url: "https://drive.google.com/file/d/1lklNv4MTdI_9bwFRMd98QQ7JLerR5gC_/view?usp=drive_link"
   - path: "models/controlnet-20datasets-e20wl100fold0bc_noi_dia_fsize_current.pt"
-    url: "https://drive.google.com/file/d/1izr52Whkk56OevNTk2QzI86eJV9TTaLk/view?usp=sharing"
+    url: "https://drive.google.com/file/d/1mLYeqeZ819_WpZPlAInhcWuCIHgn3QNT/view?usp=drive_link"
   - path: "models/mask_generation_autoencoder.pt"
-    url: "https://drive.google.com/file/d/1FzWrpv6ornYUaPiAWGOOxhRx2P9Wnynm/view?usp=drive_link"
+    url: "https://drive.google.com/file/d/19JnX-C6QAg4RfghTwpPnj4KEWhtawpCy/view?usp=drive_link"
   - path: "models/mask_generation_diffusion_unet.pt"
-    url: "https://drive.google.com/file/d/11SA9RUZ6XmCOJr5v6w6UW1kDzr6hlymw/view?usp=drive_link"
+    url: "https://drive.google.com/file/d/1yOQvlhXFGY1ZYavADM3N34vgg5AEitda/view?usp=drive_link"
   - path: "configs/candidate_masks_flexible_size_and_spacing_3000.json"
     url: "https://drive.google.com/file/d/1yMkH-lrAsn2YUGoTuVKNMpicziUmU-1J/view?usp=sharing"
   - path: "configs/all_anatomy_size_condtions.json"

diff --git a/models/maisi_ct_generative/scripts/__init__.py b/models/maisi_ct_generative/scripts/__init__.py
@@ -9,4 +9,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from . import custom_network_diffusion, custom_network_tp, sample, utils
+from . import sample, utils