diff --git a/monai/data/__init__.py b/monai/data/__init__.py index b2dbdeaaa0..e7fa2b3107 100644 --- a/monai/data/__init__.py +++ b/monai/data/__init__.py @@ -62,7 +62,6 @@ partition_dataset_classes, pickle_hashing, rectify_header_sform_qform, - rep_scalar_to_batch, resample_datalist, select_cross_validation_folds, set_rnd, diff --git a/monai/data/utils.py b/monai/data/utils.py index 1ad01976e4..e68dd387ff 100644 --- a/monai/data/utils.py +++ b/monai/data/utils.py @@ -18,7 +18,7 @@ from collections import defaultdict from copy import deepcopy from functools import reduce -from itertools import product, starmap +from itertools import product, starmap, zip_longest from pathlib import PurePath from typing import Any, Dict, Generator, Iterable, List, Mapping, Optional, Sequence, Tuple, Union @@ -73,7 +73,6 @@ "pickle_hashing", "sorted_dict", "decollate_batch", - "rep_scalar_to_batch", "pad_list_data_collate", "no_collation", "convert_tables_to_dicts", @@ -297,7 +296,32 @@ def list_data_collate(batch: Sequence): raise TypeError(re_str) from re -def decollate_batch(batch, detach: bool = True): +def _non_zipping_check(batch_data, detach, pad, fill_value): + """ + Utility function based on `decollate_batch`, to identify the largest batch size from the collated data. + returns batch_size, the list of non-iterable items, and the dictionary or list with their items decollated. + + See `decollate_batch` for more details. + """ + if isinstance(batch_data, Mapping): + _deco = {key: decollate_batch(batch_data[key], detach, pad=pad, fill_value=fill_value) for key in batch_data} + elif isinstance(batch_data, Iterable): + _deco = [decollate_batch(b, detach, pad=pad, fill_value=fill_value) for b in batch_data] + else: + raise NotImplementedError(f"Unable to de-collate: {batch_data}, type: {type(batch_data)}.") + batch_size, non_iterable = 0, [] + for k, v in _deco.items() if isinstance(_deco, Mapping) else enumerate(_deco): + if not isinstance(v, Iterable) or isinstance(v, (str, bytes)) or (isinstance(v, torch.Tensor) and v.ndim == 0): + # Not running the usual list decollate here: + # don't decollate ['test', 'test'] into [['t', 't'], ['e', 'e'], ['s', 's'], ['t', 't']] + # torch.tensor(0) is iterable but iter(torch.tensor(0)) raises TypeError: iteration over a 0-d tensor + non_iterable.append(k) + elif hasattr(v, "__len__"): + batch_size = max(batch_size, len(v)) + return batch_size, non_iterable, _deco + + +def decollate_batch(batch, detach: bool = True, pad=True, fill_value=None): """De-collate a batch of data (for example, as produced by a `DataLoader`). Returns a list of structures with the original tensor's 0-th dimension sliced into elements using `torch.unbind`. @@ -335,10 +359,23 @@ def decollate_batch(batch, detach: bool = True): print(out[0]) >>> tensor([[[4.3549e-01...43e-01]]]) + batch_data = { + "image": [1, 2, 3], "meta": [4, 5], # undetermined batch size + } + out = decollate_batch(batch_data, pad=True, fill_value=0) + print(out) + >>> [{'image': 1, 'meta': 4}, {'image': 2, 'meta': 5}, {'image': 3, 'meta': 0}] + out = decollate_batch(batch_data, pad=False) + print(out) + >>> [{'image': 1, 'meta': 4}, {'image': 2, 'meta': 5}] + Args: batch: data to be de-collated. detach: whether to detach the tensors. Scalars tensors will be detached into number types instead of torch tensors. + pad: when the items in a batch indicate different batch size, whether to pad all the sequences to the longest. + If False, the batch size will be the length of the shortest sequence. + fill_value: when `pad` is True, the `fillvalue` to use when padding, defaults to `None`. """ if batch is None: return batch @@ -353,68 +390,20 @@ def decollate_batch(batch, detach: bool = True): if out_list[0].ndim == 0 and detach: return [t.item() for t in out_list] return list(out_list) - if isinstance(batch, Mapping): - _dict_list = {key: decollate_batch(batch[key], detach) for key in batch} - return [dict(zip(_dict_list, item)) for item in zip(*_dict_list.values())] - if isinstance(batch, Iterable): - item_0 = first(batch) - if ( - not isinstance(item_0, Iterable) - or isinstance(item_0, (str, bytes)) - or (isinstance(item_0, torch.Tensor) and item_0.ndim == 0) - ): - # Not running the usual list decollate here: - # don't decollate ['test', 'test'] into [['t', 't'], ['e', 'e'], ['s', 's'], ['t', 't']] - # torch.tensor(0) is iterable but iter(torch.tensor(0)) raises TypeError: iteration over a 0-d tensor - return [decollate_batch(b, detach) for b in batch] - return [list(item) for item in zip(*(decollate_batch(b, detach) for b in batch))] - raise NotImplementedError(f"Unable to de-collate: {batch}, type: {type(batch)}.") - - -def rep_scalar_to_batch(batch_data: Union[List, Dict]) -> Union[List, Dict]: - """ - Utility tp replicate the scalar items of a list or dictionary to ensure all the items have batch dimension. - It leverages `decollate_batch(detach=False)` to filter out the scalar items. - """ - - def _detect_batch_size(batch_data: Sequence): - """ - Detect the batch size from a list of data, some items in the list have batch dim, some not. - - """ - for v in batch_data: - if isinstance(v, torch.Tensor) and v.ndim > 0: - return v.shape[0] - for v in batch_data: - if issequenceiterable(v): - warnings.warn("batch_data doesn't contain batched Tensor data, use the length of first sequence data.") - return len(v) - raise RuntimeError("failed to automatically detect the batch size.") - - if isinstance(batch_data, dict): - batch_size = _detect_batch_size(list(batch_data.values())) - dict_batch = {} - for k, v in batch_data.items(): - if decollate_batch(v, detach=False) == v and not isinstance(v, list): - # if decollating a list, the result may be the same list, so should skip this case - dict_batch[k] = [deepcopy(decollate_batch(v, detach=True)) for _ in range(batch_size)] - else: - dict_batch[k] = v - - return dict_batch - if isinstance(batch_data, list): - batch_size = _detect_batch_size(batch_data) - list_batch = [] - for b in batch_data: - if decollate_batch(b, detach=False) == b and not isinstance(b, list): - list_batch.append([deepcopy(decollate_batch(b, detach=True)) for _ in range(batch_size)]) - else: - list_batch.append(b) - - return list_batch - # if not dict or list, just return the original data - return batch_data + b, non_iterable, deco = _non_zipping_check(batch, detach, pad, fill_value) + if b <= 0: # all non-iterable, single item "batch"? {"image": 1, "label": 1} + return deco + if pad: # duplicate non-iterable items to the longest batch + for k in non_iterable: + deco[k] = [deepcopy(deco[k]) for _ in range(b)] + if isinstance(deco, Mapping): + _gen = zip_longest(*deco.values(), fillvalue=fill_value) if pad else zip(*deco.values()) + return [dict(zip(deco, item)) for item in _gen] + if isinstance(deco, Iterable): + _gen = zip_longest(*deco, fillvalue=fill_value) if pad else zip(*deco) + return [list(item) for item in _gen] + raise NotImplementedError(f"Unable to de-collate: {batch}, type: {type(batch)}.") def pad_list_data_collate( diff --git a/monai/transforms/inverse_batch_transform.py b/monai/transforms/inverse_batch_transform.py index e220d5f350..9d9d8c2a5d 100644 --- a/monai/transforms/inverse_batch_transform.py +++ b/monai/transforms/inverse_batch_transform.py @@ -17,7 +17,7 @@ from monai.config import KeysCollection from monai.data.dataloader import DataLoader -from monai.data.utils import decollate_batch, no_collation, pad_list_data_collate, rep_scalar_to_batch +from monai.data.utils import decollate_batch, no_collation, pad_list_data_collate from monai.transforms.croppad.batch import PadListDataCollate from monai.transforms.inverse import InvertibleTransform from monai.transforms.transform import MapTransform, Transform @@ -60,6 +60,8 @@ def __init__( collate_fn: Optional[Callable] = no_collation, num_workers: Optional[int] = 0, detach: bool = True, + pad_batch: bool = True, + fill_value=None, ) -> None: """ Args: @@ -73,6 +75,10 @@ def __init__( if set to `None`, use the `num_workers` of the transform data loader. detach: whether to detach the tensors. Scalars tensors will be detached into number types instead of torch tensors. + pad_batch: when the items in a batch indicate different batch size, + whether to pad all the sequences to the longest. + If False, the batch size will be the length of the shortest sequence. + fill_value: the value to fill the padded sequences when `pad_batch=True`. """ self.transform = transform @@ -80,10 +86,12 @@ def __init__( self.num_workers = loader.num_workers if num_workers is None else num_workers self.collate_fn = collate_fn self.detach = detach + self.pad_batch = pad_batch + self.fill_value = fill_value self.pad_collation_used = loader.collate_fn.__doc__ == pad_list_data_collate.__doc__ def __call__(self, data: Dict[str, Any]) -> Any: - decollated_data = decollate_batch(data, detach=self.detach) + decollated_data = decollate_batch(data, detach=self.detach, pad=self.pad_batch, fill_value=self.fill_value) inv_ds = _BatchInverseDataset(decollated_data, self.transform, self.pad_collation_used) inv_loader = DataLoader( inv_ds, batch_size=self.batch_size, num_workers=self.num_workers, collate_fn=self.collate_fn @@ -99,25 +107,36 @@ def __call__(self, data: Dict[str, Any]) -> Any: class Decollated(MapTransform): """ - Decollate a batch of data, if input a dictionary, it can also support to only decollate specified keys. - Note that unlike most MapTransforms, it will delete other keys not specified and if keys=None, will decollate - all the data in the input. - And it replicates the scalar values to every item of the decollated list. + Decollate a batch of data. If input is a dictionary, it also supports to only decollate specified keys. + Note that unlike most MapTransforms, it will delete the other keys that are not specified. + if `keys=None`, it will decollate all the data in the input. + It replicates the scalar values to every item of the decollated list. Args: keys: keys of the corresponding items to decollate, note that it will delete other keys not specified. if None, will decollate all the keys. see also: :py:class:`monai.transforms.compose.MapTransform`. detach: whether to detach the tensors. Scalars tensors will be detached into number types instead of torch tensors. + pad_batch: when the items in a batch indicate different batch size, + whether to pad all the sequences to the longest. + If False, the batch size will be the length of the shortest sequence. + fill_value: the value to fill the padded sequences when `pad_batch=True`. allow_missing_keys: don't raise exception if key is missing. """ def __init__( - self, keys: Optional[KeysCollection] = None, detach: bool = True, allow_missing_keys: bool = False + self, + keys: Optional[KeysCollection] = None, + detach: bool = True, + pad_batch: bool = True, + fill_value=None, + allow_missing_keys: bool = False, ) -> None: super().__init__(keys, allow_missing_keys) self.detach = detach + self.pad_batch = pad_batch + self.fill_value = fill_value def __call__(self, data: Union[Dict, List]): d: Union[Dict, List] @@ -131,4 +150,4 @@ def __call__(self, data: Union[Dict, List]): for key in self.key_iterator(data): d[key] = data[key] - return decollate_batch(rep_scalar_to_batch(d), detach=self.detach) + return decollate_batch(d, detach=self.detach, pad=self.pad_batch, fill_value=self.fill_value) diff --git a/tests/test_decollate.py b/tests/test_decollate.py index 763bf808d0..fc0622c8f0 100644 --- a/tests/test_decollate.py +++ b/tests/test_decollate.py @@ -75,6 +75,7 @@ [[None, None], [None, None]], [["test"], ["test"]], [[], []], + [[("ch1", "ch2"), ("ch3",)], [["ch1", "ch3"], ["ch2", None]]], # default pad None ] @@ -207,6 +208,16 @@ def test_dict_examples(self): out = decollate_batch(test_case, detach=False) self.assertEqual(out[0]["out"], "test") + test_case = { + "image": torch.tensor([[[1, 2, 3]], [[3, 4, 5]]]), + "label": torch.tensor([[[5]], [[7]]]), + "out": ["test"], + } + out = decollate_batch(test_case, detach=False, pad=False) + self.assertEqual(len(out), 1) # no padding + out = decollate_batch(test_case, detach=False, pad=True, fill_value=0) + self.assertEqual(out[1]["out"], 0) # verify padding fill_value + def test_decollated(self): test_case = { "image": torch.tensor([[[1, 2]], [[3, 4]]]), @@ -233,12 +244,16 @@ def test_decollated(self): torch.tensor([[[1, 2]], [[3, 4]]]), {"out": ["test", "test"]}, {"scl_slope": torch.Tensor((0.0, 0.0))}, + {"out2": ["test1"]}, 0.85, + [], ] - transform = Decollated(keys=None, detach=False) + transform = Decollated(keys=None, detach=False, fill_value=-1) out = transform(test_case) - # the 4th item in the list is scalar loss value - self.assertEqual(out[1][3], 0.85) + + self.assertEqual(out[0][-2], 0.85) # scalar replicates + self.assertEqual(out[1][-2], 0.85) # scalar replicates + self.assertEqual(out[1][-3], -1) # fill value for the dictionary item self.assertEqual(out[0][1]["out"], "test") self.assertEqual(out[0][2]["scl_slope"], 0.0) self.assertTrue(isinstance(out[0][2]["scl_slope"], torch.Tensor))