Repository: Lightning-Universe/lightning-flash

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


Discussion:
I can confirm this issue. Just spent a significant amount of time trying to figure out if there is anything wrong with my code. Any estimations when this is going to be fixed?
Hi, @newzealandpaul @noname202 we are sorry for this bug which seems to be very critical for segmentations... :( Would you be interested in trying to debug it and I belive that @ethanwharris could eventually give a hand... :rabbit: 

Resolving patch:
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_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)
 
 


Test patch:
diff --git a/tests/image/semantic_segm/test_data.py b/tests/image/semantic_segm/test_data.py
--- 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)