Repository: Project-MONAI/MONAI

Issue:
Support to replace original keys in post-transform
**Is your feature request related to a problem? Please describe.**
If the `output_postfix` is None, the post transform should use the original keys to save memory.



Discussion:

Patch:
diff --git a/monai/transforms/post/dictionary.py b/monai/transforms/post/dictionary.py
--- a/monai/transforms/post/dictionary.py
+++ b/monai/transforms/post/dictionary.py
@@ -71,6 +71,7 @@ def __init__(self, keys, output_postfix="act", sigmoid=False, softmax=False, oth
             output_postfix (str): the postfix string to construct keys to store converted data.
                 for example: if the keys of input data is `pred` and `label`, output_postfix is `act`,
                 the output data keys will be: `pred_act`, `label_act`.
+                if set to None, will replace the original data with the same key.
             sigmoid (bool, tuple or list of bool): whether to execute sigmoid function on model
                 output before transform.
             softmax (bool, tuple or list of bool): whether to execute softmax function on model
@@ -79,7 +80,7 @@ def __init__(self, keys, output_postfix="act", sigmoid=False, softmax=False, oth
                 for example: `other = lambda x: torch.tanh(x)`
         """
         super().__init__(keys)
-        if not isinstance(output_postfix, str):
+        if output_postfix is not None and not isinstance(output_postfix, str):
             raise ValueError("output_postfix must be a string.")
         self.output_postfix = output_postfix
         self.sigmoid = ensure_tuple_rep(sigmoid, len(self.keys))
@@ -91,7 +92,8 @@ def __call__(self, data):
         d = dict(data)
         for idx, key in enumerate(self.keys):
             ret = self.converter(d[key], self.sigmoid[idx], self.softmax[idx], self.other[idx])
-            d[f"{key}_{self.output_postfix}"] = ret
+            output_key = key if self.output_postfix is None else f"{key}_{self.output_postfix}"
+            d[output_key] = ret
         return d
 
 
@@ -117,6 +119,7 @@ def __init__(
             output_postfix (str): the postfix string to construct keys to store converted data.
                 for example: if the keys of input data is `pred` and `label`, output_postfix is `discreted`,
                 the output data keys will be: `pred_discreted`, `label_discreted`.
+                if set to None, will replace the original data with the same key.
             argmax (bool): whether to execute argmax function on input data before transform.
             to_onehot (bool): whether to convert input data into the one-hot format. Defaults to False.
             n_classes (bool): the number of classes to convert to One-Hot format.
@@ -124,7 +127,7 @@ def __init__(
             logit_thresh (float): the threshold value for thresholding operation, default is 0.5.
         """
         super().__init__(keys)
-        if not isinstance(output_postfix, str):
+        if output_postfix is not None and not isinstance(output_postfix, str):
             raise ValueError("output_postfix must be a string.")
         self.output_postfix = output_postfix
         self.argmax = ensure_tuple_rep(argmax, len(self.keys))
@@ -137,7 +140,8 @@ def __init__(
     def __call__(self, data):
         d = dict(data)
         for idx, key in enumerate(self.keys):
-            d[f"{key}_{self.output_postfix}"] = self.converter(
+            output_key = key if self.output_postfix is None else f"{key}_{self.output_postfix}"
+            d[output_key] = self.converter(
                 d[key],
                 self.argmax[idx],
                 self.to_onehot[idx],
@@ -172,9 +176,10 @@ def __init__(
             output_postfix (str): the postfix string to construct keys to store converted data.
                 for example: if the keys of input data is `label`, output_postfix is `largestcc`,
                 the output data keys will be: `label_largestcc`.
+                if set to None, will replace the original data with the same key.
         """
         super().__init__(keys)
-        if not isinstance(output_postfix, str):
+        if output_postfix is not None and not isinstance(output_postfix, str):
             raise ValueError("output_postfix must be a string.")
         self.output_postfix = output_postfix
         self.converter = KeepLargestConnectedComponent(applied_values, independent, background, connectivity)
@@ -182,7 +187,8 @@ def __init__(
     def __call__(self, data):
         d = dict(data)
         for idx, key in enumerate(self.keys):
-            d[f"{key}_{self.output_postfix}"] = self.converter(d[key])
+            output_key = key if self.output_postfix is None else f"{key}_{self.output_postfix}"
+            d[output_key] = self.converter(d[key])
         return d
 
 


Test patch:
diff --git a/tests/test_activationsd.py b/tests/test_activationsd.py
--- a/tests/test_activationsd.py
+++ b/tests/test_activationsd.py
@@ -40,10 +40,24 @@
     (1, 1, 2, 2),
 ]
 
+TEST_CASE_3 = [
+    {"keys": "pred", "output_postfix": "act", "sigmoid": False, "softmax": False, "other": lambda x: torch.tanh(x)},
+    {"pred": torch.tensor([[[[0.0, 1.0], [2.0, 3.0]]]])},
+    {"pred_act": torch.tensor([[[[0.0000, 0.7616], [0.9640, 0.9951]]]])},
+    (1, 1, 2, 2),
+]
+
+TEST_CASE_4 = [
+    {"keys": "pred", "output_postfix": None, "sigmoid": False, "softmax": False, "other": lambda x: torch.tanh(x)},
+    {"pred": torch.tensor([[[[0.0, 1.0], [2.0, 3.0]]]])},
+    {"pred": torch.tensor([[[[0.0000, 0.7616], [0.9640, 0.9951]]]])},
+    (1, 1, 2, 2),
+]
+
 
 class TestActivationsd(unittest.TestCase):
-    @parameterized.expand([TEST_CASE_1, TEST_CASE_2])
-    def test_shape(self, input_param, test_input, output, expected_shape):
+    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3])
+    def test_value_shape(self, input_param, test_input, output, expected_shape):
         result = Activationsd(**input_param)(test_input)
         torch.testing.assert_allclose(result["pred_act"], output["pred_act"])
         self.assertTupleEqual(result["pred_act"].shape, expected_shape)
@@ -51,6 +65,12 @@ def test_shape(self, input_param, test_input, output, expected_shape):
             torch.testing.assert_allclose(result["label_act"], output["label_act"])
             self.assertTupleEqual(result["label_act"].shape, expected_shape)
 
+    @parameterized.expand([TEST_CASE_4])
+    def test_none_postfix(self, input_param, test_input, output, expected_shape):
+        result = Activationsd(**input_param)(test_input)
+        torch.testing.assert_allclose(result["pred"], output["pred"])
+        self.assertTupleEqual(result["pred"].shape, expected_shape)
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/tests/test_as_discreted.py b/tests/test_as_discreted.py
--- a/tests/test_as_discreted.py
+++ b/tests/test_as_discreted.py
@@ -65,10 +65,25 @@
     (1, 2, 1, 2),
 ]
 
+TEST_CASE_4 = [
+    {
+        "keys": "pred",
+        "output_postfix": None,
+        "argmax": True,
+        "to_onehot": True,
+        "n_classes": 2,
+        "threshold_values": False,
+        "logit_thresh": 0.5,
+    },
+    {"pred": torch.tensor([[[[0.0, 1.0]], [[2.0, 3.0]]]])},
+    {"pred": torch.tensor([[[[0.0, 0.0]], [[1.0, 1.0]]]])},
+    (1, 2, 1, 2),
+]
+
 
 class TestAsDiscreted(unittest.TestCase):
     @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3])
-    def test_shape(self, input_param, test_input, output, expected_shape):
+    def test_value_shape(self, input_param, test_input, output, expected_shape):
         result = AsDiscreted(**input_param)(test_input)
         torch.testing.assert_allclose(result["pred_discrete"], output["pred_discrete"])
         self.assertTupleEqual(result["pred_discrete"].shape, expected_shape)
@@ -76,6 +91,12 @@ def test_shape(self, input_param, test_input, output, expected_shape):
             torch.testing.assert_allclose(result["label_discrete"], output["label_discrete"])
             self.assertTupleEqual(result["label_discrete"].shape, expected_shape)
 
+    @parameterized.expand([TEST_CASE_4])
+    def test_none_postfix(self, input_param, test_input, output, expected_shape):
+        result = AsDiscreted(**input_param)(test_input)
+        torch.testing.assert_allclose(result["pred"], output["pred"])
+        self.assertTupleEqual(result["pred"].shape, expected_shape)
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/tests/test_keep_largest_connected_componentd.py b/tests/test_keep_largest_connected_componentd.py
--- a/tests/test_keep_largest_connected_componentd.py
+++ b/tests/test_keep_largest_connected_componentd.py
@@ -119,6 +119,13 @@
     ),
 ]
 
+TEST_CASE_13 = [
+    "none_postfix",
+    {"keys": ["img"], "output_postfix": None, "independent": False, "applied_values": [1]},
+    grid_1,
+    torch.tensor([[[[0, 0, 1, 0, 0], [0, 2, 1, 1, 1], [0, 2, 1, 0, 0], [0, 2, 0, 1, 0], [2, 2, 0, 0, 2]]]]),
+]
+
 VALID_CASES = [
     TEST_CASE_1,
     TEST_CASE_2,
@@ -147,6 +154,13 @@ def test_correct_results(self, _, args, input_dict, expected):
             result = converter(input_dict)
             torch.allclose(result["img_largestcc"], expected)
 
+    @parameterized.expand([TEST_CASE_13])
+    def test_none_postfix(self, _, args, input_dict, expected):
+        converter = KeepLargestConnectedComponentd(**args)
+        input_dict["img"] = input_dict["img"].cpu()
+        result = converter(input_dict)
+        torch.allclose(result["img"], expected)
+
 
 if __name__ == "__main__":
     unittest.main()