Merge pull request #2040 from AdeelH/chip

Reconcile chip-sampling functionality in `RVPipelines` with `GeoDatasets` (#1496)

docs/framework/pipelines.rst

@@ -43,7 +43,7 @@ In the {{ tiny_spacenet }} example, the :class:`.SemanticSegmentationConfig` is
 
 .. literalinclude:: /../rastervision_pytorch_backend/rastervision/pytorch_backend/examples/tiny_spacenet.py
     :language: python
-    :lines: 48-53
+    :lines: 48-52
     :dedent:
 
 .. seealso:: The :class:`.ChipClassificationConfig`, :class:`.SemanticSegmentationConfig`, and :class:`.ObjectDetectionConfig` API docs have more information on configuring pipelines.

integration_tests/chip_classification/config.py

@@ -1,15 +1,16 @@
 from os.path import join, dirname, basename
 
-from rastervision.core.rv_pipeline import ChipClassificationConfig
+from rastervision.core.rv_pipeline import (ChipClassificationConfig,
+                                           ChipOptions, WindowSamplingConfig,
+                                           WindowSamplingMethod)
 from rastervision.core.data import (
     ClassConfig, ChipClassificationLabelSourceConfig,
     GeoJSONVectorSourceConfig, RasterioSourceConfig, StatsTransformerConfig,
     SceneConfig, DatasetConfig)
 from rastervision.pytorch_backend import PyTorchChipClassificationConfig
 from rastervision.pytorch_learner import (
     Backbone, SolverConfig, ClassificationModelConfig,
-    ClassificationImageDataConfig, ClassificationGeoDataConfig,
-    GeoDataWindowConfig, GeoDataWindowMethod)
+    ClassificationImageDataConfig, ClassificationGeoDataConfig)
 
 
 def get_config(runner, root_uri, data_uri=None, full_train=False,
@@ -56,13 +57,14 @@ def make_scene(img_path, label_path):
     chip_sz = 200
     img_sz = chip_sz
 
-    if nochip:
-        window_opts = GeoDataWindowConfig(
-            method=GeoDataWindowMethod.sliding, stride=chip_sz, size=chip_sz)
+    chip_options = ChipOptions(
+        sampling=WindowSamplingConfig(
+            method=WindowSamplingMethod.sliding, stride=chip_sz, size=chip_sz))
 
+    if nochip:
         data = ClassificationGeoDataConfig(
             scene_dataset=scene_dataset,
-            window_opts=window_opts,
+            sampling=chip_options.sampling,
             img_sz=img_sz,
             augmentors=[])
     else:
@@ -99,7 +101,7 @@ def make_scene(img_path, label_path):
         root_uri=root_uri,
         dataset=scene_dataset,
         backend=backend,
-        train_chip_sz=chip_sz,
+        chip_options=chip_options,
         predict_chip_sz=chip_sz)
 
     return config

integration_tests/object_detection/config.py

@@ -1,16 +1,16 @@
 from os.path import join, dirname
 
-from rastervision.core.rv_pipeline import (ObjectDetectionConfig,
-                                           ObjectDetectionChipOptions,
-                                           ObjectDetectionPredictOptions)
+from rastervision.core.rv_pipeline import (
+    ObjectDetectionConfig, ObjectDetectionChipOptions,
+    ObjectDetectionPredictOptions, ObjectDetectionWindowSamplingConfig,
+    WindowSamplingMethod)
 from rastervision.core.data import (
     ClassConfig, ObjectDetectionLabelSourceConfig, GeoJSONVectorSourceConfig,
     RasterioSourceConfig, SceneConfig, DatasetConfig)
 from rastervision.pytorch_backend import PyTorchObjectDetectionConfig
 from rastervision.pytorch_learner import (
     Backbone, SolverConfig, ObjectDetectionModelConfig,
-    ObjectDetectionImageDataConfig, ObjectDetectionGeoDataConfig,
-    ObjectDetectionGeoDataWindowConfig, GeoDataWindowMethod)
+    ObjectDetectionImageDataConfig, ObjectDetectionGeoDataConfig)
 
 
 def get_config(runner, root_uri, data_uri=None, full_train=False,
@@ -48,19 +48,18 @@ def make_scene(scene_id, img_path, label_path):
         train_scenes=scenes,
         validation_scenes=scenes)
 
-    chip_options = ObjectDetectionChipOptions(neg_ratio=1.0, ioa_thresh=1.0)
-
-    if nochip:
-        window_opts = ObjectDetectionGeoDataWindowConfig(
-            method=GeoDataWindowMethod.sliding,
+    chip_options = ObjectDetectionChipOptions(
+        sampling=ObjectDetectionWindowSamplingConfig(
+            method=WindowSamplingMethod.sliding,
             stride=chip_sz,
             size=chip_sz,
-            neg_ratio=chip_options.neg_ratio,
-            ioa_thresh=chip_options.ioa_thresh)
+            neg_ratio=1.0,
+            ioa_thresh=1.0))
 
+    if nochip:
         data = ObjectDetectionGeoDataConfig(
             scene_dataset=scene_dataset,
-            window_opts=window_opts,
+            sampling=chip_options.sampling,
             img_sz=img_sz,
             augmentors=[])
     else:
@@ -100,7 +99,6 @@ def make_scene(scene_id, img_path, label_path):
         root_uri=root_uri,
         dataset=scene_dataset,
         backend=backend,
-        train_chip_sz=chip_sz,
         predict_chip_sz=chip_sz,
         chip_options=chip_options,
         predict_options=predict_options)

integration_tests/semantic_segmentation/config.py

@@ -5,14 +5,13 @@
     SemanticSegmentationLabelStoreConfig, RasterioSourceConfig, SceneConfig,
     PolygonVectorOutputConfig, DatasetConfig, BuildingVectorOutputConfig,
     RGBClassTransformerConfig)
-from rastervision.core.rv_pipeline import (SemanticSegmentationChipOptions,
-                                           SemanticSegmentationWindowMethod,
-                                           SemanticSegmentationConfig)
+from rastervision.core.rv_pipeline import (
+    SemanticSegmentationChipOptions, SemanticSegmentationConfig,
+    WindowSamplingConfig, WindowSamplingMethod)
 from rastervision.pytorch_backend import PyTorchSemanticSegmentationConfig
 from rastervision.pytorch_learner import (
     Backbone, SolverConfig, SemanticSegmentationModelConfig,
-    SemanticSegmentationImageDataConfig, SemanticSegmentationGeoDataConfig,
-    GeoDataWindowConfig, GeoDataWindowMethod)
+    SemanticSegmentationImageDataConfig, SemanticSegmentationGeoDataConfig)
 
 
 def get_config(runner, root_uri, data_uri=None, full_train=False,
@@ -62,17 +61,14 @@ def make_scene(id, img_path, label_path):
         validation_scenes=scenes)
 
     chip_options = SemanticSegmentationChipOptions(
-        window_method=SemanticSegmentationWindowMethod.sliding, stride=chip_sz)
+        sampling=WindowSamplingConfig(
+            method=WindowSamplingMethod.sliding, stride=chip_sz, size=chip_sz))
 
     if nochip:
-        window_opts = GeoDataWindowConfig(
-            method=GeoDataWindowMethod.sliding,
-            stride=chip_options.stride,
-            size=chip_sz)
 
         data = SemanticSegmentationGeoDataConfig(
             scene_dataset=scene_dataset,
-            window_opts=window_opts,
+            sampling=chip_options.sampling,
             img_sz=img_sz,
             augmentors=[])
     else:
@@ -110,6 +106,5 @@ def make_scene(id, img_path, label_path):
         root_uri=root_uri,
         dataset=scene_dataset,
         backend=backend,
-        train_chip_sz=chip_sz,
-        predict_chip_sz=chip_sz,
-        chip_options=chip_options)
+        chip_options=chip_options,
+        predict_chip_sz=chip_sz)

rastervision_core/rastervision/core/__init__.py

@@ -2,7 +2,7 @@
 
 
 def register_plugin(registry):
-    registry.set_plugin_version('rastervision.core', 10)
+    registry.set_plugin_version('rastervision.core', 11)
     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

@@ -4,7 +4,8 @@
 
 if TYPE_CHECKING:
     from rastervision.core.data_sample import DataSample
-    from rastervision.core.data import Labels, Scene
+    from rastervision.core.data import DatasetConfig, Labels, Scene
+    from rastervision.core.rv_pipeline import ChipOptions
 
 
 class SampleWriter(AbstractContextManager):
@@ -57,3 +58,12 @@ def predict_scene(self, scene: 'Scene', chip_sz: int,
         Return:
             Labels object containing predictions
         """
+
+    @abstractmethod
+    def chip_dataset(self, dataset: 'DatasetConfig',
+                     chip_options: 'ChipOptions') -> None:
+        """Create and write chips for scenes in a :class:`.DatasetConfig`.
+
+        Args:
+            scenes: Scenes to chip.
+        """

rastervision_core/rastervision/core/data/label_source/chip_classification_label_source_config.py

@@ -55,7 +55,7 @@ class ChipClassificationLabelSourceConfig(LabelSourceConfig):
         None,
         description=
         ('Size of a cell to use in pixels. If None, and this Config is part '
-         'of an RVPipeline, this field will be set from RVPipeline.train_chip_sz.'
+         'of an RVPipeline, this field will be set from RVPipeline.chip_options.'
          ))
     lazy: bool = Field(
         False,
@@ -107,6 +107,6 @@ def build(self, class_config, crs_transformer, bbox=None,
     def update(self, pipeline=None, scene=None):
         super().update(pipeline, scene)
         if self.cell_sz is None and pipeline is not None:
-            self.cell_sz = pipeline.train_chip_sz
+            self.cell_sz = pipeline.chip_options.get_chip_sz(scene.id)
         if self.vector_source is not None:
             self.vector_source.update(pipeline, scene)

rastervision_core/rastervision/core/data/label_source/semantic_segmentation_label_source.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Any, List, Optional
+from typing import TYPE_CHECKING, Any, Optional
 
 import numpy as np
 
@@ -36,28 +36,6 @@ def __init__(self,
         if bbox is not None:
             self.set_bbox(bbox)
 
-    def enough_target_pixels(self, window: Box, target_count_threshold: int,
-                             target_classes: List[int]) -> bool:
-        """Check if window contains enough pixels of the given target classes.
-
-        Args:
-             window: The larger window from-which the sub-window will
-                  be clipped.
-             target_count_threshold:  Minimum number of target pixels.
-             target_classes: The classes of interest.  The given
-                  window is examined to make sure that it contains a
-                  sufficient number of target pixels.
-        Returns:
-             True (the window does contain interesting pixels) or False.
-        """
-        label_arr = self.get_label_arr(window)
-
-        target_count = 0
-        for class_id in target_classes:
-            target_count += (label_arr == class_id).sum()
-
-        return target_count >= target_count_threshold
-
     def get_labels(self,
                    window: Optional[Box] = None) -> SemanticSegmentationLabels:
         """Get labels for a window.

rastervision_core/rastervision/core/data/utils/__init__.py

@@ -6,3 +6,4 @@
 from rastervision.core.data.utils.raster import *
 from rastervision.core.data.utils.rasterio import *
 from rastervision.core.data.utils.vectorization import *
+from rastervision.core.data.utils.aoi_sampler import *

rastervision_core/rastervision/core/data_sample.py

@@ -1,17 +1,16 @@
-from pydantic import BaseModel
+from typing import Any, Literal, Optional
+from dataclasses import dataclass
+
 from numpy import ndarray
 
 from rastervision.core.box import Box
-from rastervision.core.data import Labels
 
 
-class DataSample(BaseModel):
+@dataclass
+class DataSample:
     """A chip and labels along with metadata."""
     chip: ndarray
-    window: Box
-    labels: Labels
-    scene_id: str = 'default'
-    is_train: bool = True
-
-    class Config:
-        arbitrary_types_allowed = True
+    label: Optional[Any] = None
+    split: Optional[Literal['train', 'valid', 'test']] = None
+    scene_id: Optional[str] = None
+    window: Optional[Box] = None

rastervision_core/rastervision/core/rv_pipeline/__init__.py

@@ -3,6 +3,7 @@
 TRAIN = 'train'
 VALIDATION = 'validation'
 
+from rastervision.core.rv_pipeline.chip_options import *
 from rastervision.core.rv_pipeline.rv_pipeline import *
 from rastervision.core.rv_pipeline.rv_pipeline_config import *
 from rastervision.core.rv_pipeline.chip_classification import *
@@ -22,10 +23,12 @@
     SemanticSegmentationConfig.__name__,
     SemanticSegmentationChipOptions.__name__,
     SemanticSegmentationPredictOptions.__name__,
-    SemanticSegmentationWindowMethod.__name__,
     ObjectDetection.__name__,
     ObjectDetectionConfig.__name__,
-    ObjectDetectionChipOptions.__name__,
+    ObjectDetectionWindowSamplingConfig.__name__,
     ObjectDetectionChipOptions.__name__,
     ObjectDetectionPredictOptions.__name__,
+    ChipOptions.__name__,
+    WindowSamplingConfig.__name__,
+    WindowSamplingMethod.__name__,
 ]

rastervision_core/rastervision/core/rv_pipeline/chip_classification.py

@@ -1,42 +1,5 @@
-from typing import List
-import logging
-
 from rastervision.core.rv_pipeline.rv_pipeline import RVPipeline
-from rastervision.core.rv_pipeline.utils import nodata_below_threshold
-from rastervision.core.box import Box
-from rastervision.core.data import Scene
-
-log = logging.getLogger(__name__)
-
-
-def get_train_windows(scene: Scene,
-                      chip_size: int,
-                      chip_nodata_threshold: float = 1.) -> List[Box]:
-    train_windows = []
-    extent = scene.raster_source.extent
-    stride = chip_size
-    windows = extent.get_windows(chip_size, stride)
-
-    total_windows = len(windows)
-    if scene.aoi_polygons_bbox_coords:
-        windows = Box.filter_by_aoi(windows, scene.aoi_polygons_bbox_coords)
-        log.info(f'AOI filtering: {len(windows)}/{total_windows} '
-                 'chips accepted')
-    for window in windows:
-        chip = scene.raster_source.get_chip(window)
-        if nodata_below_threshold(chip, chip_nodata_threshold, nodata_val=0):
-            train_windows.append(window)
-    log.info('NODATA filtering: '
-             f'{len(train_windows)}/{len(windows)} chips accepted')
-    return train_windows
 
 
 class ChipClassification(RVPipeline):
-    def get_train_windows(self, scene: Scene) -> List[Box]:
-        return get_train_windows(
-            scene,
-            self.config.train_chip_sz,
-            chip_nodata_threshold=self.config.chip_nodata_threshold)
-
-    def get_train_labels(self, window: Box, scene: Scene):
-        return scene.label_source.get_labels(window=window)
+    pass