Repository: Lightning-Universe/lightning-flash

Issue:
BUG: cannot use `faster_rcnn` in ObjectDetector
## 🐛 Bug
```
Traceback (most recent call last):
  File "/mnt/data/buduanban_detection/test.py", line 25, in <module>
    trainer.finetune(model, datamodule=datamodule, strategy="freeze")
  File "/data/miniconda3/envs/torch/lib/python3.9/site-packages/lightning_flash-0.8.0.dev0-py3.9.egg/flash/core/trainer.py", line 161, in finetune
    return super().fit(model, train_dataloader, val_dataloaders, datamodule)
  File "/data/miniconda3/envs/torch/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 740, in fit
    self._call_and_handle_interrupt(
  File "/data/miniconda3/envs/torch/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 685, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "/data/miniconda3/envs/torch/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 777, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/data/miniconda3/envs/torch/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1199, in _run
    self._dispatch()
  File "/data/miniconda3/envs/torch/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1279, in _dispatch
    self.training_type_plugin.start_training(self)
  File "/data/miniconda3/envs/torch/lib/python3.9/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py", line 202, in start_training
    self._results = trainer.run_stage()
  File "/data/miniconda3/envs/torch/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1289, in run_stage
    return self._run_train()
  File "/data/miniconda3/envs/torch/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1311, in _run_train
    self._run_sanity_check(self.lightning_module)
  File "/data/miniconda3/envs/torch/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1375, in _run_sanity_check
    self._evaluation_loop.run()
  File "/data/miniconda3/envs/torch/lib/python3.9/site-packages/pytorch_lightning/loops/base.py", line 145, in run
    self.advance(*args, **kwargs)
  File "/data/miniconda3/envs/torch/lib/python3.9/site-packages/pytorch_lightning/loops/dataloader/evaluation_loop.py", line 110, in advance
    dl_outputs = self.epoch_loop.run(dataloader, dataloader_idx, dl_max_batches, self.num_dataloaders)
  File "/data/miniconda3/envs/torch/lib/python3.9/site-packages/pytorch_lightning/loops/base.py", line 145, in run
    self.advance(*args, **kwargs)
  File "/data/miniconda3/envs/torch/lib/python3.9/site-packages/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py", line 122, in advance
    output = self._evaluation_step(batch, batch_idx, dataloader_idx)
  File "/data/miniconda3/envs/torch/lib/python3.9/site-packages/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py", line 217, in _evaluation_step
    output = self.trainer.accelerator.validation_step(step_kwargs)
  File "/data/miniconda3/envs/torch/lib/python3.9/site-packages/pytorch_lightning/accelerators/accelerator.py", line 239, in validation_step
    return self.training_type_plugin.validation_step(*step_kwargs.values())
  File "/data/miniconda3/envs/torch/lib/python3.9/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py", line 219, in validation_step
    return self.model.validation_step(*args, **kwargs)
  File "/data/miniconda3/envs/torch/lib/python3.9/site-packages/lightning_flash-0.8.0.dev0-py3.9.egg/flash/core/adapter.py", line 113, in validation_step
    return self.adapter.validation_step(batch, batch_idx)
  File "/data/miniconda3/envs/torch/lib/python3.9/site-packages/lightning_flash-0.8.0.dev0-py3.9.egg/flash/core/integrations/icevision/adapter.py", line 233, in validation_step
    return self.icevision_adapter.validation_step(batch[DataKeys.INPUT], batch_idx)
  File "/data/miniconda3/envs/torch/lib/python3.9/site-packages/icevision/models/torchvision/lightning/model_adapter.py", line 41, in validation_step
    train_preds = self(xb, yb)
  File "/data/miniconda3/envs/torch/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/data/miniconda3/envs/torch/lib/python3.9/site-packages/icevision/models/torchvision/lightning/model_adapter.py", line 26, in forward
    return self.model(*args, **kwargs)
  File "/data/miniconda3/envs/torch/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/data/miniconda3/envs/torch/lib/python3.9/site-packages/torchvision/models/detection/generalized_rcnn.py", line 96, in forward
    proposals, proposal_losses = self.rpn(images, features, targets)
  File "/data/miniconda3/envs/torch/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/data/miniconda3/envs/torch/lib/python3.9/site-packages/torchvision/models/detection/rpn.py", line 344, in forward
    anchors = self.anchor_generator(images, features)
  File "/data/miniconda3/envs/torch/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/data/miniconda3/envs/torch/lib/python3.9/site-packages/torchvision/models/detection/anchor_utils.py", line 125, in forward
    anchors_over_all_feature_maps = self.grid_anchors(grid_sizes, strides)
  File "/data/miniconda3/envs/torch/lib/python3.9/site-packages/torchvision/models/detection/anchor_utils.py", line 86, in grid_anchors
    raise ValueError("Anchors should be Tuple[Tuple[int]] because each feature "
ValueError: Anchors should be Tuple[Tuple[int]] because each feature map could potentially have different sizes and aspect ratios. There needs to be a match between the number of feature maps passed and the number of sizes / aspect ratios specified.
```
### To Reproduce
Change the head and backbone of [example in the doc](https://lightning-flash.readthedocs.io/en/latest/reference/object_detection.html) to  `head="faster_rcnn", backbone="mobilenet" `
#### Code sample
```python
import flash
from flash.core.data.utils import download_data
from flash.image import ObjectDetectionData, ObjectDetector

# 1. Create the DataModule
# Dataset Credit: https://www.kaggle.com/ultralytics/coco128
# download_data("https://github.com/zhiqwang/yolov5-rt-stack/releases/download/v0.3.0/coco128.zip", "data/")

datamodule = ObjectDetectionData.from_coco(
    train_folder="data/coco128/images/train2017/",
    train_ann_file="data/coco128/annotations/instances_train2017.json",
    val_split=0.1,
    transform_kwargs={"image_size": 512},
    batch_size=4,
)

# 2. Build the task
model = ObjectDetector(head="faster_rcnn",   #  <----  use faster_rcnn
                       backbone="mobilenet",     #  <----  use mobilenet
                       num_classes=datamodule.num_classes, 
                       image_size=512)

# 3. Create the trainer and finetune the model
trainer = flash.Trainer(max_epochs=1)
trainer.finetune(model, datamodule=datamodule, strategy="freeze")

# 4. Detect objects in a few images!
datamodule = ObjectDetectionData.from_files(
    predict_files=[
        "data/coco128/images/train2017/000000000625.jpg",
        "data/coco128/images/train2017/000000000626.jpg",
        "data/coco128/images/train2017/000000000629.jpg",
    ],
    transform_kwargs={"image_size": 512},
    batch_size=4,
)
predictions = trainer.predict(model, datamodule=datamodule)
print(predictions)

# 5. Save the model!
trainer.save_checkpoint("object_detection_model.pt")
```

### Expected behavior

<!-- A clear and concise description of what you expected to happen. -->

### Environment

 - OS (e.g., Linux):
 - Python version:
 - PyTorch/Lightning/Flash Version (e.g., 1.10/1.5/0.7):
 - GPU models and configuration:
 - Any other relevant information:

### Additional context

<!-- Add any other context about the problem here. -->



Discussion:
BTW, all backbones in `faster_rcnn` has the same bug.
Hey @IncubatorShokuhou Thanks for reporting this! I can reproduce the error, looking in to it 😃 
@IncubatorShokuhou Just to give you a bit of an update, the issue only relates to non-fpn backbones. So any backbone with `fpn` in the name should work. E.g. `resnet18_fpn` is working fine for me. Working on a fix for the other backbones 😃 

Resolving patch:
diff --git a/flash/core/integrations/icevision/backbones.py b/flash/core/integrations/icevision/backbones.py
--- a/flash/core/integrations/icevision/backbones.py
+++ b/flash/core/integrations/icevision/backbones.py
@@ -51,6 +51,10 @@ def get_backbones(model_type):
     _BACKBONES = FlashRegistry("backbones")
 
     for backbone_name, backbone_config in getmembers(model_type.backbones, lambda x: isinstance(x, BackboneConfig)):
+        # Only torchvision backbones with an FPN are supported
+        if "torchvision" in model_type.__name__ and "fpn" not in backbone_name:
+            continue
+
         _BACKBONES(
             backbone_config,
             name=backbone_name,


Test patch:
diff --git a/tests/image/detection/test_data_model_integration.py b/tests/image/detection/test_data_model_integration.py
--- a/tests/image/detection/test_data_model_integration.py
+++ b/tests/image/detection/test_data_model_integration.py
@@ -41,7 +41,7 @@
 
 @pytest.mark.skipif(not _IMAGE_AVAILABLE, reason="image libraries aren't installed.")
 @pytest.mark.skipif(not _ICEVISION_AVAILABLE, reason="IceVision is not installed for testing")
-@pytest.mark.parametrize(["head", "backbone"], [("retinanet", "resnet18_fpn")])
+@pytest.mark.parametrize(["head", "backbone"], [("retinanet", "resnet18_fpn"), ("faster_rcnn", "resnet18_fpn")])
 def test_detection(tmpdir, head, backbone):
 
     train_folder, coco_ann_path = _create_synth_coco_dataset(tmpdir)
diff --git a/tests/image/instance_segmentation/test_model.py b/tests/image/instance_segmentation/test_model.py
--- a/tests/image/instance_segmentation/test_model.py
+++ b/tests/image/instance_segmentation/test_model.py
@@ -11,88 +11,110 @@
 # 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 collections
+import json
 import os
-from functools import partial
-from pathlib import Path
+from unittest import mock
+
+import numpy as np
+import pytest
 
-import flash
-from flash.core.utilities.imports import _ICEDATA_AVAILABLE, _ICEVISION_AVAILABLE, _IMAGE_AVAILABLE
+from flash import Trainer
+from flash.__main__ import main
+from flash.core.utilities.imports import _ICEVISION_AVAILABLE, _IMAGE_AVAILABLE
 from flash.image import InstanceSegmentation, InstanceSegmentationData
 
-if _ICEDATA_AVAILABLE:
-    import icedata
-if _ICEVISION_AVAILABLE:
-    import icevision
+if _IMAGE_AVAILABLE:
+    from PIL import Image
 
-from unittest import mock
+COCODataConfig = collections.namedtuple("COCODataConfig", "train_folder train_ann_file predict_folder")
 
-import pytest
 
-from flash.__main__ import main
+@pytest.fixture
+def coco_instances(tmpdir):
+    rand_image = Image.fromarray(np.random.randint(0, 255, (64, 64, 3), dtype="uint8"))
+    os.makedirs(tmpdir / "train_folder", exist_ok=True)
+    os.makedirs(tmpdir / "predict_folder", exist_ok=True)
 
+    train_folder = tmpdir / "train_folder"
+    train_ann_file = tmpdir / "train_annotations.json"
+    predict_folder = tmpdir / "predict_folder"
 
-@pytest.mark.skipif(not _IMAGE_AVAILABLE, reason="image libraries aren't installed.")
-@pytest.mark.skipif(not _ICEVISION_AVAILABLE, reason="IceVision is not installed for testing")
-def test_cli():
-    cli_args = ["flash", "instance_segmentation", "--trainer.fast_dev_run", "True"]
-    with mock.patch("sys.argv", cli_args):
-        try:
-            main()
-        except SystemExit:
-            pass
+    _ = [rand_image.save(str(train_folder / f"image_{i}.png")) for i in range(1, 4)]
+    _ = [rand_image.save(str(predict_folder / f"predict_image_{i}.png")) for i in range(1, 4)]
+    annotations = {
+        "annotations": [
+            {
+                "area": 50,
+                "bbox": [10, 20, 5, 10],
+                "category_id": 1,
+                "id": 1,
+                "image_id": 1,
+                "iscrowd": 0,
+                "segmentation": [[10, 20, 15, 20, 15, 30, 10, 30]],
+            },
+            {
+                "area": 100,
+                "bbox": [20, 30, 10, 10],
+                "category_id": 2,
+                "id": 2,
+                "image_id": 2,
+                "iscrowd": 0,
+                "segmentation": [[20, 30, 30, 30, 30, 40, 20, 40]],
+            },
+            {
+                "area": 125,
+                "bbox": [10, 20, 5, 25],
+                "category_id": 1,
+                "id": 3,
+                "image_id": 3,
+                "iscrowd": 0,
+                "segmentation": [[10, 20, 15, 20, 15, 45, 10, 45]],
+            },
+        ],
+        "categories": [
+            {"id": 1, "name": "cat", "supercategory": "annimal"},
+            {"id": 2, "name": "dog", "supercategory": "annimal"},
+        ],
+        "images": [
+            {"file_name": "image_1.png", "height": 64, "width": 64, "id": 1},
+            {"file_name": "image_2.png", "height": 64, "width": 64, "id": 2},
+            {"file_name": "image_3.png", "height": 64, "width": 64, "id": 3},
+        ],
+    }
+    with open(train_ann_file, "w") as annotation_file:
+        json.dump(annotations, annotation_file)
+
+    return COCODataConfig(train_folder, train_ann_file, predict_folder)
 
 
-# todo: this test takes around 25s because of the icedata download, can we speed it up?
 @pytest.mark.skipif(not _IMAGE_AVAILABLE, reason="image libraries aren't installed.")
 @pytest.mark.skipif(not _ICEVISION_AVAILABLE, reason="IceVision is not installed for testing")
-def test_instance_segmentation_inference(tmpdir):
-    """Test to ensure that inference runs with instance segmentation from input paths."""
-    # modify the root path to use 'data' where our CI caches datasets
-    icevision.utils.data_dir.data_dir = Path("data/icevision/")
-    icevision.utils.data_dir.data_dir.mkdir(exist_ok=True, parents=True)
-    data_dir = icedata.pets.load_data()
-
-    datamodule = InstanceSegmentationData.from_icedata(
-        train_folder=data_dir,
-        val_split=0.1,
-        parser=partial(icedata.pets.parser, mask=True),
-        batch_size=1,
+@pytest.mark.parametrize("backbone, head", [("resnet18_fpn", "mask_rcnn")])
+def test_model(coco_instances, backbone, head):
+    datamodule = InstanceSegmentationData.from_coco(
+        train_folder=coco_instances.train_folder,
+        train_ann_file=coco_instances.train_ann_file,
+        predict_folder=coco_instances.predict_folder,
+        transform_kwargs=dict(image_size=(128, 128)),
+        batch_size=2,
     )
 
-    model = InstanceSegmentation(
-        head="mask_rcnn",
-        backbone="resnet18_fpn",
-        num_classes=datamodule.num_classes,
-    )
+    assert datamodule.num_classes == 3
+    assert datamodule.labels == ["background", "cat", "dog"]
 
-    # 3. Create the trainer and finetune the model
-    trainer = flash.Trainer(max_epochs=1, fast_dev_run=True)
-    trainer.finetune(model, datamodule=datamodule, strategy="freeze")
+    model = InstanceSegmentation(num_classes=datamodule.num_classes, backbone=backbone, head=head)
+    trainer = Trainer(fast_dev_run=True)
+    trainer.fit(model, datamodule=datamodule)
+    trainer.predict(model, datamodule=datamodule)
 
-    datamodule = InstanceSegmentationData.from_files(
-        predict_files=[
-            str(data_dir / "images/yorkshire_terrier_9.jpg"),
-            str(data_dir / "images/yorkshire_terrier_12.jpg"),
-            str(data_dir / "images/yorkshire_terrier_13.jpg"),
-        ],
-        batch_size=4,
-    )
-    predictions = trainer.predict(model, datamodule=datamodule)
-    assert len(predictions[0]) == 3
-    assert len(list(predictions[0][0].keys())) == 4
-
-    model_path = os.path.join(tmpdir, "model.pt")
-    trainer.save_checkpoint(model_path)
-    InstanceSegmentation.load_from_checkpoint(model_path)
 
-    datamodule = InstanceSegmentationData.from_files(
-        predict_files=[
-            str(data_dir / "images/yorkshire_terrier_9.jpg"),
-            str(data_dir / "images/yorkshire_terrier_12.jpg"),
-            str(data_dir / "images/yorkshire_terrier_13.jpg"),
-        ],
-        batch_size=4,
-    )
-    predictions = trainer.predict(model, datamodule=datamodule)
-    assert len(predictions[0]) == 3
-    assert len(list(predictions[0][0].keys())) == 4
+@pytest.mark.skipif(not _IMAGE_AVAILABLE, reason="image libraries aren't installed.")
+@pytest.mark.skipif(not _ICEVISION_AVAILABLE, reason="IceVision is not installed for testing")
+def test_cli():
+    cli_args = ["flash", "instance_segmentation", "--trainer.fast_dev_run", "True"]
+    with mock.patch("sys.argv", cli_args):
+        try:
+            main()
+        except SystemExit:
+            pass
diff --git a/tests/image/keypoint_detection/test_model.py b/tests/image/keypoint_detection/test_model.py
--- a/tests/image/keypoint_detection/test_model.py
+++ b/tests/image/keypoint_detection/test_model.py
@@ -11,12 +11,105 @@
 # 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 collections
+import json
+import os
 from unittest import mock
 
+import numpy as np
 import pytest
 
+from flash import Trainer
 from flash.__main__ import main
 from flash.core.utilities.imports import _ICEVISION_AVAILABLE, _IMAGE_AVAILABLE
+from flash.image import KeypointDetectionData, KeypointDetector
+
+if _IMAGE_AVAILABLE:
+    from PIL import Image
+
+COCODataConfig = collections.namedtuple("COCODataConfig", "train_folder train_ann_file predict_folder")
+
+
+@pytest.fixture
+def coco_keypoints(tmpdir):
+    rand_image = Image.fromarray(np.random.randint(0, 255, (64, 64, 3), dtype="uint8"))
+    os.makedirs(tmpdir / "train_folder", exist_ok=True)
+    os.makedirs(tmpdir / "predict_folder", exist_ok=True)
+
+    train_folder = tmpdir / "train_folder"
+    train_ann_file = tmpdir / "train_annotations.json"
+    predict_folder = tmpdir / "predict_folder"
+
+    _ = [rand_image.save(str(train_folder / f"image_{i}.png")) for i in range(1, 4)]
+    _ = [rand_image.save(str(predict_folder / f"predict_image_{i}.png")) for i in range(1, 4)]
+    annotations = {
+        "annotations": [
+            {
+                "area": 50,
+                "bbox": [10, 20, 5, 10],
+                "num_keypoints": 2,
+                "keypoints": [10, 15, 2, 20, 30, 2],
+                "category_id": 1,
+                "id": 1,
+                "image_id": 1,
+                "iscrowd": 0,
+            },
+            {
+                "area": 100,
+                "bbox": [20, 30, 10, 10],
+                "num_keypoints": 2,
+                "keypoints": [20, 30, 2, 30, 40, 2],
+                "category_id": 2,
+                "id": 2,
+                "image_id": 2,
+                "iscrowd": 0,
+            },
+            {
+                "area": 125,
+                "bbox": [10, 20, 5, 25],
+                "num_keypoints": 2,
+                "keypoints": [10, 15, 2, 20, 45, 2],
+                "category_id": 1,
+                "id": 3,
+                "image_id": 3,
+                "iscrowd": 0,
+            },
+        ],
+        "categories": [
+            {"id": 1, "name": "cat", "supercategory": "cat", "keypoints": ["left ear", "right ear"]},
+            {"id": 2, "name": "dog", "supercategory": "dog", "keypoints": ["left ear", "right ear"]},
+        ],
+        "images": [
+            {"file_name": "image_1.png", "height": 64, "width": 64, "id": 1},
+            {"file_name": "image_2.png", "height": 64, "width": 64, "id": 2},
+            {"file_name": "image_3.png", "height": 64, "width": 64, "id": 3},
+        ],
+    }
+    with open(train_ann_file, "w") as annotation_file:
+        json.dump(annotations, annotation_file)
+
+    return COCODataConfig(train_folder, train_ann_file, predict_folder)
+
+
+@pytest.mark.skipif(not _IMAGE_AVAILABLE, reason="image libraries aren't installed.")
+@pytest.mark.skipif(not _ICEVISION_AVAILABLE, reason="IceVision is not installed for testing")
+@pytest.mark.parametrize("backbone, head", [("resnet18_fpn", "keypoint_rcnn")])
+def test_model(coco_keypoints, backbone, head):
+    datamodule = KeypointDetectionData.from_coco(
+        train_folder=coco_keypoints.train_folder,
+        train_ann_file=coco_keypoints.train_ann_file,
+        predict_folder=coco_keypoints.predict_folder,
+        transform_kwargs=dict(image_size=(128, 128)),
+        batch_size=2,
+    )
+
+    assert datamodule.num_classes == 3
+    assert datamodule.labels == ["background", "cat", "dog"]
+
+    model = KeypointDetector(2, num_classes=datamodule.num_classes, backbone=backbone, head=head)
+    trainer = Trainer(fast_dev_run=True)
+    trainer.fit(model, datamodule=datamodule)
+    trainer.predict(model, datamodule=datamodule)
 
 
 @pytest.mark.skipif(not _IMAGE_AVAILABLE, reason="image libraries aren't installed.")