IceVision integration

.github/workflows/ci-testing.yml

@@ -137,7 +137,7 @@ jobs:
       run: |
         sudo apt-get install libsndfile1
         pip install matplotlib
-        pip install '.[image]' --pre --upgrade
+        pip install '.[image]' --pre --upgrade torchaudio
 
     - name: Cache datasets
       uses: actions/cache@v2

.gitignore

@@ -161,5 +161,5 @@ jigsaw_toxic_comments
 flash_examples/serve/tabular_classification/data
 logs/cache/*
 flash_examples/data
-flash_examples/cli/*/data
+flash_examples/checkpoints
 timit/

CHANGELOG.md

@@ -34,12 +34,20 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Added Flash Zero, a zero code command line ML platform built with flash ([#611](https://github.com/PyTorchLightning/lightning-flash/pull/611))
 
+- Added integration with IceVision for the `ObjectDetector` ([#608](https://github.com/PyTorchLightning/lightning-flash/pull/608))
+
+- Added keypoint detection task ([#608](https://github.com/PyTorchLightning/lightning-flash/pull/608))
+
+- Added instance segmentation task ([#608](https://github.com/PyTorchLightning/lightning-flash/pull/608))
+
 ### Changed
 
 - Changed how pretrained flag works for loading weights for ImageClassifier task ([#560](https://github.com/PyTorchLightning/lightning-flash/pull/560))
 
 - Removed bolts pretrained weights for SSL from ImageClassifier task ([#560](https://github.com/PyTorchLightning/lightning-flash/pull/560))
 
+- Changed arguments to `ObjectDetector`, use `head` instead of `model` and append `_fpn` to the backbone name instead of the `fpn` argument ([#608](https://github.com/PyTorchLightning/lightning-flash/pull/608))
+
 ### Fixed
 
 - Fixed a bug where serve sanity checking would not be triggered using the latest PyTorchLightning version ([#493](https://github.com/PyTorchLightning/lightning-flash/pull/493))

docs/source/api/image.rst

@@ -31,8 +31,8 @@ ______________
     classification.transforms.default_transforms
     classification.transforms.train_default_transforms
 
-Detection
-_________
+Object Detection
+________________
 
 .. autosummary::
     :toctree: generated/
@@ -42,21 +42,37 @@ _________
     ~detection.model.ObjectDetector
     ~detection.data.ObjectDetectionData
 
-    detection.data.COCODataSource
+    detection.data.FiftyOneParser
     detection.data.ObjectDetectionFiftyOneDataSource
     detection.data.ObjectDetectionPreprocess
-    detection.finetuning.ObjectDetectionFineTuning
-    detection.model.ObjectDetector
     detection.serialization.DetectionLabels
     detection.serialization.FiftyOneDetectionLabels
 
+Keypoint Detection
+__________________
+
 .. autosummary::
     :toctree: generated/
     :nosignatures:
-    :template:
+    :template: classtemplate.rst
+
+    ~keypoint_detection.model.KeypointDetector
+    ~keypoint_detection.data.KeypointDetectionData
+
+    keypoint_detection.data.KeypointDetectionPreprocess
+
+Instance Segmentation
+_____________________
+
+.. autosummary::
+    :toctree: generated/
+    :nosignatures:
+    :template: classtemplate.rst
+
+    ~instance_segmentation.model.InstanceSegmentation
+    ~instance_segmentation.data.InstanceSegmentationData
 
-    detection.transforms.collate
-    detection.transforms.default_transforms
+    instance_segmentation.data.InstanceSegmentationPreprocess
 
 Embedding
 _________

docs/source/index.rst

@@ -37,6 +37,8 @@ Lightning Flash
    reference/image_classification_multi_label
    reference/image_embedder
    reference/object_detection
+   reference/keypoint_detection
+   reference/instance_segmentation
    reference/semantic_segmentation
    reference/style_transfer
    reference/video_classification

docs/source/reference/instance_segmentation.rst

@@ -0,0 +1,31 @@
+
+.. _instance_segmentation:
+
+#####################
+Instance Segmentation
+#####################
+
+********
+The Task
+********
+
+Instance segmentation is the task of segmenting objects images and determining their associated classes.
+
+The :class:`~flash.image.instance_segmentation.model.InstanceSegmentation` and :class:`~flash.image.instance_segmentation.data.InstanceSegmentationData` classes internally rely on `IceVision <https://airctic.com/>`_.
+
+------
+
+*******
+Example
+*******
+
+Let's look at instance segmentation with `The Oxford-IIIT Pet Dataset <https://www.robots.ox.ac.uk/~vgg/data/pets/>`_ from `IceData <https://github.com/airctic/icedata>`_.
+Once we've downloaded the data, we can create the :class:`~flash.image.instance_segmentation.data.InstanceSegmentationData`.
+We select a ``mask_rcnn`` with a ``resnet18_fpn`` backbone to use for our :class:`~flash.image.instance_segmentation.model.InstanceSegmentation` and fine-tune on the pets data.
+We then use the trained :class:`~flash.image.instance_segmentation.model.InstanceSegmentation` for inference.
+Finally, we save the model.
+Here's the full example:
+
+.. literalinclude:: ../../../flash_examples/instance_segmentation.py
+    :language: python
+    :lines: 14-

docs/source/reference/keypoint_detection.rst

@@ -0,0 +1,31 @@
+
+.. _keypoint_detection:
+
+##################
+Keypoint Detection
+##################
+
+********
+The Task
+********
+
+Keypoint detection is the task of identifying keypoints in images and their associated classes.
+
+The :class:`~flash.image.keypoint_detection.model.KeypointDetector` and :class:`~flash.image.keypoint_detection.data.KeypointDetectionData` classes internally rely on `IceVision <https://airctic.com/>`_.
+
+------
+
+*******
+Example
+*******
+
+Let's look at keypoint detection with `BIWI Sample Keypoints (center of face) <https://www.kaggle.com/kmader/biwi-kinect-head-pose-database>`_ from `IceData <https://github.com/airctic/icedata>`_.
+Once we've downloaded the data, we can create the :class:`~flash.image.keypoint_detection.data.KeypointDetectionData`.
+We select a ``keypoint_rcnn`` with a ``resnet18_fpn`` backbone to use for our :class:`~flash.image.keypoint_detection.model.KeypointDetector` and fine-tune on the BIWI data.
+We then use the trained :class:`~flash.image.keypoint_detection.model.KeypointDetector` for inference.
+Finally, we save the model.
+Here's the full example:
+
+.. literalinclude:: ../../../flash_examples/keypoint_detection.py
+    :language: python
+    :lines: 14-

docs/source/reference/object_detection.rst

@@ -11,6 +11,8 @@ The Task
 
 Object detection is the task of identifying objects in images and their associated classes and bounding boxes.
 
+The :class:`~flash.image.detection.model.ObjectDetector` and :class:`~flash.image.detection.data.ObjectDetectionData` classes internally rely on `IceVision <https://airctic.com/>`_.
+
 ------
 
 *******

flash/__about__.py

@@ -1,4 +1,4 @@
-__version__ = "0.4.1dev"
+__version__ = "0.5.0dev"
 __author__ = "PyTorchLightning et al."
 __author_email__ = "[email protected]"
 __license__ = "Apache-2.0"

flash/core/adapter.py

@@ -0,0 +1,148 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 abc import abstractmethod
+from typing import Any, Callable, Optional
+
+from torch import nn
+from torch.utils.data import DataLoader, Sampler
+
+import flash
+from flash.core.data.auto_dataset import BaseAutoDataset
+from flash.core.model import DatasetProcessor, Task, Wrapper
+
+
+class Adapter(DatasetProcessor, Wrapper, nn.Module):
+    @classmethod
+    @abstractmethod
+    def from_task(cls, task: "flash.Task", **kwargs) -> "Adapter":
+        pass
+
+    def forward(self, x: Any) -> Any:
+        pass
+
+    def training_step(self, batch: Any, batch_idx: int) -> Any:
+        pass
+
+    def validation_step(self, batch: Any, batch_idx: int) -> None:
+        pass
+
+    def test_step(self, batch: Any, batch_idx: int) -> None:
+        pass
+
+    def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> Any:
+        pass
+
+    def training_epoch_end(self, outputs) -> None:
+        pass
+
+    def validation_epoch_end(self, outputs) -> None:
+        pass
+
+    def test_epoch_end(self, outputs) -> None:
+        pass
+
+
+class AdapterTask(Task):
+    def __init__(self, adapter: Adapter, **kwargs):
+        super().__init__(**kwargs)
+
+        self.adapter = adapter
+
+    @property
+    def backbone(self) -> nn.Module:
+        return self.adapter.backbone
+
+    def forward(self, x: Any) -> Any:
+        return self.adapter.forward(x)
+
+    def training_step(self, batch: Any, batch_idx: int) -> Any:
+        return self.adapter.training_step(batch, batch_idx)
+
+    def validation_step(self, batch: Any, batch_idx: int) -> None:
+        return self.adapter.validation_step(batch, batch_idx)
+
+    def test_step(self, batch: Any, batch_idx: int) -> None:
+        return self.adapter.test_step(batch, batch_idx)
+
+    def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> Any:
+        return self.adapter.predict_step(batch, batch_idx, dataloader_idx=dataloader_idx)
+
+    def training_epoch_end(self, outputs) -> None:
+        return self.adapter.training_epoch_end(outputs)
+
+    def validation_epoch_end(self, outputs) -> None:
+        return self.adapter.validation_epoch_end(outputs)
+
+    def test_epoch_end(self, outputs) -> None:
+        return self.adapter.test_epoch_end(outputs)
+
+    def process_train_dataset(
+        self,
+        dataset: BaseAutoDataset,
+        batch_size: int,
+        num_workers: int,
+        pin_memory: bool,
+        collate_fn: Callable,
+        shuffle: bool = False,
+        drop_last: bool = True,
+        sampler: Optional[Sampler] = None,
+    ) -> DataLoader:
+        return self.adapter.process_train_dataset(
+            dataset, batch_size, num_workers, pin_memory, collate_fn, shuffle, drop_last, sampler
+        )
+
+    def process_val_dataset(
+        self,
+        dataset: BaseAutoDataset,
+        batch_size: int,
+        num_workers: int,
+        pin_memory: bool,
+        collate_fn: Callable,
+        shuffle: bool = False,
+        drop_last: bool = False,
+        sampler: Optional[Sampler] = None,
+    ) -> DataLoader:
+        return self.adapter.process_val_dataset(
+            dataset, batch_size, num_workers, pin_memory, collate_fn, shuffle, drop_last, sampler
+        )
+
+    def process_test_dataset(
+        self,
+        dataset: BaseAutoDataset,
+        batch_size: int,
+        num_workers: int,
+        pin_memory: bool,
+        collate_fn: Callable,
+        shuffle: bool = False,
+        drop_last: bool = False,
+        sampler: Optional[Sampler] = None,
+    ) -> DataLoader:
+        return self.adapter.process_test_dataset(
+            dataset, batch_size, num_workers, pin_memory, collate_fn, shuffle, drop_last, sampler
+        )
+
+    def process_predict_dataset(
+        self,
+        dataset: BaseAutoDataset,
+        batch_size: int = 1,
+        num_workers: int = 0,
+        pin_memory: bool = False,
+        collate_fn: Callable = lambda x: x,
+        shuffle: bool = False,
+        drop_last: bool = True,
+        sampler: Optional[Sampler] = None,
+    ) -> DataLoader:
+        return self.adapter.process_predict_dataset(
+            dataset, batch_size, num_workers, pin_memory, collate_fn, shuffle, drop_last, sampler
+        )

flash/core/data/data_module.py

@@ -361,7 +361,7 @@ def _predict_dataloader(self) -> DataLoader:
         pin_memory = True
 
         if isinstance(getattr(self, "trainer", None), pl.Trainer):
-            return self.trainer.lightning_module.process_test_dataset(
+            return self.trainer.lightning_module.process_predict_dataset(
                 predict_ds,
                 batch_size=batch_size,
                 num_workers=self.num_workers,

flash/core/data/data_pipeline.py

@@ -164,8 +164,10 @@ def _identity(samples: Sequence[Any]) -> Sequence[Any]:
     def deserialize_processor(self) -> _DeserializeProcessor:
         return self._create_collate_preprocessors(RunningStage.PREDICTING)[0]
 
-    def worker_preprocessor(self, running_stage: RunningStage, is_serving: bool = False) -> _Preprocessor:
-        return self._create_collate_preprocessors(running_stage, is_serving=is_serving)[1]
+    def worker_preprocessor(
+        self, running_stage: RunningStage, collate_fn: Optional[Callable] = None, is_serving: bool = False
+    ) -> _Preprocessor:
+        return self._create_collate_preprocessors(running_stage, collate_fn=collate_fn, is_serving=is_serving)[1]
 
     def device_preprocessor(self, running_stage: RunningStage) -> _Preprocessor:
         return self._create_collate_preprocessors(running_stage)[2]