3370 option of zip_longest to decollate

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,

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(

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,17 +75,23 @@ 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
         self.batch_size = loader.batch_size
         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)

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))