Merge branch 'Project-MONAI:dev' into 8328-nnunet-bundle-integration

Author: SimoneBendazzoli93

README.md

@@ -18,12 +18,13 @@
 
 MONAI is a [PyTorch](https://pytorch.org/)-based, [open-source](https://github.com/Project-MONAI/MONAI/blob/dev/LICENSE) framework for deep learning in healthcare imaging, part of the [PyTorch Ecosystem](https://pytorch.org/ecosystem/).
 Its ambitions are as follows:
+
 - Developing a community of academic, industrial and clinical researchers collaborating on a common foundation;
 - Creating state-of-the-art, end-to-end training workflows for healthcare imaging;
 - Providing researchers with the optimized and standardized way to create and evaluate deep learning models.
 
-
 ## Features
+
 > _Please see [the technical highlights](https://docs.monai.io/en/latest/highlights.html) and [What's New](https://docs.monai.io/en/latest/whatsnew.html) of the milestone releases._
 
 - flexible pre-processing for multi-dimensional medical imaging data;
@@ -32,7 +33,6 @@ Its ambitions are as follows:
 - customizable design for varying user expertise;
 - multi-GPU multi-node data parallelism support.
 
-
 ## Installation
 
 To install [the current release](https://pypi.org/project/monai/), you can simply run:
@@ -53,30 +53,34 @@ Technical documentation is available at [docs.monai.io](https://docs.monai.io).
 
 ## Citation
 
-If you have used MONAI in your research, please cite us! The citation can be exported from: https://arxiv.org/abs/2211.02701.
+If you have used MONAI in your research, please cite us! The citation can be exported from: <https://arxiv.org/abs/2211.02701>.
 
 ## Model Zoo
+
 [The MONAI Model Zoo](https://github.com/Project-MONAI/model-zoo) is a place for researchers and data scientists to share the latest and great models from the community.
 Utilizing [the MONAI Bundle format](https://docs.monai.io/en/latest/bundle_intro.html) makes it easy to [get started](https://github.com/Project-MONAI/tutorials/tree/main/model_zoo) building workflows with MONAI.
 
 ## Contributing
+
 For guidance on making a contribution to MONAI, see the [contributing guidelines](https://github.com/Project-MONAI/MONAI/blob/dev/CONTRIBUTING.md).
 
 ## Community
+
 Join the conversation on Twitter/X [@ProjectMONAI](https://twitter.com/ProjectMONAI) or join our [Slack channel](https://forms.gle/QTxJq3hFictp31UM9).
 
 Ask and answer questions over on [MONAI's GitHub Discussions tab](https://github.com/Project-MONAI/MONAI/discussions).
 
 ## Links
-- Website: https://monai.io/
-- API documentation (milestone): https://docs.monai.io/
-- API documentation (latest dev): https://docs.monai.io/en/latest/
-- Code: https://github.com/Project-MONAI/MONAI
-- Project tracker: https://github.com/Project-MONAI/MONAI/projects
-- Issue tracker: https://github.com/Project-MONAI/MONAI/issues
-- Wiki: https://github.com/Project-MONAI/MONAI/wiki
-- Test status: https://github.com/Project-MONAI/MONAI/actions
-- PyPI package: https://pypi.org/project/monai/
-- conda-forge: https://anaconda.org/conda-forge/monai
-- Weekly previews: https://pypi.org/project/monai-weekly/
-- Docker Hub: https://hub.docker.com/r/projectmonai/monai
+
+- Website: <https://monai.io/>
+- API documentation (milestone): <https://docs.monai.io/>
+- API documentation (latest dev): <https://docs.monai.io/en/latest/>
+- Code: <https://github.com/Project-MONAI/MONAI>
+- Project tracker: <https://github.com/Project-MONAI/MONAI/projects>
+- Issue tracker: <https://github.com/Project-MONAI/MONAI/issues>
+- Wiki: <https://github.com/Project-MONAI/MONAI/wiki>
+- Test status: <https://github.com/Project-MONAI/MONAI/actions>
+- PyPI package: <https://pypi.org/project/monai/>
+- conda-forge: <https://anaconda.org/conda-forge/monai>
+- Weekly previews: <https://pypi.org/project/monai-weekly/>
+- Docker Hub: <https://hub.docker.com/r/projectmonai/monai>

docs/source/handlers.rst

@@ -53,6 +53,12 @@ ROC AUC metrics handler
     :members:
 
 
+Average Precision metric handler
+--------------------------------
+.. autoclass:: AveragePrecision
+    :members:
+
+
 Confusion matrix metrics handler
 --------------------------------
 .. autoclass:: ConfusionMatrix

docs/source/metrics.rst

@@ -80,6 +80,13 @@ Metrics
 .. autoclass:: ROCAUCMetric
     :members:
 
+`Average Precision`
+-------------------
+.. autofunction:: compute_average_precision
+
+.. autoclass:: AveragePrecisionMetric
+    :members:
+
 `Confusion matrix`
 ------------------
 .. autofunction:: get_confusion_matrix

monai/handlers/__init__.py

@@ -11,6 +11,7 @@
 
 from __future__ import annotations
 
+from .average_precision import AveragePrecision
 from .checkpoint_loader import CheckpointLoader
 from .checkpoint_saver import CheckpointSaver
 from .classification_saver import ClassificationSaver

monai/handlers/average_precision.py

@@ -0,0 +1,53 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+from monai.handlers.ignite_metric import IgniteMetricHandler
+from monai.metrics import AveragePrecisionMetric
+from monai.utils import Average
+
+
+class AveragePrecision(IgniteMetricHandler):
+    """
+    Computes Average Precision (AP).
+    accumulating predictions and the ground-truth during an epoch and applying `compute_average_precision`.
+
+    Args:
+        average: {``"macro"``, ``"weighted"``, ``"micro"``, ``"none"``}
+            Type of averaging performed if not binary classification. Defaults to ``"macro"``.
+
+            - ``"macro"``: calculate metrics for each label, and find their unweighted mean.
+                This does not take label imbalance into account.
+            - ``"weighted"``: calculate metrics for each label, and find their average,
+                weighted by support (the number of true instances for each label).
+            - ``"micro"``: calculate metrics globally by considering each element of the label
+                indicator matrix as a label.
+            - ``"none"``: the scores for each class are returned.
+
+        output_transform: callable to extract `y_pred` and `y` from `ignite.engine.state.output` then
+            construct `(y_pred, y)` pair, where `y_pred` and `y` can be `batch-first` Tensors or
+            lists of `channel-first` Tensors. the form of `(y_pred, y)` is required by the `update()`.
+            `engine.state` and `output_transform` inherit from the ignite concept:
+            https://pytorch.org/ignite/concepts.html#state, explanation and usage example are in the tutorial:
+            https://github.com/Project-MONAI/tutorials/blob/master/modules/batch_output_transform.ipynb.
+
+    Note:
+        Average Precision expects y to be comprised of 0's and 1's.
+        y_pred must either be probability estimates or confidence values.
+
+    """
+
+    def __init__(self, average: Average | str = Average.MACRO, output_transform: Callable = lambda x: x) -> None:
+        metric_fn = AveragePrecisionMetric(average=Average(average))
+        super().__init__(metric_fn=metric_fn, output_transform=output_transform, save_details=False)

monai/inferers/inferer.py

@@ -1202,15 +1202,16 @@ def sample(  # type: ignore[override]
 
         if self.autoencoder_latent_shape is not None:
             latent = torch.stack([self.autoencoder_resizer(i) for i in decollate_batch(latent)], 0)
-            latent_intermediates = [
-                torch.stack([self.autoencoder_resizer(i) for i in decollate_batch(l)], 0) for l in latent_intermediates
-            ]
+            if save_intermediates:
+                latent_intermediates = [
+                    torch.stack([self.autoencoder_resizer(i) for i in decollate_batch(l)], 0)
+                    for l in latent_intermediates
+                ]
 
         decode = autoencoder_model.decode_stage_2_outputs
         if isinstance(autoencoder_model, SPADEAutoencoderKL):
             decode = partial(autoencoder_model.decode_stage_2_outputs, seg=seg)
         image = decode(latent / self.scale_factor)
-
         if save_intermediates:
             intermediates = []
             for latent_intermediate in latent_intermediates:
@@ -1333,13 +1334,15 @@ def __call__(  # type: ignore[override]
             raise NotImplementedError(f"{mode} condition is not supported")
 
         noisy_image = self.scheduler.add_noise(original_samples=inputs, noise=noise, timesteps=timesteps)
-        down_block_res_samples, mid_block_res_sample = controlnet(
-            x=noisy_image, timesteps=timesteps, controlnet_cond=cn_cond
-        )
+
         if mode == "concat" and condition is not None:
             noisy_image = torch.cat([noisy_image, condition], dim=1)
             condition = None
 
+        down_block_res_samples, mid_block_res_sample = controlnet(
+            x=noisy_image, timesteps=timesteps, controlnet_cond=cn_cond, context=condition
+        )
+
         diffuse = diffusion_model
         if isinstance(diffusion_model, SPADEDiffusionModelUNet):
             diffuse = partial(diffusion_model, seg=seg)
@@ -1395,17 +1398,21 @@ def sample(  # type: ignore[override]
             progress_bar = iter(scheduler.timesteps)
         intermediates = []
         for t in progress_bar:
-            # 1. ControlNet forward
-            down_block_res_samples, mid_block_res_sample = controlnet(
-                x=image, timesteps=torch.Tensor((t,)).to(input_noise.device), controlnet_cond=cn_cond
-            )
-            # 2. predict noise model_output
             diffuse = diffusion_model
             if isinstance(diffusion_model, SPADEDiffusionModelUNet):
                 diffuse = partial(diffusion_model, seg=seg)
 
             if mode == "concat" and conditioning is not None:
+                # 1. Conditioning
                 model_input = torch.cat([image, conditioning], dim=1)
+                # 2. ControlNet forward
+                down_block_res_samples, mid_block_res_sample = controlnet(
+                    x=model_input,
+                    timesteps=torch.Tensor((t,)).to(input_noise.device),
+                    controlnet_cond=cn_cond,
+                    context=None,
+                )
+                # 3. predict noise model_output
                 model_output = diffuse(
                     model_input,
                     timesteps=torch.Tensor((t,)).to(input_noise.device),
@@ -1414,6 +1421,12 @@ def sample(  # type: ignore[override]
                     mid_block_additional_residual=mid_block_res_sample,
                 )
             else:
+                down_block_res_samples, mid_block_res_sample = controlnet(
+                    x=image,
+                    timesteps=torch.Tensor((t,)).to(input_noise.device),
+                    controlnet_cond=cn_cond,
+                    context=conditioning,
+                )
                 model_output = diffuse(
                     image,
                     timesteps=torch.Tensor((t,)).to(input_noise.device),
@@ -1484,16 +1497,16 @@ def get_likelihood(  # type: ignore[override]
         for t in progress_bar:
             timesteps = torch.full(inputs.shape[:1], t, device=inputs.device).long()
             noisy_image = self.scheduler.add_noise(original_samples=inputs, noise=noise, timesteps=timesteps)
-            down_block_res_samples, mid_block_res_sample = controlnet(
-                x=noisy_image, timesteps=torch.Tensor((t,)).to(inputs.device), controlnet_cond=cn_cond
-            )
 
             diffuse = diffusion_model
             if isinstance(diffusion_model, SPADEDiffusionModelUNet):
                 diffuse = partial(diffusion_model, seg=seg)
 
             if mode == "concat" and conditioning is not None:
                 noisy_image = torch.cat([noisy_image, conditioning], dim=1)
+                down_block_res_samples, mid_block_res_sample = controlnet(
+                    x=noisy_image, timesteps=torch.Tensor((t,)).to(inputs.device), controlnet_cond=cn_cond, context=None
+                )
                 model_output = diffuse(
                     noisy_image,
                     timesteps=timesteps,
@@ -1502,6 +1515,12 @@ def get_likelihood(  # type: ignore[override]
                     mid_block_additional_residual=mid_block_res_sample,
                 )
             else:
+                down_block_res_samples, mid_block_res_sample = controlnet(
+                    x=noisy_image,
+                    timesteps=torch.Tensor((t,)).to(inputs.device),
+                    controlnet_cond=cn_cond,
+                    context=conditioning,
+                )
                 model_output = diffuse(
                     x=noisy_image,
                     timesteps=timesteps,
@@ -1727,9 +1746,11 @@ def sample(  # type: ignore[override]
 
         if self.autoencoder_latent_shape is not None:
             latent = torch.stack([self.autoencoder_resizer(i) for i in decollate_batch(latent)], 0)
-            latent_intermediates = [
-                torch.stack([self.autoencoder_resizer(i) for i in decollate_batch(l)], 0) for l in latent_intermediates
-            ]
+            if save_intermediates:
+                latent_intermediates = [
+                    torch.stack([self.autoencoder_resizer(i) for i in decollate_batch(l)], 0)
+                    for l in latent_intermediates
+                ]
 
         decode = autoencoder_model.decode_stage_2_outputs
         if isinstance(autoencoder_model, SPADEAutoencoderKL):

monai/metrics/__init__.py

@@ -12,6 +12,7 @@
 from __future__ import annotations
 
 from .active_learning_metrics import LabelQualityScore, VarianceMetric, compute_variance, label_quality_score
+from .average_precision import AveragePrecisionMetric, compute_average_precision
 from .confusion_matrix import ConfusionMatrixMetric, compute_confusion_matrix_metric, get_confusion_matrix
 from .cumulative_average import CumulativeAverage
 from .f_beta_score import FBetaScore