Repository: Lightning-Universe/lightning-flash

Issue:
SemanticSegmentationFilesInput does not accept .npy files
## 🐛 Bug

It seems that `SemanticSegmentationFilesInput` does not accept `.npy` files.

```python
class SemanticSegmentationFilesInput(SemanticSegmentationInput):
    def load_data(
        self,
        files: Union[PATH_TYPE, List[PATH_TYPE]],
        mask_files: Optional[Union[PATH_TYPE, List[PATH_TYPE]]] = None,
        num_classes: Optional[int] = None,
        labels_map: Optional[Dict[int, Tuple[int, int, int]]] = None,
    ) -> List[Dict[str, Any]]:
        self.load_labels_map(num_classes, labels_map)
        if mask_files is None:
            files = filter_valid_files(files, valid_extensions=IMG_EXTENSIONS)
        else:
            files, mask_files = filter_valid_files(files, mask_files, valid_extensions=IMG_EXTENSIONS)
        return to_samples(files, mask_files)
```

`files, mask_files = filter_valid_files(files, mask_files, valid_extensions=IMG_EXTENSIONS)` 
should be 
`files, mask_files = filter_valid_files(files, mask_files, valid_extensions=IMG_EXTENSIONS + NP_EXTENSIONS)` ?





Discussion:
Hi, @chris-clem - Thank you for creating this issue! Could you please share a sample `npy` file with us? I would love to consider this if it's not supported already!

Also, yes, I agree that they should be supported. Maybe we just didn't have anyone demanding this, so we went with the image extensions initially. :)
@chris-clem mind sending a quick PR?

Resolving patch:
diff --git a/flash/core/data/utilities/loading.py b/flash/core/data/utilities/loading.py
--- a/flash/core/data/utilities/loading.py
+++ b/flash/core/data/utilities/loading.py
@@ -70,7 +70,7 @@ def _load_image_from_image(file, drop_alpha: bool = True):
 
 
 def _load_image_from_numpy(file):
-    return Image.fromarray(np.load(file).astype("uint8"), "RGB")
+    return Image.fromarray(np.load(file).astype("uint8")).convert("RGB")
 
 
 def _load_spectrogram_from_image(file):
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
@@ -115,9 +115,9 @@ def from_files(
 
             >>> from PIL import Image
             >>> rand_image = Image.fromarray(np.random.randint(0, 255, (64, 64, 3), dtype="uint8"))
-            >>> rand_mask= Image.fromarray(np.random.randint(0, 10, (64, 64), dtype="uint8"))
+            >>> rand_mask= np.random.randint(0, 10, (64, 64), dtype="uint8")
             >>> _ = [rand_image.save(f"image_{i}.png") for i in range(1, 4)]
-            >>> _ = [rand_mask.save(f"mask_{i}.png") for i in range(1, 4)]
+            >>> _ = [np.save(f"mask_{i}.npy", rand_mask) for i in range(1, 4)]
             >>> _ = [rand_image.save(f"predict_image_{i}.png") for i in range(1, 4)]
 
         .. doctest::
@@ -126,7 +126,7 @@ def from_files(
             >>> from flash.image import SemanticSegmentation, SemanticSegmentationData
             >>> datamodule = SemanticSegmentationData.from_files(
             ...     train_files=["image_1.png", "image_2.png", "image_3.png"],
-            ...     train_targets=["mask_1.png", "mask_2.png", "mask_3.png"],
+            ...     train_targets=["mask_1.npy", "mask_2.npy", "mask_3.npy"],
             ...     predict_files=["predict_image_1.png", "predict_image_2.png", "predict_image_3.png"],
             ...     transform_kwargs=dict(image_size=(128, 128)),
             ...     num_classes=10,
@@ -145,7 +145,7 @@ def from_files(
 
             >>> import os
             >>> _ = [os.remove(f"image_{i}.png") for i in range(1, 4)]
-            >>> _ = [os.remove(f"mask_{i}.png") for i in range(1, 4)]
+            >>> _ = [os.remove(f"mask_{i}.npy") for i in range(1, 4)]
             >>> _ = [os.remove(f"predict_image_{i}.png") for i in range(1, 4)]
         """
 
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
@@ -17,7 +17,7 @@
 import torch
 
 from flash.core.data.io.input import DataKeys, Input
-from flash.core.data.utilities.loading import IMG_EXTENSIONS, load_image
+from flash.core.data.utilities.loading import IMG_EXTENSIONS, load_image, NP_EXTENSIONS
 from flash.core.data.utilities.paths import filter_valid_files, PATH_TYPE
 from flash.core.data.utilities.samples import to_samples
 from flash.core.integrations.fiftyone.utils import FiftyOneLabelUtilities
@@ -98,9 +98,9 @@ def load_data(
     ) -> List[Dict[str, Any]]:
         self.load_labels_map(num_classes, labels_map)
         if mask_files is None:
-            files = filter_valid_files(files, valid_extensions=IMG_EXTENSIONS)
+            files = filter_valid_files(files, valid_extensions=IMG_EXTENSIONS + NP_EXTENSIONS)
         else:
-            files, mask_files = filter_valid_files(files, mask_files, valid_extensions=IMG_EXTENSIONS)
+            files, mask_files = filter_valid_files(files, mask_files, valid_extensions=IMG_EXTENSIONS + NP_EXTENSIONS)
         return to_samples(files, mask_files)
 
     def load_sample(self, sample: Dict[str, Any]) -> Dict[str, Any]:


Test patch: