[feat] Update Segmentation Task in preparation to adding resizing

flash/__init__.py

@@ -18,6 +18,7 @@
 
 _PACKAGE_ROOT = os.path.dirname(__file__)
 PROJECT_ROOT = os.path.dirname(_PACKAGE_ROOT)
+_IS_TESTING = os.getenv("FLASH_TESTING", "0") == "1"
 
 from flash.core.model import Task  # noqa: E402
 from flash.core.trainer import Trainer  # noqa: E402

flash/data/data_pipeline.py

@@ -61,6 +61,9 @@ def get_state(self, state_type: Type[ProcessState]) -> Optional[ProcessState]:
         else:
             return None
 
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(initialized={self._initialized}, state={self._state})"
+
 
 class DataPipeline:
     """

flash/data/data_source.py

@@ -44,6 +44,12 @@ class LabelsState(ProcessState):
     labels: Optional[Sequence[str]]
 
 
+@dataclass(unsafe_hash=True, frozen=True)
+class ImageLabelsMap(ProcessState):
+
+    labels_map: Optional[Dict[int, Tuple[int, int, int]]]
+
+
 class DefaultDataSources(LightningEnum):
     """The ``DefaultDataSources`` enum contains the data source names used by all of the default ``from_*`` methods in
     :class:`~flash.data.data_module.DataModule`."""
@@ -65,6 +71,7 @@ class DefaultDataKeys(LightningEnum):
 
     INPUT = "input"
     TARGET = "target"
+    METADATA = "metadata"
 
     # TODO: Create a FlashEnum class???
     def __hash__(self) -> int:

flash/data/splits.py

@@ -38,6 +38,11 @@ def __init__(self, dataset: Any, indices: List[int] = [], use_duplicated_indices
         self.dataset = dataset
         self.indices = indices
 
+    def __getattr__(self, key: str):
+        if key in ("dataset", "indices", "data"):
+            return getattr(self, key)
+        return getattr(self.dataset, key)
+
     def __getitem__(self, index: int) -> Any:
         return self.dataset[self.indices[index]]
 

flash/vision/segmentation/data.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import os
+from dataclasses import dataclass
 from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple, Union
 
 import numpy as np
@@ -23,14 +24,18 @@
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from torchvision.datasets.folder import has_file_allowed_extension, IMG_EXTENSIONS
 
+import flash
+from flash.data.auto_dataset import BaseAutoDataset
 from flash.data.base_viz import BaseVisualization  # for viz
 from flash.data.callback import BaseDataFetcher
 from flash.data.data_module import DataModule
 from flash.data.data_source import (
     DefaultDataKeys,
     DefaultDataSources,
+    ImageLabelsMap,
     NumpyDataSource,
     PathsDataSource,
+    SEQUENCE_DATA_TYPE,
     TensorDataSource,
 )
 from flash.data.process import Preprocess
@@ -44,19 +49,38 @@
     plt = None
 
 
+def create_random_labels_map(num_classes: int) -> Dict[int, Tuple[int, int, int]]:
+    labels_map: Dict[int, Tuple[int, int, int]] = {}
+    for i in range(num_classes):
+        labels_map[i] = torch.randint(0, 255, (3, ))
+    return labels_map
+
+
 class SemanticSegmentationNumpyDataSource(NumpyDataSource):
 
+    def __init__(self, num_classes: int):
+        self.num_classes = num_classes
+
+    def load_data(self, data: Tuple[Sequence[SEQUENCE_DATA_TYPE], Optional[Sequence]],
+                  dataset: Optional[Any]) -> Sequence[Mapping[str, Any]]:
+        data = super().load_data(data, dataset=dataset)
+        if self.training:
+            dataset.num_classes = self.num_classes
+        return data
+
     def load_sample(self, sample: Dict[str, Any], dataset: Optional[Any] = None) -> Dict[str, Any]:
         sample[DefaultDataKeys.INPUT] = torch.from_numpy(sample[DefaultDataKeys.INPUT]).float()
         return sample
 
 
 class SemanticSegmentationPathsDataSource(PathsDataSource):
 
-    def __init__(self):
+    def __init__(self, num_classes: int):
         super().__init__(IMG_EXTENSIONS)
+        self.num_classes = num_classes
 
-    def load_data(self, data: Union[Tuple[str, str], Tuple[List[str], List[str]]]) -> Sequence[Mapping[str, Any]]:
+    def load_data(self, data: Union[Tuple[str, str], Tuple[List[str], List[str]]],
+                  dataset: BaseAutoDataset) -> Sequence[Mapping[str, Any]]:
         input_data, target_data = data
 
         if self.isdir(input_data) and self.isdir(target_data):
@@ -93,6 +117,9 @@ def load_data(self, data: Union[Tuple[str, str], Tuple[List[str], List[str]]]) -
             zip(input_data, target_data),
         )
 
+        if self.training:
+            dataset.num_classes = self.num_classes
+
         return [{DefaultDataKeys.INPUT: input, DefaultDataKeys.TARGET: target} for input, target in data]
 
     def predict_load_data(self, data: Union[str, List[str]]):
@@ -108,7 +135,11 @@ def load_sample(self, sample: Mapping[str, Any]) -> Mapping[str, torch.Tensor]:
         img_labels: torch.Tensor = torchvision.io.read_image(img_labels_path)  # CxHxW
         img_labels = img_labels[0]  # HxW
 
-        return {DefaultDataKeys.INPUT: img.float(), DefaultDataKeys.TARGET: img_labels.float()}
+        return {
+            DefaultDataKeys.INPUT: img.float(),
+            DefaultDataKeys.TARGET: img_labels.float(),
+            DefaultDataKeys.TARGET: img_labels.float()
+        }
 
     def predict_load_sample(self, sample: Mapping[str, Any]) -> Mapping[str, Any]:
         return {DefaultDataKeys.INPUT: torchvision.io.read_image(sample[DefaultDataKeys.INPUT]).float()}
@@ -123,6 +154,8 @@ def __init__(
         test_transform: Optional[Dict[str, Callable]] = None,
         predict_transform: Optional[Dict[str, Callable]] = None,
         image_size: Tuple[int, int] = (196, 196),
+        num_classes: int = None,
+        labels_map: Dict[int, Tuple[int, int, int]] = None,
     ) -> None:
         """Preprocess pipeline for semantic segmentation tasks.
 
@@ -133,25 +166,36 @@ def __init__(
             predict_transform: Dictionary with the set of transforms to apply during prediction.
             image_size: A tuple with the expected output image size.
         """
+        if not num_classes or not isinstance(num_classes, int):
+            raise MisconfigurationException("`num_classes` should be provided for instantiation.")
+
+        if not labels_map:
+            raise MisconfigurationException("`labels_map` should be provided for instantiation.")
+
         self.image_size = image_size
+        self.num_classes = num_classes
+        self.labels_map = labels_map
 
         super().__init__(
             train_transform=train_transform,
             val_transform=val_transform,
             test_transform=test_transform,
             predict_transform=predict_transform,
             data_sources={
-                DefaultDataSources.PATHS: SemanticSegmentationPathsDataSource(),
+                DefaultDataSources.PATHS: SemanticSegmentationPathsDataSource(num_classes),
                 DefaultDataSources.TENSOR: TensorDataSource(),
-                DefaultDataSources.NUMPY: SemanticSegmentationNumpyDataSource(),
+                DefaultDataSources.NUMPY: SemanticSegmentationNumpyDataSource(num_classes),
             },
             default_data_source=DefaultDataSources.PATHS,
         )
 
+        self.set_state(ImageLabelsMap(labels_map))
+
     def get_state_dict(self) -> Dict[str, Any]:
         return {
-            **self.transforms,
-            "image_size": self.image_size,
+            **self.transforms, "image_size": self.image_size,
+            "num_classes": self.num_classes,
+            "labels_map": self.labels_map
         }
 
     @classmethod
@@ -184,9 +228,6 @@ class SemanticSegmentationData(DataModule):
     def configure_data_fetcher(*args, **kwargs) -> BaseDataFetcher:
         return SegmentationMatplotlibVisualization(*args, **kwargs)
 
-    def set_labels_map(self, labels_map: Dict[int, Tuple[int, int, int]]):
-        self.data_fetcher.labels_map = labels_map
-
     def set_block_viz_window(self, value: bool) -> None:
         """Setter method to switch on/off matplotlib to pop up windows."""
         self.data_fetcher.block_viz_window = value
@@ -210,7 +251,9 @@ def from_folders(
         val_split: Optional[float] = None,
         batch_size: int = 4,
         num_workers: Optional[int] = None,
-        **preprocess_kwargs: Any,
+        num_classes: Optional[int] = None,
+        labels_map: Dict[int, Tuple[int, int, int]] = None,
+        **preprocess_kwargs,
     ) -> 'DataModule':
         """Creates a :class:`~flash.vision.segmentation.data.SemanticSegmentationData` object from the given data
         folders and corresponding target folders.
@@ -242,6 +285,8 @@ def from_folders(
             val_split: The ``val_split`` argument to pass to the :class:`~flash.data.data_module.DataModule`.
             batch_size: The ``batch_size`` argument to pass to the :class:`~flash.data.data_module.DataModule`.
             num_workers: The ``num_workers`` argument to pass to the :class:`~flash.data.data_module.DataModule`.
+            num_classes: Number of classes within the segmentation mask.
+            labels_map: Mapping between a class_id and its corresponding color.
             preprocess_kwargs: Additional keyword arguments to use when constructing the preprocess. Will only be used
                 if ``preprocess = None``.
 
@@ -255,6 +300,17 @@ def from_folders(
                 train_target_folder="train_masks",
             )
         """
+
+        if not num_classes or not isinstance(num_classes, int):
+            raise MisconfigurationException("`num_classes` should be provided during instantiation.")
+
+        labels_map = labels_map or create_random_labels_map(num_classes)
+
+        data_fetcher = data_fetcher or cls.configure_data_fetcher(labels_map)
+
+        if flash._IS_TESTING:
+            data_fetcher.block_viz_window = True
+
         return cls.from_data_source(
             DefaultDataSources.PATHS,
             (train_folder, train_target_folder),
@@ -270,6 +326,8 @@ def from_folders(
             val_split=val_split,
             batch_size=batch_size,
             num_workers=num_workers,
+            num_classes=num_classes,
+            labels_map=labels_map,
             **preprocess_kwargs,
         )
 
@@ -278,11 +336,12 @@ class SegmentationMatplotlibVisualization(BaseVisualization):
     """Process and show the image batch and its associated label using matplotlib.
     """
 
-    def __init__(self):
-        super().__init__(self)
+    def __init__(self, labels_map: Dict[int, Tuple[int, int, int]]):
+        super().__init__()
+
         self.max_cols: int = 4  # maximum number of columns we accept
         self.block_viz_window: bool = True  # parameter to allow user to block visualisation windows
-        self.labels_map: Dict[int, Tuple[int, int, int]] = {}
+        self.labels_map: Dict[int, Tuple[int, int, int]] = labels_map
 
     @staticmethod
     def _to_numpy(img: Union[torch.Tensor, Image.Image]) -> np.ndarray:

flash/vision/segmentation/serialization.py

@@ -16,6 +16,8 @@
 
 import torch
 
+import flash
+from flash.data.data_source import ImageLabelsMap
 from flash.data.process import Serializer
 from flash.utils.imports import _KORNIA_AVAILABLE, _MATPLOTLIB_AVAILABLE
 
@@ -58,24 +60,14 @@ def labels_to_image(img_labels: torch.Tensor, labels_map: Dict[int, Tuple[int, i
                 out[i].masked_fill_(mask, label_val[i])
         return out
 
-    @staticmethod
-    def create_random_labels_map(num_classes: int) -> Dict[int, Tuple[int, int, int]]:
-        labels_map: Dict[int, Tuple[int, int, int]] = {}
-        for i in range(num_classes):
-            labels_map[i] = torch.randint(0, 255, (3, ))
-        return labels_map
-
     def serialize(self, sample: torch.Tensor) -> torch.Tensor:
         assert len(sample.shape) == 3, sample.shape
         labels = torch.argmax(sample, dim=-3)  # HxW
-        if self.visualize and os.getenv("FLASH_TESTING", "0") == "0":
+
+        if self.visualize and not flash._IS_TESTING:
             if self.labels_map is None:
-                # create random colors map
-                num_classes = sample.shape[-3]
-                labels_map = self.create_random_labels_map(num_classes)
-            else:
-                labels_map = self.labels_map
-            labels_vis = self.labels_to_image(labels, labels_map)
+                self.labels_map = self.get_state(ImageLabelsMap).labels_map
+            labels_vis = self.labels_to_image(labels, self.labels_map)
             labels_vis = K.utils.tensor_to_image(labels_vis)
             plt.imshow(labels_vis)
             plt.show()

flash_examples/finetuning/semantic_segmentation.py

@@ -31,17 +31,17 @@
     batch_size=4,
     val_split=0.3,
     image_size=(200, 200),  # (600, 800)
+    num_classes=21,
 )
 
 # 2.2 Visualise the samples
-labels_map = SegmentationLabels.create_random_labels_map(num_classes=21)
-datamodule.set_labels_map(labels_map)
 datamodule.show_train_batch(["load_sample", "post_tensor_transform"])
 
 # 3. Build the model
 model = SemanticSegmentation(
     backbone="torchvision/fcn_resnet50",
-    num_classes=21,
+    num_classes=datamodule.num_classes,
+    serializer=SegmentationLabels(visualize=True)
 )
 
 # 4. Create the trainer.
@@ -53,9 +53,6 @@
 # 5. Train the model
 trainer.finetune(model, datamodule=datamodule, strategy="freeze")
 
-# 6. Predict what's on a few images!
-model.serializer = SegmentationLabels(labels_map, visualize=True)
-
 predictions = model.predict([
     "data/CameraRGB/F61-1.png",
     "data/CameraRGB/F62-1.png",

flash_examples/predict/semantic_segmentation.py

@@ -24,9 +24,7 @@
 )
 
 # 2. Load the model from a checkpoint
-model = SemanticSegmentation.load_from_checkpoint(
-    "https://flash-weights.s3.amazonaws.com/semantic_segmentation_model.pt"
-)
+model = SemanticSegmentation.load_from_checkpoint("semantic_segmentation_model.pt")
 model.serializer = SegmentationLabels(visualize=True)
 
 # 3. Predict what's on a few images and visualize!

tests/data/test_split_dataset.py

@@ -37,3 +37,18 @@ def test_split_dataset(tmpdir):
 
     with pytest.raises(MisconfigurationException, match="[0, 99]"):
         SplitDataset(list(range(50)) + list(range(50)), indices=[-1], use_duplicated_indices=True)
+
+    class Dataset:
+
+        def __init__(self):
+            self.data = [0, 1, 2]
+            self.name = "something"
+
+        def __getitem__(self, index):
+            return self.data[index]
+
+        def __len__(self):
+            return len(self.data)
+
+    split_dataset = SplitDataset(Dataset(), indices=[0])
+    assert split_dataset.name == "something"