Repository: Project-MONAI/MONAI

Issue:
Promote Child Items to Top Level
The MONAO dict-based transforms acts on the top level keys, for instance `data={"image": ..., "pred": ...}`; however, when the output is nested (like in the hovernet model), the actual data will lie one level below like `data={"image": ..., "pred": {"a":..., "b":...}}`.

Currently, if user want to apply transform on "a" or "b", it can use `Lambdad` and apply that transfrom on `data["pred"]` instead and use "a" and "b" as keys. However, if the user want to apply a transform on "image" and "a" for instance, I didn't find any transform that can enable it.

Here a transform that can promote these children items ("a" and "b") to the same level of "image" can solve the problem.



Discussion:
Hi, I'm not familiar with this model, but i'm just curious  at which place do we get the output in this form `"pred": {"a":..., "b":...}}` during loading pipeline or post-processing after prediction? I don't think any loading, augmentation transforms produce such format.   or if the model  produces such output, can we simple reformat the output of the model? thank you
Hi @myron, we got this in the output of HoVerNet model.
```
{"pred": {"horizontal_vertical": ..., "nucleus_prediction": ..., "type_prediction": ...}
```
https://github.com/Project-MONAI/MONAI/blob/77651add61b50f2a31fcc3c8b510d3179daa078c/monai/networks/nets/hovernet.py#L581-L586


Changing that will cause backward incompatibilities internally and externally.

I see, thank you.  then if you have access to model output in the python code,  you could directly add those keys "horizontal_vertical", 'nucleus_prediction".. to you main dict 
e.g. `batch_data.update(output["pred])`
this would save you from writing another monai transform. 
but, I guess you don't have access to it in python code, then yes, you'd need a new transfor . I think there is a similar in-spirit transform  **splitdim** https://docs.monai.io/en/stable/transforms.html#splitdim, but that only split a tensor dimension to create new key, value pairs. 


Resolving patch:
diff --git a/monai/transforms/__init__.py b/monai/transforms/__init__.py
--- a/monai/transforms/__init__.py
+++ b/monai/transforms/__init__.py
@@ -545,6 +545,9 @@
     FgBgToIndicesd,
     FgBgToIndicesD,
     FgBgToIndicesDict,
+    FlattenSubKeysd,
+    FlattenSubKeysD,
+    FlattenSubKeysDict,
     Identityd,
     IdentityD,
     IdentityDict,
diff --git a/monai/transforms/utility/dictionary.py b/monai/transforms/utility/dictionary.py
--- a/monai/transforms/utility/dictionary.py
+++ b/monai/transforms/utility/dictionary.py
@@ -127,6 +127,9 @@
     "MapLabelValueD",
     "MapLabelValueDict",
     "MapLabelValued",
+    "FlattenSubKeysd",
+    "FlattenSubKeysD",
+    "FlattenSubKeysDict",
     "RandCuCIMd",
     "RandCuCIMD",
     "RandCuCIMDict",
@@ -770,7 +773,7 @@ def _delete_item(keys, d, use_re: bool = False):
 class SelectItemsd(MapTransform):
     """
     Select only specified items from data dictionary to release memory.
-    It will copy the selected key-values and construct and new dictionary.
+    It will copy the selected key-values and construct a new dictionary.
     """
 
     backend = [TransformBackends.TORCH, TransformBackends.NUMPY]
@@ -779,6 +782,55 @@ def __call__(self, data):
         return {key: data[key] for key in self.key_iterator(data)}
 
 
+class FlattenSubKeysd(MapTransform):
+    """
+    If an item is dictionary, it flatten the item by moving the sub-items (defined by sub-keys) to the top level.
+    {"pred": {"a": ..., "b", ... }} --> {"a": ..., "b", ... }
+
+    Args:
+        keys: keys of the corresponding items to be flatten
+        sub_keys: the sub-keys of items to be flatten. If not provided all the sub-keys are flattened.
+        delete_keys: whether to delete the key of the items that their sub-keys are flattened. Default to True.
+        prefix: optional prefix to be added to the sub-keys when moving to the top level.
+            By default no prefix will be added.
+    """
+
+    backend = [TransformBackends.TORCH, TransformBackends.NUMPY]
+
+    def __init__(
+        self,
+        keys: KeysCollection,
+        sub_keys: Optional[KeysCollection] = None,
+        delete_keys: bool = True,
+        prefix: Optional[str] = None,
+    ) -> None:
+        super().__init__(keys)
+        self.sub_keys = sub_keys
+        self.delete_keys = delete_keys
+        self.prefix = prefix
+
+    def __call__(self, data):
+        d = dict(data)
+        for key in self.key_iterator(d):
+            # set the sub-keys for the specified key
+            sub_keys = d[key].keys() if self.sub_keys is None else self.sub_keys
+
+            # move all the sub-keys to the top level
+            for sk in sub_keys:
+                # set the top-level key for the sub-key
+                sk_top = f"{self.prefix}_{sk}" if self.prefix else sk
+                if sk_top in d:
+                    raise ValueError(
+                        f"'{sk_top}' already exists in the top-level keys. Please change `prefix` to avoid duplicity."
+                    )
+                d[sk_top] = d[key][sk]
+
+            # delete top level key that is flattened
+            if self.delete_keys:
+                del d[key]
+        return d
+
+
 class SqueezeDimd(MapTransform):
     """
     Dictionary-based wrapper of :py:class:`monai.transforms.SqueezeDim`.
@@ -1742,3 +1794,4 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N
 CuCIMD = CuCIMDict = CuCIMd
 RandCuCIMD = RandCuCIMDict = RandCuCIMd
 AddCoordinateChannelsD = AddCoordinateChannelsDict = AddCoordinateChannelsd
+FlattenSubKeysD = FlattenSubKeysDict = FlattenSubKeysd


Test patch:
diff --git a/tests/test_flatten_sub_keysd.py b/tests/test_flatten_sub_keysd.py
new file mode 100644
--- /dev/null
+++ b/tests/test_flatten_sub_keysd.py
@@ -0,0 +1,62 @@
+# 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.
+
+import unittest
+
+import torch
+from parameterized import parameterized
+
+from monai.transforms import FlattenSubKeysd
+
+A = torch.randn(2, 2)
+B = torch.randn(3, 3)
+C = torch.randn(1, 3)
+I = torch.randn(2, 3)
+D1 = {"a": A, "b": B}
+D2 = {"a": A, "b": B, "c": C}
+
+
+TEST_CASE_0 = [{"keys": "pred"}, {"image": I, "pred": D1}, {"a": A, "b": B, "image": I}]
+TEST_CASE_1 = [{"keys": "pred"}, {"image": I, "pred": D2}, {"a": A, "b": B, "c": C, "image": I}]
+TEST_CASE_2 = [{"keys": "pred", "sub_keys": ["a", "b"]}, {"image": I, "pred": D1}, {"a": A, "b": B, "image": I}]
+TEST_CASE_3 = [{"keys": "pred", "sub_keys": ["a", "b"]}, {"image": I, "pred": D2}, {"a": A, "b": B, "image": I}]
+TEST_CASE_4 = [
+    {"keys": "pred", "sub_keys": ["a", "b"], "delete_keys": False},
+    {"image": I, "pred": D1},
+    {"a": A, "b": B, "image": I, "pred": D1},
+]
+TEST_CASE_5 = [
+    {"keys": "pred", "sub_keys": ["a", "b"], "prefix": "new"},
+    {"image": I, "pred": D2},
+    {"new_a": A, "new_b": B, "image": I},
+]
+TEST_CASE_ERROR_1 = [  # error for duplicate key
+    {"keys": "pred", "sub_keys": ["a", "b"]},
+    {"image": I, "pred": D2, "a": None},
+]
+
+
+class TestFlattenSubKeysd(unittest.TestCase):
+    @parameterized.expand([TEST_CASE_0, TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5])
+    def test_dict(self, params, input_data, expected):
+        result = FlattenSubKeysd(**params)(input_data)
+        self.assertSetEqual(set(result.keys()), set(expected.keys()))
+        for k in expected:
+            self.assertEqual(id(result[k]), id(expected[k]))
+
+    @parameterized.expand([TEST_CASE_ERROR_1])
+    def test_error(self, params, input_data):
+        with self.assertRaises(ValueError):
+            FlattenSubKeysd(**params)(input_data)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_module_list.py b/tests/test_module_list.py
--- a/tests/test_module_list.py
+++ b/tests/test_module_list.py
@@ -39,7 +39,7 @@ def test_transform_api(self):
         """monai subclasses of MapTransforms must have alias names ending with 'd', 'D', 'Dict'"""
         to_exclude = {"MapTransform"}  # except for these transforms
         to_exclude_docs = {"Decollate", "Ensemble", "Invert", "SaveClassification", "RandTorchVision", "RandCrop"}
-        to_exclude_docs.update({"DeleteItems", "SelectItems", "CopyItems", "ConcatItems"})
+        to_exclude_docs.update({"DeleteItems", "SelectItems", "FlattenSubKeys", "CopyItems", "ConcatItems"})
         to_exclude_docs.update({"ToMetaTensor", "FromMetaTensor"})
         xforms = {
             name: obj