Repository: Lightning-Universe/lightning-flash

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.



Discussion:
Hey @Nico995 Thanks for the suggestion! Your changes look great, would you be willing to open a PR for this?

Resolving patch:
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 @@ def remove_extra_dimensions(batch: Dict[str, Any]):
 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


Test patch: