standardize prediction options

rastervision_core/rastervision/core/__init__.py

@@ -2,7 +2,7 @@
 
 
 def register_plugin(registry):
-    registry.set_plugin_version('rastervision.core', 11)
+    registry.set_plugin_version('rastervision.core', 12)
     from rastervision.core.cli import predict, predict_scene
     registry.add_plugin_command(predict)
     registry.add_plugin_command(predict_scene)

rastervision_core/rastervision/core/backend/backend.py

@@ -5,7 +5,7 @@
 if TYPE_CHECKING:
     from rastervision.core.data_sample import DataSample
     from rastervision.core.data import DatasetConfig, Labels, Scene
-    from rastervision.core.rv_pipeline import ChipOptions
+    from rastervision.core.rv_pipeline import ChipOptions, PredictOptions
 
 
 class SampleWriter(AbstractContextManager):
@@ -48,12 +48,13 @@ def load_model(self, uri: Optional[str] = None):
         """
 
     @abstractmethod
-    def predict_scene(self, scene: 'Scene', chip_sz: int,
-                      stride: int) -> 'Labels':
+    def predict_scene(self, scene: 'Scene',
+                      predict_options: 'PredictOptions') -> 'Labels':
         """Return predictions for an entire scene using the model.
 
         Args:
-            scene (Scene): Scene to run inference on.
+            scene: Scene to run inference on.
+            predict_options: Prediction options.
 
         Return:
             Labels object containing predictions

rastervision_core/rastervision/core/rv_pipeline/__init__.py

@@ -31,4 +31,5 @@
     ChipOptions.__name__,
     WindowSamplingConfig.__name__,
     WindowSamplingMethod.__name__,
+    PredictOptions.__name__,
 ]

rastervision_core/rastervision/core/rv_pipeline/object_detection.py

@@ -1,33 +1,5 @@
-from typing import TYPE_CHECKING
-import logging
-
 from rastervision.core.rv_pipeline import RVPipeline
-from rastervision.core.data.label import ObjectDetectionLabels
-
-if TYPE_CHECKING:
-    from rastervision.core.data import Labels, Scene
-
-log = logging.getLogger(__name__)
 
 
 class ObjectDetection(RVPipeline):
-    def predict_scene(self, scene: 'Scene') -> 'Labels':
-        if self.backend is None:
-            self.build_backend()
-
-        # Use strided windowing to ensure that each object is fully visible (ie. not
-        # cut off) within some window. This means prediction takes 4x longer for object
-        # detection :(
-        chip_sz = self.config.predict_chip_sz
-        stride = chip_sz // 2
-        labels = self.backend.predict_scene(
-            scene, chip_sz=chip_sz, stride=stride)
-        labels = self.post_process_predictions(labels, scene)
-        return labels
-
-    def post_process_predictions(self, labels: ObjectDetectionLabels,
-                                 scene: 'Scene') -> ObjectDetectionLabels:
-        return ObjectDetectionLabels.prune_duplicates(
-            labels,
-            score_thresh=self.config.predict_options.score_thresh,
-            merge_thresh=self.config.predict_options.merge_thresh)
+    pass

rastervision_core/rastervision/core/rv_pipeline/object_detection_config.py

@@ -1,6 +1,6 @@
 from typing import Optional
 
-from rastervision.pipeline.config import register_config, Field
+from rastervision.pipeline.config import Field, register_config, validator
 from rastervision.core.rv_pipeline import (
     ChipOptions, RVPipelineConfig, PredictOptions, WindowSamplingConfig)
 from rastervision.core.data.label_store import ObjectDetectionGeoJSONStoreConfig
@@ -42,6 +42,10 @@ class ObjectDetectionChipOptions(ChipOptions):
 
 @register_config('object_detection_predict_options')
 class ObjectDetectionPredictOptions(PredictOptions):
+    stride: Optional[int] = Field(
+        None,
+        description='Stride of the sliding window for generating chips. '
+        'Defaults to half of ``chip_sz``.')
     merge_thresh: float = Field(
         0.5,
         description=
@@ -55,15 +59,20 @@ class ObjectDetectionPredictOptions(PredictOptions):
         ('Predicted boxes are only output if their score is above score_thresh.'
          ))
 
+    @validator('stride', always=True)
+    def validate_stride(cls, v: Optional[int], values: dict) -> dict:
+        if v is None:
+            chip_sz: int = values['chip_sz']
+            return chip_sz // 2
+        return v
+
 
 @register_config('object_detection')
 class ObjectDetectionConfig(RVPipelineConfig):
     """Configure an :class:`.ObjectDetection` pipeline."""
 
-    chip_options: Optional[ObjectDetectionChipOptions] = Field(
-        None, description='Config for chip stage.')
-    predict_options: Optional[
-        ObjectDetectionPredictOptions] = ObjectDetectionPredictOptions()
+    chip_options: Optional[ObjectDetectionChipOptions]
+    predict_options: Optional[ObjectDetectionPredictOptions]
 
     def build(self, tmp_dir):
         from rastervision.core.rv_pipeline.object_detection import ObjectDetection

rastervision_core/rastervision/core/rv_pipeline/rv_pipeline.py

@@ -162,10 +162,8 @@ def predict(self, split_ind=0, num_splits=1):
     def predict_scene(self, scene: Scene) -> Labels:
         if self.backend is None:
             self.build_backend()
-        chip_sz = self.config.predict_chip_sz
-        stride = chip_sz
         labels = self.backend.predict_scene(
-            scene, chip_sz=chip_sz, stride=stride)
+            scene, predict_options=self.config.predict_options)
         labels = self.post_process_predictions(labels, scene)
         return labels
 

rastervision_core/rastervision/core/rv_pipeline/rv_pipeline_config.py

@@ -10,16 +10,30 @@
 from rastervision.core.analyzer import AnalyzerConfig
 from rastervision.core.rv_pipeline.chip_options import (ChipOptions,
                                                         WindowSamplingConfig)
-from rastervision.pipeline.config import (Config, Field, register_config)
+from rastervision.pipeline.config import (Config, Field, register_config,
+                                          validator)
 
 if TYPE_CHECKING:
     from rastervision.core.backend.backend import Backend  # noqa
 
 
 @register_config('predict_options')
 class PredictOptions(Config):
-    # TODO: predict_chip_sz and predict_batch_sz should probably be moved here
-    pass
+    chip_sz: int = Field(
+        300, description='Size of predictions chips in pixels.')
+    stride: Optional[int] = Field(
+        None,
+        description='Stride of the sliding window for generating chips.'
+        'Defaults to ``chip_sz``.')
+    batch_sz: int = Field(
+        8, description='Batch size to use during prediction.')
+
+    @validator('stride', always=True)
+    def validate_stride(cls, v: Optional[int], values: dict) -> dict:
+        if v is None:
+            chip_sz: int = values['chip_sz']
+            return chip_sz
+        return v
 
 
 def rv_pipeline_config_upgrader(cfg_dict: dict, version: int) -> dict:
@@ -33,6 +47,15 @@ def rv_pipeline_config_upgrader(cfg_dict: dict, version: int) -> dict:
         else:
             cfg_dict['chip_options']['sampling']['size'] = train_chip_sz
             cfg_dict['chip_options']['nodata_threshold'] = nodata_threshold
+    elif version == 11:
+        predict_chip_sz = cfg_dict.pop('predict_chip_sz', 300)
+        predict_batch_sz = cfg_dict.pop('predict_batch_sz', 8)
+        if 'predict_options' not in cfg_dict:
+            cfg_dict['predict_options'] = PredictOptions(
+                chip_sz=predict_chip_sz, batch_sz=predict_batch_sz)
+        else:
+            cfg_dict['predict_options']['chip_sz'] = predict_chip_sz
+            cfg_dict['predict_options']['batch_sz'] = predict_batch_sz
     return cfg_dict
 
 
@@ -57,13 +80,6 @@ class RVPipelineConfig(PipelineConfig):
         ('Analyzers to run during analyzer command. A StatsAnalyzer will be added '
          'automatically if any scenes have a RasterTransformer.'))
 
-    chip_options: Optional[ChipOptions] = Field(
-        None, description='Config for chip stage.')
-    predict_chip_sz: int = Field(
-        300, description='Size of predictions chips in pixels.')
-    predict_batch_sz: int = Field(
-        8, description='Batch size to use during prediction.')
-
     analyze_uri: Optional[str] = Field(
         None,
         description=
@@ -91,6 +107,11 @@ class RVPipelineConfig(PipelineConfig):
         description='If provided, the model will be loaded from this bundle '
         'for the train stage. Useful for fine-tuning.')
 
+    chip_options: Optional[ChipOptions] = Field(
+        None, description='Config for chip stage.')
+    predict_options: Optional[PredictOptions] = Field(
+        None, description='Config for predict stage.')
+
     def update(self):
         super().update()
 

rastervision_core/rastervision/core/rv_pipeline/semantic_segmentation.py

@@ -1,53 +1,5 @@
-from typing import TYPE_CHECKING
-import logging
-
-import numpy as np
-
 from rastervision.core.rv_pipeline import RVPipeline
 
-if TYPE_CHECKING:
-    from rastervision.core.data import (
-        Labels,
-        Scene,
-    )
-    from rastervision.core.rv_pipeline.semantic_segmentation_config import (
-        SemanticSegmentationConfig)
-
-log = logging.getLogger(__name__)
-
 
 class SemanticSegmentation(RVPipeline):
-    def post_process_batch(self, windows, chips, labels):
-        # Fill in null class for any NODATA pixels.
-        null_class_id = self.config.dataset.class_config.null_class_id
-        for window, chip in zip(windows, chips):
-            nodata_mask = np.sum(chip, axis=2) == 0
-            labels.mask_fill(window, nodata_mask, fill_value=null_class_id)
-
-        return labels
-
-    def predict_scene(self, scene: 'Scene') -> 'Labels':
-        if self.backend is None:
-            self.build_backend()
-
-        cfg: 'SemanticSegmentationConfig' = self.config
-        chip_sz = cfg.predict_chip_sz
-        stride = cfg.predict_options.stride
-        crop_sz = cfg.predict_options.crop_sz
-
-        if stride is None:
-            stride = chip_sz
-
-        if crop_sz == 'auto':
-            overlap_sz = chip_sz - stride
-            if overlap_sz % 2 == 1:
-                log.warning(
-                    'Using crop_sz="auto" but overlap size (chip_sz minus '
-                    'stride) is odd. This means that one pixel row/col will '
-                    'still overlap after cropping.')
-            crop_sz = overlap_sz // 2
-
-        labels = self.backend.predict_scene(
-            scene, chip_sz=chip_sz, stride=stride, crop_sz=crop_sz)
-        labels = self.post_process_predictions(labels, scene)
-        return labels
+    pass

rastervision_core/rastervision/core/rv_pipeline/semantic_segmentation_config.py

@@ -1,17 +1,19 @@
 from typing import (List, Literal, Optional, Union)
 from pydantic import conint
+import logging
 
 import numpy as np
 
-from rastervision.pipeline.config import (ConfigError, register_config, Field,
-                                          validator)
+from rastervision.pipeline.config import (register_config, Field, validator)
 from rastervision.core.rv_pipeline.rv_pipeline_config import (PredictOptions,
                                                               RVPipelineConfig)
 from rastervision.core.rv_pipeline.chip_options import (ChipOptions,
                                                         WindowSamplingConfig)
 from rastervision.core.data import SemanticSegmentationLabelStoreConfig
 from rastervision.core.evaluation import SemanticSegmentationEvaluatorConfig
 
+log = logging.getLogger(__name__)
+
 
 def ss_chip_options_upgrader(cfg_dict: dict, version: int) -> dict:
     if version == 10:
@@ -76,10 +78,9 @@ def enough_target_pixels(self, label_arr: np.ndarray) -> bool:
 class SemanticSegmentationPredictOptions(PredictOptions):
     stride: Optional[int] = Field(
         None,
-        description=
-        'Stride of windows across image. Allows aggregating multiple '
-        'predictions for each pixel if less than the chip size. '
-        'Defaults to predict_chip_sz.')
+        description='Stride of the sliding window for generating chips. '
+        'Allows aggregating multiple predictions for each pixel if less than '
+        'the chip size. Defaults to ``chip_sz``.')
     crop_sz: Optional[Union[conint(gt=0), Literal['auto']]] = Field(
         None,
         description=
@@ -93,12 +94,17 @@ class SemanticSegmentationPredictOptions(PredictOptions):
     def validate_crop_sz(cls,
                          v: Optional[Union[conint(gt=0), Literal['auto']]],
                          values: dict) -> dict:
-        stride: Optional[int] = values.get('stride')
         crop_sz = v
-
-        if stride is None and crop_sz is not None:
-            raise ConfigError('Cannot use crop_sz if stride is None.')
-
+        if crop_sz == 'auto':
+            chip_sz: int = values['chip_sz']
+            stride: int = values['stride']
+            overlap_sz = chip_sz - stride
+            if overlap_sz % 2 == 1:
+                log.warning(
+                    'Using crop_sz="auto" but overlap size (chip_sz minus '
+                    'stride) is odd. This means that one pixel row/col will '
+                    'still overlap after cropping.')
+            crop_sz = overlap_sz // 2
         return crop_sz
 
 
@@ -118,10 +124,8 @@ def ss_config_upgrader(cfg_dict: dict, version: int) -> dict:
 class SemanticSegmentationConfig(RVPipelineConfig):
     """Configure a :class:`.SemanticSegmentation` pipeline."""
 
-    chip_options: SemanticSegmentationChipOptions = Field(
-        None, description='Config for chip stage.')
-    predict_options: SemanticSegmentationPredictOptions = \
-        SemanticSegmentationPredictOptions()
+    chip_options: Optional[SemanticSegmentationChipOptions]
+    predict_options: Optional[SemanticSegmentationPredictOptions]
 
     def build(self, tmp_dir):
         from rastervision.core.rv_pipeline.semantic_segmentation import (

rastervision_pytorch_backend/rastervision/pytorch_backend/examples/chip_classification/spacenet_rio.py

@@ -1,9 +1,9 @@
 import os
 from os.path import join
 
-from rastervision.core.rv_pipeline import (ChipClassificationConfig,
-                                           ChipOptions, WindowSamplingConfig,
-                                           WindowSamplingMethod)
+from rastervision.core.rv_pipeline import (
+    ChipClassificationConfig, ChipOptions, PredictOptions,
+    WindowSamplingConfig, WindowSamplingMethod)
 from rastervision.core.data import (
     ChipClassificationLabelSourceConfig, ClassConfig,
     ClassInferenceTransformerConfig, DatasetConfig, GeoJSONVectorSourceConfig,
@@ -192,6 +192,6 @@ def make_scene(scene_info) -> SceneConfig:
         dataset=scene_dataset,
         backend=backend,
         chip_options=chip_options,
-        predict_chip_sz=chip_sz)
+        predict_options=PredictOptions(chip_sz=chip_sz))
 
     return pipeline

rastervision_pytorch_backend/rastervision/pytorch_backend/examples/object_detection/cowc_potsdam.py

@@ -196,14 +196,13 @@ def make_scene(id: str) -> SceneConfig:
     )
 
     predict_options = ObjectDetectionPredictOptions(
-        merge_thresh=0.5, score_thresh=0.9)
+        chip_sz=chip_sz, merge_thresh=0.5, score_thresh=0.9)
 
     pipeline = ObjectDetectionConfig(
         root_uri=root_uri,
         dataset=scene_dataset,
         backend=backend,
         chip_options=chip_options,
-        predict_chip_sz=chip_sz,
         predict_options=predict_options)
 
     return pipeline

rastervision_pytorch_backend/rastervision/pytorch_backend/examples/object_detection/xview.py

@@ -107,9 +107,6 @@ def make_scene(scene_info):
     else:
         data = ObjectDetectionImageDataConfig(img_sz=img_sz, num_workers=4)
 
-    predict_options = ObjectDetectionPredictOptions(
-        merge_thresh=0.1, score_thresh=0.5)
-
     backend = PyTorchObjectDetectionConfig(
         data=data,
         model=ObjectDetectionModelConfig(backbone=Backbone.resnet50),
@@ -123,10 +120,12 @@ def make_scene(scene_info):
         run_tensorboard=False,
     )
 
+    predict_options = ObjectDetectionPredictOptions(
+        chip_sz=chip_sz, merge_thresh=0.1, score_thresh=0.5)
+
     return ObjectDetectionConfig(
         root_uri=root_uri,
         dataset=scene_dataset,
         backend=backend,
-        predict_chip_sz=chip_sz,
         chip_options=chip_options,
         predict_options=predict_options)