# Copyright The PyTorch Lightning team.
#
# 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 os
from typing import Any, Dict, List, Optional, Tuple, Union

import numpy as np

from flash.core.data.io.input import DataKeys, Input
from flash.core.data.utilities.loading import IMG_EXTENSIONS, NP_EXTENSIONS, load_image
from flash.core.data.utilities.paths import PATH_TYPE, filter_valid_files
from flash.core.data.utilities.samples import to_samples
from flash.core.integrations.fiftyone.utils import FiftyOneLabelUtilities
from flash.core.utilities.imports import _FIFTYONE_AVAILABLE, lazy_import
from flash.image.data import ImageFilesInput, ImageNumpyInput, ImageTensorInput
from flash.image.segmentation.output import SegmentationLabelsOutput

if _FIFTYONE_AVAILABLE:
    fo = lazy_import("fiftyone")
    SampleCollection = "fiftyone.core.collections.SampleCollection"
else:
    fo = None
    SampleCollection = None


class SemanticSegmentationInput(Input):
    num_classes: int
    labels_map: Dict[int, Tuple[int, int, int]]

    def load_labels_map(
        self, num_classes: Optional[int] = None, labels_map: Optional[Dict[int, Tuple[int, int, int]]] = None
    ) -> None:
        if num_classes is not None:
            self.num_classes = num_classes
            labels_map = labels_map or SegmentationLabelsOutput.create_random_labels_map(num_classes)

        if labels_map is not None:
            self.labels_map = labels_map


class SemanticSegmentationTensorInput(SemanticSegmentationInput, ImageTensorInput):
    def load_data(
        self,
        tensor: Any,
        masks: Any = None,
        num_classes: Optional[int] = None,
        labels_map: Optional[Dict[int, Tuple[int, int, int]]] = None,
    ) -> List[Dict[str, Any]]:
        self.load_labels_map(num_classes, labels_map)
        return to_samples(tensor, masks)

    def load_sample(self, sample: Dict[str, Any]) -> Dict[str, Any]:
        if DataKeys.TARGET in sample:
            sample[DataKeys.TARGET] = sample[DataKeys.TARGET].numpy()
        return super().load_sample(sample)


class SemanticSegmentationNumpyInput(SemanticSegmentationInput, ImageNumpyInput):
    def load_data(
        self,
        array: Any,
        masks: Any = None,
        num_classes: Optional[int] = None,
        labels_map: Optional[Dict[int, Tuple[int, int, int]]] = None,
    ) -> List[Dict[str, Any]]:
        self.load_labels_map(num_classes, labels_map)
        return to_samples(array, masks)


class SemanticSegmentationFilesInput(SemanticSegmentationInput, ImageFilesInput):
    def load_data(
        self,
        files: Union[PATH_TYPE, List[PATH_TYPE]],
        mask_files: Optional[Union[PATH_TYPE, List[PATH_TYPE]]] = None,
        num_classes: Optional[int] = None,
        labels_map: Optional[Dict[int, Tuple[int, int, int]]] = None,
    ) -> List[Dict[str, Any]]:
        self.load_labels_map(num_classes, labels_map)
        if mask_files is None:
            files = filter_valid_files(files, valid_extensions=IMG_EXTENSIONS + NP_EXTENSIONS)
        else:
            files, mask_files = filter_valid_files(files, mask_files, valid_extensions=IMG_EXTENSIONS + NP_EXTENSIONS)
        return to_samples(files, mask_files)

    def load_sample(self, sample: Dict[str, Any]) -> Dict[str, Any]:
        if DataKeys.TARGET in sample:
            sample[DataKeys.TARGET] = np.array(load_image(sample[DataKeys.TARGET])).transpose((2, 0, 1))[:, :, 0]
        return super().load_sample(sample)


class SemanticSegmentationFolderInput(SemanticSegmentationFilesInput):
    def load_data(
        self,
        folder: PATH_TYPE,
        mask_folder: Optional[PATH_TYPE] = None,
        num_classes: Optional[int] = None,
        labels_map: Optional[Dict[int, Tuple[int, int, int]]] = None,
    ) -> List[Dict[str, Any]]:
        self.load_labels_map(num_classes, labels_map)
        files = os.listdir(folder)
        files.sort()
        if mask_folder is not None:
            mask_files = {os.path.splitext(file)[0]: file for file in os.listdir(mask_folder)}
            file_names = [os.path.splitext(file)[0] for file in files]

            if len(set(file_names) - mask_files.keys()) != 0:
                raise ValueError(
                    f"Found inconsistent files in input folder: {folder} and mask folder: {mask_folder}. All input "
                    f"files must have a corresponding mask file with the same name."
                )

            files = [os.path.join(folder, file) for file in files]
            mask_files = [os.path.join(mask_folder, mask_files[file_name]) for file_name in file_names]
            return super().load_data(files, mask_files)
        return super().load_data([os.path.join(folder, file) for file in files])


class SemanticSegmentationFiftyOneInput(SemanticSegmentationFilesInput):
    label_field: str

    def load_data(
        self,
        sample_collection: SampleCollection,
        label_field: str = "ground_truth",
        num_classes: Optional[int] = None,
        labels_map: Optional[Dict[int, Tuple[int, int, int]]] = None,
    ) -> List[Dict[str, Any]]:
        self.load_labels_map(num_classes, labels_map)

        self.label_field = label_field
        label_utilities = FiftyOneLabelUtilities(label_field, fo.Segmentation)
        label_utilities.validate(sample_collection)

        self._fo_dataset_name = sample_collection.name
        return to_samples(sample_collection.values("filepath"))

    def predict_load_data(
        self,
        sample_collection: SampleCollection,
    ) -> List[Dict[str, Any]]:
        return to_samples(sample_collection.values("filepath"))

    def load_sample(self, sample: Dict[str, Any]) -> Dict[str, Any]:
        filepath = sample[DataKeys.INPUT]
        sample = super().load_sample(sample)
        if not self.predicting:
            fo_dataset = fo.load_dataset(self._fo_dataset_name)
            fo_sample = fo_dataset[filepath]
            sample[DataKeys.TARGET] = fo_sample[self.label_field].mask
        return sample

src/flash/image/segmentation/input.py
# Copyright The PyTorch Lightning team.
#
# 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 os
from typing import Any, Dict, List, Optional, Tuple, Union

import numpy as np

from flash.core.data.io.input import DataKeys, Input
from flash.core.data.utilities.loading import IMG_EXTENSIONS, NP_EXTENSIONS, load_image
from flash.core.data.utilities.paths import PATH_TYPE, filter_valid_files
from flash.core.data.utilities.samples import to_samples
from flash.core.integrations.fiftyone.utils import FiftyOneLabelUtilities
from flash.core.utilities.imports import _FIFTYONE_AVAILABLE, lazy_import
from flash.image.data import ImageFilesInput, ImageNumpyInput, ImageTensorInput
from flash.image.segmentation.output import SegmentationLabelsOutput

if _FIFTYONE_AVAILABLE:
    fo = lazy_import("fiftyone")
    SampleCollection = "fiftyone.core.collections.SampleCollection"
else:
    fo = None
    SampleCollection = None


class SemanticSegmentationInput(Input):
    num_classes: int
    labels_map: Dict[int, Tuple[int, int, int]]

    def load_labels_map(
        self, num_classes: Optional[int] = None, labels_map: Optional[Dict[int, Tuple[int, int, int]]] = None
    ) -> None:
        if num_classes is not None:
            self.num_classes = num_classes
            labels_map = labels_map or SegmentationLabelsOutput.create_random_labels_map(num_classes)

        if labels_map is not None:
            self.labels_map = labels_map


class SemanticSegmentationTensorInput(SemanticSegmentationInput, ImageTensorInput):
    def load_data(
        self,
        tensor: Any,
        masks: Any = None,
        num_classes: Optional[int] = None,
        labels_map: Optional[Dict[int, Tuple[int, int, int]]] = None,
    ) -> List[Dict[str, Any]]:
        self.load_labels_map(num_classes, labels_map)
        return to_samples(tensor, masks)

    def load_sample(self, sample: Dict[str, Any]) -> Dict[str, Any]:
        if DataKeys.TARGET in sample:
            sample[DataKeys.TARGET] = sample[DataKeys.TARGET].numpy()
        return super().load_sample(sample)


class SemanticSegmentationNumpyInput(SemanticSegmentationInput, ImageNumpyInput):
    def load_data(
        self,
        array: Any,
        masks: Any = None,
        num_classes: Optional[int] = None,
        labels_map: Optional[Dict[int, Tuple[int, int, int]]] = None,
    ) -> List[Dict[str, Any]]:
        self.load_labels_map(num_classes, labels_map)
        return to_samples(array, masks)


class SemanticSegmentationFilesInput(SemanticSegmentationInput, ImageFilesInput):
    def load_data(
        self,
        files: Union[PATH_TYPE, List[PATH_TYPE]],
        mask_files: Optional[Union[PATH_TYPE, List[PATH_TYPE]]] = None,
        num_classes: Optional[int] = None,
        labels_map: Optional[Dict[int, Tuple[int, int, int]]] = None,
    ) -> List[Dict[str, Any]]:
        self.load_labels_map(num_classes, labels_map)
        if mask_files is None:
            files = filter_valid_files(files, valid_extensions=IMG_EXTENSIONS + NP_EXTENSIONS)
        else:
            files, mask_files = filter_valid_files(files, mask_files, valid_extensions=IMG_EXTENSIONS + NP_EXTENSIONS)
        return to_samples(files, mask_files)

    def load_sample(self, sample: Dict[str, Any]) -> Dict[str, Any]:
        if DataKeys.TARGET in sample:
            sample[DataKeys.TARGET] = np.array(load_image(sample[DataKeys.TARGET])).transpose((2, 0, 1))[0, :, :]
        return super().load_sample(sample)


class SemanticSegmentationFolderInput(SemanticSegmentationFilesInput):
    def load_data(
        self,
        folder: PATH_TYPE,
        mask_folder: Optional[PATH_TYPE] = None,
        num_classes: Optional[int] = None,
        labels_map: Optional[Dict[int, Tuple[int, int, int]]] = None,
    ) -> List[Dict[str, Any]]:
        self.load_labels_map(num_classes, labels_map)
        files = os.listdir(folder)
        files.sort()
        if mask_folder is not None:
            mask_files = {os.path.splitext(file)[0]: file for file in os.listdir(mask_folder)}
            file_names = [os.path.splitext(file)[0] for file in files]

            if len(set(file_names) - mask_files.keys()) != 0:
                raise ValueError(
                    f"Found inconsistent files in input folder: {folder} and mask folder: {mask_folder}. All input "
                    f"files must have a corresponding mask file with the same name."
                )

            files = [os.path.join(folder, file) for file in files]
            mask_files = [os.path.join(mask_folder, mask_files[file_name]) for file_name in file_names]
            return super().load_data(files, mask_files)
        return super().load_data([os.path.join(folder, file) for file in files])


class SemanticSegmentationFiftyOneInput(SemanticSegmentationFilesInput):
    label_field: str

    def load_data(
        self,
        sample_collection: SampleCollection,
        label_field: str = "ground_truth",
        num_classes: Optional[int] = None,
        labels_map: Optional[Dict[int, Tuple[int, int, int]]] = None,
    ) -> List[Dict[str, Any]]:
        self.load_labels_map(num_classes, labels_map)

        self.label_field = label_field
        label_utilities = FiftyOneLabelUtilities(label_field, fo.Segmentation)
        label_utilities.validate(sample_collection)

        self._fo_dataset_name = sample_collection.name
        return to_samples(sample_collection.values("filepath"))

    def predict_load_data(
        self,
        sample_collection: SampleCollection,
    ) -> List[Dict[str, Any]]:
        return to_samples(sample_collection.values("filepath"))

    def load_sample(self, sample: Dict[str, Any]) -> Dict[str, Any]:
        filepath = sample[DataKeys.INPUT]
        sample = super().load_sample(sample)
        if not self.predicting:
            fo_dataset = fo.load_dataset(self._fo_dataset_name)
            fo_sample = fo_dataset[filepath]
            sample[DataKeys.TARGET] = fo_sample[self.label_field].mask
        return sample

src/flash/image/segmentation/input.py
diff --git a/src/flash/image/segmentation/input.py b/src/flash/image/segmentation/input.py
index 4662e1c986..02068c96be 100644
--- a/src/flash/image/segmentation/input.py
+++ b/src/flash/image/segmentation/input.py
@@ -94,7 +94,7 @@ def load_data(
 
     def load_sample(self, sample: Dict[str, Any]) -> Dict[str, Any]:
         if DataKeys.TARGET in sample:
-            sample[DataKeys.TARGET] = np.array(load_image(sample[DataKeys.TARGET])).transpose((2, 0, 1))[:, :, 0]
+            sample[DataKeys.TARGET] = np.array(load_image(sample[DataKeys.TARGET])).transpose((2, 0, 1))[0, :, :]
         return super().load_sample(sample)
 
 
diff --git a/tests/image/semantic_segm/test_data.py b/tests/image/semantic_segm/test_data.py
index 0c07e99fb0..aae5129c90 100644
--- a/tests/image/semantic_segm/test_data.py
+++ b/tests/image/semantic_segm/test_data.py
@@ -1,6 +1,6 @@
 import os
 from pathlib import Path
-from typing import Dict, List, Tuple
+from typing import Callable, Dict, List, Tuple
 
 import numpy as np
 import pytest
@@ -8,6 +8,8 @@
 
 from flash import Trainer
 from flash.core.data.io.input import DataKeys
+from flash.core.data.io.input_transform import InputTransform
+from flash.core.data.transforms import ApplyToKeys
 from flash.core.utilities.imports import (
     _FIFTYONE_AVAILABLE,
     _MATPLOTLIB_AVAILABLE,
@@ -43,12 +45,22 @@ def _rand_labels(size: Tuple[int, int], num_classes: int):
     return Image.fromarray(data.astype(np.uint8))
 
 
-def create_random_data(image_files: List[str], label_files: List[str], size: Tuple[int, int], num_classes: int):
+def create_random_data(
+    image_files: List[str], label_files: List[str], size: Tuple[int, int], num_classes: int
+) -> Tuple[List[Image.Image], List[Image.Image]]:
+    imgs = []
     for img_file in image_files:
-        _rand_image(size).save(img_file)
+        img = _rand_image(size)
+        img.save(img_file)
+        imgs.append(img)
 
+    labels = []
     for label_file in label_files:
-        _rand_labels(size, num_classes).save(label_file)
+        label = _rand_labels(size, num_classes)
+        label.save(label_file)
+        labels.append(label)
+
+    return imgs, labels
 
 
 @pytest.mark.skipif(not _TOPIC_IMAGE_AVAILABLE, reason="image libraries aren't installed.")
@@ -58,6 +70,56 @@ def test_smoke():
         dm = SemanticSegmentationData(batch_size=1)
         assert dm is not None
 
+    @staticmethod
+    @pytest.mark.skipif(not _TOPIC_IMAGE_AVAILABLE, reason="image libraries aren't installed.")
+    def test_identity(tmpdir):
+        class IdentityTransform(InputTransform):
+            def per_sample_transform(self) -> Callable:
+                return ApplyToKeys(
+                    DataKeys.INPUT,
+                    np.array,
+                )
+
+            def per_batch_transform(self) -> Callable:
+                return lambda x: x
+
+        tmp_dir = Path(tmpdir)
+
+        # create random dummy data
+
+        os.makedirs(str(tmp_dir / "images"))
+        os.makedirs(str(tmp_dir / "targets"))
+
+        images = [str(tmp_dir / "images" / "img1.png")]
+
+        targets = [str(tmp_dir / "targets" / "img1.png")]
+
+        num_classes: int = 2
+        img_size: Tuple[int, int] = (128, 128)
+        images_data, targets_data = create_random_data(images, targets, img_size, num_classes)
+
+        # instantiate the data module
+
+        dm = SemanticSegmentationData.from_files(
+            test_files=images,
+            test_targets=targets,
+            batch_size=1,
+            num_workers=0,
+            num_classes=num_classes,
+            transform=IdentityTransform(),
+        )
+
+        assert dm is not None
+        assert dm.test_dataloader() is not None
+
+        # check test data
+        data = next(iter(dm.test_dataloader()))
+        imgs, labels = data[DataKeys.INPUT], data[DataKeys.TARGET]
+        assert imgs.shape == (1, 128, 128, 3)
+        assert labels.shape == (1, 128, 128)
+        assert torch.allclose(imgs, torch.from_numpy(np.array(images_data[0])))
+        assert torch.allclose(labels, torch.from_numpy(np.array(targets_data[0]))[:, :, 0])
+
     @staticmethod
     def test_from_folders(tmpdir):
         tmp_dir = Path(tmpdir)

Semantic Segmentation target masks broken >0.7.5
## 🐛 Bug

The switch to albumentation in newer releases of lightning-flash seem to have broken transformation of segmentation targets.

This is what I expect masks to look like (screenshot showing below code sample running on 0.7.5):

![2022-11-25 25-10-10-49--259_chrome](https://user-images.githubusercontent.com/9672/203894021-42a3b66f-ab8a-4216-9c66-7cd03e58a595.png)

This is what it looks like on the latest release (0.8.1):

![2022-11-25 25-13-13-32--468_chrome](https://user-images.githubusercontent.com/9672/203894038-e4d3b942-b6f9-4932-ba44-b30b6506edd4.png)

### To Reproduce

Run the below sample with lightning-flash=0.7.5 and lightning-flash=0.8.1 and compare behavior. 

#### Code sample

```py
import torch

import flash
from flash.core.data.utils import download_data
from flash.image import SemanticSegmentation, SemanticSegmentationData

import matplotlib.pyplot as plt
import numpy as np

# 1. Create the DataModule
# The data was generated with the  CARLA self-driving simulator as part of the Kaggle Lyft Udacity Challenge.
# More info here: https://www.kaggle.com/kumaresanmanickavelu/lyft-udacity-challenge
# download_data(
#     "https://github.com/ongchinkiat/LyftPerceptionChallenge/releases/download/v0.1/carla-capture-20180513A.zip",
#     "./data",
# )

datamodule = SemanticSegmentationData.from_folders(
    train_folder="data/CameraRGB",
    train_target_folder="data/CameraSeg",
    val_split=0.1,
    transform_kwargs=dict(image_size=(256, 256)),
    num_classes=21,
    batch_size=4,
)

# 2. Build the task
model = SemanticSegmentation(
    backbone="mobilenetv3_large_100",
    head="fpn",
    num_classes=datamodule.num_classes,
)

n = 3
fig, axarr = plt.subplots(ncols=2, nrows=n, figsize=(8, 4*n))

for batch in datamodule.train_dataloader():
    print(batch.keys())
    for i in range(n):
        segm = batch['target'][i]
        print(segm.shape)
        img = np.rollaxis(batch['input'][i].numpy(), 0, 3)
        axarr[i, 0].imshow(img)
        axarr[i, 1].imshow(segm)
    break
```

### Environment

 - OS: Ubuntu WSL2
 - Python version: 3.10.8
 - GPU model: RTX 3080
 - CUDA Version: 11.6
