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
 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,95 @@ def from_fiftyone(
             num_workers=num_workers,
             **preprocess_kwargs,
         )
+
+    @classmethod
+    def from_labelstudio(
+        cls,
+        export_json: str = None,
+        img_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,
+        **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
+            img_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',
+                img_folder='label-studio/media/upload',
+                val_split=0.8,
+            )
+        """
+        # loading export data
+        with open(export_json) as f:
+            js = json.load(f)
+        # loading data sets
+        full_dataset = LabelStudioDataset(js, img_folder)
+        val_dataset = LabelStudioDataset(js, img_folder, val=True)
+        # creating splitting params
+        l = len(full_dataset)
+        prop = int(l * val_split)
+        # splitting full data set
+        train_dataset, test_dataset = random_split(full_dataset, [prop, l - prop])
+
+        preprocess = preprocess or cls.preprocess_cls(
+            train_transform,
+            val_transform,
+            test_transform,
+            predict_transform,
+            **preprocess_kwargs,
+        )
+        data_source = preprocess.data_source_of_name(DefaultDataSources.FOLDERS)
+        data = cls(
+            train_dataset,
+            val_dataset,
+            test_dataset,
+            None,
+            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

flash/core/data/data_source.py

@@ -38,6 +38,7 @@
 from pytorch_lightning.utilities.enums import LightningEnum
 from torch.nn import Module
 from torch.utils.data.dataset import Dataset
+from torchvision.datasets.folder import default_loader
 
 from flash.core.data.auto_dataset import AutoDataset, BaseAutoDataset, IterableAutoDataset
 from flash.core.data.properties import ProcessState, Properties
@@ -156,6 +157,7 @@ class DefaultDataSources(LightningEnum):
     JSON = "json"
     DATASET = "dataset"
     FIFTYONE = "fiftyone"
+    LABELSTUDIO = "labelstudio"
 
     # TODO: Create a FlashEnum class???
     def __hash__(self) -> int:
@@ -543,3 +545,65 @@ def _get_classes(self, data):
             classes = data.distinct(label_path)
 
         return classes
+
+
+class LabelStudioDataset(Dataset):
+    r"""Dataset wrapping label studio annotations.
+
+    Each sample will be retrieved by checking result field of the annotation.
+
+    Args:
+        *js: json of export file
+        *img_folder: path to image folder of label studio
+    """
+
+    def __init__(self, js, img_folder, val=False):
+        self._raw_data = js
+        self._img_folder = img_folder
+        self.results = list()
+        self.classes = set()
+        # iterate through all tasks in exported data
+        for task in self._raw_data:
+            for annotation in task['annotations']:
+                # Adding ground_truth annotation to separate dataset
+                result = annotation['result']
+                for res in result:
+                    t = res['type']
+                    for label in res['value'][t]:
+                        # check if labeling result is a list of labels
+                        if isinstance(label, list):
+                            for sublabel in label:
+                                self.classes.add(sublabel)
+                                temp = dict()
+                                temp['file_upload'] = task['file_upload']
+                                temp['label'] = sublabel
+                                if annotation['ground_truth'] & val:
+                                    self.results.append(temp)
+                                elif not annotation['ground_truth'] or not val:
+                                    self.results.append(temp)
+                        else:
+                            self.classes.add(label)
+                            temp = dict()
+                            temp['file_upload'] = task['file_upload']
+                            temp['label'] = label
+                            if annotation['ground_truth'] & val:
+                                self.results.append(temp)
+                            elif not annotation['ground_truth'] or not val:
+                                self.results.append(temp)
+        self.num_classes = len(self.classes)
+
+    def __getitem__(self, idx):
+        r = self.results[idx]
+        # extracting path to file
+        p = os.path.join(self._img_folder, r['file_upload'])
+        # loading image
+        sample = default_loader(p)
+        # casting to list and sorting classes
+        sorted_labels = sorted(list(self.classes))
+        # checking index of class
+        label = sorted_labels.index(r['label'])
+        result = {DefaultDataKeys.INPUT: sample, DefaultDataKeys.TARGET: label}
+        return result
+
+    def __len__(self):
+        return len(self.results)

flash_examples/integrations/labelstudio/image_classification.py

@@ -0,0 +1,40 @@
+from itertools import chain
+
+import flash
+from flash.core.classification import Labels
+from flash.core.data.utils import download_data
+from flash.core.finetuning import FreezeUnfreeze
+from flash.image import ImageClassificationData, ImageClassifier
+
+# 1 Download data
+download_data("https://label-studio-testdata.s3.us-east-2.amazonaws.com/lightning-flash/data.zip")
+
+# 1. Load export data
+datamodule = ImageClassificationData.from_labelstudio(
+    export_json='data/project.json',
+    img_folder='data/upload/',
+    val_split=0.8,
+)
+
+# 2. Fine tune a model
+model = ImageClassifier(
+    backbone="resnet18",
+    num_classes=datamodule.num_classes,
+)
+trainer = flash.Trainer(max_epochs=3)
+
+trainer.finetune(
+    model,
+    datamodule=datamodule,
+    strategy=FreezeUnfreeze(unfreeze_epoch=1),
+)
+trainer.save_checkpoint("image_classification_model.pt")
+
+# 3. Predict from checkpoint
+model = ImageClassifier.load_from_checkpoint("image_classification_model.pt")
+model.serializer = Labels()
+
+predictions = model.predict([
+    "data/test/1.jpg",
+    "data/test/2.jpg",
+])