We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
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

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `src/flash/image/segmentation/input.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 import os
15 from typing import Any, Dict, List, Optional, Tuple, Union
16 
17 import numpy as np
18 
19 from flash.core.data.io.input import DataKeys, Input
20 from flash.core.data.utilities.loading import IMG_EXTENSIONS, NP_EXTENSIONS, load_image
21 from flash.core.data.utilities.paths import PATH_TYPE, filter_valid_files
22 from flash.core.data.utilities.samples import to_samples
23 from flash.core.integrations.fiftyone.utils import FiftyOneLabelUtilities
24 from flash.core.utilities.imports import _FIFTYONE_AVAILABLE, lazy_import
25 from flash.image.data import ImageFilesInput, ImageNumpyInput, ImageTensorInput
26 from flash.image.segmentation.output import SegmentationLabelsOutput
27 
28 if _FIFTYONE_AVAILABLE:
29     fo = lazy_import("fiftyone")
30     SampleCollection = "fiftyone.core.collections.SampleCollection"
31 else:
32     fo = None
33     SampleCollection = None
34 
35 
36 class SemanticSegmentationInput(Input):
37     num_classes: int
38     labels_map: Dict[int, Tuple[int, int, int]]
39 
40     def load_labels_map(
41         self, num_classes: Optional[int] = None, labels_map: Optional[Dict[int, Tuple[int, int, int]]] = None
42     ) -> None:
43         if num_classes is not None:
44             self.num_classes = num_classes
45             labels_map = labels_map or SegmentationLabelsOutput.create_random_labels_map(num_classes)
46 
47         if labels_map is not None:
48             self.labels_map = labels_map
49 
50 
51 class SemanticSegmentationTensorInput(SemanticSegmentationInput, ImageTensorInput):
52     def load_data(
53         self,
54         tensor: Any,
55         masks: Any = None,
56         num_classes: Optional[int] = None,
57         labels_map: Optional[Dict[int, Tuple[int, int, int]]] = None,
58     ) -> List[Dict[str, Any]]:
59         self.load_labels_map(num_classes, labels_map)
60         return to_samples(tensor, masks)
61 
62     def load_sample(self, sample: Dict[str, Any]) -> Dict[str, Any]:
63         if DataKeys.TARGET in sample:
64             sample[DataKeys.TARGET] = sample[DataKeys.TARGET].numpy()
65         return super().load_sample(sample)
66 
67 
68 class SemanticSegmentationNumpyInput(SemanticSegmentationInput, ImageNumpyInput):
69     def load_data(
70         self,
71         array: Any,
72         masks: Any = None,
73         num_classes: Optional[int] = None,
74         labels_map: Optional[Dict[int, Tuple[int, int, int]]] = None,
75     ) -> List[Dict[str, Any]]:
76         self.load_labels_map(num_classes, labels_map)
77         return to_samples(array, masks)
78 
79 
80 class SemanticSegmentationFilesInput(SemanticSegmentationInput, ImageFilesInput):
81     def load_data(
82         self,
83         files: Union[PATH_TYPE, List[PATH_TYPE]],
84         mask_files: Optional[Union[PATH_TYPE, List[PATH_TYPE]]] = None,
85         num_classes: Optional[int] = None,
86         labels_map: Optional[Dict[int, Tuple[int, int, int]]] = None,
87     ) -> List[Dict[str, Any]]:
88         self.load_labels_map(num_classes, labels_map)
89         if mask_files is None:
90             files = filter_valid_files(files, valid_extensions=IMG_EXTENSIONS + NP_EXTENSIONS)
91         else:
92             files, mask_files = filter_valid_files(files, mask_files, valid_extensions=IMG_EXTENSIONS + NP_EXTENSIONS)
93         return to_samples(files, mask_files)
94 
95     def load_sample(self, sample: Dict[str, Any]) -> Dict[str, Any]:
96         if DataKeys.TARGET in sample:
97             sample[DataKeys.TARGET] = np.array(load_image(sample[DataKeys.TARGET])).transpose((2, 0, 1))[:, :, 0]
98         return super().load_sample(sample)
99 
100 
101 class SemanticSegmentationFolderInput(SemanticSegmentationFilesInput):
102     def load_data(
103         self,
104         folder: PATH_TYPE,
105         mask_folder: Optional[PATH_TYPE] = None,
106         num_classes: Optional[int] = None,
107         labels_map: Optional[Dict[int, Tuple[int, int, int]]] = None,
108     ) -> List[Dict[str, Any]]:
109         self.load_labels_map(num_classes, labels_map)
110         files = os.listdir(folder)
111         files.sort()
112         if mask_folder is not None:
113             mask_files = {os.path.splitext(file)[0]: file for file in os.listdir(mask_folder)}
114             file_names = [os.path.splitext(file)[0] for file in files]
115 
116             if len(set(file_names) - mask_files.keys()) != 0:
117                 raise ValueError(
118                     f"Found inconsistent files in input folder: {folder} and mask folder: {mask_folder}. All input "
119                     f"files must have a corresponding mask file with the same name."
120                 )
121 
122             files = [os.path.join(folder, file) for file in files]
123             mask_files = [os.path.join(mask_folder, mask_files[file_name]) for file_name in file_names]
124             return super().load_data(files, mask_files)
125         return super().load_data([os.path.join(folder, file) for file in files])
126 
127 
128 class SemanticSegmentationFiftyOneInput(SemanticSegmentationFilesInput):
129     label_field: str
130 
131     def load_data(
132         self,
133         sample_collection: SampleCollection,
134         label_field: str = "ground_truth",
135         num_classes: Optional[int] = None,
136         labels_map: Optional[Dict[int, Tuple[int, int, int]]] = None,
137     ) -> List[Dict[str, Any]]:
138         self.load_labels_map(num_classes, labels_map)
139 
140         self.label_field = label_field
141         label_utilities = FiftyOneLabelUtilities(label_field, fo.Segmentation)
142         label_utilities.validate(sample_collection)
143 
144         self._fo_dataset_name = sample_collection.name
145         return to_samples(sample_collection.values("filepath"))
146 
147     def predict_load_data(
148         self,
149         sample_collection: SampleCollection,
150     ) -> List[Dict[str, Any]]:
151         return to_samples(sample_collection.values("filepath"))
152 
153     def load_sample(self, sample: Dict[str, Any]) -> Dict[str, Any]:
154         filepath = sample[DataKeys.INPUT]
155         sample = super().load_sample(sample)
156         if not self.predicting:
157             fo_dataset = fo.load_dataset(self._fo_dataset_name)
158             fo_sample = fo_dataset[filepath]
159             sample[DataKeys.TARGET] = fo_sample[self.label_field].mask
160         return sample
161
```
--- 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()

```

diff --git a/src/flash/image/segmentation/input.py b/src/flash/image/segmentation/input.py
--- a/src/flash/image/segmentation/input.py
+++ b/src/flash/image/segmentation/input.py
@@ -94,7 +94,7 @@
 
     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)

{"golden_diff": "diff --git a/src/flash/image/segmentation/input.py b/src/flash/image/segmentation/input.py\n--- a/src/flash/image/segmentation/input.py\n+++ b/src/flash/image/segmentation/input.py\n@@ -94,7 +94,7 @@\n \n     def load_sample(self, sample: Dict[str, Any]) -> Dict[str, Any]:\n         if DataKeys.TARGET in sample:\n-            sample[DataKeys.TARGET] = np.array(load_image(sample[DataKeys.TARGET])).transpose((2, 0, 1))[:, :, 0]\n+            sample[DataKeys.TARGET] = np.array(load_image(sample[DataKeys.TARGET])).transpose((2, 0, 1))[0, :, :]\n         return super().load_sample(sample)\n", "issue": "Semantic Segmentation target masks broken >0.7.5\n## \ud83d\udc1b Bug\r\n\r\nThe switch to albumentation in newer releases of lightning-flash seem to have broken transformation of segmentation targets.\r\n\r\nThis is what I expect masks to look like (screenshot showing below code sample running on 0.7.5):\r\n\r\n![2022-11-25 25-10-10-49--259_chrome](https://user-images.githubusercontent.com/9672/203894021-42a3b66f-ab8a-4216-9c66-7cd03e58a595.png)\r\n\r\nThis is what it looks like on the latest release (0.8.1):\r\n\r\n![2022-11-25 25-13-13-32--468_chrome](https://user-images.githubusercontent.com/9672/203894038-e4d3b942-b6f9-4932-ba44-b30b6506edd4.png)\r\n\r\n### To Reproduce\r\n\r\nRun the below sample with lightning-flash=0.7.5 and lightning-flash=0.8.1 and compare behavior. \r\n\r\n#### Code sample\r\n\r\n```py\r\nimport torch\r\n\r\nimport flash\r\nfrom flash.core.data.utils import download_data\r\nfrom flash.image import SemanticSegmentation, SemanticSegmentationData\r\n\r\nimport matplotlib.pyplot as plt\r\nimport numpy as np\r\n\r\n# 1. Create the DataModule\r\n# The data was generated with the  CARLA self-driving simulator as part of the Kaggle Lyft Udacity Challenge.\r\n# More info here: https://www.kaggle.com/kumaresanmanickavelu/lyft-udacity-challenge\r\n# download_data(\r\n#     \"https://github.com/ongchinkiat/LyftPerceptionChallenge/releases/download/v0.1/carla-capture-20180513A.zip\",\r\n#     \"./data\",\r\n# )\r\n\r\ndatamodule = SemanticSegmentationData.from_folders(\r\n    train_folder=\"data/CameraRGB\",\r\n    train_target_folder=\"data/CameraSeg\",\r\n    val_split=0.1,\r\n    transform_kwargs=dict(image_size=(256, 256)),\r\n    num_classes=21,\r\n    batch_size=4,\r\n)\r\n\r\n# 2. Build the task\r\nmodel = SemanticSegmentation(\r\n    backbone=\"mobilenetv3_large_100\",\r\n    head=\"fpn\",\r\n    num_classes=datamodule.num_classes,\r\n)\r\n\r\nn = 3\r\nfig, axarr = plt.subplots(ncols=2, nrows=n, figsize=(8, 4*n))\r\n\r\nfor batch in datamodule.train_dataloader():\r\n    print(batch.keys())\r\n    for i in range(n):\r\n        segm = batch['target'][i]\r\n        print(segm.shape)\r\n        img = np.rollaxis(batch['input'][i].numpy(), 0, 3)\r\n        axarr[i, 0].imshow(img)\r\n        axarr[i, 1].imshow(segm)\r\n    break\r\n```\r\n\r\n### Environment\r\n\r\n - OS: Ubuntu WSL2\r\n - Python version: 3.10.8\r\n - GPU model: RTX 3080\r\n - CUDA Version: 11.6\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.\nimport os\nfrom typing import Any, Dict, List, Optional, Tuple, Union\n\nimport numpy as np\n\nfrom flash.core.data.io.input import DataKeys, Input\nfrom flash.core.data.utilities.loading import IMG_EXTENSIONS, NP_EXTENSIONS, load_image\nfrom flash.core.data.utilities.paths import PATH_TYPE, filter_valid_files\nfrom flash.core.data.utilities.samples import to_samples\nfrom flash.core.integrations.fiftyone.utils import FiftyOneLabelUtilities\nfrom flash.core.utilities.imports import _FIFTYONE_AVAILABLE, lazy_import\nfrom flash.image.data import ImageFilesInput, ImageNumpyInput, ImageTensorInput\nfrom flash.image.segmentation.output import SegmentationLabelsOutput\n\nif _FIFTYONE_AVAILABLE:\n    fo = lazy_import(\"fiftyone\")\n    SampleCollection = \"fiftyone.core.collections.SampleCollection\"\nelse:\n    fo = None\n    SampleCollection = None\n\n\nclass SemanticSegmentationInput(Input):\n    num_classes: int\n    labels_map: Dict[int, Tuple[int, int, int]]\n\n    def load_labels_map(\n        self, num_classes: Optional[int] = None, labels_map: Optional[Dict[int, Tuple[int, int, int]]] = None\n    ) -> None:\n        if num_classes is not None:\n            self.num_classes = num_classes\n            labels_map = labels_map or SegmentationLabelsOutput.create_random_labels_map(num_classes)\n\n        if labels_map is not None:\n            self.labels_map = labels_map\n\n\nclass SemanticSegmentationTensorInput(SemanticSegmentationInput, ImageTensorInput):\n    def load_data(\n        self,\n        tensor: Any,\n        masks: Any = None,\n        num_classes: Optional[int] = None,\n        labels_map: Optional[Dict[int, Tuple[int, int, int]]] = None,\n    ) -> List[Dict[str, Any]]:\n        self.load_labels_map(num_classes, labels_map)\n        return to_samples(tensor, masks)\n\n    def load_sample(self, sample: Dict[str, Any]) -> Dict[str, Any]:\n        if DataKeys.TARGET in sample:\n            sample[DataKeys.TARGET] = sample[DataKeys.TARGET].numpy()\n        return super().load_sample(sample)\n\n\nclass SemanticSegmentationNumpyInput(SemanticSegmentationInput, ImageNumpyInput):\n    def load_data(\n        self,\n        array: Any,\n        masks: Any = None,\n        num_classes: Optional[int] = None,\n        labels_map: Optional[Dict[int, Tuple[int, int, int]]] = None,\n    ) -> List[Dict[str, Any]]:\n        self.load_labels_map(num_classes, labels_map)\n        return to_samples(array, masks)\n\n\nclass SemanticSegmentationFilesInput(SemanticSegmentationInput, ImageFilesInput):\n    def load_data(\n        self,\n        files: Union[PATH_TYPE, List[PATH_TYPE]],\n        mask_files: Optional[Union[PATH_TYPE, List[PATH_TYPE]]] = None,\n        num_classes: Optional[int] = None,\n        labels_map: Optional[Dict[int, Tuple[int, int, int]]] = None,\n    ) -> List[Dict[str, Any]]:\n        self.load_labels_map(num_classes, labels_map)\n        if mask_files is None:\n            files = filter_valid_files(files, valid_extensions=IMG_EXTENSIONS + NP_EXTENSIONS)\n        else:\n            files, mask_files = filter_valid_files(files, mask_files, valid_extensions=IMG_EXTENSIONS + NP_EXTENSIONS)\n        return to_samples(files, mask_files)\n\n    def load_sample(self, sample: Dict[str, Any]) -> Dict[str, Any]:\n        if DataKeys.TARGET in sample:\n            sample[DataKeys.TARGET] = np.array(load_image(sample[DataKeys.TARGET])).transpose((2, 0, 1))[:, :, 0]\n        return super().load_sample(sample)\n\n\nclass SemanticSegmentationFolderInput(SemanticSegmentationFilesInput):\n    def load_data(\n        self,\n        folder: PATH_TYPE,\n        mask_folder: Optional[PATH_TYPE] = None,\n        num_classes: Optional[int] = None,\n        labels_map: Optional[Dict[int, Tuple[int, int, int]]] = None,\n    ) -> List[Dict[str, Any]]:\n        self.load_labels_map(num_classes, labels_map)\n        files = os.listdir(folder)\n        files.sort()\n        if mask_folder is not None:\n            mask_files = {os.path.splitext(file)[0]: file for file in os.listdir(mask_folder)}\n            file_names = [os.path.splitext(file)[0] for file in files]\n\n            if len(set(file_names) - mask_files.keys()) != 0:\n                raise ValueError(\n                    f\"Found inconsistent files in input folder: {folder} and mask folder: {mask_folder}. All input \"\n                    f\"files must have a corresponding mask file with the same name.\"\n                )\n\n            files = [os.path.join(folder, file) for file in files]\n            mask_files = [os.path.join(mask_folder, mask_files[file_name]) for file_name in file_names]\n            return super().load_data(files, mask_files)\n        return super().load_data([os.path.join(folder, file) for file in files])\n\n\nclass SemanticSegmentationFiftyOneInput(SemanticSegmentationFilesInput):\n    label_field: str\n\n    def load_data(\n        self,\n        sample_collection: SampleCollection,\n        label_field: str = \"ground_truth\",\n        num_classes: Optional[int] = None,\n        labels_map: Optional[Dict[int, Tuple[int, int, int]]] = None,\n    ) -> List[Dict[str, Any]]:\n        self.load_labels_map(num_classes, labels_map)\n\n        self.label_field = label_field\n        label_utilities = FiftyOneLabelUtilities(label_field, fo.Segmentation)\n        label_utilities.validate(sample_collection)\n\n        self._fo_dataset_name = sample_collection.name\n        return to_samples(sample_collection.values(\"filepath\"))\n\n    def predict_load_data(\n        self,\n        sample_collection: SampleCollection,\n    ) -> List[Dict[str, Any]]:\n        return to_samples(sample_collection.values(\"filepath\"))\n\n    def load_sample(self, sample: Dict[str, Any]) -> Dict[str, Any]:\n        filepath = sample[DataKeys.INPUT]\n        sample = super().load_sample(sample)\n        if not self.predicting:\n            fo_dataset = fo.load_dataset(self._fo_dataset_name)\n            fo_sample = fo_dataset[filepath]\n            sample[DataKeys.TARGET] = fo_sample[self.label_field].mask\n        return sample\n", "path": "src/flash/image/segmentation/input.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.\nimport os\nfrom typing import Any, Dict, List, Optional, Tuple, Union\n\nimport numpy as np\n\nfrom flash.core.data.io.input import DataKeys, Input\nfrom flash.core.data.utilities.loading import IMG_EXTENSIONS, NP_EXTENSIONS, load_image\nfrom flash.core.data.utilities.paths import PATH_TYPE, filter_valid_files\nfrom flash.core.data.utilities.samples import to_samples\nfrom flash.core.integrations.fiftyone.utils import FiftyOneLabelUtilities\nfrom flash.core.utilities.imports import _FIFTYONE_AVAILABLE, lazy_import\nfrom flash.image.data import ImageFilesInput, ImageNumpyInput, ImageTensorInput\nfrom flash.image.segmentation.output import SegmentationLabelsOutput\n\nif _FIFTYONE_AVAILABLE:\n    fo = lazy_import(\"fiftyone\")\n    SampleCollection = \"fiftyone.core.collections.SampleCollection\"\nelse:\n    fo = None\n    SampleCollection = None\n\n\nclass SemanticSegmentationInput(Input):\n    num_classes: int\n    labels_map: Dict[int, Tuple[int, int, int]]\n\n    def load_labels_map(\n        self, num_classes: Optional[int] = None, labels_map: Optional[Dict[int, Tuple[int, int, int]]] = None\n    ) -> None:\n        if num_classes is not None:\n            self.num_classes = num_classes\n            labels_map = labels_map or SegmentationLabelsOutput.create_random_labels_map(num_classes)\n\n        if labels_map is not None:\n            self.labels_map = labels_map\n\n\nclass SemanticSegmentationTensorInput(SemanticSegmentationInput, ImageTensorInput):\n    def load_data(\n        self,\n        tensor: Any,\n        masks: Any = None,\n        num_classes: Optional[int] = None,\n        labels_map: Optional[Dict[int, Tuple[int, int, int]]] = None,\n    ) -> List[Dict[str, Any]]:\n        self.load_labels_map(num_classes, labels_map)\n        return to_samples(tensor, masks)\n\n    def load_sample(self, sample: Dict[str, Any]) -> Dict[str, Any]:\n        if DataKeys.TARGET in sample:\n            sample[DataKeys.TARGET] = sample[DataKeys.TARGET].numpy()\n        return super().load_sample(sample)\n\n\nclass SemanticSegmentationNumpyInput(SemanticSegmentationInput, ImageNumpyInput):\n    def load_data(\n        self,\n        array: Any,\n        masks: Any = None,\n        num_classes: Optional[int] = None,\n        labels_map: Optional[Dict[int, Tuple[int, int, int]]] = None,\n    ) -> List[Dict[str, Any]]:\n        self.load_labels_map(num_classes, labels_map)\n        return to_samples(array, masks)\n\n\nclass SemanticSegmentationFilesInput(SemanticSegmentationInput, ImageFilesInput):\n    def load_data(\n        self,\n        files: Union[PATH_TYPE, List[PATH_TYPE]],\n        mask_files: Optional[Union[PATH_TYPE, List[PATH_TYPE]]] = None,\n        num_classes: Optional[int] = None,\n        labels_map: Optional[Dict[int, Tuple[int, int, int]]] = None,\n    ) -> List[Dict[str, Any]]:\n        self.load_labels_map(num_classes, labels_map)\n        if mask_files is None:\n            files = filter_valid_files(files, valid_extensions=IMG_EXTENSIONS + NP_EXTENSIONS)\n        else:\n            files, mask_files = filter_valid_files(files, mask_files, valid_extensions=IMG_EXTENSIONS + NP_EXTENSIONS)\n        return to_samples(files, mask_files)\n\n    def load_sample(self, sample: Dict[str, Any]) -> Dict[str, Any]:\n        if DataKeys.TARGET in sample:\n            sample[DataKeys.TARGET] = np.array(load_image(sample[DataKeys.TARGET])).transpose((2, 0, 1))[0, :, :]\n        return super().load_sample(sample)\n\n\nclass SemanticSegmentationFolderInput(SemanticSegmentationFilesInput):\n    def load_data(\n        self,\n        folder: PATH_TYPE,\n        mask_folder: Optional[PATH_TYPE] = None,\n        num_classes: Optional[int] = None,\n        labels_map: Optional[Dict[int, Tuple[int, int, int]]] = None,\n    ) -> List[Dict[str, Any]]:\n        self.load_labels_map(num_classes, labels_map)\n        files = os.listdir(folder)\n        files.sort()\n        if mask_folder is not None:\n            mask_files = {os.path.splitext(file)[0]: file for file in os.listdir(mask_folder)}\n            file_names = [os.path.splitext(file)[0] for file in files]\n\n            if len(set(file_names) - mask_files.keys()) != 0:\n                raise ValueError(\n                    f\"Found inconsistent files in input folder: {folder} and mask folder: {mask_folder}. All input \"\n                    f\"files must have a corresponding mask file with the same name.\"\n                )\n\n            files = [os.path.join(folder, file) for file in files]\n            mask_files = [os.path.join(mask_folder, mask_files[file_name]) for file_name in file_names]\n            return super().load_data(files, mask_files)\n        return super().load_data([os.path.join(folder, file) for file in files])\n\n\nclass SemanticSegmentationFiftyOneInput(SemanticSegmentationFilesInput):\n    label_field: str\n\n    def load_data(\n        self,\n        sample_collection: SampleCollection,\n        label_field: str = \"ground_truth\",\n        num_classes: Optional[int] = None,\n        labels_map: Optional[Dict[int, Tuple[int, int, int]]] = None,\n    ) -> List[Dict[str, Any]]:\n        self.load_labels_map(num_classes, labels_map)\n\n        self.label_field = label_field\n        label_utilities = FiftyOneLabelUtilities(label_field, fo.Segmentation)\n        label_utilities.validate(sample_collection)\n\n        self._fo_dataset_name = sample_collection.name\n        return to_samples(sample_collection.values(\"filepath\"))\n\n    def predict_load_data(\n        self,\n        sample_collection: SampleCollection,\n    ) -> List[Dict[str, Any]]:\n        return to_samples(sample_collection.values(\"filepath\"))\n\n    def load_sample(self, sample: Dict[str, Any]) -> Dict[str, Any]:\n        filepath = sample[DataKeys.INPUT]\n        sample = super().load_sample(sample)\n        if not self.predicting:\n            fo_dataset = fo.load_dataset(self._fo_dataset_name)\n            fo_sample = fo_dataset[filepath]\n            sample[DataKeys.TARGET] = fo_sample[self.label_field].mask\n        return sample\n", "path": "src/flash/image/segmentation/input.py"}]}