Merge pull request #67 from djhoese/feature-overview-resampling

Add option for geotiff overview resampling and auto-levels

doc/conf.py

@@ -251,4 +251,5 @@ def __getattr__(cls, name):
     'scipy': ('https://docs.scipy.org/doc/scipy/reference', None),
     'xarray': ('https://xarray.pydata.org/en/stable', None),
     'dask': ('https://dask.pydata.org/en/latest', None),
+    'rasterio': ('https://rasterio.readthedocs.io/en/latest', None),
 }

trollimage/tests/test_image.py

@@ -1216,6 +1216,30 @@ def test_save_overviews(self):
             with rio.open(tmp.name) as f:
                 self.assertEqual(len(f.overviews(1)), 2)
 
+        # auto-levels
+        data = np.zeros(25*25*3, dtype=np.uint8).reshape(25, 25, 3)
+        data = xr.DataArray(data, dims=[
+            'y', 'x', 'bands'], coords={'bands': ['R', 'G', 'B']})
+        img = xrimage.XRImage(data)
+        self.assertTrue(np.issubdtype(img.data.dtype, np.integer))
+        with NamedTemporaryFile(suffix='.tif') as tmp:
+            img.save(tmp.name, overviews=[], overviews_minsize=2)
+            with rio.open(tmp.name) as f:
+                self.assertEqual(len(f.overviews(1)), 4)
+
+        # auto-levels and resampling
+        data = np.zeros(25*25*3, dtype=np.uint8).reshape(25, 25, 3)
+        data = xr.DataArray(data, dims=[
+            'y', 'x', 'bands'], coords={'bands': ['R', 'G', 'B']})
+        img = xrimage.XRImage(data)
+        self.assertTrue(np.issubdtype(img.data.dtype, np.integer))
+        with NamedTemporaryFile(suffix='.tif') as tmp:
+            img.save(tmp.name, overviews=[], overviews_minsize=2,
+                     overviews_resampling='average')
+            with rio.open(tmp.name) as f:
+                # no way to check resampling method from the file
+                self.assertEqual(len(f.overviews(1)), 4)
+
     @unittest.skipIf(sys.platform.startswith('win'), "'NamedTemporaryFile' not supported on Windows")
     def test_save_tags(self):
         """Test saving geotiffs with tags."""

trollimage/xrimage.py

@@ -47,6 +47,7 @@
 
 try:
     import rasterio
+    from rasterio.enums import Resampling
 except ImportError:
     rasterio = None
 
@@ -74,6 +75,16 @@ def __init__(self, path, mode='w', **kwargs):
         self.rfile = None
         self.lock = threading.Lock()
 
+    @property
+    def width(self):
+        """Width of the band images."""
+        return self.kwargs['width']
+
+    @property
+    def height(self):
+        """Height of the band images."""
+        return self.kwargs['height']
+
     @property
     def closed(self):
         """Check if the file is closed."""
@@ -160,10 +171,15 @@ def close(self):
 class RIODataset:
     """A wrapper for a rasterio dataset."""
 
-    def __init__(self, rfile, overviews=None):
+    def __init__(self, rfile, overviews=None, overviews_resampling=None,
+                 overviews_minsize=256):
         """Init the rasterio dataset."""
         self.rfile = rfile
         self.overviews = overviews
+        if overviews_resampling is None:
+            overviews_resampling = 'nearest'
+        self.overviews_resampling = Resampling[overviews_resampling]
+        self.overviews_minsize = overviews_minsize
 
     def __setitem__(self, key, item):
         """Put the data chunk in the image."""
@@ -190,9 +206,19 @@ def __setitem__(self, key, item):
 
     def close(self):
         """Close the file."""
-        if self.overviews:
-            logger.debug('Building overviews %s', str(self.overviews))
-            self.rfile.build_overviews(self.overviews)
+        if self.overviews is not None:
+            overviews = self.overviews
+            # it's an empty list
+            if len(overviews) == 0:
+                from rasterio.rio.overview import get_maximum_overview_level
+                width = self.rfile.width
+                height = self.rfile.height
+                max_level = get_maximum_overview_level(
+                    width, height, self.overviews_minsize)
+                overviews = [2 ** j for j in range(1, max_level + 1)]
+            logger.debug('Building overviews %s with %s resampling',
+                         str(overviews), self.overviews_resampling.name)
+            self.rfile.build_overviews(overviews, resampling=self.overviews_resampling)
 
         return self.rfile.close()
 
@@ -401,6 +427,7 @@ def save(self, filename, fformat=None, fill_value=None, compute=True,
     def rio_save(self, filename, fformat=None, fill_value=None,
                  dtype=np.uint8, compute=True, tags=None,
                  keep_palette=False, cmap=None, overviews=None,
+                 overviews_minsize=256, overviews_resampling=None,
                  include_scale_offset_tags=False,
                  **format_kwargs):
         """Save the image using rasterio.
@@ -424,6 +451,20 @@ def rio_save(self, filename, fformat=None, fill_value=None,
 
                     img.rio_save('myfile.tif', overviews=[2, 4, 8, 16])
 
+                If provided as an empty list, then levels will be
+                computed as powers of two until the last level has less
+                pixels than `overviews_minsize`.
+                Default is to not add overviews.
+            overviews_minsize (int): Minimum number of pixels for the smallest
+                overview size generated when `overviews` is auto-generated.
+                Defaults to 256.
+            overviews_resampling (str): Resampling method
+                to use when generating overviews. This must be the name of an
+                enum value from :class:`rasterio.enums.Resampling` and
+                only takes effect if the `overviews` keyword argument is
+                provided. Common values include `nearest` (default),
+                `bilinear`, `average`, and many others. See the rasterio
+                documentation for more information.
             include_scale_offset_tags (bool): Whether or not (default) to
                 include a ``scale`` and an ``offset`` tag in the data that would
                 help retrieving original data values from pixel values.
@@ -534,7 +575,9 @@ def rio_save(self, filename, fformat=None, fill_value=None,
                 continue
 
         r_file.rfile.update_tags(**tags)
-        r_dataset = RIODataset(r_file, overviews)
+        r_dataset = RIODataset(r_file, overviews,
+                               overviews_resampling=overviews_resampling,
+                               overviews_minsize=overviews_minsize)
 
         to_store = (data.data, r_dataset)
         if da_tags: