Repository: Lightning-Universe/lightning-flash

Issue:
Onboard audio inputs to new Input object
Tracking issue #955 


Discussion:

Resolving patch:
diff --git a/flash/audio/classification/data.py b/flash/audio/classification/data.py
--- a/flash/audio/classification/data.py
+++ b/flash/audio/classification/data.py
@@ -11,24 +11,28 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Callable, Dict, Optional, Tuple
+import os
+from typing import Any, Callable, Collection, Dict, List, Optional, Sequence, Tuple, Union
 
 import numpy as np
+import pandas as pd
+import torch
 
 from flash.audio.classification.transforms import default_transforms, train_default_transforms
 from flash.core.data.callback import BaseDataFetcher
 from flash.core.data.data_module import DataModule
-from flash.core.data.io.input import (
-    DataKeys,
-    has_file_allowed_extension,
-    InputFormat,
-    LoaderDataFrameInput,
-    NumpyInput,
-    PathsInput,
-)
+from flash.core.data.data_pipeline import DataPipelineState
+from flash.core.data.io.classification_input import ClassificationInput, ClassificationState
+from flash.core.data.io.input import DataKeys, has_file_allowed_extension, InputFormat
 from flash.core.data.io.input_transform import InputTransform
 from flash.core.data.process import Deserializer
+from flash.core.data.utilities.classification import TargetMode
+from flash.core.data.utilities.data_frame import read_csv, resolve_files, resolve_targets
+from flash.core.data.utilities.paths import filter_valid_files, make_dataset, PATH_TYPE
+from flash.core.data.utilities.samples import to_samples
 from flash.core.data.utils import image_default_loader
+from flash.core.utilities.imports import requires
+from flash.core.utilities.stages import RunningStage
 from flash.image.classification.data import MatplotlibVisualization
 from flash.image.data import ImageDeserializer, IMG_EXTENSIONS, NP_EXTENSIONS
 
@@ -42,26 +46,104 @@ def spectrogram_loader(filepath: str):
     return data
 
 
-class AudioClassificationNumpyInput(NumpyInput):
-    def load_sample(self, sample: Dict[str, Any], dataset: Optional[Any] = None) -> Dict[str, Any]:
+class AudioClassificationInput(ClassificationInput):
+    @requires("audio")
+    def load_sample(self, sample: Dict[str, Any]) -> Dict[str, Any]:
+        h, w = sample[DataKeys.INPUT].shape[-2:]  # H x W
+        if DataKeys.METADATA not in sample:
+            sample[DataKeys.METADATA] = {}
+        sample[DataKeys.METADATA]["size"] = (h, w)
+        if DataKeys.TARGET in sample:
+            sample[DataKeys.TARGET] = self.format_target(sample[DataKeys.TARGET])
+        return sample
+
+
+class AudioClassificationFilesInput(AudioClassificationInput):
+    def load_data(
+        self,
+        files: List[PATH_TYPE],
+        targets: Optional[List[Any]] = None,
+    ) -> List[Dict[str, Any]]:
+        if targets is None:
+            files = filter_valid_files(files, valid_extensions=IMG_EXTENSIONS + NP_EXTENSIONS)
+            return to_samples(files)
+        files, targets = filter_valid_files(files, targets, valid_extensions=IMG_EXTENSIONS + NP_EXTENSIONS)
+        self.load_target_metadata(targets)
+        return to_samples(files, targets)
+
+    def load_sample(self, sample: Dict[str, Any]) -> Dict[str, Any]:
+        filepath = sample[DataKeys.INPUT]
+        sample[DataKeys.INPUT] = spectrogram_loader(filepath)
+        sample = super().load_sample(sample)
+        sample[DataKeys.METADATA]["filepath"] = filepath
+        return sample
+
+
+class AudioClassificationFolderInput(AudioClassificationFilesInput):
+    def load_data(self, folder: PATH_TYPE) -> List[Dict[str, Any]]:
+        files, targets = make_dataset(folder, extensions=IMG_EXTENSIONS + NP_EXTENSIONS)
+        return super().load_data(files, targets)
+
+
+class AudioClassificationNumpyInput(AudioClassificationInput):
+    def load_data(self, array: Any, targets: Optional[List[Any]] = None) -> List[Dict[str, Any]]:
+        if targets is not None:
+            self.load_target_metadata(targets)
+        return to_samples(array, targets)
+
+    def load_sample(self, sample: Dict[str, Any]) -> Dict[str, Any]:
         sample[DataKeys.INPUT] = np.transpose(sample[DataKeys.INPUT], (1, 2, 0))
         return sample
 
 
 class AudioClassificationTensorInput(AudioClassificationNumpyInput):
-    def load_sample(self, sample: Dict[str, Any], dataset: Optional[Any] = None) -> Dict[str, Any]:
+    def load_data(self, tensor: Any, targets: Optional[List[Any]] = None) -> List[Dict[str, Any]]:
+        if targets is not None:
+            self.load_target_metadata(targets)
+        return to_samples(tensor, targets)
+
+    def load_sample(self, sample: Dict[str, Any]) -> Dict[str, Any]:
         sample[DataKeys.INPUT] = sample[DataKeys.INPUT].numpy()
-        return super().load_sample(sample, dataset=dataset)
+        return super().load_sample(sample)
+
+
+class AudioClassificationDataFrameInput(AudioClassificationFilesInput):
+    def load_data(
+        self,
+        data_frame: pd.DataFrame,
+        input_key: str,
+        target_keys: Optional[Union[str, List[str]]] = None,
+        root: Optional[PATH_TYPE] = None,
+        resolver: Optional[Callable[[Optional[PATH_TYPE], Any], PATH_TYPE]] = None,
+    ) -> List[Dict[str, Any]]:
+        files = resolve_files(data_frame, input_key, root, resolver)
+        if target_keys is not None:
+            targets = resolve_targets(data_frame, target_keys)
+        else:
+            targets = None
+        result = super().load_data(files, targets)
 
+        # If we had binary multi-class targets then we also know the labels (column names)
+        if self.training and self.target_mode is TargetMode.MULTI_BINARY and isinstance(target_keys, List):
+            classification_state = self.get_state(ClassificationState)
+            self.set_state(ClassificationState(target_keys, classification_state.num_classes))
 
-class AudioClassificationPathsInput(PathsInput):
-    def __init__(self):
-        super().__init__(loader=spectrogram_loader, extensions=IMG_EXTENSIONS + NP_EXTENSIONS)
+        return result
 
 
-class AudioClassificationDataFrameInput(LoaderDataFrameInput):
-    def __init__(self):
-        super().__init__(spectrogram_loader)
+class AudioClassificationCSVInput(AudioClassificationDataFrameInput):
+    def load_data(
+        self,
+        csv_file: PATH_TYPE,
+        input_key: str,
+        target_keys: Optional[Union[str, List[str]]] = None,
+        root: Optional[PATH_TYPE] = None,
+        resolver: Optional[Callable[[Optional[PATH_TYPE], Any], PATH_TYPE]] = None,
+    ) -> List[Dict[str, Any]]:
+        data_frame = read_csv(csv_file)
+        if root is None:
+            root = os.path.dirname(csv_file)
+        return super().load_data(data_frame, input_key, target_keys, root, resolver)
 
 
 class AudioClassificationInputTransform(InputTransform):
@@ -86,12 +168,12 @@ def __init__(
             test_transform=test_transform,
             predict_transform=predict_transform,
             inputs={
-                InputFormat.FILES: AudioClassificationPathsInput(),
-                InputFormat.FOLDERS: AudioClassificationPathsInput(),
-                "data_frame": AudioClassificationDataFrameInput(),
-                InputFormat.CSV: AudioClassificationDataFrameInput(),
-                InputFormat.NUMPY: AudioClassificationNumpyInput(),
-                InputFormat.TENSORS: AudioClassificationTensorInput(),
+                InputFormat.FILES: AudioClassificationFilesInput,
+                InputFormat.FOLDERS: AudioClassificationFolderInput,
+                InputFormat.DATAFRAME: AudioClassificationDataFrameInput,
+                InputFormat.CSV: AudioClassificationDataFrameInput,
+                InputFormat.NUMPY: AudioClassificationNumpyInput,
+                InputFormat.TENSORS: AudioClassificationTensorInput,
             },
             deserializer=deserializer or ImageDeserializer(),
             default_input=InputFormat.FILES,
@@ -121,6 +203,255 @@ class AudioClassificationData(DataModule):
 
     input_transform_cls = AudioClassificationInputTransform
 
+    @classmethod
+    def from_files(
+        cls,
+        train_files: Optional[Sequence[str]] = None,
+        train_targets: Optional[Sequence[Any]] = None,
+        val_files: Optional[Sequence[str]] = None,
+        val_targets: Optional[Sequence[Any]] = None,
+        test_files: Optional[Sequence[str]] = None,
+        test_targets: Optional[Sequence[Any]] = None,
+        predict_files: Optional[Sequence[str]] = None,
+        train_transform: Optional[Dict[str, Callable]] = None,
+        val_transform: Optional[Dict[str, Callable]] = None,
+        test_transform: Optional[Dict[str, Callable]] = None,
+        predict_transform: Optional[Dict[str, Callable]] = None,
+        spectrogram_size: Tuple[int, int] = (128, 128),
+        time_mask_param: Optional[int] = None,
+        freq_mask_param: Optional[int] = None,
+        **data_module_kwargs: Any,
+    ) -> "AudioClassificationData":
+        dataset_kwargs = dict(data_pipeline_state=DataPipelineState())
+
+        return cls(
+            AudioClassificationFilesInput(RunningStage.TRAINING, train_files, train_targets, **dataset_kwargs),
+            AudioClassificationFilesInput(RunningStage.VALIDATING, val_files, val_targets, **dataset_kwargs),
+            AudioClassificationFilesInput(RunningStage.TESTING, test_files, test_targets, **dataset_kwargs),
+            AudioClassificationFilesInput(RunningStage.PREDICTING, predict_files, **dataset_kwargs),
+            input_transform=cls.input_transform_cls(
+                train_transform,
+                val_transform,
+                test_transform,
+                predict_transform,
+                spectrogram_size=spectrogram_size,
+                time_mask_param=time_mask_param,
+                freq_mask_param=freq_mask_param,
+            ),
+            **data_module_kwargs,
+        )
+
+    @classmethod
+    def from_folders(
+        cls,
+        train_folder: Optional[str] = None,
+        val_folder: Optional[str] = None,
+        test_folder: Optional[str] = None,
+        predict_folder: Optional[str] = None,
+        train_transform: Optional[Dict[str, Callable]] = None,
+        val_transform: Optional[Dict[str, Callable]] = None,
+        test_transform: Optional[Dict[str, Callable]] = None,
+        predict_transform: Optional[Dict[str, Callable]] = None,
+        spectrogram_size: Tuple[int, int] = (128, 128),
+        time_mask_param: Optional[int] = None,
+        freq_mask_param: Optional[int] = None,
+        **data_module_kwargs: Any,
+    ) -> "AudioClassificationData":
+        dataset_kwargs = dict(data_pipeline_state=DataPipelineState())
+
+        return cls(
+            AudioClassificationFolderInput(RunningStage.TRAINING, train_folder, **dataset_kwargs),
+            AudioClassificationFolderInput(RunningStage.VALIDATING, val_folder, **dataset_kwargs),
+            AudioClassificationFolderInput(RunningStage.TESTING, test_folder, **dataset_kwargs),
+            AudioClassificationFolderInput(RunningStage.PREDICTING, predict_folder, **dataset_kwargs),
+            input_transform=cls.input_transform_cls(
+                train_transform,
+                val_transform,
+                test_transform,
+                predict_transform,
+                spectrogram_size=spectrogram_size,
+                time_mask_param=time_mask_param,
+                freq_mask_param=freq_mask_param,
+            ),
+            **data_module_kwargs,
+        )
+
+    @classmethod
+    def from_numpy(
+        cls,
+        train_data: Optional[Collection[np.ndarray]] = None,
+        train_targets: Optional[Collection[Any]] = None,
+        val_data: Optional[Collection[np.ndarray]] = None,
+        val_targets: Optional[Sequence[Any]] = None,
+        test_data: Optional[Collection[np.ndarray]] = None,
+        test_targets: Optional[Sequence[Any]] = None,
+        predict_data: Optional[Collection[np.ndarray]] = None,
+        train_transform: Optional[Dict[str, Callable]] = None,
+        val_transform: Optional[Dict[str, Callable]] = None,
+        test_transform: Optional[Dict[str, Callable]] = None,
+        predict_transform: Optional[Dict[str, Callable]] = None,
+        spectrogram_size: Tuple[int, int] = (128, 128),
+        time_mask_param: Optional[int] = None,
+        freq_mask_param: Optional[int] = None,
+        **data_module_kwargs: Any,
+    ) -> "AudioClassificationData":
+        dataset_kwargs = dict(data_pipeline_state=DataPipelineState())
+
+        return cls(
+            AudioClassificationNumpyInput(RunningStage.TRAINING, train_data, train_targets, **dataset_kwargs),
+            AudioClassificationNumpyInput(RunningStage.VALIDATING, val_data, val_targets, **dataset_kwargs),
+            AudioClassificationNumpyInput(RunningStage.TESTING, test_data, test_targets, **dataset_kwargs),
+            AudioClassificationNumpyInput(RunningStage.PREDICTING, predict_data, **dataset_kwargs),
+            input_transform=cls.input_transform_cls(
+                train_transform,
+                val_transform,
+                test_transform,
+                predict_transform,
+                spectrogram_size=spectrogram_size,
+                time_mask_param=time_mask_param,
+                freq_mask_param=freq_mask_param,
+            ),
+            **data_module_kwargs,
+        )
+
+    @classmethod
+    def from_tensors(
+        cls,
+        train_data: Optional[Collection[torch.Tensor]] = None,
+        train_targets: Optional[Collection[Any]] = None,
+        val_data: Optional[Collection[torch.Tensor]] = None,
+        val_targets: Optional[Sequence[Any]] = None,
+        test_data: Optional[Collection[torch.Tensor]] = None,
+        test_targets: Optional[Sequence[Any]] = None,
+        predict_data: Optional[Collection[torch.Tensor]] = None,
+        train_transform: Optional[Dict[str, Callable]] = None,
+        val_transform: Optional[Dict[str, Callable]] = None,
+        test_transform: Optional[Dict[str, Callable]] = None,
+        predict_transform: Optional[Dict[str, Callable]] = None,
+        spectrogram_size: Tuple[int, int] = (128, 128),
+        time_mask_param: Optional[int] = None,
+        freq_mask_param: Optional[int] = None,
+        **data_module_kwargs: Any,
+    ) -> "AudioClassificationData":
+        dataset_kwargs = dict(data_pipeline_state=DataPipelineState())
+
+        return cls(
+            AudioClassificationTensorInput(RunningStage.TRAINING, train_data, train_targets, **dataset_kwargs),
+            AudioClassificationTensorInput(RunningStage.VALIDATING, val_data, val_targets, **dataset_kwargs),
+            AudioClassificationTensorInput(RunningStage.TESTING, test_data, test_targets, **dataset_kwargs),
+            AudioClassificationTensorInput(RunningStage.PREDICTING, predict_data, **dataset_kwargs),
+            input_transform=cls.input_transform_cls(
+                train_transform,
+                val_transform,
+                test_transform,
+                predict_transform,
+                spectrogram_size=spectrogram_size,
+                time_mask_param=time_mask_param,
+                freq_mask_param=freq_mask_param,
+            ),
+            **data_module_kwargs,
+        )
+
+    @classmethod
+    def from_data_frame(
+        cls,
+        input_field: str,
+        target_fields: Optional[Union[str, Sequence[str]]] = None,
+        train_data_frame: Optional[pd.DataFrame] = None,
+        train_images_root: Optional[str] = None,
+        train_resolver: Optional[Callable[[str, str], str]] = None,
+        val_data_frame: Optional[pd.DataFrame] = None,
+        val_images_root: Optional[str] = None,
+        val_resolver: Optional[Callable[[str, str], str]] = None,
+        test_data_frame: Optional[pd.DataFrame] = None,
+        test_images_root: Optional[str] = None,
+        test_resolver: Optional[Callable[[str, str], str]] = None,
+        predict_data_frame: Optional[pd.DataFrame] = None,
+        predict_images_root: Optional[str] = None,
+        predict_resolver: Optional[Callable[[str, str], str]] = None,
+        train_transform: Optional[Union[Callable, List, Dict[str, Callable]]] = None,
+        val_transform: Optional[Union[Callable, List, Dict[str, Callable]]] = None,
+        test_transform: Optional[Union[Callable, List, Dict[str, Callable]]] = None,
+        predict_transform: Optional[Dict[str, Callable]] = None,
+        spectrogram_size: Tuple[int, int] = (128, 128),
+        time_mask_param: Optional[int] = None,
+        freq_mask_param: Optional[int] = None,
+        **data_module_kwargs: Any,
+    ) -> "AudioClassificationData":
+        dataset_kwargs = dict(data_pipeline_state=DataPipelineState())
+
+        train_data = (train_data_frame, input_field, target_fields, train_images_root, train_resolver)
+        val_data = (val_data_frame, input_field, target_fields, val_images_root, val_resolver)
+        test_data = (test_data_frame, input_field, target_fields, test_images_root, test_resolver)
+        predict_data = (predict_data_frame, input_field, None, predict_images_root, predict_resolver)
+
+        return cls(
+            AudioClassificationDataFrameInput(RunningStage.TRAINING, *train_data, **dataset_kwargs),
+            AudioClassificationDataFrameInput(RunningStage.VALIDATING, *val_data, **dataset_kwargs),
+            AudioClassificationDataFrameInput(RunningStage.TESTING, *test_data, **dataset_kwargs),
+            AudioClassificationDataFrameInput(RunningStage.PREDICTING, *predict_data, **dataset_kwargs),
+            input_transform=cls.input_transform_cls(
+                train_transform,
+                val_transform,
+                test_transform,
+                predict_transform,
+                spectrogram_size=spectrogram_size,
+                time_mask_param=time_mask_param,
+                freq_mask_param=freq_mask_param,
+            ),
+            **data_module_kwargs,
+        )
+
+    @classmethod
+    def from_csv(
+        cls,
+        input_field: str,
+        target_fields: Optional[Union[str, List[str]]] = None,
+        train_file: Optional[PATH_TYPE] = None,
+        train_images_root: Optional[PATH_TYPE] = None,
+        train_resolver: Optional[Callable[[PATH_TYPE, Any], PATH_TYPE]] = None,
+        val_file: Optional[PATH_TYPE] = None,
+        val_images_root: Optional[PATH_TYPE] = None,
+        val_resolver: Optional[Callable[[PATH_TYPE, Any], PATH_TYPE]] = None,
+        test_file: Optional[str] = None,
+        test_images_root: Optional[str] = None,
+        test_resolver: Optional[Callable[[PATH_TYPE, Any], PATH_TYPE]] = None,
+        predict_file: Optional[str] = None,
+        predict_images_root: Optional[str] = None,
+        predict_resolver: Optional[Callable[[PATH_TYPE, Any], PATH_TYPE]] = None,
+        train_transform: Optional[Union[Callable, List, Dict[str, Callable]]] = None,
+        val_transform: Optional[Union[Callable, List, Dict[str, Callable]]] = None,
+        test_transform: Optional[Union[Callable, List, Dict[str, Callable]]] = None,
+        predict_transform: Optional[Dict[str, Callable]] = None,
+        spectrogram_size: Tuple[int, int] = (128, 128),
+        time_mask_param: Optional[int] = None,
+        freq_mask_param: Optional[int] = None,
+        **data_module_kwargs: Any,
+    ) -> "AudioClassificationData":
+        dataset_kwargs = dict(data_pipeline_state=DataPipelineState())
+
+        train_data = (train_file, input_field, target_fields, train_images_root, train_resolver)
+        val_data = (val_file, input_field, target_fields, val_images_root, val_resolver)
+        test_data = (test_file, input_field, target_fields, test_images_root, test_resolver)
+        predict_data = (predict_file, input_field, None, predict_images_root, predict_resolver)
+
+        return cls(
+            AudioClassificationCSVInput(RunningStage.TRAINING, *train_data, **dataset_kwargs),
+            AudioClassificationCSVInput(RunningStage.VALIDATING, *val_data, **dataset_kwargs),
+            AudioClassificationCSVInput(RunningStage.TESTING, *test_data, **dataset_kwargs),
+            AudioClassificationCSVInput(RunningStage.PREDICTING, *predict_data, **dataset_kwargs),
+            input_transform=cls.input_transform_cls(
+                train_transform,
+                val_transform,
+                test_transform,
+                predict_transform,
+                spectrogram_size=spectrogram_size,
+                time_mask_param=time_mask_param,
+                freq_mask_param=freq_mask_param,
+            ),
+            **data_module_kwargs,
+        )
+
     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


Test patch: