From 6a14ba7a1409c1b7e1eafc20a04fb43d799c697d Mon Sep 17 00:00:00 2001 From: Nic Ma Date: Thu, 22 Apr 2021 02:15:27 +0800 Subject: [PATCH 1/3] [DLMED] save inverted meta dict Signed-off-by: Nic Ma --- monai/handlers/segmentation_saver.py | 14 +++++++------- monai/handlers/transform_inverter.py | 21 +++++++++++++++------ 2 files changed, 22 insertions(+), 13 deletions(-) diff --git a/monai/handlers/segmentation_saver.py b/monai/handlers/segmentation_saver.py index c3eca3b9c8..24dc7cf523 100644 --- a/monai/handlers/segmentation_saver.py +++ b/monai/handlers/segmentation_saver.py @@ -10,11 +10,9 @@ # limitations under the License. import logging -import warnings from typing import TYPE_CHECKING, Callable, Optional, Union import numpy as np - from monai.config import DtypeLike from monai.transforms import SaveImage from monai.utils import GridSampleMode, GridSamplePadMode, InterpolateMode, exact_version, optional_import @@ -121,7 +119,6 @@ def __init__( squeeze_end_dims=squeeze_end_dims, data_root_dir=data_root_dir, ) - self.resample = resample self.batch_transform = batch_transform self.output_transform = output_transform @@ -150,12 +147,15 @@ def __call__(self, engine: Engine) -> None: engine_output = self.output_transform(engine.state.output) if isinstance(engine_output, (tuple, list)): # if a list of data in shape: [channel, H, W, [D]], save every item separately - if self.resample: - warnings.warn("if saving inverted data, please set `resample=False` as it's already resampled.") - self._saver.save_batch = False for i, d in enumerate(engine_output): - self._saver(d, {k: meta_data[k][i] for k in meta_data} if meta_data is not None else None) + if isinstance(meta_data, dict): + meta_ = {k: meta_data[k][i] for k in meta_data} + elif isinstance(meta_data, (list, tuple)): + meta_ = meta_data[i] + else: + meta_ = meta_data + self._saver(d, meta_) else: # if the data is in shape: [batch, channel, H, W, [D]] self._saver.save_batch = True diff --git a/monai/handlers/transform_inverter.py b/monai/handlers/transform_inverter.py index 5f5c141189..c14cf94566 100644 --- a/monai/handlers/transform_inverter.py +++ b/monai/handlers/transform_inverter.py @@ -33,7 +33,10 @@ class TransformInverter: """ Ignite handler to automatically invert `transforms`. It takes `engine.state.output` as the input data and uses the transforms information from `engine.state.batch`. - The inverted results are stored in `engine.state.output` with key: "{output_key}_{postfix}". + The inverted data are stored in `engine.state.output` with key: "{output_key}_{postfix}". + And the inverted meta dict will be stored in `engine.state.batch` + with key: "{output_key}_{postfix}_{meta_key_postfix}". + """ def __init__( @@ -136,8 +139,14 @@ def __call__(self, engine: Engine) -> None: segs_dict[meta_dict_key] = engine.state.batch[meta_dict_key] with allow_missing_keys_mode(self.transform): # type: ignore - inverted_key = f"{output_key}_{self.postfix}" - engine.state.output[inverted_key] = [ - post_func(self._totensor(i[batch_key]).to(device) if to_tensor else i[batch_key]) - for i in self.inverter(segs_dict) - ] + inverted = self.inverter(segs_dict) + + # save the inverted data into state.output + inverted_key = f"{output_key}_{self.postfix}" + engine.state.output[inverted_key] = [ + post_func(self._totensor(i[batch_key]).to(device) if to_tensor else i[batch_key]) for i in inverted + ] + + # save the inverted meta dict into state.batch + if meta_dict_key in engine.state.batch: + engine.state.batch[f"{inverted_key}_{self.meta_key_postfix}"] = [i[meta_dict_key] for i in inverted] From c98a01ec467878ad6b37ef6cd9dd3be88c16bd75 Mon Sep 17 00:00:00 2001 From: monai-bot Date: Wed, 21 Apr 2021 18:23:44 +0000 Subject: [PATCH 2/3] [MONAI] python code formatting Signed-off-by: monai-bot --- monai/handlers/segmentation_saver.py | 1 + 1 file changed, 1 insertion(+) diff --git a/monai/handlers/segmentation_saver.py b/monai/handlers/segmentation_saver.py index 24dc7cf523..c04f9e1701 100644 --- a/monai/handlers/segmentation_saver.py +++ b/monai/handlers/segmentation_saver.py @@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Callable, Optional, Union import numpy as np + from monai.config import DtypeLike from monai.transforms import SaveImage from monai.utils import GridSampleMode, GridSamplePadMode, InterpolateMode, exact_version, optional_import From 41559bd3848381653e506e1746c239c8331e34aa Mon Sep 17 00:00:00 2001 From: Wenqi Li Date: Thu, 22 Apr 2021 00:14:14 +0100 Subject: [PATCH 3/3] safe key get Signed-off-by: Wenqi Li --- monai/handlers/transform_inverter.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monai/handlers/transform_inverter.py b/monai/handlers/transform_inverter.py index c14cf94566..651a47c39c 100644 --- a/monai/handlers/transform_inverter.py +++ b/monai/handlers/transform_inverter.py @@ -149,4 +149,4 @@ def __call__(self, engine: Engine) -> None: # save the inverted meta dict into state.batch if meta_dict_key in engine.state.batch: - engine.state.batch[f"{inverted_key}_{self.meta_key_postfix}"] = [i[meta_dict_key] for i in inverted] + engine.state.batch[f"{inverted_key}_{self.meta_key_postfix}"] = [i.get(meta_dict_key) for i in inverted]