prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  Enable input normalization in SemanticSegmentationData module
  ## 🚀 Feature
  Add the possibility to normalize Input images in SemanticSegmentationData module
  
  ### Motivation
  Enable effortless normalization, as already implemented by ImageClassificationData: optionally configurable by doing: 
  ```python
  dm = SemanticSegmentationData.from_folders(
      # ...
      args_transforms=dict(mean=mean,std=std)
  )
  ```
  
  ### Pitch
  Change [/flash/image/segmentation/input_transform.py:43](https://github.com/Lightning-AI/lightning-flash/blob/master/flash/image/segmentation/input_transform.py#L43)
  
  ```python
  
  @dataclass
  class SemanticSegmentationInputTransform(InputTransform):
  
      image_size: Tuple[int, int] = (128, 128)
  
      def train_per_sample_transform(self) -> Callable:
          return ApplyToKeys(
              [DataKeys.INPUT, DataKeys.TARGET],
              KorniaParallelTransforms(
                  K.geometry.Resize(self.image_size, interpolation="nearest"), K.augmentation.RandomHorizontalFlip(p=0.5)
              ),
          )
  
      def per_sample_transform(self) -> Callable:
          return ApplyToKeys(
              [DataKeys.INPUT, DataKeys.TARGET],
              KorniaParallelTransforms(K.geometry.Resize(self.image_size, interpolation="nearest")),
          )
  
      def predict_per_sample_transform(self) -> Callable:
          return ApplyToKeys(DataKeys.INPUT, K.geometry.Resize(self.image_size, interpolation="nearest"))
  ```
  
  into this
  
  ```python
  @dataclass
  class SemanticSegmentationInputTransform(InputTransform):
  
      image_size: Tuple[int, int] = (128, 128)
      mean: Union[float, Tuple[float, float, float]] = (0.485, 0.456, 0.406)
      std: Union[float, Tuple[float, float, float]] = (0.229, 0.224, 0.225)
  
  
      def train_per_sample_transform(self) -> Callable:
          return T.Compose(
              [
                  ApplyToKeys(
                      [DataKeys.INPUT, DataKeys.TARGET],
                      KorniaParallelTransforms(
                          K.geometry.Resize(self.image_size, interpolation="nearest"),
                      )
                  ),
                  ApplyToKeys(
                      [DataKeys.INPUT],
                      K.augmentation.Normalize(mean=mean, std=std)
                      
                  ),
              ]
          )
  
      def per_sample_transform(self) -> Callable:
          return T.Compose(
              [
                  ApplyToKeys(
                      [DataKeys.INPUT, DataKeys.TARGET],
                      KorniaParallelTransforms(
                          K.geometry.Resize(self.image_size, interpolation="nearest"),
                      )
                  ),
                  ApplyToKeys(
                      [DataKeys.INPUT],
                      K.augmentation.Normalize(mean=mean, std=std)
                      
                  ),
              ]
          )
  
      def predict_per_sample_transform(self) -> Callable: 
          return ApplyToKeys(
                      DataKeys.INPUT, 
                      K.geometry.Resize(self.image_size, interpolation="nearest"), 
                      K.augmentation.Normalize(mean=mean, std=std)
                  )
  
  ```
  
  ### Alternatives
  The alternative is to write a custom InputTransform object every time.
  
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `flash/image/segmentation/input_transform.py`
  Content:
  ```
  1 # Copyright The PyTorch Lightning team.
  2 #
  3 # Licensed under the Apache License, Version 2.0 (the "License");
  4 # you may not use this file except in compliance with the License.
  5 # You may obtain a copy of the License at
  6 #
  7 #     http://www.apache.org/licenses/LICENSE-2.0
  8 #
  9 # Unless required by applicable law or agreed to in writing, software
  10 # distributed under the License is distributed on an "AS IS" BASIS,
  11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12 # See the License for the specific language governing permissions and
  13 # limitations under the License.
  14 from dataclasses import dataclass
  15 from typing import Any, Callable, Dict, Tuple
  16 
  17 from flash.core.data.io.input import DataKeys
  18 from flash.core.data.io.input_transform import InputTransform
  19 from flash.core.data.transforms import ApplyToKeys, kornia_collate, KorniaParallelTransforms
  20 from flash.core.utilities.imports import _KORNIA_AVAILABLE, _TORCHVISION_AVAILABLE
  21 
  22 if _KORNIA_AVAILABLE:
  23     import kornia as K
  24 
  25 if _TORCHVISION_AVAILABLE:
  26     from torchvision import transforms as T
  27 
  28 
  29 def prepare_target(batch: Dict[str, Any]) -> Dict[str, Any]:
  30     """Convert the target mask to long and remove the channel dimension."""
  31     if DataKeys.TARGET in batch:
  32         batch[DataKeys.TARGET] = batch[DataKeys.TARGET].long().squeeze(1)
  33     return batch
  34 
  35 
  36 def remove_extra_dimensions(batch: Dict[str, Any]):
  37     if isinstance(batch[DataKeys.INPUT], list):
  38         assert len(batch[DataKeys.INPUT]) == 1
  39         batch[DataKeys.INPUT] = batch[DataKeys.INPUT][0]
  40     return batch
  41 
  42 
  43 @dataclass
  44 class SemanticSegmentationInputTransform(InputTransform):
  45 
  46     image_size: Tuple[int, int] = (128, 128)
  47 
  48     def train_per_sample_transform(self) -> Callable:
  49         return ApplyToKeys(
  50             [DataKeys.INPUT, DataKeys.TARGET],
  51             KorniaParallelTransforms(
  52                 K.geometry.Resize(self.image_size, interpolation="nearest"), K.augmentation.RandomHorizontalFlip(p=0.5)
  53             ),
  54         )
  55 
  56     def per_sample_transform(self) -> Callable:
  57         return ApplyToKeys(
  58             [DataKeys.INPUT, DataKeys.TARGET],
  59             KorniaParallelTransforms(K.geometry.Resize(self.image_size, interpolation="nearest")),
  60         )
  61 
  62     def predict_per_sample_transform(self) -> Callable:
  63         return ApplyToKeys(DataKeys.INPUT, K.geometry.Resize(self.image_size, interpolation="nearest"))
  64 
  65     def collate(self) -> Callable:
  66         return kornia_collate
  67 
  68     def per_batch_transform(self) -> Callable:
  69         return T.Compose([prepare_target, remove_extra_dimensions])
  70
  ```
  --- END FILES ---
  
  Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.
  
  Here is an example:
  
  ```diff
  diff --git a/examples/server_async.py b/examples/server_async.py
  --- a/examples/server_async.py
  +++ b/examples/server_async.py
  @@ -313,4 +313,4 @@
   
   
   if __name__ == "__main__":
  -    asyncio.run(run_async_server("."), debug=True)
  +    asyncio.run(run_async_server(), debug=True)
  diff --git a/examples/server_sync.py b/examples/server_sync.py
  --- a/examples/server_sync.py
  +++ b/examples/server_sync.py
  @@ -313,5 +313,5 @@
   
   
   if __name__ == "__main__":
  -    server = run_sync_server(".")
  +    server = run_sync_server()
       server.shutdown()
  
  ```
  
golden_diff: diff --git a/flash/image/segmentation/input_transform.py b/flash/image/segmentation/input_transform.py
  --- a/flash/image/segmentation/input_transform.py
  +++ b/flash/image/segmentation/input_transform.py
  @@ -12,7 +12,7 @@
   # See the License for the specific language governing permissions and
   # limitations under the License.
   from dataclasses import dataclass
  -from typing import Any, Callable, Dict, Tuple
  +from typing import Any, Callable, Dict, Tuple, Union
   
   from flash.core.data.io.input import DataKeys
   from flash.core.data.io.input_transform import InputTransform
  @@ -44,23 +44,43 @@
   class SemanticSegmentationInputTransform(InputTransform):
   
       image_size: Tuple[int, int] = (128, 128)
  +    mean: Union[float, Tuple[float, float, float]] = (0.485, 0.456, 0.406)
  +    std: Union[float, Tuple[float, float, float]] = (0.229, 0.224, 0.225)
   
       def train_per_sample_transform(self) -> Callable:
  -        return ApplyToKeys(
  -            [DataKeys.INPUT, DataKeys.TARGET],
  -            KorniaParallelTransforms(
  -                K.geometry.Resize(self.image_size, interpolation="nearest"), K.augmentation.RandomHorizontalFlip(p=0.5)
  -            ),
  +        return T.Compose(
  +            [
  +                ApplyToKeys(
  +                    [DataKeys.INPUT, DataKeys.TARGET],
  +                    KorniaParallelTransforms(
  +                        K.geometry.Resize(self.image_size, interpolation="nearest"),
  +                        K.augmentation.RandomHorizontalFlip(p=0.5),
  +                    ),
  +                ),
  +                ApplyToKeys([DataKeys.INPUT], K.augmentation.Normalize(mean=self.mean, std=self.std)),
  +            ]
           )
   
       def per_sample_transform(self) -> Callable:
  -        return ApplyToKeys(
  -            [DataKeys.INPUT, DataKeys.TARGET],
  -            KorniaParallelTransforms(K.geometry.Resize(self.image_size, interpolation="nearest")),
  +        return T.Compose(
  +            [
  +                ApplyToKeys(
  +                    [DataKeys.INPUT, DataKeys.TARGET],
  +                    KorniaParallelTransforms(K.geometry.Resize(self.image_size, interpolation="nearest")),
  +                ),
  +                ApplyToKeys([DataKeys.INPUT], K.augmentation.Normalize(mean=self.mean, std=self.std)),
  +            ]
           )
   
       def predict_per_sample_transform(self) -> Callable:
  -        return ApplyToKeys(DataKeys.INPUT, K.geometry.Resize(self.image_size, interpolation="nearest"))
  +        return ApplyToKeys(
  +            DataKeys.INPUT,
  +            K.geometry.Resize(
  +                self.image_size,
  +                interpolation="nearest",
  +            ),
  +            K.augmentation.Normalize(mean=self.mean, std=self.std),
  +        )
   
       def collate(self) -> Callable:
           return kornia_collate
  
verification_info: {"golden_diff": "diff --git a/flash/image/segmentation/input_transform.py b/flash/image/segmentation/input_transform.py\n--- a/flash/image/segmentation/input_transform.py\n+++ b/flash/image/segmentation/input_transform.py\n@@ -12,7 +12,7 @@\n # See the License for the specific language governing permissions and\n # limitations under the License.\n from dataclasses import dataclass\n-from typing import Any, Callable, Dict, Tuple\n+from typing import Any, Callable, Dict, Tuple, Union\n \n from flash.core.data.io.input import DataKeys\n from flash.core.data.io.input_transform import InputTransform\n@@ -44,23 +44,43 @@\n class SemanticSegmentationInputTransform(InputTransform):\n \n     image_size: Tuple[int, int] = (128, 128)\n+    mean: Union[float, Tuple[float, float, float]] = (0.485, 0.456, 0.406)\n+    std: Union[float, Tuple[float, float, float]] = (0.229, 0.224, 0.225)\n \n     def train_per_sample_transform(self) -> Callable:\n-        return ApplyToKeys(\n-            [DataKeys.INPUT, DataKeys.TARGET],\n-            KorniaParallelTransforms(\n-                K.geometry.Resize(self.image_size, interpolation=\"nearest\"), K.augmentation.RandomHorizontalFlip(p=0.5)\n-            ),\n+        return T.Compose(\n+            [\n+                ApplyToKeys(\n+                    [DataKeys.INPUT, DataKeys.TARGET],\n+                    KorniaParallelTransforms(\n+                        K.geometry.Resize(self.image_size, interpolation=\"nearest\"),\n+                        K.augmentation.RandomHorizontalFlip(p=0.5),\n+                    ),\n+                ),\n+                ApplyToKeys([DataKeys.INPUT], K.augmentation.Normalize(mean=self.mean, std=self.std)),\n+            ]\n         )\n \n     def per_sample_transform(self) -> Callable:\n-        return ApplyToKeys(\n-            [DataKeys.INPUT, DataKeys.TARGET],\n-            KorniaParallelTransforms(K.geometry.Resize(self.image_size, interpolation=\"nearest\")),\n+        return T.Compose(\n+            [\n+                ApplyToKeys(\n+                    [DataKeys.INPUT, DataKeys.TARGET],\n+                    KorniaParallelTransforms(K.geometry.Resize(self.image_size, interpolation=\"nearest\")),\n+                ),\n+                ApplyToKeys([DataKeys.INPUT], K.augmentation.Normalize(mean=self.mean, std=self.std)),\n+            ]\n         )\n \n     def predict_per_sample_transform(self) -> Callable:\n-        return ApplyToKeys(DataKeys.INPUT, K.geometry.Resize(self.image_size, interpolation=\"nearest\"))\n+        return ApplyToKeys(\n+            DataKeys.INPUT,\n+            K.geometry.Resize(\n+                self.image_size,\n+                interpolation=\"nearest\",\n+            ),\n+            K.augmentation.Normalize(mean=self.mean, std=self.std),\n+        )\n \n     def collate(self) -> Callable:\n         return kornia_collate\n", "issue": "Enable input normalization in SemanticSegmentationData module\n## \ud83d\ude80 Feature\r\nAdd the possibility to normalize Input images in SemanticSegmentationData module\r\n\r\n### Motivation\r\nEnable effortless normalization, as already implemented by ImageClassificationData: optionally configurable by doing: \r\n```python\r\ndm = SemanticSegmentationData.from_folders(\r\n    # ...\r\n    args_transforms=dict(mean=mean,std=std)\r\n)\r\n```\r\n\r\n### Pitch\r\nChange [/flash/image/segmentation/input_transform.py:43](https://github.com/Lightning-AI/lightning-flash/blob/master/flash/image/segmentation/input_transform.py#L43)\r\n\r\n```python\r\n\r\n@dataclass\r\nclass SemanticSegmentationInputTransform(InputTransform):\r\n\r\n    image_size: Tuple[int, int] = (128, 128)\r\n\r\n    def train_per_sample_transform(self) -> Callable:\r\n        return ApplyToKeys(\r\n            [DataKeys.INPUT, DataKeys.TARGET],\r\n            KorniaParallelTransforms(\r\n                K.geometry.Resize(self.image_size, interpolation=\"nearest\"), K.augmentation.RandomHorizontalFlip(p=0.5)\r\n            ),\r\n        )\r\n\r\n    def per_sample_transform(self) -> Callable:\r\n        return ApplyToKeys(\r\n            [DataKeys.INPUT, DataKeys.TARGET],\r\n            KorniaParallelTransforms(K.geometry.Resize(self.image_size, interpolation=\"nearest\")),\r\n        )\r\n\r\n    def predict_per_sample_transform(self) -> Callable:\r\n        return ApplyToKeys(DataKeys.INPUT, K.geometry.Resize(self.image_size, interpolation=\"nearest\"))\r\n```\r\n\r\ninto this\r\n\r\n```python\r\n@dataclass\r\nclass SemanticSegmentationInputTransform(InputTransform):\r\n\r\n    image_size: Tuple[int, int] = (128, 128)\r\n    mean: Union[float, Tuple[float, float, float]] = (0.485, 0.456, 0.406)\r\n    std: Union[float, Tuple[float, float, float]] = (0.229, 0.224, 0.225)\r\n\r\n\r\n    def train_per_sample_transform(self) -> Callable:\r\n        return T.Compose(\r\n            [\r\n                ApplyToKeys(\r\n                    [DataKeys.INPUT, DataKeys.TARGET],\r\n                    KorniaParallelTransforms(\r\n                        K.geometry.Resize(self.image_size, interpolation=\"nearest\"),\r\n                    )\r\n                ),\r\n                ApplyToKeys(\r\n                    [DataKeys.INPUT],\r\n                    K.augmentation.Normalize(mean=mean, std=std)\r\n                    \r\n                ),\r\n            ]\r\n        )\r\n\r\n    def per_sample_transform(self) -> Callable:\r\n        return T.Compose(\r\n            [\r\n                ApplyToKeys(\r\n                    [DataKeys.INPUT, DataKeys.TARGET],\r\n                    KorniaParallelTransforms(\r\n                        K.geometry.Resize(self.image_size, interpolation=\"nearest\"),\r\n                    )\r\n                ),\r\n                ApplyToKeys(\r\n                    [DataKeys.INPUT],\r\n                    K.augmentation.Normalize(mean=mean, std=std)\r\n                    \r\n                ),\r\n            ]\r\n        )\r\n\r\n    def predict_per_sample_transform(self) -> Callable: \r\n        return ApplyToKeys(\r\n                    DataKeys.INPUT, \r\n                    K.geometry.Resize(self.image_size, interpolation=\"nearest\"), \r\n                    K.augmentation.Normalize(mean=mean, std=std)\r\n                )\r\n\r\n```\r\n\r\n### Alternatives\r\nThe alternative is to write a custom InputTransform object every time.\r\n\n", "before_files": [{"content": "# Copyright The PyTorch Lightning team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom dataclasses import dataclass\nfrom typing import Any, Callable, Dict, Tuple\n\nfrom flash.core.data.io.input import DataKeys\nfrom flash.core.data.io.input_transform import InputTransform\nfrom flash.core.data.transforms import ApplyToKeys, kornia_collate, KorniaParallelTransforms\nfrom flash.core.utilities.imports import _KORNIA_AVAILABLE, _TORCHVISION_AVAILABLE\n\nif _KORNIA_AVAILABLE:\n    import kornia as K\n\nif _TORCHVISION_AVAILABLE:\n    from torchvision import transforms as T\n\n\ndef prepare_target(batch: Dict[str, Any]) -> Dict[str, Any]:\n    \"\"\"Convert the target mask to long and remove the channel dimension.\"\"\"\n    if DataKeys.TARGET in batch:\n        batch[DataKeys.TARGET] = batch[DataKeys.TARGET].long().squeeze(1)\n    return batch\n\n\ndef remove_extra_dimensions(batch: Dict[str, Any]):\n    if isinstance(batch[DataKeys.INPUT], list):\n        assert len(batch[DataKeys.INPUT]) == 1\n        batch[DataKeys.INPUT] = batch[DataKeys.INPUT][0]\n    return batch\n\n\n@dataclass\nclass SemanticSegmentationInputTransform(InputTransform):\n\n    image_size: Tuple[int, int] = (128, 128)\n\n    def train_per_sample_transform(self) -> Callable:\n        return ApplyToKeys(\n            [DataKeys.INPUT, DataKeys.TARGET],\n            KorniaParallelTransforms(\n                K.geometry.Resize(self.image_size, interpolation=\"nearest\"), K.augmentation.RandomHorizontalFlip(p=0.5)\n            ),\n        )\n\n    def per_sample_transform(self) -> Callable:\n        return ApplyToKeys(\n            [DataKeys.INPUT, DataKeys.TARGET],\n            KorniaParallelTransforms(K.geometry.Resize(self.image_size, interpolation=\"nearest\")),\n        )\n\n    def predict_per_sample_transform(self) -> Callable:\n        return ApplyToKeys(DataKeys.INPUT, K.geometry.Resize(self.image_size, interpolation=\"nearest\"))\n\n    def collate(self) -> Callable:\n        return kornia_collate\n\n    def per_batch_transform(self) -> Callable:\n        return T.Compose([prepare_target, remove_extra_dimensions])\n", "path": "flash/image/segmentation/input_transform.py"}], "after_files": [{"content": "# Copyright The PyTorch Lightning team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom dataclasses import dataclass\nfrom typing import Any, Callable, Dict, Tuple, Union\n\nfrom flash.core.data.io.input import DataKeys\nfrom flash.core.data.io.input_transform import InputTransform\nfrom flash.core.data.transforms import ApplyToKeys, kornia_collate, KorniaParallelTransforms\nfrom flash.core.utilities.imports import _KORNIA_AVAILABLE, _TORCHVISION_AVAILABLE\n\nif _KORNIA_AVAILABLE:\n    import kornia as K\n\nif _TORCHVISION_AVAILABLE:\n    from torchvision import transforms as T\n\n\ndef prepare_target(batch: Dict[str, Any]) -> Dict[str, Any]:\n    \"\"\"Convert the target mask to long and remove the channel dimension.\"\"\"\n    if DataKeys.TARGET in batch:\n        batch[DataKeys.TARGET] = batch[DataKeys.TARGET].long().squeeze(1)\n    return batch\n\n\ndef remove_extra_dimensions(batch: Dict[str, Any]):\n    if isinstance(batch[DataKeys.INPUT], list):\n        assert len(batch[DataKeys.INPUT]) == 1\n        batch[DataKeys.INPUT] = batch[DataKeys.INPUT][0]\n    return batch\n\n\n@dataclass\nclass SemanticSegmentationInputTransform(InputTransform):\n\n    image_size: Tuple[int, int] = (128, 128)\n    mean: Union[float, Tuple[float, float, float]] = (0.485, 0.456, 0.406)\n    std: Union[float, Tuple[float, float, float]] = (0.229, 0.224, 0.225)\n\n    def train_per_sample_transform(self) -> Callable:\n        return T.Compose(\n            [\n                ApplyToKeys(\n                    [DataKeys.INPUT, DataKeys.TARGET],\n                    KorniaParallelTransforms(\n                        K.geometry.Resize(self.image_size, interpolation=\"nearest\"),\n                        K.augmentation.RandomHorizontalFlip(p=0.5),\n                    ),\n                ),\n                ApplyToKeys([DataKeys.INPUT], K.augmentation.Normalize(mean=self.mean, std=self.std)),\n            ]\n        )\n\n    def per_sample_transform(self) -> Callable:\n        return T.Compose(\n            [\n                ApplyToKeys(\n                    [DataKeys.INPUT, DataKeys.TARGET],\n                    KorniaParallelTransforms(K.geometry.Resize(self.image_size, interpolation=\"nearest\")),\n                ),\n                ApplyToKeys([DataKeys.INPUT], K.augmentation.Normalize(mean=self.mean, std=self.std)),\n            ]\n        )\n\n    def predict_per_sample_transform(self) -> Callable:\n        return ApplyToKeys(\n            DataKeys.INPUT,\n            K.geometry.Resize(\n                self.image_size,\n                interpolation=\"nearest\",\n            ),\n            K.augmentation.Normalize(mean=self.mean, std=self.std),\n        )\n\n    def collate(self) -> Callable:\n        return kornia_collate\n\n    def per_batch_transform(self) -> Callable:\n        return T.Compose([prepare_target, remove_extra_dimensions])\n", "path": "flash/image/segmentation/input_transform.py"}]}