Repository: voxel51/fiftyone

Issue:
[BUG] Lightning Flash integration doesn't work with current Flash version
**tl;dr** our Flash integration needs to be updated to support the latest Flash version

Our [Lightning Flash integration](https://voxel51.com/docs/fiftyone/integrations/lightning_flash.html) was originally written for these versions of Flash/Pytorch Lightning:

```
pip install lightning-flash==0.4.0
pip install pytorch-lightning==1.3.6
```

With the above versions of those packages installed, I am still able to run the detection and embeddings examples from our integration docs. However, I am no longer able to load the pre-trained image checkpoint:

```
rm -rf ~/.cache/torch/hub/checkpoints
```

```py
from flash.image import ImageClassifier

model = ImageClassifier.load_from_checkpoint(
    "https://flash-weights.s3.amazonaws.com/image_classification_model.pt"
)
```

```
AttributeError: Can't get attribute 'image_loader' on <module 'flash.image.data' from '/Users/Brian/dev/env/fifty3/lib/python3.6/site-packages/flash/image/data.py'>
```

If I upgrade to the latest version of Lightning Flash (0.6.0 as of this writing) then our examples no longer work.

Example code I tried running:

```py
from flash.image import ImageClassifier, ObjectDetector

import fiftyone as fo
import fiftyone.zoo as foz

dataset = foz.load_zoo_dataset("quickstart", max_samples=5)

model = ObjectDetector.load_from_checkpoint(
    "https://flash-weights.s3.amazonaws.com/object_detection_model.pt"
)
dataset.apply_model(model, label_field="flash_detector")

model = ImageClassifier.load_from_checkpoint(
    "https://flash-weights.s3.amazonaws.com/image_classification_model.pt"
)
dataset.apply_model(model, label_field="flash_classifier")
```



Discussion:
Would love to see a fix for this.
I've tried updating the imports but there seem to be some other issues. With a little bit of direction I'd be happy to make a PR :)
@MrBlenny Thanks for offering your help! There are multiple things that need to be updated, including some of the changes to the integration on the Lightning Flash side. I am working out what all needs to be changed and I'll let you know if there is a way for you to assist! :hugs: 
I also quite interested in a new working version of integrationof the "semantic segmentation" example / use case. In the meanwhile up to the time the fix is available maybe adding a comment that newer flash version are not supported yet, would be nice. Thank you!

Resolving patch:
diff --git a/fiftyone/core/collections.py b/fiftyone/core/collections.py
--- a/fiftyone/core/collections.py
+++ b/fiftyone/core/collections.py
@@ -1526,7 +1526,7 @@ def apply_model(
         batch_size=None,
         num_workers=None,
         skip_failures=True,
-        **trainer_kwargs,
+        **kwargs,
     ):
         """Applies the :class:`FiftyOne model <fiftyone.core.models.Model>` or
         :class:`Lightning Flash model <flash:flash.core.model.Task>` to the
@@ -1563,10 +1563,8 @@ def apply_model(
                 raising an error if predictions cannot be generated for a
                 sample. Only applicable to :class:`fiftyone.core.models.Model`
                 instances
-            **trainer_kwargs: optional keyword arguments used to initialize the
-                :mod:`Trainer <flash:flash.core.trainer>` when using Flash
-                models. These can be used to, for example, configure the number
-                of GPUs to use and other distributed inference parameters
+            **kwargs: optional model-specific keyword arguments passed through
+                to the underlying inference implementation
         """
         fomo.apply_model(
             self,
@@ -1577,7 +1575,7 @@ def apply_model(
             batch_size=batch_size,
             num_workers=num_workers,
             skip_failures=skip_failures,
-            **trainer_kwargs,
+            **kwargs,
         )
 
     def compute_embeddings(
@@ -1587,7 +1585,7 @@ def compute_embeddings(
         batch_size=None,
         num_workers=None,
         skip_failures=True,
-        **trainer_kwargs,
+        **kwargs,
     ):
         """Computes embeddings for the samples in the collection using the
         given :class:`FiftyOne model <fiftyone.core.models.Model>` or
@@ -1626,10 +1624,8 @@ def compute_embeddings(
                 raising an error if embeddings cannot be generated for a
                 sample. Only applicable to :class:`fiftyone.core.models.Model`
                 instances
-            **trainer_kwargs: optional keyword arguments used to initialize the
-                :mod:`Trainer <flash:flash.core.trainer>` when using Flash
-                models. These can be used to, for example, configure the number
-                of GPUs to use and other distributed inference parameters
+            **kwargs: optional model-specific keyword arguments passed through
+                to the underlying inference implementation
 
         Returns:
             one of the following:
@@ -1657,7 +1653,7 @@ def compute_embeddings(
             batch_size=batch_size,
             num_workers=num_workers,
             skip_failures=skip_failures,
-            **trainer_kwargs,
+            **kwargs,
         )
 
     def compute_patch_embeddings(
diff --git a/fiftyone/core/models.py b/fiftyone/core/models.py
--- a/fiftyone/core/models.py
+++ b/fiftyone/core/models.py
@@ -53,7 +53,7 @@ def apply_model(
     batch_size=None,
     num_workers=None,
     skip_failures=True,
-    **trainer_kwargs,
+    **kwargs,
 ):
     """Applies the :class:`FiftyOne model <Model>` or
     :class:`Lightning Flash model <flash:flash.core.model.Task>` to the samples
@@ -85,10 +85,8 @@ def apply_model(
         skip_failures (True): whether to gracefully continue without raising an
             error if predictions cannot be generated for a sample. Only
             applicable to :class:`Model` instances
-        **trainer_kwargs: optional keyword arguments used to initialize the
-            :mod:`Trainer <flash:flash.core.trainer>` when using Flash models.
-            These can be used to, for example, configure the number of GPUs to
-            use and other distributed inference parameters
+        **kwargs: optional model-specific keyword arguments passed through
+            to the underlying inference implementation
     """
     if _is_flash_model(model):
         return fouf.apply_flash_model(
@@ -99,7 +97,7 @@ def apply_model(
             store_logits=store_logits,
             batch_size=batch_size,
             num_workers=num_workers,
-            **trainer_kwargs,
+            **kwargs,
         )
 
     if not isinstance(model, Model):
@@ -543,7 +541,7 @@ def compute_embeddings(
     batch_size=None,
     num_workers=None,
     skip_failures=True,
-    **trainer_kwargs,
+    **kwargs,
 ):
     """Computes embeddings for the samples in the collection using the given
     :class:`FiftyOne model <Model>` or
@@ -579,10 +577,8 @@ def compute_embeddings(
         skip_failures (True): whether to gracefully continue without raising an
             error if embeddings cannot be generated for a sample. Only
             applicable to :class:`Model` instances
-        **trainer_kwargs: optional keyword arguments used to initialize the
-            :mod:`Trainer <flash:flash.core.trainer>` when using Flash models.
-            These can be used to, for example, configure the number of GPUs to
-            use and other distributed inference parameters
+        **kwargs: optional model-specific keyword arguments passed through
+            to the underlying inference implementation
 
     Returns:
         one of the following:
@@ -609,7 +605,7 @@ def compute_embeddings(
             embeddings_field=embeddings_field,
             batch_size=batch_size,
             num_workers=num_workers,
-            **trainer_kwargs,
+            **kwargs,
         )
 
     if not isinstance(model, Model):
diff --git a/fiftyone/utils/flash.py b/fiftyone/utils/flash.py
--- a/fiftyone/utils/flash.py
+++ b/fiftyone/utils/flash.py
@@ -11,15 +11,17 @@
 
 import numpy as np
 
-import fiftyone.core.metadata as fom
+import fiftyone.core.labels as fol
+import fiftyone.core.media as fom
 import fiftyone.core.utils as fou
 
-fou.ensure_import("flash>=0.4")
+fou.ensure_import("flash>=0.7dev")
 import flash
 import flash.core.classification as fc
 import flash.image as fi
-import flash.image.detection.serialization as fds
-import flash.image.segmentation.serialization as fss
+import flash.image.detection.output as fdo
+import flash.image.segmentation.output as fso
+import flash.video as fv
 
 
 _SUPPORTED_MODELS = (
@@ -28,7 +30,7 @@
     fi.SemanticSegmentation,
 )
 
-_SUPPORTED_EMBEDDERS = (fi.ImageEmbedder,)
+_SUPPORTED_EMBEDDERS = fi.ImageEmbedder
 
 
 def apply_flash_model(
@@ -39,7 +41,8 @@ def apply_flash_model(
     store_logits=False,
     batch_size=None,
     num_workers=None,
-    **trainer_kwargs,
+    transform_kwargs=None,
+    trainer_kwargs=None,
 ):
     """Applies the given
     :class:`Lightning Flash model <flash:flash.core.model.Task>` to the samples
@@ -59,35 +62,37 @@ def apply_flash_model(
         batch_size (None): an optional batch size to use. If not provided, a
             default batch size is used
         num_workers (None): the number of workers for the data loader to use
-        **trainer_kwargs: optional keyword arguments used to initialize the
+        transform_kwargs (None): an optional dict of transform kwargs to pass
+            into the created data module used by some models
+        trainer_kwargs (None): an optional dict of kwargs used to initialize the
             :mod:`Trainer <flash:flash.core.trainer>`. These can be used to,
             for example, configure the number of GPUs to use and other
             distributed inference parameters
     """
-    serializer = _get_serializer(model, confidence_thresh, store_logits)
+    output = _get_output(model, confidence_thresh, store_logits)
 
-    # Running `Trainer().predict()` seems to cause `model.data_pipeline` to be
-    # garbage-collected after inference, so we restore it to its original value
-    data_pipeline = model.data_pipeline
+    if trainer_kwargs is None:
+        trainer_kwargs = {}
 
-    with fou.SetAttributes(
-        model, data_pipeline=data_pipeline, serializer=serializer
-    ):
-        data_kwargs = dict(
-            preprocess=model.preprocess, num_workers=num_workers,
-        )
-        if batch_size is not None:
-            data_kwargs["batch_size"] = batch_size
+    data_kwargs = {
+        "num_workers": num_workers or 1,
+        "batch_size": batch_size or 1,
+    }
+    if transform_kwargs:
+        data_kwargs["transform_kwargs"] = transform_kwargs
 
-        datamodule = fi.ImageClassificationData.from_fiftyone(
-            predict_dataset=samples, **data_kwargs
-        )
-        predictions = flash.Trainer(**trainer_kwargs).predict(
-            model, datamodule=datamodule
-        )
-        predictions = list(itertools.chain.from_iterable(predictions))
+    datamodule_cls = _MODEL_TO_DATAMODULE_MAP[type(model)]
 
-        samples.set_values(label_field, predictions)
+    datamodule = datamodule_cls.from_fiftyone(
+        predict_dataset=samples, **data_kwargs
+    )
+    predictions = flash.Trainer(**trainer_kwargs).predict(
+        model, datamodule=datamodule, output=output
+    )
+    predictions = list(itertools.chain.from_iterable(predictions))
+    predictions = {p["filepath"]: p["predictions"] for p in predictions}
+
+    samples.set_values(label_field, predictions, key_field="filepath")
 
 
 def compute_flash_embeddings(
@@ -96,7 +101,8 @@ def compute_flash_embeddings(
     embeddings_field=None,
     batch_size=None,
     num_workers=None,
-    **trainer_kwargs,
+    transform_kwargs=None,
+    trainer_kwargs=None,
 ):
     """Computes embeddings for the samples in the collection using the given
     :class:`Lightning Flash model <flash:flash.core.model.Task>`.
@@ -115,7 +121,9 @@ def compute_flash_embeddings(
         batch_size (None): an optional batch size to use. If not provided, a
             default batch size is used
         num_workers (None): the number of workers for the data loader to use
-        **trainer_kwargs: optional keyword arguments used to initialize the
+        transform_kwargs (None): an optional dict of transform kwargs to pass
+            into the created data module used by some models
+        trainer_kwargs (None): an optional dict of kwargs used to initialize the
             :mod:`Trainer <flash:flash.core.trainer>`. These can be used to,
             for example, configure the number of GPUs to use and other
             distributed inference parameters
@@ -133,35 +141,38 @@ def compute_flash_embeddings(
             % (type(model), _SUPPORTED_EMBEDDERS)
         )
 
-    # Running `Trainer().predict()` seems to cause `model.data_pipeline` to be
-    # garbage-collected after inference, so we restore it to its original value
-    data_pipeline = model.data_pipeline
+    if trainer_kwargs is None:
+        trainer_kwargs = {}
 
-    with fou.SetAttributes(model, data_pipeline=data_pipeline):
-        data_kwargs = dict(
-            preprocess=model.preprocess, num_workers=num_workers,
-        )
-        if batch_size is not None:
-            data_kwargs["batch_size"] = batch_size
+    data_kwargs = {
+        "num_workers": num_workers or 1,
+        "batch_size": batch_size or 1,
+    }
+    if transform_kwargs:
+        data_kwargs["transform_kwargs"] = transform_kwargs
 
-        datamodule = fi.ImageClassificationData.from_fiftyone(
-            predict_dataset=samples, **data_kwargs
-        )
-        embeddings = flash.Trainer(**trainer_kwargs).predict(
-            model, datamodule=datamodule
-        )
-        embeddings = list(itertools.chain.from_iterable(embeddings))
+    datamodule = fi.ImageClassificationData.from_fiftyone(
+        predict_dataset=samples, **data_kwargs
+    )
+    embeddings = flash.Trainer(**trainer_kwargs).predict(
+        model, datamodule=datamodule
+    )
+    embeddings = np.stack(sum(embeddings, []))
 
-        if embeddings_field is not None:
-            samples.set_values(embeddings_field, embeddings)
-            return
+    if embeddings_field is not None:
+        samples.set_values(embeddings_field, embeddings)
+        return
 
-        return np.stack(embeddings)
+    return embeddings
 
 
-def _get_serializer(model, confidence_thresh, store_logits):
+def _get_output(model, confidence_thresh, store_logits):
     if isinstance(model, fi.ImageClassifier):
-        prev_args = dict(inspect.getmembers(model.serializer))
+        prev_args = {}
+        try:
+            prev_args = dict(inspect.getmembers(model.output))
+        except AttributeError as e:
+            pass
 
         kwargs = {
             "multi_label": prev_args.get("multi_label", False),
@@ -174,15 +185,23 @@ def _get_serializer(model, confidence_thresh, store_logits):
         if confidence_thresh is not None:
             kwargs["threshold"] = confidence_thresh
 
-        return fc.FiftyOneLabels(**kwargs)
+        return fc.FiftyOneLabelsOutput(**kwargs)
 
     if isinstance(model, fi.ObjectDetector):
-        return fds.FiftyOneDetectionLabels(threshold=confidence_thresh)
+        return fdo.FiftyOneDetectionLabelsOutput(threshold=confidence_thresh)
 
     if isinstance(model, fi.SemanticSegmentation):
-        return fss.FiftyOneSegmentationLabels()
+        return fso.FiftyOneSegmentationLabelsOutput()
 
     raise ValueError(
         "Unsupported model type %s. Supported model types are %s"
         % (type(model), _SUPPORTED_MODELS)
     )
+
+
+_MODEL_TO_DATAMODULE_MAP = {
+    fi.ObjectDetector: fi.ObjectDetectionData,
+    fi.ImageClassifier: fi.ImageClassificationData,
+    fi.SemanticSegmentation: fi.SemanticSegmentationData,
+    fv.VideoClassifier: fv.VideoClassificationData,
+}


Test patch:
diff --git a/tests/intensive/lightning_flash_tests.py b/tests/intensive/lightning_flash_tests.py
new file mode 100644
--- /dev/null
+++ b/tests/intensive/lightning_flash_tests.py
@@ -0,0 +1,438 @@
+"""
+Tests for the :mod:`fiftyone.utils.flash` module.
+
+You must run these tests interactively as follows::
+
+    python tests/intensive/lightning_flash_tests.py
+
+| Copyright 2017-2022, Voxel51, Inc.
+| `voxel51.com <https://voxel51.com/>`_
+|
+"""
+import numpy as np
+import unittest
+from itertools import chain
+
+from flash.core.data.utils import download_data
+from flash.image.detection.output import FiftyOneDetectionLabelsOutput
+from flash.core.classification import FiftyOneLabelsOutput
+from flash.image.segmentation.output import FiftyOneSegmentationLabelsOutput
+from flash.image import ImageClassificationData, ImageClassifier, ImageEmbedder
+from flash.image import ObjectDetectionData, ObjectDetector
+from flash.image import SemanticSegmentation, SemanticSegmentationData
+from flash import Trainer
+from flash.video import VideoClassificationData, VideoClassifier
+
+import fiftyone as fo
+import fiftyone.brain as fob
+import fiftyone.utils.splits as fous
+import fiftyone.zoo as foz
+
+
+class LightningFlashTests(unittest.TestCase):
+    def test_apply_model(self):
+        # Load your dataset
+        dataset = foz.load_zoo_dataset("quickstart", max_samples=5).clone()
+
+        num_classes = len(dataset.distinct("ground_truth.detections.label"))
+        # Load your Flash model
+        cls_model = ImageClassifier(
+            backbone="resnet18", num_classes=num_classes
+        )
+
+        det_model = ObjectDetector(
+            head="efficientdet",
+            backbone="d0",
+            num_classes=num_classes,
+            image_size=512,
+        )
+
+        # Predict!
+        dataset.apply_model(
+            cls_model, label_field="flash_classifications",
+        )
+
+        transform_kwargs = {"image_size": 512}
+        dataset.apply_model(
+            det_model,
+            label_field="flash_detections",
+            transform_kwargs=transform_kwargs,
+        )
+
+    def test_image_classifier(self):
+        # 1 Load your FiftyOne dataset
+        dataset = foz.load_zoo_dataset(
+            "cifar10", split="test", max_samples=300
+        ).clone()
+        dataset.untag_samples("test")
+
+        # Create splits from the dataset
+        splits = {"train": 0.7, "test": 0.1, "val": 0.1, "pred": 0.1}
+        fous.random_split(dataset, splits)
+
+        # Here we use views into one dataset,
+        # but you can also use a different dataset for each split
+        train_dataset = dataset.match_tags("train")
+        test_dataset = dataset.match_tags("test")
+        val_dataset = dataset.match_tags("val")
+        predict_dataset = dataset.match_tags("pred")
+
+        # 2 Create the Datamodule
+        datamodule = ImageClassificationData.from_fiftyone(
+            train_dataset=train_dataset,
+            test_dataset=test_dataset,
+            val_dataset=val_dataset,
+            predict_dataset=predict_dataset,
+            label_field="ground_truth",
+            batch_size=4,
+            num_workers=4,
+        )
+
+        # 3 Build the model
+        model = ImageClassifier(backbone="resnet18", labels=datamodule.labels)
+
+        # 4 Create the trainer
+        trainer = Trainer(
+            max_epochs=1, limit_train_batches=10, limit_val_batches=10,
+        )
+
+        # 5 Finetune the model
+        trainer.finetune(model, datamodule=datamodule)
+
+        # 6 Save it!
+        trainer.save_checkpoint("/tmp/image_classification_model.pt")
+
+        # 7 Generate predictions
+        predictions = trainer.predict(
+            model,
+            datamodule=datamodule,
+            output=FiftyOneLabelsOutput(labels=datamodule.labels),
+        )
+        predictions = list(chain.from_iterable(predictions))  # flatten batches
+
+        # Map filepaths to predictions
+        predictions = {p["filepath"]: p["predictions"] for p in predictions}
+
+        # Add predictions to FiftyOne dataset
+        predict_dataset.set_values(
+            "flash_predictions", predictions, key_field="filepath",
+        )
+
+    def test_object_detector(self):
+        # 1 Load your FiftyOne dataset
+        dataset = foz.load_zoo_dataset(
+            "coco-2017",
+            split="validation",
+            max_samples=100,
+            classes=["person"],
+        ).clone()
+
+        # Create splits from the dataset
+        splits = {"train": 0.7, "test": 0.1, "val": 0.1}
+        fous.random_split(dataset, splits)
+
+        # Here we use views into one dataset,
+        # but you can also use a different dataset for each split
+        train_dataset = dataset.match_tags("train")
+        test_dataset = dataset.match_tags("test")
+        val_dataset = dataset.match_tags("val")
+        predict_dataset = train_dataset.take(5)
+
+        # Remove background label, gets added by datamodule
+        dataset.default_classes.pop(0)
+
+        # 2 Create the Datamodule
+        datamodule = ObjectDetectionData.from_fiftyone(
+            train_dataset=train_dataset,
+            test_dataset=test_dataset,
+            val_dataset=val_dataset,
+            predict_dataset=predict_dataset,
+            label_field="ground_truth",
+            transform_kwargs={"image_size": 512},
+            batch_size=4,
+        )
+
+        # 3 Build the model
+        model = ObjectDetector(
+            head="efficientdet",
+            backbone="d0",
+            num_classes=datamodule.num_classes,
+            image_size=512,
+        )
+
+        # 4 Create the trainer
+        trainer = Trainer(max_epochs=1, limit_train_batches=10)
+
+        # 5 Finetune the model
+        trainer.finetune(model, datamodule=datamodule, strategy="freeze")
+
+        # 6 Save it!
+        trainer.save_checkpoint("/tmp/object_detection_model.pt")
+
+        # 7 Generate predictions
+        predictions = trainer.predict(
+            model,
+            datamodule=datamodule,
+            output=FiftyOneDetectionLabelsOutput(labels=datamodule.labels),
+        )
+        predictions = list(chain.from_iterable(predictions))  # flatten batches
+
+        # Map filepaths to predictions
+        predictions = {p["filepath"]: p["predictions"] for p in predictions}
+
+        # Add predictions to FiftyOne dataset
+        dataset.set_values(
+            "flash_predictions", predictions, key_field="filepath",
+        )
+
+    def test_semantic_segmentation(self):
+        # 1 Load your FiftyOne dataset
+        # source: https://www.kaggle.com/kumaresanmanickavelu/lyft-udacity-challenge
+        download_data(
+            "https://github.com/ongchinkiat/LyftPerceptionChallenge/releases/download/v0.1/carla-capture-20180513A.zip",
+            "/tmp/carla_data/",
+        )
+
+        dataset = fo.Dataset.from_dir(
+            dataset_dir="/tmp/carla_data",
+            dataset_type=fo.types.ImageSegmentationDirectory,
+            data_path="CameraRGB",
+            labels_path="CameraSeg",
+            force_grayscale=True,
+            shuffle=True,
+        )
+
+        # Just test and val on train dataset for this example
+        predict_dataset = dataset.take(5)
+
+        # 2 Create the Datamodule
+        datamodule = SemanticSegmentationData.from_fiftyone(
+            train_dataset=dataset,
+            test_dataset=dataset,
+            val_dataset=dataset,
+            predict_dataset=predict_dataset,
+            label_field="ground_truth",
+            transform_kwargs=dict(image_size=(256, 256)),
+            num_classes=21,
+            batch_size=4,
+        )
+
+        # 3 Build the model
+        model = SemanticSegmentation(
+            backbone="mobilenetv3_large_100",
+            head="fpn",
+            num_classes=datamodule.num_classes,
+        )
+
+        # 4 Create the trainer
+        trainer = Trainer(
+            max_epochs=1, limit_train_batches=10, limit_val_batches=5
+        )
+
+        # 5 Finetune the model
+        trainer.finetune(model, datamodule=datamodule, strategy="freeze")
+
+        # 6 Save it!
+        trainer.save_checkpoint("/tmp/semantic_segmentation_model.pt")
+
+        # 7 Generate predictions
+        predictions = trainer.predict(
+            model,
+            datamodule=datamodule,
+            output=FiftyOneSegmentationLabelsOutput(),
+        )
+        predictions = list(chain.from_iterable(predictions))  # flatten batches
+
+        # Map filepaths to predictions
+        predictions = {p["filepath"]: p["predictions"] for p in predictions}
+
+        # Add predictions to FiftyOne dataset
+        dataset.set_values(
+            "flash_predictions", predictions, key_field="filepath",
+        )
+
+        # Test segmentation apply_model
+        predict_dataset.apply_model(model, "seg_apply_model")
+
+    def test_video_classification(self):
+        # 1 Load the data
+        dataset = foz.load_zoo_dataset(
+            "kinetics-700-2020",
+            split="validation",
+            max_samples=15,
+            shuffle=True,
+        ).clone()
+        dataset.untag_samples("validation")
+
+        # Replace spaces in class names with underscore
+        labels = dataset.distinct("ground_truth.label")
+        labels_map = {l: l.replace(" ", "_") for l in labels}
+        dataset = dataset.map_labels("ground_truth", labels_map).clone()
+
+        # Get list of labels in dataset
+        labels = dataset.distinct("ground_truth.label")
+
+        # Create splits from the dataset
+        splits = {"train": 0.7, "pred": 0.3}
+        fous.random_split(dataset, splits)
+
+        # Here we use views into one dataset,
+        # but you can also use a different dataset for each split
+        train_dataset = dataset.match_tags("train")
+        predict_dataset = dataset.match_tags("pred")
+
+        # 2 Create the Datamodule
+        datamodule = VideoClassificationData.from_fiftyone(
+            train_dataset=dataset,
+            predict_dataset=predict_dataset,
+            label_field="ground_truth",
+            batch_size=1,
+            clip_sampler="uniform",
+            clip_duration=1,
+            decode_audio=False,
+        )
+
+        # 3 Build the model
+        model = VideoClassifier(
+            backbone="x3d_xs", labels=datamodule.labels, pretrained=False,
+        )
+
+        # 4 Create the trainer
+        trainer = Trainer(max_epochs=1, limit_train_batches=5)
+
+        # 5 Finetune the model
+        trainer.finetune(model, datamodule=datamodule, strategy="freeze")
+
+        # 6 Save it!
+        trainer.save_checkpoint("/tmp/video_classification.pt")
+
+        # 7 Generate predictions
+        predictions = trainer.predict(
+            model,
+            datamodule=datamodule,
+            output=FiftyOneLabelsOutput(labels=datamodule.labels),
+        )
+        predictions = list(chain.from_iterable(predictions))  # flatten batches
+
+        # Map filepaths to predictions
+        predictions = {p["filepath"]: p["predictions"] for p in predictions}
+
+        # Add predictions to FiftyOne dataset
+        predict_dataset.set_values(
+            "flash_predictions", predictions, key_field="filepath",
+        )
+
+    def test_manually_adding_predictions(self):
+        # Load your dataset
+        dataset = (
+            foz.load_zoo_dataset("quickstart", max_samples=5)
+            .select_fields("ground_truth")
+            .clone()
+        )
+        labels = dataset.distinct("ground_truth.detections.label")
+
+        # Load your Flash model
+        model = ImageClassifier(labels=labels)
+
+        # Create prediction datamodule
+        datamodule = ImageClassificationData.from_fiftyone(
+            predict_dataset=dataset, batch_size=1,
+        )
+
+        # Output FiftyOne format
+        output = FiftyOneLabelsOutput(return_filepath=False, labels=labels)
+        # Predict with trainer
+        predictions = Trainer().predict(
+            model, datamodule=datamodule, output=output
+        )
+
+        predictions = list(chain.from_iterable(predictions))  # flatten batches
+
+        # Predictions is a list of Label objects since ``return_filepath=False``
+        # Order corresponds to order of the ``predict_dataset``
+
+        # Add predictions to dataset
+        dataset.set_values("flash_predictions", predictions)
+
+    def test_specifying_class_names(self):
+        # Load your dataset
+        dataset = foz.load_zoo_dataset("quickstart", max_samples=5).clone()
+
+        datamodule = ImageClassificationData.from_fiftyone(
+            predict_dataset=dataset, batch_size=1
+        )
+
+        # Load your Flash model
+        num_classes = 100
+        model = ImageClassifier(backbone="resnet18", num_classes=num_classes)
+
+        # Configure output with class labels
+        labels = [
+            "label_" + str(i) for i in range(num_classes)
+        ]  # example class labels
+        output = FiftyOneLabelsOutput(labels=labels)  # output FiftyOne format
+
+        # Predict with model
+        trainer = Trainer()
+        predictions = trainer.predict(
+            model, datamodule=datamodule, output=output
+        )
+
+        predictions = list(chain.from_iterable(predictions))  # flatten batches
+
+        # Map filepaths to predictions
+        predictions = {p["filepath"]: p["predictions"] for p in predictions}
+
+        # Add predictions to dataset
+        dataset.set_values(
+            "flash_predictions", predictions, key_field="filepath"
+        )
+
+        print(dataset.distinct("flash_predictions.label"))
+        # ['label_57', 'label_60']
+
+    def test_image_embedder(self):
+        ## 1 Download data
+        download_data(
+            "https://pl-flash-data.s3.amazonaws.com/hymenoptera_data.zip",
+            "/tmp",
+        )
+
+        # 2 Load data into FiftyOne
+        dataset = fo.Dataset.from_dir(
+            "/tmp/hymenoptera_data/test/",
+            fo.types.ImageClassificationDirectoryTree,
+        )
+        datamodule = ImageClassificationData.from_fiftyone(
+            predict_dataset=dataset, batch_size=1,
+        )
+
+        # 3 Load model
+        # embedder = ImageEmbedder(backbone="resnet", head=None,
+        #        pretraining_transform=None, training_strategy="default")
+        embedder = ImageEmbedder(
+            backbone="vision_transformer",
+            training_strategy="barlow_twins",
+            head="barlow_twins_head",
+            pretraining_transform="barlow_twins_transform",
+            training_strategy_kwargs={"latent_embedding_dim": 128},
+            pretraining_transform_kwargs={"size_crops": [32]},
+        )
+
+        # 4 Generate embeddings
+        trainer = Trainer()
+        embedding_batches = trainer.predict(embedder, datamodule=datamodule)
+        embeddings = np.stack(sum(embedding_batches, []))
+
+        # 5 Visualize in FiftyOne App
+        results = fob.compute_visualization(dataset, embeddings=embeddings)
+        plot = results.visualize(labels="ground_truth.label")
+        plot.show()
+
+        embeddings = dataset.compute_embeddings(embedder)
+        dataset.compute_embeddings(embedder, embeddings_field="embeddings")
+
+
+if __name__ == "__main__":
+    fo.config.show_progress_bars = False
+    unittest.main(verbosity=2)