Repository: Lightning-Universe/lightning-flash

Issue:
flash.image.SemanticSegmentationData dataloaders to convert greyscale/RGBA to RGB for train and val source images.
## 🚀 Feature
Have The dataloaders convert any image type to rgb for use in segmentation.

### Motivation

greyscale or RGBA image types crash a training run



### Pitch

Most generic segmentation models work with RGB images [3 x length x width], some sources of images may contain either greyscale or RGBA image types. This simple modification would make the Flash segmentation more robust and easier to use. 

### Alternatives

Check image formats and give a clear error as to the the problem, current stack trace is not indicative of the problem.

### Additional context

<!-- Add any other context or screenshots about the feature request here. -->



Discussion:
Hi @kurtjcu Thanks for the issue!

I think the problem is here: https://github.com/PyTorchLightning/lightning-flash/blob/113efabc3e28e80ef03f9b3ed92b11156a4a3ca5/flash/image/segmentation/data.py#L146
We should be loading the image with `torchvision.datasets.folder.default_loader` and then using `torchvision.transforms.ToTensor()` to convert it to a tensor. Those methods in torchvision will ensure that the final tensors always have three channels even if the image was greyscale.

Would you be interested in making a PR that implements this fix?
This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.


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
@@ -63,7 +63,8 @@
 
 if _TORCHVISION_AVAILABLE:
     import torchvision
-    from torchvision.datasets.folder import has_file_allowed_extension, IMG_EXTENSIONS
+    import torchvision.transforms.functional as FT
+    from torchvision.datasets.folder import default_loader, has_file_allowed_extension, IMG_EXTENSIONS
 else:
     IMG_EXTENSIONS = None
 
@@ -148,7 +149,7 @@ def load_sample(self, sample: Mapping[str, Any]) -> Mapping[str, Union[torch.Ten
         img_labels_path = sample[DefaultDataKeys.TARGET]
 
         # load images directly to torch tensors
-        img: torch.Tensor = torchvision.io.read_image(img_path)  # CxHxW
+        img: torch.Tensor = FT.to_tensor(default_loader(img_path))  # CxHxW
         img_labels: torch.Tensor = torchvision.io.read_image(img_labels_path)  # CxHxW
         img_labels = img_labels[0]  # HxW
 
@@ -163,7 +164,7 @@ def load_sample(self, sample: Mapping[str, Any]) -> Mapping[str, Union[torch.Ten
     @staticmethod
     def predict_load_sample(sample: Mapping[str, Any]) -> Mapping[str, Any]:
         img_path = sample[DefaultDataKeys.INPUT]
-        img = torchvision.io.read_image(img_path).float()
+        img = FT.to_tensor(default_loader(img_path)).float()
 
         sample[DefaultDataKeys.INPUT] = img
         sample[DefaultDataKeys.METADATA] = {
@@ -195,7 +196,7 @@ def load_sample(self, sample: Mapping[str, str]) -> Mapping[str, Union[torch.Ten
         img_path = sample[DefaultDataKeys.INPUT]
         fo_sample = _fo_dataset[img_path]
 
-        img: torch.Tensor = torchvision.io.read_image(img_path)  # CxHxW
+        img: torch.Tensor = FT.to_tensor(default_loader(img_path))  # CxHxW
         img_labels: torch.Tensor = torch.from_numpy(fo_sample[self.label_field].mask)  # HxW
 
         sample[DefaultDataKeys.INPUT] = img.float()
@@ -209,7 +210,7 @@ def load_sample(self, sample: Mapping[str, str]) -> Mapping[str, Union[torch.Ten
     @staticmethod
     def predict_load_sample(sample: Mapping[str, Any]) -> Mapping[str, Any]:
         img_path = sample[DefaultDataKeys.INPUT]
-        img = torchvision.io.read_image(img_path).float()
+        img = FT.to_tensor(default_loader(img_path)).float()
 
         sample[DefaultDataKeys.INPUT] = img
         sample[DefaultDataKeys.METADATA] = {


Test patch: