IceVision integration

.gitignore

@@ -161,3 +161,4 @@ jigsaw_toxic_comments
 flash_examples/serve/tabular_classification/data
 logs/cache/*
 flash_examples/data
+flash_examples/checkpoints

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

@@ -36,6 +36,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/core/data/data_module.py

@@ -363,12 +363,13 @@ 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,
                 pin_memory=pin_memory,
-                collate_fn=collate_fn
+                collate_fn=collate_fn,
+                convert_to_dataloader=True,
             )
 
         return DataLoader(

flash/core/data/data_pipeline.py

@@ -223,6 +223,7 @@ def _create_collate_preprocessors(
         prefix: str = _STAGES_PREFIX[stage]
 
         if collate_fn is not None:
+            preprocess._original_default_collate = preprocess._default_collate
             preprocess._default_collate = collate_fn
 
         func_names: Dict[str, str] = {
@@ -486,6 +487,10 @@ def _detach_preprocessing_from_model(self, model: 'Task', stage: Optional[Runnin
         elif isinstance(stage, RunningStage):
             stages = [stage]
 
+        self._preprocess_pipeline._default_collate = getattr(
+            self._preprocess_pipeline, "_original_default_collate", self._preprocess_pipeline._default_collate
+        )
+
         for stage in stages:
 
             device_collate = None

flash/core/integrations/icevision/backbones.py

@@ -0,0 +1,67 @@
+# 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 inspect import getmembers
+
+from torch import nn
+
+from flash.core.registry import FlashRegistry
+from flash.core.utilities.imports import _ICEVISION_AVAILABLE
+
+if _ICEVISION_AVAILABLE:
+    from icevision.backbones import BackboneConfig
+
+OBJECT_DETECTION_HEADS = FlashRegistry("heads")
+
+
+def icevision_model_adapter(model_type):
+
+    class IceVisionModelAdapter(model_type.lightning.ModelAdapter):
+
+        def log(self, name, value, **kwargs):
+            if "prog_bar" not in kwargs:
+                kwargs["prog_bar"] = True
+            return super().log(name.split("/")[-1], value, **kwargs)
+
+    return IceVisionModelAdapter
+
+
+def load_icevision(adapter, model_type, backbone, num_classes, **kwargs):
+    model = model_type.model(backbone=backbone, num_classes=num_classes, **kwargs)
+
+    backbone = nn.Module()
+    params = model.param_groups()[0]
+    for i, param in enumerate(params):
+        backbone.register_parameter(f"backbone_{i}", param)
+
+    return model_type, model, adapter(model_type), backbone
+
+
+def load_icevision_ignore_image_size(adapter, model_type, backbone, num_classes, image_size=None, **kwargs):
+    return load_icevision(adapter, model_type, backbone, num_classes, **kwargs)
+
+
+def load_icevision_with_image_size(adapter, model_type, backbone, num_classes, image_size=None, **kwargs):
+    kwargs["img_size"] = image_size
+    return load_icevision(adapter, model_type, backbone, num_classes, **kwargs)
+
+
+def get_backbones(model_type):
+    _BACKBONES = FlashRegistry("backbones")
+
+    for backbone_name, backbone_config in getmembers(model_type.backbones, lambda x: isinstance(x, BackboneConfig)):
+        _BACKBONES(
+            backbone_config,
+            name=backbone_name,
+        )
+    return _BACKBONES

flash/core/integrations/icevision/data.py

@@ -0,0 +1,79 @@
+# 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 typing import Any, Callable, Dict, Optional, Sequence, Tuple, Type
+
+import numpy as np
+
+from flash.core.data.data_source import DefaultDataKeys
+from flash.core.utilities.imports import _ICEVISION_AVAILABLE
+from flash.image.data import ImagePathsDataSource
+
+if _ICEVISION_AVAILABLE:
+    from icevision.core import BaseRecord, ClassMapRecordComponent, ImageRecordComponent, tasks
+    from icevision.data import SingleSplitSplitter
+    from icevision.parsers import Parser
+
+
+class IceVisionPathsDataSource(ImagePathsDataSource):
+
+    def predict_load_data(self, data: Tuple[str, str], dataset: Optional[Any] = None) -> Sequence[Dict[str, Any]]:
+        return super().predict_load_data(data, dataset)
+
+    def load_sample(self, sample: Dict[str, Any]) -> Dict[str, Any]:
+        return sample[DefaultDataKeys.INPUT].load()
+
+    def predict_load_sample(self, sample: Dict[str, Any]) -> Dict[str, Any]:
+        sample = super().load_sample(sample)
+        image = np.array(sample[DefaultDataKeys.INPUT])
+        record = BaseRecord([ImageRecordComponent()])
+
+        record.set_img(image)
+        record.add_component(ClassMapRecordComponent(task=tasks.detection))
+        return record
+
+
+class IceVisionParserDataSource(IceVisionPathsDataSource):
+
+    def __init__(self, parser: Optional[Type['Parser']] = None):
+        super().__init__()
+        self.parser = parser
+
+    def load_data(self, data: Tuple[str, str], dataset: Optional[Any] = None) -> Sequence[Dict[str, Any]]:
+        root, ann_file = data
+
+        if self.parser is not None:
+            parser = self.parser(ann_file, root)
+            dataset.num_classes = len(parser.class_map)
+            records = parser.parse(data_splitter=SingleSplitSplitter())
+            return [{DefaultDataKeys.INPUT: record} for record in records[0]]
+        else:
+            raise ValueError("The parser type must be provided")
+
+
+class IceDataParserDataSource(IceVisionPathsDataSource):
+
+    def __init__(self, parser: Optional[Callable] = None):
+        super().__init__()
+        self.parser = parser
+
+    def load_data(self, data: Tuple[str, str], dataset: Optional[Any] = None) -> Sequence[Dict[str, Any]]:
+        root = data
+
+        if self.parser is not None:
+            parser = self.parser(root)
+            dataset.num_classes = len(parser.class_map)
+            records = parser.parse(data_splitter=SingleSplitSplitter())
+            return [{DefaultDataKeys.INPUT: record} for record in records[0]]
+        else:
+            raise ValueError("The parser must be provided")