Repository: Lightning-Universe/lightning-flash

Issue:
Support targets file names with different file extension than x file names 
## 🚀 Feature
<!-- A clear and concise description of the feature proposal -->
The feature consists in supporting the use case where we have train/val/test images with an extension and train/val/test targets with a different extension.

eg: 
### train_folder
- File extension : `.jpg`  

### train_target_folder  
- File extension : `.tiff`  

### val_folder  
- File extension : `.jpg`  

### val_target_folder  
- File extension : `.tiff`  

### Motivation
This is a common use case that should be supported in order to open Lightning Flash to a broader audience.  
This also reduces the barrier to try out Flash by eliminating the need to convert your data to a matching file extension.  

Currently this is not supported if your target file paths have a different file extension than the X file paths : 
```
data_module = SemanticSegmentationData.from_folders(
    train_folder=train_images.path,
    train_target_folder=train_masks.path,
    val_folder=valid_images.path,
    val_target_folder=valid_masks.path,
    num_classes=number_of_classes,
    batch_size=batch_size
)
```
Actual Outcome : `ValueError: not enough values to unpack (expected 2, got 0)`
Expected Outcome : `No error message and it works`
<!-- Please outline the motivation for the proposal. Is your feature request related to a problem? e.g., I'm always frustrated when [...]. If this is related to another GitHub issue, please link here too -->

### Additional context
See slack discussion that led to this issue https://pytorch-lightning.slack.com/archives/CRBLFHY79/p1642626657155300
<!-- Add any other context or screenshots about the feature request here. -->



Discussion:

Resolving patch:
diff --git a/flash/image/segmentation/data.py b/flash/image/segmentation/data.py
--- a/flash/image/segmentation/data.py
+++ b/flash/image/segmentation/data.py
@@ -203,7 +203,8 @@ def from_folders(
             ├── image_3.png
             ...
 
-        your ``train_masks`` folder (passed to the ``train_target_folder`` argument) would need to look like this:
+        your ``train_masks`` folder (passed to the ``train_target_folder`` argument) would need to look like this
+        (although the file extensions could be different):
 
         .. code-block::
 
diff --git a/flash/image/segmentation/input.py b/flash/image/segmentation/input.py
--- a/flash/image/segmentation/input.py
+++ b/flash/image/segmentation/input.py
@@ -15,7 +15,6 @@
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import torch
-from pytorch_lightning.utilities import rank_zero_warn
 
 from flash.core.data.io.input import DataKeys, ImageLabelsMap, Input
 from flash.core.data.utilities.paths import filter_valid_files, PATH_TYPE
@@ -33,8 +32,7 @@
     SampleCollection = None
 
 if _TORCHVISION_AVAILABLE:
-    import torchvision
-    import torchvision.transforms.functional as FT
+    from torchvision.transforms.functional import to_tensor
 
 
 class SemanticSegmentationInput(Input):
@@ -104,9 +102,9 @@ def load_data(
 
     def load_sample(self, sample: Dict[str, Any]) -> Dict[str, Any]:
         filepath = sample[DataKeys.INPUT]
-        sample[DataKeys.INPUT] = FT.to_tensor(image_loader(filepath))
+        sample[DataKeys.INPUT] = to_tensor(image_loader(filepath))
         if DataKeys.TARGET in sample:
-            sample[DataKeys.TARGET] = torchvision.io.read_image(sample[DataKeys.TARGET])[0]
+            sample[DataKeys.TARGET] = (to_tensor(image_loader(sample[DataKeys.TARGET])) * 255).long()[0]
         sample = super().load_sample(sample)
         sample[DataKeys.METADATA]["filepath"] = filepath
         return sample
@@ -124,20 +122,17 @@ def load_data(
         files = os.listdir(folder)
         files.sort()
         if mask_folder is not None:
-            mask_files = os.listdir(mask_folder)
-
-            all_files = set(files).intersection(set(mask_files))
-            if len(all_files) != len(files) or len(all_files) != len(mask_files):
-                rank_zero_warn(
-                    f"Found inconsistent files in input folder: {folder} and mask folder: {mask_folder}. Some files"
-                    " have been dropped.",
-                    UserWarning,
+            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 all_files]
-            mask_files = [os.path.join(mask_folder, file) for file in all_files]
-            files.sort()
-            mask_files.sort()
+            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])
 
@@ -172,6 +167,6 @@ def load_sample(self, sample: Dict[str, Any]) -> Dict[str, Any]:
 class SemanticSegmentationDeserializer(ImageDeserializer):
     def serve_load_sample(self, data: str) -> Dict[str, Any]:
         result = super().serve_load_sample(data)
-        result[DataKeys.INPUT] = FT.to_tensor(result[DataKeys.INPUT])
+        result[DataKeys.INPUT] = to_tensor(result[DataKeys.INPUT])
         result[DataKeys.METADATA] = {"size": result[DataKeys.INPUT].shape[-2:]}
         return result


Test patch:
diff --git a/tests/image/segmentation/test_data.py b/tests/image/segmentation/test_data.py
--- a/tests/image/segmentation/test_data.py
+++ b/tests/image/segmentation/test_data.py
@@ -126,7 +126,69 @@ def test_from_folders(tmpdir):
 
     @staticmethod
     @pytest.mark.skipif(not _IMAGE_TESTING, reason="image libraries aren't installed.")
-    def test_from_folders_warning(tmpdir):
+    def test_from_folders_different_extensions(tmpdir):
+        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"),
+            str(tmp_dir / "images" / "img2.png"),
+            str(tmp_dir / "images" / "img3.png"),
+        ]
+
+        targets = [
+            str(tmp_dir / "targets" / "img1.tiff"),
+            str(tmp_dir / "targets" / "img2.tiff"),
+            str(tmp_dir / "targets" / "img3.tiff"),
+        ]
+
+        num_classes: int = 2
+        img_size: Tuple[int, int] = (128, 128)
+        create_random_data(images, targets, img_size, num_classes)
+
+        # instantiate the data module
+
+        dm = SemanticSegmentationData.from_folders(
+            train_folder=str(tmp_dir / "images"),
+            train_target_folder=str(tmp_dir / "targets"),
+            val_folder=str(tmp_dir / "images"),
+            val_target_folder=str(tmp_dir / "targets"),
+            test_folder=str(tmp_dir / "images"),
+            test_target_folder=str(tmp_dir / "targets"),
+            batch_size=2,
+            num_workers=0,
+            num_classes=num_classes,
+        )
+        assert dm is not None
+        assert dm.train_dataloader() is not None
+        assert dm.val_dataloader() is not None
+        assert dm.test_dataloader() is not None
+
+        # check training data
+        data = next(iter(dm.train_dataloader()))
+        imgs, labels = data[DataKeys.INPUT], data[DataKeys.TARGET]
+        assert imgs.shape == (2, 3, 128, 128)
+        assert labels.shape == (2, 128, 128)
+
+        # check val data
+        data = next(iter(dm.val_dataloader()))
+        imgs, labels = data[DataKeys.INPUT], data[DataKeys.TARGET]
+        assert imgs.shape == (2, 3, 128, 128)
+        assert labels.shape == (2, 128, 128)
+
+        # check test data
+        data = next(iter(dm.test_dataloader()))
+        imgs, labels = data[DataKeys.INPUT], data[DataKeys.TARGET]
+        assert imgs.shape == (2, 3, 128, 128)
+        assert labels.shape == (2, 128, 128)
+
+    @staticmethod
+    @pytest.mark.skipif(not _IMAGE_TESTING, reason="image libraries aren't installed.")
+    def test_from_folders_error(tmpdir):
         tmp_dir = Path(tmpdir)
 
         # create random dummy data
@@ -150,22 +212,14 @@ def test_from_folders_warning(tmpdir):
 
         # instantiate the data module
 
-        with pytest.warns(UserWarning, match="Found inconsistent files"):
-            dm = SemanticSegmentationData.from_folders(
+        with pytest.raises(ValueError, match="Found inconsistent files"):
+            SemanticSegmentationData.from_folders(
                 train_folder=str(tmp_dir / "images"),
                 train_target_folder=str(tmp_dir / "targets"),
                 batch_size=1,
                 num_workers=0,
                 num_classes=num_classes,
             )
-        assert dm is not None
-        assert dm.train_dataloader() is not None
-
-        # check training data
-        data = next(iter(dm.train_dataloader()))
-        imgs, labels = data[DataKeys.INPUT], data[DataKeys.TARGET]
-        assert imgs.shape == (1, 3, 128, 128)
-        assert labels.shape == (1, 128, 128)
 
     @staticmethod
     @pytest.mark.skipif(not _IMAGE_TESTING, reason="image libraries aren't installed.")