Adding integration with Label Studio

flash/core/data/data_module.py

@@ -11,6 +11,7 @@
 # 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.
+import json
 import os
 import platform
 from typing import Any, Callable, Collection, Dict, Iterable, List, Optional, Sequence, Tuple, TYPE_CHECKING, Union
@@ -21,14 +22,14 @@
 from pytorch_lightning.trainer.states import RunningStage
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from torch.utils.data import DataLoader, Dataset
-from torch.utils.data.dataset import IterableDataset, Subset
+from torch.utils.data.dataset import IterableDataset, random_split, Subset
 from torch.utils.data.sampler import Sampler
 
 from flash.core.data.auto_dataset import BaseAutoDataset, IterableAutoDataset
 from flash.core.data.base_viz import BaseVisualization
 from flash.core.data.callback import BaseDataFetcher
 from flash.core.data.data_pipeline import DataPipeline, DefaultPreprocess, Postprocess, Preprocess
-from flash.core.data.data_source import DataSource, DefaultDataSources
+from flash.core.data.data_source import DataSource, DefaultDataSources, LabelStudioDataset, LabelStudioDataSource
 from flash.core.data.splits import SplitDataset
 from flash.core.data.utils import _STAGES_PREFIX
 from flash.core.utilities.imports import _FIFTYONE_AVAILABLE
@@ -1152,3 +1153,101 @@ def from_fiftyone(
             num_workers=num_workers,
             **preprocess_kwargs,
         )
+
+    @classmethod
+    def from_labelstudio(
+        cls,
+        export_json: str = None,
+        data_folder: 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,
+        data_fetcher: Optional[BaseDataFetcher] = None,
+        preprocess: Optional[Preprocess] = None,
+        val_split: Optional[float] = None,
+        batch_size: int = 4,
+        num_workers: Optional[int] = None,
+        sampler: Optional[Sampler] = None,
+        data_type: str = "other",
+        **preprocess_kwargs: Any,
+    ) -> "DataModule":
+        """Creates a :class:`~flash.core.data.data_module.DataModule` object
+        from the given export file and data directory using the
+        :class:`~flash.core.data.data_source.DataSource` of name
+        :attr:`~flash.core.data.data_source.DefaultDataSources.FOLDERS`
+        from the passed or constructed :class:`~flash.core.data.process.Preprocess`.
+
+        Args:
+            export_json: path to label studio export file
+            data_folder: path to label studio data folder
+            train_transform: The dictionary of transforms to use during training which maps
+                :class:`~flash.core.data.process.Preprocess` hook names to callable transforms.
+            val_transform: The dictionary of transforms to use during validation which maps
+                :class:`~flash.core.data.process.Preprocess` hook names to callable transforms.
+            test_transform: The dictionary of transforms to use during testing which maps
+                :class:`~flash.core.data.process.Preprocess` hook names to callable transforms.
+            predict_transform: The dictionary of transforms to use during predicting which maps
+                :class:`~flash.core.data.process.Preprocess` hook names to callable transforms.
+            data_fetcher: The :class:`~flash.core.data.callback.BaseDataFetcher` to pass to the
+                :class:`~flash.core.data.data_module.DataModule`.
+            preprocess: The :class:`~flash.core.data.data.Preprocess` to pass to the
+                :class:`~flash.core.data.data_module.DataModule`. If ``None``, ``cls.preprocess_cls``
+                will be constructed and used.
+            val_split: The ``val_split`` argument to pass to the :class:`~flash.core.data.data_module.DataModule`.
+            batch_size: The ``batch_size`` argument to pass to the :class:`~flash.core.data.data_module.DataModule`.
+            num_workers: The ``num_workers`` argument to pass to the :class:`~flash.core.data.data_module.DataModule`.
+            sampler: The ``sampler`` argument to pass to the :class:`~flash.core.data.data_module.DataModule`.
+            preprocess_kwargs: Additional keyword arguments to use when constructing the preprocess. Will only be used
+                if ``preprocess = None``.
+
+        Returns:
+            The constructed data module.
+
+        Examples::
+
+            data_module = DataModule.from_labelstudio(
+                export_json='project.json',
+                data_folder='label-studio/media/upload',
+                val_split=0.8,
+            )
+        """
+        # create preprocess
+        preprocess = preprocess or cls.preprocess_cls(
+            train_transform,
+            val_transform,
+            test_transform,
+            predict_transform,
+            **preprocess_kwargs,
+        )
+        # loading export data
+        data_source = LabelStudioDataSource(
+            data_folder=data_folder,
+            export_json=export_json,
+            data_type=data_type,
+            backbone=preprocess_kwargs.get("backbone"),
+            max_length=preprocess_kwargs.get("max_length", 128),
+            clip_sampler=preprocess_kwargs.get("clip_sampler"),
+            clip_duration=preprocess_kwargs.get("clip_duration", 1),
+            decode_audio=preprocess_kwargs.get("decode_audio", False),
+            decoder=preprocess_kwargs.get("decoder", "pyav"),
+            multi_label=preprocess_kwargs.get("multi_label", False),
+        )
+        data_source.load_data()
+        # loading data sets
+        train_dataset, val_dataset, test_dataset, predict_dataset = data_source.to_datasets()
+
+        data = cls(
+            train_dataset,
+            val_dataset,
+            test_dataset,
+            predict_dataset,
+            data_source=data_source,
+            preprocess=preprocess,
+            data_fetcher=data_fetcher,
+            val_split=val_split,
+            batch_size=batch_size,
+            num_workers=num_workers,
+            sampler=sampler,
+        )
+        return data