ObjectDetection/InstanceSegmentationTask: fix support for non-RGB images

tests/trainers/test_detection.py

@@ -5,6 +5,7 @@
 from typing import Any
 
 import pytest
+import torch
 from lightning.pytorch import Trainer
 from pytest import MonkeyPatch
 
@@ -122,3 +123,17 @@ def test_freeze_backbone(self, model_name: str) -> None:
             model=model_name, backbone='resnet18', freeze_backbone=True
         )
         assert not all([param.requires_grad for param in model.model.parameters()])
+
+    @pytest.mark.parametrize('model_name', ['faster-rcnn', 'fcos', 'retinanet'])
+    @pytest.mark.parametrize('in_channels', [1, 4])
+    def test_multispectral_support(self, model_name: str, in_channels: int) -> None:
+        model = ObjectDetectionTask(
+            model=model_name,
+            backbone='resnet18',
+            num_classes=2,
+            in_channels=in_channels,
+        )
+        model.eval()
+        sample = [torch.randn(in_channels, 224, 224)]
+        with torch.inference_mode():
+            model(sample)

tests/trainers/test_instance_segmentation.py

@@ -5,6 +5,7 @@
 from typing import Any
 
 import pytest
+import torch
 from lightning.pytorch import Trainer
 from pytest import MonkeyPatch
 
@@ -123,3 +124,11 @@ def test_freeze_backbone(self) -> None:
         for head in ['rpn', 'roi_heads']:
             for param in getattr(task.model, head).parameters():
                 assert param.requires_grad is True
+
+    @pytest.mark.parametrize('in_channels', [1, 4])
+    def test_multispectral_support(self, in_channels: int) -> None:
+        model = InstanceSegmentationTask(in_channels=in_channels, num_classes=2)
+        model.eval()
+        sample = [torch.randn(in_channels, 224, 224)]
+        with torch.inference_mode():
+            model(sample)

torchgeo/trainers/detection.py

@@ -151,6 +151,10 @@ def configure_models(self) -> None:
                 num_classes,
                 rpn_anchor_generator=anchor_generator,
                 box_roi_pool=roi_pooler,
+                min_size=1,
+                max_size=4096,
+                image_mean=[0],
+                image_std=[1],
             )
         elif model == 'fcos':
             kwargs['extra_blocks'] = feature_pyramid_network.LastLevelP6P7(256, 256)
@@ -169,7 +173,13 @@ def configure_models(self) -> None:
                     param.requires_grad = False
 
             self.model = torchvision.models.detection.FCOS(
-                model_backbone, num_classes, anchor_generator=anchor_generator
+                model_backbone,
+                num_classes,
+                anchor_generator=anchor_generator,
+                min_size=1,
+                max_size=4096,
+                image_mean=[0],
+                image_std=[1],
             )
         elif model == 'retinanet':
             kwargs['extra_blocks'] = feature_pyramid_network.LastLevelP6P7(
@@ -204,6 +214,10 @@ def configure_models(self) -> None:
                 num_classes,
                 anchor_generator=anchor_generator,
                 head=head,
+                min_size=1,
+                max_size=4096,
+                image_mean=[0],
+                image_std=[1],
             )
         else:
             raise ValueError(f"Model type '{model}' is not valid.")

torchgeo/trainers/instance_segmentation.py

@@ -87,6 +87,10 @@ def configure_models(self) -> None:
                     weights=weights,
                     num_classes=num_classes,
                     weights_backbone=weights_backbone,
+                    min_size=1,
+                    max_size=4096,
+                    image_mean=[0],
+                    image_std=[1],
                 )
             else:
                 msg = f"Invalid backbone type '{backbone}'. Supported backbone: 'resnet50'"