Repository: azavea/raster-vision

Issue:
Handle edge case: reading a window that overflows the RasterSource's extent that has been artificially cropped using extent_crop can cause data leakage
If, for example, there are two `RasterSource`'s with cropped extents such that they both respectively read from the left and right halves of the same underlying image, there will be instances at the boundary where windows straddle both halves. This is not an issue if the two halves are both in the training split or both in the validation split, but if they are in different splits, there will be data leakage.

**Possible solution**
`RasterSource` should fill the overflowing regions in the window with NODATA values if `extent_crop` is set. This can be a configurable option so that it can be enabled only when the image is part of scenes from different data splits.


Discussion:

Resolving patch:
diff --git a/rastervision_core/rastervision/core/data/raster_source/rasterio_source.py b/rastervision_core/rastervision/core/data/raster_source/rasterio_source.py
--- a/rastervision_core/rastervision/core/data/raster_source/rasterio_source.py
+++ b/rastervision_core/rastervision/core/data/raster_source/rasterio_source.py
@@ -73,6 +73,26 @@ def load_window(image_dataset, window=None, is_masked=False):
     return im
 
 
+def fill_overflow(extent: Box,
+                  window: Box,
+                  arr: np.ndarray,
+                  fill_value: int = 0) -> np.ndarray:
+    """Given a window and corresponding array of values, if the window
+    overflows the extent, fill the overflowing regions with fill_value.
+    """
+    top_overflow = max(0, extent.ymin - window.ymin)
+    bottom_overflow = max(0, window.ymax - extent.ymax)
+    left_overflow = max(0, extent.xmin - window.xmin)
+    right_overflow = max(0, window.xmax - extent.xmax)
+
+    h, w = arr.shape[:2]
+    arr[:top_overflow] = fill_value
+    arr[h - bottom_overflow:] = fill_value
+    arr[:, :left_overflow] = fill_value
+    arr[:, w - right_overflow:] = fill_value
+    return arr
+
+
 class RasterioSource(ActivateMixin, RasterSource):
     def __init__(self,
                  uris,
@@ -175,14 +195,17 @@ def get_dtype(self):
         """Return the numpy.dtype of this scene"""
         return self.dtype
 
-    def _get_chip(self, window):
+    def _get_chip(self, window: Box) -> np.ndarray:
         if self.image_dataset is None:
             raise ActivationError('RasterSource must be activated before use')
         shifted_window = self._get_shifted_window(window)
-        return load_window(
+        chip = load_window(
             self.image_dataset,
             window=shifted_window.rasterio_format(),
             is_masked=self.is_masked)
+        if self.extent_crop is not None:
+            chip = fill_overflow(self.get_extent(), window, chip)
+        return chip
 
     def _activate(self):
         # Download images to temporary directory and delete them when done.


Test patch:
diff --git a/tests/core/data/raster_source/test_rasterio_source.py b/tests/core/data/raster_source/test_rasterio_source.py
--- a/tests/core/data/raster_source/test_rasterio_source.py
+++ b/tests/core/data/raster_source/test_rasterio_source.py
@@ -1,15 +1,18 @@
 import unittest
 from os.path import join
 from pydantic import ValidationError
+from tempfile import NamedTemporaryFile
 
 import numpy as np
 import rasterio
 from rasterio.enums import ColorInterp
 
 from rastervision.core import (RasterStats)
+from rastervision.core.box import Box
 from rastervision.core.utils.misc import save_img
 from rastervision.core.data import (ChannelOrderError, RasterioSourceConfig,
-                                    StatsTransformerConfig, CropOffsets)
+                                    StatsTransformerConfig, CropOffsets,
+                                    fill_overflow)
 from rastervision.pipeline import rv_config
 
 from tests import data_file_path
@@ -306,13 +309,13 @@ def test_extent_crop(self):
         self.assertRaises(
             ValidationError,
             lambda: RasterioSourceConfig(uris=[img_path],
-                                         extent_crop=extent_crop))
+                                            extent_crop=extent_crop))
 
         extent_crop = CropOffsets(skip_left=.5, skip_right=.5)
         self.assertRaises(
             ValidationError,
             lambda: RasterioSourceConfig(uris=[img_path],
-                                         extent_crop=extent_crop))
+                                            extent_crop=extent_crop))
 
         # test extent_crop=None
         try:
@@ -320,6 +323,48 @@ def test_extent_crop(self):
         except Exception:
             self.fail('extent_crop=None caused an error.')
 
+    def test_fill_overflow(self):
+        extent = Box(10, 10, 90, 90)
+        window = Box(0, 0, 100, 100)
+        arr = np.ones((100, 100), dtype=np.uint8)
+        out = fill_overflow(extent, window, arr)
+        mask = np.zeros_like(arr).astype(np.bool)
+        mask[10:90, 10:90] = 1
+        self.assertTrue(np.all(out[mask] == 1))
+        self.assertTrue(np.all(out[~mask] == 0))
+
+        window = Box(0, 0, 80, 100)
+        arr = np.ones((80, 100), dtype=np.uint8)
+        out = fill_overflow(extent, window, arr)
+        mask = np.zeros((80, 100), dtype=np.bool)
+        mask[10:90, 10:90] = 1
+        self.assertTrue(np.all(out[mask] == 1))
+        self.assertTrue(np.all(out[~mask] == 0))
+
+    def test_extent_crop_overflow(self):
+        f = 1 / 10
+        arr = np.ones((100, 100), dtype=np.uint8)
+        mask = np.zeros_like(arr).astype(np.bool)
+        mask[10:90, 10:90] = 1
+        with NamedTemporaryFile('wb') as fp:
+            uri = fp.name
+            with rasterio.open(
+                    uri,
+                    'w',
+                    driver='GTiff',
+                    height=100,
+                    width=100,
+                    count=1,
+                    dtype=np.uint8) as ds:
+                ds.write_band(1, arr)
+            cfg = RasterioSourceConfig(uris=[uri], extent_crop=(f, f, f, f))
+            rs = cfg.build(tmp_dir=self.tmp_dir)
+            with rs.activate():
+                out = rs.get_chip(Box(0, 0, 100, 100))[..., 0]
+
+        self.assertTrue(np.all(out[mask] == 1))
+        self.assertTrue(np.all(out[~mask] == 0))
+
 
 if __name__ == '__main__':
     unittest.main()