Filter fpn backbones for torchvision

CHANGELOG.md

@@ -50,6 +50,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Fixed a bug where using `ReduceLROnPlateau` would raise an error ([#1251](https://github.com/PyTorchLightning/lightning-flash/pull/1251))
 
+- Fixed a bug where some backbones were incorrectly listed as available for the `ObjectDetector`, `InstanceSegmentation`, and `KeypointDetector` ([#1267](https://github.com/PyTorchLightning/lightning-flash/pull/1267))
+
 ## [0.7.0] - 2022-02-15
 
 ### Added

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,

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)

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

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.")