add windows_within_aoi arg to SlidingWindowGeoDataset

azavea · Mar 29, 2024 · ba75659 · ba75659
1 parent 53de47b
commit ba75659
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/rastervision_pytorch_learner/rastervision/pytorch_learner/dataset/dataset.py b/rastervision_pytorch_learner/rastervision/pytorch_learner/dataset/dataset.py
@@ -172,6 +172,7 @@ def __init__(self,
                  padding: Optional[Union[NonNegInt, Tuple[NonNegInt,
                                                           NonNegInt]]] = None,
                  pad_direction: Literal['both', 'start', 'end'] = 'end',
+                 windows_within_aoi: bool = True,
                  transform: Optional[A.BasicTransform] = None,
                  transform_type: Optional[TransformType] = None,
                  normalize: bool = True,
@@ -192,6 +193,9 @@ def __init__(self,
                 ymax and xmax (bottom and right). If 'start', only pad ymin and
                 xmin (top and left). If 'both', pad all sides. Has no effect if
                 paddiong is zero. Defaults to 'end'.
+            windows_within_aoi: If True, discard windows that are not fully
+                within the AOI. If False, keep windows that intersect AOI.
+                Defaults to True.
             transform (Optional[A.BasicTransform], optional): Albumentations
                 transform to apply to the windows. Defaults to None.
                 Each transform in Albumentations takes images of type uint8, and
@@ -221,6 +225,7 @@ def __init__(self,
         self.stride = _to_tuple(stride)
         self.padding = padding
         self.pad_direction = pad_direction
+        self.windows_within_aoi = windows_within_aoi
         self.init_windows()
 
     def init_windows(self) -> None:
@@ -231,8 +236,10 @@ def init_windows(self) -> None:
             padding=self.padding,
             pad_direction=self.pad_direction)
         if len(self.scene.aoi_polygons_bbox_coords) > 0:
-            windows = Box.filter_by_aoi(windows,
-                                        self.scene.aoi_polygons_bbox_coords)
+            windows = Box.filter_by_aoi(
+                windows,
+                self.scene.aoi_polygons_bbox_coords,
+                within=self.windows_within_aoi)
         self.windows = windows
 
     def __getitem__(self, idx: int):