Add semantic segmentation task

flash/vision/segmentation/data.py

@@ -62,10 +62,20 @@ def load_data(self, data: Union[Tuple[str, str], Tuple[List[str], List[str]]]) -
 
             input_data = input_files
             target_data = target_files
+        else:
+            if not isinstance(input_data, list) and not isinstance(target_data, list):
+                input_data = [input_data]
+                target_data = [target_data]
 
-        if not isinstance(input_data, list) and not isinstance(target_data, list):
-            input_data = [input_data]
-            target_data = [target_data]
+            if len(input_data) != len(target_data):
+                rank_zero_warn(
+                    f"The number of input files ({len(input_data)}) and number of target files ({len(target_data)}) are"
+                    " different. Some files have been dropped.",
+                    UserWarning,
+                )
+                length = min(len(input_data), len(target_data))
+                input_data = input_data[:length]
+                target_data = target_data[:length]
 
         data = filter(
             lambda sample: (

tests/vision/segmentation/test_data.py

@@ -140,36 +140,18 @@ def test_from_folders_warning(self, tmpdir):
             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=1,
                 num_workers=0,
             )
         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[DefaultDataKeys.INPUT], data[DefaultDataKeys.TARGET]
         assert imgs.shape == (1, 3, 196, 196)
         assert labels.shape == (1, 196, 196)
 
-        # check val data
-        data = next(iter(dm.val_dataloader()))
-        imgs, labels = data[DefaultDataKeys.INPUT], data[DefaultDataKeys.TARGET]
-        assert imgs.shape == (1, 3, 196, 196)
-        assert labels.shape == (1, 196, 196)
-
-        # check test data
-        data = next(iter(dm.test_dataloader()))
-        imgs, labels = data[DefaultDataKeys.INPUT], data[DefaultDataKeys.TARGET]
-        assert imgs.shape == (1, 3, 196, 196)
-        assert labels.shape == (1, 196, 196)
-
     def test_from_files(self, tmpdir):
         tmp_dir = Path(tmpdir)
 
@@ -226,6 +208,45 @@ def test_from_files(self, tmpdir):
         assert imgs.shape == (2, 3, 196, 196)
         assert labels.shape == (2, 196, 196)
 
+    def test_from_files_warning(self, tmpdir):
+        tmp_dir = Path(tmpdir)
+
+        # create random dummy data
+
+        images = [
+            str(tmp_dir / "img1.png"),
+            str(tmp_dir / "img2.png"),
+            str(tmp_dir / "img3.png"),
+        ]
+
+        targets = [
+            str(tmp_dir / "labels_img1.png"),
+            str(tmp_dir / "labels_img2.png"),
+            str(tmp_dir / "labels_img3.png"),
+        ]
+
+        num_classes: int = 2
+        img_size: Tuple[int, int] = (196, 196)
+        create_random_data(images, targets, img_size, num_classes)
+
+        # instantiate the data module
+
+        with pytest.warns(UserWarning, match="The number of input files"):
+            dm = SemanticSegmentationData.from_files(
+                train_files=images,
+                train_targets=targets + [str(tmp_dir / "labels_img4.png")],
+                batch_size=2,
+                num_workers=0,
+            )
+        assert dm is not None
+        assert dm.train_dataloader() is not None
+
+        # check training data
+        data = next(iter(dm.train_dataloader()))
+        imgs, labels = data[DefaultDataKeys.INPUT], data[DefaultDataKeys.TARGET]
+        assert imgs.shape == (2, 3, 196, 196)
+        assert labels.shape == (2, 196, 196)
+
     def test_map_labels(self, tmpdir):
         tmp_dir = Path(tmpdir)