Porting of MedNIST DDPM model-zoo from MONAI Generative (#721)

### Description Initial porting from the Generative model-zoo of the mednist-ddpm model (https://github.com/Project-MONAI/GenerativeModels/tree/main/model-zoo/models/mednist_ddpm/bundle). The bundle files are tested within the docs/2d_ddpm_bundle_tutorial.ipynb ### Status **Ready** ### Please ensure all the checkboxes:  - [X] Update `version` and `changelog` in `metadata.json` if changing an existing bundle. - [X] Please ensure the naming rules in config files meet our requirements (please refer to: `CONTRIBUTING.md`). - [X] Ensure versions of packages such as `monai`, `pytorch` and `numpy` are correct in `metadata.json`. - [X] Avoid using path that contains personal information within config files (such as use `/home/your_name/` for `"bundle_root"`). --------- Signed-off-by: Eric Kerfoot <[email protected]> Signed-off-by: Eric Kerfoot <[email protected]> Co-authored-by: virginia <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Eric Kerfoot <[email protected]> Co-authored-by: Eric Kerfoot <[email protected]>
Project-MONAI · Feb 5, 2025 · 713255b · 713255b
1 parent c664109
commit 713255b
Show file tree

Hide file tree

Showing 13 changed files with 1,168 additions and 0 deletions.
diff --git a/ci/bundle_custom_data.py b/ci/bundle_custom_data.py
@@ -15,6 +15,7 @@
 # If a bundle does not need to be tested, please add the bundle name into the list.
 exclude_verify_shape_list = [
     "mednist_gan",
+    "mednist_ddpm",
     "lung_nodule_ct_detection",
     "pathology_nuclei_segmentation_classification",
     "brats_mri_generative_diffusion",
@@ -42,6 +43,7 @@
     "vista3d",
     "maisi_ct_generative",
     "vista2d",
+    "mednist_ddpm",
     "cxr_image_synthesis_latent_diffusion_model",
 ]
 

diff --git a/models/mednist_ddpm/LICENSE b/models/mednist_ddpm/LICENSE
@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 MONAI Consortium
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
diff --git a/models/mednist_ddpm/configs/inference.yaml b/models/mednist_ddpm/configs/inference.yaml
@@ -0,0 +1,99 @@
+# This defines an inference script for generating a random image to a Pytorch file
+imports:
+- $import os
+- $import datetime
+- $import torch
+- $import scripts
+- $import monai
+- $import torch.distributed as dist
+- $import operator
+
+# Common elements to all yaml files
+-
+image: $monai.utils.CommonKeys.IMAGE
+label: $monai.utils.CommonKeys.LABEL
+pred: $monai.utils.CommonKeys.PRED
+
+is_dist: '$dist.is_initialized()'
+rank: '$dist.get_rank() if @is_dist else 0'
+is_not_rank0: '$@rank > 0'
+device: '$torch.device(f"cuda:{@rank}" if torch.cuda.is_available() else "cpu")'
+
+network_def:
+  _target_: monai.networks.nets.DiffusionModelUNet
+  spatial_dims: 2
+  in_channels: 1
+  out_channels: 1
+  channels: [64, 128, 128]
+  attention_levels: [false, true, true]
+  num_res_blocks: 1
+  num_head_channels: 128
+
+network: $@network_def.to(@device)
+bundle_root: .
+ckpt_path: $@bundle_root + '/models/model.pt'
+use_amp: true
+image_dim: 64
+image_size: [1, '@image_dim', '@image_dim']
+num_train_timesteps: 1000
+
+base_transforms:
+- _target_: LoadImaged
+  keys: '@image'
+  image_only: true
+- _target_: EnsureChannelFirstd
+  keys: '@image'
+- _target_: ScaleIntensityRanged
+  keys: '@image'
+  a_min: 0.0
+  a_max: 255.0
+  b_min: 0.0
+  b_max: 1.0
+  clip: true
+
+scheduler:
+  _target_: monai.networks.schedulers.DDPMScheduler
+  num_train_timesteps: '@num_train_timesteps'
+
+inferer:
+  _target_: monai.inferers.DiffusionInferer
+  scheduler: '@scheduler'
+
+# Inference-specific
+
+batch_size: 1
+num_workers: 0
+
+noise: $torch.rand(1,1,@image_dim,@image_dim)  # create a random image every time this program is run
+
+out_file: ""  # where to save the tensor to
+
+# using a lambda this defines a simple sampling function used below
+sample: '$lambda x: @inferer.sample(input_noise=x, diffusion_model=@network, scheduler=@scheduler)'
+
+load_state: '[email protected]_state_dict(torch.load(@ckpt_path, weights_only = True))'  # command to load the saved model weights
+
+save_trans:
+  _target_: Compose
+  transforms:
+  - _target_: ScaleIntensity
+    minv: 0.0
+    maxv: 255.0
+  - _target_: ToTensor
+    track_meta: false
+  - _target_: SaveImage
+    output_ext: "jpg"
+    resample: false
+    output_dtype: '$torch.uint8'
+    separate_folder: false
+    output_postfix: '@out_file'
+
+# program to load the model weights, run `sample`, and store results to `out_file`
+testing:
+- '@load_state'
+- '$torch.save(@sample(@noise.to(@device)), @out_file)'
+
+#alternative version which saves to a jpg file
+testing_jpg:
+- '@load_state'
+- '$@save_trans(@sample(@noise.to(@device))[0])'
diff --git a/models/mednist_ddpm/configs/logging.conf b/models/mednist_ddpm/configs/logging.conf
@@ -0,0 +1,21 @@
+[loggers]
+keys=root
+
+[handlers]
+keys=consoleHandler
+
+[formatters]
+keys=fullFormatter
+
+[logger_root]
+level=INFO
+handlers=consoleHandler
+
+[handler_consoleHandler]
+class=StreamHandler
+level=INFO
+formatter=fullFormatter
+args=(sys.stdout,)
+
+[formatter_fullFormatter]
+format=%(asctime)s - %(name)s - %(levelname)s - %(message)s
diff --git a/models/mednist_ddpm/configs/metadata.json b/models/mednist_ddpm/configs/metadata.json
@@ -0,0 +1,59 @@
+{
+    "schema": "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/meta_schema_20240725.json",
+    "version": "1.0.0",
+    "changelog": {
+        "1.0.0": "Initial release"
+    },
+    "monai_version": "1.4.0",
+    "pytorch_version": "2.5.1",
+    "numpy_version": "1.26.4",
+    "required_packages_version": {},
+    "task": "MedNIST Hand Generation",
+    "description": "",
+    "authors": "Walter Hugo Lopez Pinaya, Mark Graham, and Eric Kerfoot",
+    "copyright": "Copyright (c) KCL",
+    "references": [],
+    "intended_use": "This is suitable for research purposes only.",
+    "image_classes": "Single channel magnitude data.",
+    "data_source": "MedNIST",
+    "network_data_format": {
+        "inputs": {
+            "image": {
+                "type": "image",
+                "format": "magnitude",
+                "modality": "xray",
+                "num_channels": 1,
+                "spatial_shape": [
+                    1,
+                    64,
+                    64
+                ],
+                "dtype": "float32",
+                "value_range": [],
+                "is_patch_data": false,
+                "channel_def": {
+                    "0": "image"
+                }
+            }
+        },
+        "outputs": {
+            "pred": {
+                "type": "image",
+                "format": "magnitude",
+                "modality": "xray",
+                "num_channels": 1,
+                "spatial_shape": [
+                    1,
+                    64,
+                    64
+                ],
+                "dtype": "float32",
+                "value_range": [],
+                "is_patch_data": false,
+                "channel_def": {
+                    "0": "image"
+                }
+            }
+        }
+    }
+}