Merge pull request #1079 from AdeelH/overflow

Zero out overflowing regions when reading a window from a RasterioSource with a cropped extent

Author: AdeelH

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.

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,20 +309,62 @@ 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:
             _ = RasterioSourceConfig(uris=[img_path], extent_crop=None)  # noqa
         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()