consolidate prediction options into PredictOptions for all tasks (#2055)

Changes:
- Move `predict_chip_sz` and `predict_batch_sz` to `PredictOptions`. Also add `stride`.
- Make `predict_options` a field in `RVPipelineConfig`. Previously, this was only defined in the SS and OD subclasses.
- Make `Backend.predict_scene()` take `PredictOptions` instead of `chip_sz`, `stride` etc.
- Move default `stride`, `crop_sz` initialization to pydantic validators.
- Move OD prediction post-processing to the OD PyTorch `Backend`.
- Remove unused SS post processing.
- Remove support for `RASTERVISION_PREDICT_BATCH_SIZE` `RVConfig` param that was used by `Learner.predict_dataset()`.
- Update usage in examples.
- Update unit and integration tests.

integration_tests/chip_classification/config.py

@@ -1,8 +1,8 @@
 from os.path import join, dirname, basename
 
-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 (
     ClassConfig, ChipClassificationLabelSourceConfig,
     GeoJSONVectorSourceConfig, RasterioSourceConfig, StatsTransformerConfig,
@@ -102,6 +102,6 @@ def make_scene(img_path, label_path):
         dataset=scene_dataset,
         backend=backend,
         chip_options=chip_options,
-        predict_chip_sz=chip_sz)
+        predict_options=PredictOptions(chip_sz=chip_sz))
 
     return config

integration_tests/object_detection/config.py

@@ -93,12 +93,11 @@ def make_scene(scene_id, img_path, label_path):
         run_tensorboard=False)
 
     predict_options = ObjectDetectionPredictOptions(
-        merge_thresh=0.1, score_thresh=0.5)
+        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)

integration_tests/semantic_segmentation/config.py

@@ -7,7 +7,8 @@
     RGBClassTransformerConfig)
 from rastervision.core.rv_pipeline import (
     SemanticSegmentationChipOptions, SemanticSegmentationConfig,
-    WindowSamplingConfig, WindowSamplingMethod)
+    SemanticSegmentationPredictOptions, WindowSamplingConfig,
+    WindowSamplingMethod)
 from rastervision.pytorch_backend import PyTorchSemanticSegmentationConfig
 from rastervision.pytorch_learner import (
     Backbone, SolverConfig, SemanticSegmentationModelConfig,
@@ -101,10 +102,11 @@ def make_scene(id, img_path, label_path):
         solver=solver,
         log_tensorboard=False,
         run_tensorboard=False)
+    predict_options = SemanticSegmentationPredictOptions(chip_sz=chip_sz)
 
     return SemanticSegmentationConfig(
         root_uri=root_uri,
         dataset=scene_dataset,
         backend=backend,
         chip_options=chip_options,
-        predict_chip_sz=chip_sz)
+        predict_options=predict_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

@@ -8,31 +8,51 @@
 from rastervision.core.backend import BackendConfig
 from rastervision.core.evaluation import EvaluatorConfig
 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.core.rv_pipeline.chip_options import ChipOptions
+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:
     if version == 10:
         train_chip_sz = cfg_dict.pop('train_chip_sz', 300)
-        nodata_threshold = cfg_dict.pop('chip_nodata_threshold')
-        if 'chip_options' not in cfg_dict:
-            cfg_dict['chip_options'] = ChipOptions(
-                sampling=WindowSamplingConfig(size=train_chip_sz),
-                nodata_threshold=nodata_threshold)
-        else:
-            cfg_dict['chip_options']['sampling']['size'] = train_chip_sz
-            cfg_dict['chip_options']['nodata_threshold'] = nodata_threshold
+        nodata_threshold = cfg_dict.pop('chip_nodata_threshold', 1.)
+        chip_options: dict = cfg_dict.get('chip_options', {})
+        method = chip_options.pop('method', 'sliding')
+        if method != 'sliding':
+            method = 'random'
+        chip_options['sampling'] = dict(size=train_chip_sz, method=method)
+        chip_options['nodata_threshold'] = nodata_threshold
+        cfg_dict['chip_options'] = chip_options
+    elif version == 11:
+        predict_chip_sz = cfg_dict.pop('predict_chip_sz', 300)
+        predict_batch_sz = cfg_dict.pop('predict_batch_sz', 8)
+        predict_options = cfg_dict.get('predict_options', {})
+        predict_options['chip_sz'] = predict_chip_sz
+        predict_options['batch_sz'] = predict_batch_sz
+        cfg_dict['predict_options'] = predict_options
     return cfg_dict
 
 
@@ -57,13 +77,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 +104,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)

rastervision_pytorch_backend/rastervision/pytorch_backend/examples/semantic_segmentation/isprs_potsdam.py

@@ -5,7 +5,8 @@
 
 from rastervision.core.rv_pipeline import (
     SemanticSegmentationConfig, SemanticSegmentationChipOptions,
-    WindowSamplingConfig, WindowSamplingMethod)
+    SemanticSegmentationPredictOptions, WindowSamplingConfig,
+    WindowSamplingMethod)
 from rastervision.core.data import (
     ClassConfig, DatasetConfig, PolygonVectorOutputConfig,
     RasterioSourceConfig, RGBClassTransformerConfig, SceneConfig,
@@ -238,11 +239,13 @@ def make_scene(id) -> SceneConfig:
         run_tensorboard=False,
     )
 
+    predict_options = SemanticSegmentationPredictOptions(chip_sz=chip_sz)
+
     pipeline = SemanticSegmentationConfig(
         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/semantic_segmentation/isprs_potsdam_multi_source.py

@@ -3,7 +3,8 @@
 
 from rastervision.core.rv_pipeline import (
     SceneConfig, DatasetConfig, SemanticSegmentationChipOptions,
-    SemanticSegmentationConfig, WindowSamplingConfig, WindowSamplingMethod)
+    SemanticSegmentationConfig, SemanticSegmentationPredictOptions,
+    WindowSamplingConfig, WindowSamplingMethod)
 
 from rastervision.core.data import (
     ClassConfig, RasterioSourceConfig, MultiRasterSourceConfig,
@@ -161,15 +162,17 @@ def get_config(runner,
         run_tensorboard=RUN_TENSORBOARD,
     )
 
+    predict_options = SemanticSegmentationPredictOptions(chip_sz=CHIP_SIZE)
+
     # -----------------------------------------------
     # Pass configurations to the pipeline config
     # -----------------------------------------------
     pipeline_config = SemanticSegmentationConfig(
         root_uri=root_uri,
-        predict_chip_sz=CHIP_SIZE,
-        chip_options=chip_options,
         dataset=dataset_config,
-        backend=backend_config)
+        backend=backend_config,
+        chip_options=chip_options,
+        predict_options=predict_options)
 
     return pipeline_config
 

rastervision_pytorch_backend/rastervision/pytorch_backend/examples/semantic_segmentation/spacenet_vegas.py

@@ -7,7 +7,8 @@
 from rastervision.pipeline.file_system.utils import list_paths
 from rastervision.core.rv_pipeline import (
     SemanticSegmentationConfig, SemanticSegmentationChipOptions,
-    WindowSamplingConfig, WindowSamplingMethod)
+    SemanticSegmentationPredictOptions, WindowSamplingConfig,
+    WindowSamplingMethod)
 from rastervision.core.data import (
     BufferTransformerConfig, ClassConfig, ClassInferenceTransformerConfig,
     DatasetConfig, GeoJSONVectorSourceConfig, PolygonVectorOutputConfig,
@@ -222,9 +223,11 @@ def get_config(runner,
         run_tensorboard=False,
     )
 
+    predict_options = SemanticSegmentationPredictOptions(chip_sz=chip_sz)
+
     return SemanticSegmentationConfig(
         root_uri=root_uri,
         dataset=scene_dataset,
         backend=backend,
-        predict_chip_sz=chip_sz,
-        chip_options=chip_options)
+        chip_options=chip_options,
+        predict_options=predict_options)

rastervision_pytorch_backend/rastervision/pytorch_backend/examples/tiny_spacenet.py

@@ -49,7 +49,7 @@ def get_config(runner) -> SemanticSegmentationConfig:
         root_uri=output_root_uri,
         dataset=scene_dataset,
         backend=backend,
-        predict_chip_sz=chip_sz)
+        predict_options=SemanticSegmentationPredictOptions(chip_sz=chip_sz))
 
 
 def make_scene(scene_id: str, image_uri: str, label_uri: str,

rastervision_pytorch_backend/rastervision/pytorch_backend/pytorch_chip_classification.py

@@ -1,9 +1,7 @@
-from typing import TYPE_CHECKING, Iterator, Optional
+from typing import TYPE_CHECKING, Iterator
 from os.path import join
 import uuid
 
-import numpy as np
-
 from rastervision.pipeline.file_system import make_dir
 from rastervision.core.data_sample import DataSample
 from rastervision.pytorch_backend.pytorch_learner_backend import (
@@ -14,8 +12,9 @@
 from rastervision.core.data import ChipClassificationLabels
 
 if TYPE_CHECKING:
+    import numpy as np
     from rastervision.core.data import DatasetConfig, Scene
-    from rastervision.core.rv_pipeline import ChipOptions
+    from rastervision.core.rv_pipeline import ChipOptions, PredictOptions
 
 
 class PyTorchChipClassificationSampleWriter(PyTorchLearnerSampleWriter):
@@ -58,17 +57,16 @@ def chip_dataset(self,
         dataloader_kw = dict(**dataloader_kw, collate_fn=chip_collate_fn_cc)
         return super().chip_dataset(dataset, chip_options, dataloader_kw)
 
-    def predict_scene(self,
-                      scene: 'Scene',
-                      chip_sz: int,
-                      stride: Optional[int] = None
+    def predict_scene(self, scene: 'Scene', predict_options: 'PredictOptions'
                       ) -> 'ChipClassificationLabels':
-        if stride is None:
-            stride = chip_sz
 
         if self.learner is None:
             self.load_model()
 
+        chip_sz = predict_options.chip_sz
+        stride = predict_options.stride
+        batch_sz = predict_options.batch_sz
+
         # Important to use self.learner.cfg.data instead of
         # self.learner_cfg.data because of the updates
         # Learner.from_model_bundle() makes to the custom transforms.
@@ -80,6 +78,7 @@ def predict_scene(self,
             ds,
             raw_out=True,
             numpy_out=True,
+            dataloader_kw=dict(batch_size=batch_sz),
             progress_bar=True,
             progress_bar_kw=dict(desc=f'Making predictions on {scene.id}'))
 

rastervision_pytorch_backend/rastervision/pytorch_backend/pytorch_object_detection.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Dict, Iterator, Optional
+from typing import TYPE_CHECKING, Dict, Iterator
 from os.path import join, basename
 import uuid
 
@@ -17,7 +17,8 @@
 
 if TYPE_CHECKING:
     from rastervision.core.data import DatasetConfig, Scene
-    from rastervision.core.rv_pipeline import ChipOptions
+    from rastervision.core.rv_pipeline import (ChipOptions,
+                                               ObjectDetectionPredictOptions)
     from rastervision.pytorch_learner.object_detection_utils import BoxList
 
 
@@ -113,12 +114,13 @@ def chip_dataset(self,
         dataloader_kw = dict(**dataloader_kw, collate_fn=chip_collate_fn_od)
         return super().chip_dataset(dataset, chip_options, dataloader_kw)
 
-    def predict_scene(self,
-                      scene: 'Scene',
-                      chip_sz: int,
-                      stride: Optional[int] = None) -> ObjectDetectionLabels:
-        if stride is None:
-            stride = chip_sz
+    def predict_scene(self, scene: 'Scene',
+                      predict_options: 'ObjectDetectionPredictOptions'
+                      ) -> ObjectDetectionLabels:
+
+        chip_sz = predict_options.chip_sz
+        stride = predict_options.stride
+        batch_sz = predict_options.batch_sz
 
         if self.learner is None:
             self.load_model()
@@ -136,12 +138,17 @@ def predict_scene(self,
                 raw_out=True,
                 numpy_out=True,
                 predict_kw=dict(out_shape=(chip_sz, chip_sz)),
+                dataloader_kw=dict(batch_size=batch_sz),
                 progress_bar=True,
                 progress_bar_kw=dict(desc=f'Making predictions on {scene.id}'))
         )
 
         labels = ObjectDetectionLabels.from_predictions(
             ds.windows, predictions)
+        labels = ObjectDetectionLabels.prune_duplicates(
+            labels,
+            score_thresh=predict_options.score_thresh,
+            merge_thresh=predict_options.merge_thresh)
 
         return labels
 

rastervision_pytorch_backend/rastervision/pytorch_backend/pytorch_semantic_segmentation.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Iterator, Optional
+from typing import TYPE_CHECKING, Iterator
 from os.path import join
 import uuid
 
@@ -17,7 +17,8 @@
 if TYPE_CHECKING:
     from rastervision.core.data import (DatasetConfig, Scene,
                                         SemanticSegmentationLabelStore)
-    from rastervision.core.rv_pipeline import ChipOptions
+    from rastervision.core.rv_pipeline import (
+        ChipOptions, SemanticSegmentationPredictOptions)
 
 
 class PyTorchSemanticSegmentationSampleWriter(PyTorchLearnerSampleWriter):
@@ -67,23 +68,22 @@ def chip_dataset(self,
         dataloader_kw = dict(**dataloader_kw, collate_fn=chip_collate_fn_ss)
         return super().chip_dataset(dataset, chip_options, dataloader_kw)
 
-    def predict_scene(
-            self,
-            scene: 'Scene',
-            chip_sz: int,
-            stride: Optional[int] = None,
-            crop_sz: Optional[int] = None) -> 'SemanticSegmentationLabels':
+    def predict_scene(self, scene: 'Scene',
+                      predict_options: 'SemanticSegmentationPredictOptions'
+                      ) -> 'SemanticSegmentationLabels':
 
         if scene.label_store is None:
             raise ValueError(
                 f'Scene.label_store is not set for scene {scene.id}')
 
-        if stride is None:
-            stride = chip_sz
-
         if self.learner is None:
             self.load_model()
 
+        chip_sz = predict_options.chip_sz
+        stride = predict_options.stride
+        crop_sz = predict_options.crop_sz
+        batch_sz = predict_options.batch_sz
+
         label_store: 'SemanticSegmentationLabelStore' = scene.label_store
         raw_out = label_store.smooth_output
 
@@ -104,6 +104,7 @@ def predict_scene(
             raw_out=raw_out,
             numpy_out=True,
             predict_kw=dict(out_shape=(chip_sz, chip_sz)),
+            dataloader_kw=dict(batch_size=batch_sz),
             progress_bar=True,
             progress_bar_kw=dict(desc=f'Making predictions on {scene.id}'))
 

rastervision_pytorch_learner/rastervision/pytorch_learner/learner.py

@@ -776,12 +776,10 @@ def predict_dataset(self,
             'rastervision',
             'PREDICT_NUM_WORKERS',
             default=cfg.data.num_workers)
-        batch_size = rv_config.get_namespace_option(
-            'rastervision', 'PREDICT_BATCH_SIZE', default=cfg.solver.batch_sz)
 
         dl_kw = dict(
             collate_fn=self.get_collate_fn(),
-            batch_size=int(batch_size),
+            batch_size=cfg.solver.batch_sz,
             num_workers=int(num_workers),
             shuffle=False,
             pin_memory=True)

tests/core/rv_pipeline/test_object_detection_config.py

@@ -0,0 +1,15 @@
+import unittest
+
+from rastervision.core.rv_pipeline import (ObjectDetectionPredictOptions)
+
+
+class TestObjectDetectionPredictOptions(unittest.TestCase):
+    def test_stride_validator(self):
+        cfg = ObjectDetectionPredictOptions(chip_sz=10)
+        self.assertEqual(cfg.stride, 5)
+        cfg = ObjectDetectionPredictOptions(chip_sz=11)
+        self.assertEqual(cfg.stride, 5)
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/core/rv_pipeline/test_rv_pipeline_config.py

@@ -0,0 +1,53 @@
+import unittest
+
+from rastervision.core.data import (ClassConfig, DatasetConfig)
+from rastervision.core.backend import (BackendConfig)
+from rastervision.core.rv_pipeline.rv_pipeline_config import (
+    PredictOptions, RVPipelineConfig, rv_pipeline_config_upgrader)
+
+
+class TestPredictOptions(unittest.TestCase):
+    def test_stride_validator(self):
+        cfg = PredictOptions(chip_sz=10)
+        self.assertEqual(cfg.stride, 10)
+
+
+class TestRVPipelineConfig(unittest.TestCase):
+    def test_upgrader(self):
+        cfg_dict = dict(
+            dataset=DatasetConfig(
+                class_config=ClassConfig(names=[]),
+                train_scenes=[],
+                validation_scenes=[]),
+            backend=BackendConfig(),
+            train_chip_sz=20,
+            chip_nodata_threshold=0.5,
+            predict_chip_sz=20,
+            predict_batch_sz=8)
+        cfg_dict = rv_pipeline_config_upgrader(cfg_dict, 10)
+        cfg_dict = rv_pipeline_config_upgrader(cfg_dict, 11)
+        cfg = RVPipelineConfig(**cfg_dict)
+
+        cfg_dict = dict(
+            dataset=DatasetConfig(
+                class_config=ClassConfig(names=[]),
+                train_scenes=[],
+                validation_scenes=[]),
+            backend=BackendConfig(),
+            train_chip_sz=20,
+            chip_nodata_threshold=0.5,
+            chip_options=dict(method='random'),
+            predict_chip_sz=20,
+            predict_batch_sz=8,
+            predict_options=dict())
+        cfg_dict = rv_pipeline_config_upgrader(cfg_dict, 10)
+        cfg_dict = rv_pipeline_config_upgrader(cfg_dict, 11)
+        cfg = RVPipelineConfig(**cfg_dict)
+        self.assertEqual(cfg.chip_options.get_chip_sz(), 20)
+        self.assertEqual(cfg.chip_options.nodata_threshold, 0.5)
+        self.assertEqual(cfg.predict_options.chip_sz, 20)
+        self.assertEqual(cfg.predict_options.batch_sz, 8)
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/core/rv_pipeline/test_semantic_segmentation_config.py

@@ -1,7 +1,6 @@
 from typing import Callable
 import unittest
 
-from pydantic import ValidationError
 import numpy as np
 
 from rastervision.core.data import (ClassConfig, DatasetConfig)
@@ -48,27 +47,14 @@ def test_upgrader(self):
 
 
 class TestSemanticSegmentationPredictOptions(unittest.TestCase):
-    def assertNoError(self, fn: Callable, msg: str = ''):
-        try:
-            fn()
-        except Exception:
-            self.fail(msg)
-
-    def test_upgrader(self):
-        args = dict(stride=None, crop_sz=None)
-        self.assertNoError(lambda: SemanticSegmentationPredictOptions(**args))
-
-        args = dict(stride=None, crop_sz='auto')
-        self.assertRaises(ValidationError,
-                          lambda: SemanticSegmentationPredictOptions(**args))
-
-        args = dict(stride=None, crop_sz=10)
-        self.assertRaises(ValidationError,
-                          lambda: SemanticSegmentationPredictOptions(**args))
-
-        args = dict(stride=10, crop_sz=0)
-        self.assertRaises(ValidationError,
-                          lambda: SemanticSegmentationPredictOptions(**args))
+    def test_crop_sz_validator(self):
+        args = dict(chip_sz=10, stride=4, crop_sz='auto')
+        cfg = SemanticSegmentationPredictOptions(**args)
+        self.assertEqual(cfg.crop_sz, 3)
+
+        args = dict(chip_sz=10, stride=5, crop_sz='auto')
+        cfg = SemanticSegmentationPredictOptions(**args)
+        self.assertEqual(cfg.crop_sz, 2)
 
 
 class TestSemanticSegmentationChipOptions(unittest.TestCase):

tests/pytorch_backend/test_pytorch_chip_classification.py

@@ -4,8 +4,8 @@
 from rastervision.pipeline.config import save_pipeline_config
 from rastervision.core.data import (ClassConfig, DatasetConfig)
 from rastervision.core.rv_pipeline import (
-    ChipOptions, ChipClassificationConfig, WindowSamplingConfig,
-    WindowSamplingMethod)
+    ChipOptions, ChipClassificationConfig, PredictOptions,
+    WindowSamplingConfig, WindowSamplingMethod)
 from rastervision.pytorch_backend import PyTorchChipClassificationConfig
 from rastervision.pytorch_learner import (ClassificationModelConfig,
                                           SolverConfig)
@@ -48,7 +48,8 @@ def make_pipeline(tmp_dir: str, num_channels: int, nochip: bool = False):
         root_uri=tmp_dir,
         dataset=dataset_cfg,
         backend=backend_cfg,
-        chip_options=chip_options)
+        chip_options=chip_options,
+        predict_options=PredictOptions(chip_sz=100))
     pipeline_cfg.update()
     save_pipeline_config(pipeline_cfg, pipeline_cfg.get_config_uri())
     pipeline = pipeline_cfg.build(tmp_dir)

tests/pytorch_backend/test_pytorch_object_detection.py

@@ -5,7 +5,8 @@
 from rastervision.core.data import (ClassConfig, DatasetConfig)
 from rastervision.core.rv_pipeline import (
     ObjectDetectionConfig, ObjectDetectionChipOptions,
-    ObjectDetectionWindowSamplingConfig, WindowSamplingMethod)
+    ObjectDetectionPredictOptions, ObjectDetectionWindowSamplingConfig,
+    WindowSamplingMethod)
 from rastervision.pytorch_backend import PyTorchObjectDetectionConfig
 from rastervision.pytorch_learner import (ObjectDetectionModelConfig,
                                           SolverConfig)
@@ -48,7 +49,8 @@ def make_pipeline(tmp_dir: str, num_channels: int, nochip: bool = False):
         root_uri=tmp_dir,
         dataset=dataset_cfg,
         backend=backend_cfg,
-        chip_options=chip_options)
+        chip_options=chip_options,
+        predict_options=ObjectDetectionPredictOptions(chip_sz=100, stride=50))
     pipeline_cfg.update()
     save_pipeline_config(pipeline_cfg, pipeline_cfg.get_config_uri())
     pipeline = pipeline_cfg.build(tmp_dir)

tests/pytorch_backend/test_pytorch_semantic_segmentation.py

@@ -5,7 +5,8 @@
 from rastervision.core.data import (ClassConfig, DatasetConfig)
 from rastervision.core.rv_pipeline import (
     SemanticSegmentationConfig, SemanticSegmentationChipOptions,
-    WindowSamplingConfig, WindowSamplingMethod)
+    SemanticSegmentationPredictOptions, WindowSamplingConfig,
+    WindowSamplingMethod)
 from rastervision.pytorch_backend import PyTorchSemanticSegmentationConfig
 from rastervision.pytorch_learner import (SemanticSegmentationModelConfig,
                                           SolverConfig)
@@ -34,7 +35,7 @@ def make_pipeline(tmp_dir: str, num_channels: int, nochip: bool = False):
         test_scenes=[])
     chip_options = SemanticSegmentationChipOptions(
         sampling=WindowSamplingConfig(
-            method=WindowSamplingMethod.random, size=20, max_windows=8))
+            method=WindowSamplingMethod.random, size=100, max_windows=8))
     if nochip:
         data_cfg = SemanticSegmentationGeoDataConfig(
             scene_dataset=dataset_cfg,
@@ -51,7 +52,9 @@ def make_pipeline(tmp_dir: str, num_channels: int, nochip: bool = False):
         root_uri=tmp_dir,
         dataset=dataset_cfg,
         backend=backend_cfg,
-        chip_options=chip_options)
+        chip_options=chip_options,
+        predict_options=SemanticSegmentationPredictOptions(
+            chip_sz=100, stride=50, crop_sz='auto'))
     pipeline_cfg.update()
     save_pipeline_config(pipeline_cfg, pipeline_cfg.get_config_uri())
     pipeline = pipeline_cfg.build(tmp_dir)

tests/pytorch_learner/test_classification_learner.py

@@ -122,14 +122,6 @@ def _test_learner(self,
             learner.plot_predictions(split='valid')
             learner.save_model_bundle()
 
-            learner = None
-            backend.learner = None
-            backend.load_model()
-
-            pred_scene = dataset_cfg.validation_scenes[0].build(
-                class_config, tmp_dir)
-            _ = backend.predict_scene(pred_scene, chip_sz=100)
-
 
 if __name__ == '__main__':
     unittest.main()

tests/pytorch_learner/test_object_detection_learner.py

@@ -124,14 +124,6 @@ def _test_learner(self,
             learner.plot_predictions(split='valid')
             learner.save_model_bundle()
 
-            learner = None
-            backend.learner = None
-            backend.load_model()
-
-            pred_scene = dataset_cfg.validation_scenes[0].build(
-                class_config, tmp_dir)
-            _ = backend.predict_scene(pred_scene, chip_sz=100)
-
 
 if __name__ == '__main__':
     unittest.main()

tests/pytorch_learner/test_semantic_segmentation_learner.py

@@ -128,14 +128,6 @@ def _test_learner(self,
             learner.plot_predictions(split='valid')
             learner.save_model_bundle()
 
-            learner = None
-            backend.learner = None
-            backend.load_model()
-
-            pred_scene = dataset_cfg.validation_scenes[0].build(
-                class_config, tmp_dir)
-            _ = backend.predict_scene(pred_scene, chip_sz=100)
-
 
 if __name__ == '__main__':
     unittest.main()