Repository: opendatacube/odc-geo
Pull request #146 (2024-04-30 07:50:10)

Add tag writing support including scale, offset and units special cases
Added a simple test.

Open to feedback on a better method, but I think this works ok and hopefully handles all cases!

github-actions[bot]: <!-- NETLIFY DEPLOY COMMENT GENERATED BY ACTIONS_NETLIFY - APP ID SHA256: 2516d918dc82767c089255cf7129f75d9f959e53c07db8d3368972c4656b3dce -->
🚀 Deployed on https://6630a3e4198d27938a651018--odc-geo-docs.netlify.app
alexgleith: So, I checked, and all these tags and special metadata items are written to the GDAL section in the tif metadata:

```
  GDAL Metadata: <GDALMetadata>
  <Item name="random_tag">random_value</Item>
  <Item name="OFFSET" sample="0" role="offset">0.989999999999999991</Item>
  <Item name="SCALE" sample="0" role="scale">100</Item>
  <Item name="UNITTYPE" sample="0" role="unittype">microwaves</Item>
</GDALMetadata>
```
Kirill888: This is all fine, I do have issue with the name `tags`: since this is TIFF making function, my first assumption is that `tags` refers to TIFF tags and not to "GDAL" tags, aka gdal metadata items, or general "annotation" of some sort. On the other hand [that's the nomenclature used by `rioxarray`](https://github.com/corteva/rioxarray/blob/13a3a99b3779f86e06a5130210a835af937cfa99/rioxarray/raster_writer.py#L88) also, so...
Kirill888: @alexgleith do you want to rebase this into one clean commit and maybe add a note to docs clarifying that tags refer to "GDAL tags/metadata". As a person who actually read TIFF spec I'm probably biased in my interpretation of `tags`, but it won't hurt being more precise about it.

Patch:
diff --git a/odc/geo/cog/_rio.py b/odc/geo/cog/_rio.py
index 132fe51..6176a60 100644
--- a/odc/geo/cog/_rio.py
+++ b/odc/geo/cog/_rio.py
@@ -83,6 +83,40 @@ def _norm_compression_opts(
     return compression
 
 
+def _get_gdal_metadata(
+    xx: xr.DataArray | list[xr.DataArray],
+    tags: Optional[Dict[str, Any]],
+) -> Dict[str, Any]:
+    """
+    Extract metadata from xarray for GDAL.
+
+    :param xx: xarray.DataArray
+    :return: Dictionary of metadata
+    """
+    meta: Dict[str, Any] = {"tags": tags}
+    if not isinstance(xx, list):
+        xx = [xx]
+
+    # The logic for this is to fill the lists with
+    # None values if the attribute is not present
+    # as only some bands may have scale and offset
+    # and we need to know which ones...
+    scales, offsets, units = [], [], []
+    for x in xx:
+        scales.append(x.attrs.get("scales", None))
+        offsets.append(x.attrs.get("offsets", None))
+        units.append(x.attrs.get("units", None))
+
+    if any(v is not None for v in scales):
+        meta["scales"] = scales
+    if any(v is not None for v in offsets):
+        meta["offsets"] = offsets
+    if any(v is not None for v in units):
+        meta["units"] = units
+
+    return meta
+
+
 def _write_cog(
     pix: np.ndarray,
     geobox: GeoBox,
@@ -95,6 +129,7 @@ def _write_cog(
     ovr_blocksize: Optional[int] = None,
     use_windowed_writes: bool = False,
     intermediate_compression: Union[bool, str, Dict[str, Any]] = False,
+    gdal_metadata: Optional[Dict[str, Any]] = None,
     **extra_rio_opts,
 ) -> Union[Path, bytes]:
     if blocksize is None:
@@ -158,15 +193,33 @@ def _write_cog(
     def _write(pix, band, dst):
         if not use_windowed_writes:
             dst.write(pix, band)
-            return
+        else:
+            for _, win in dst.block_windows():
+                if pix.ndim == 2:
+                    block = pix[win.toslices()]
+                else:
+                    block = pix[(slice(None),) + win.toslices()]
 
-        for _, win in dst.block_windows():
-            if pix.ndim == 2:
-                block = pix[win.toslices()]
-            else:
-                block = pix[(slice(None),) + win.toslices()]
+                dst.write(block, indexes=band, window=win)
+
+        if gdal_metadata is not None:
+            # Add scales, offsets and units if present
+            scales = gdal_metadata.get("scales", None)
+            if scales is not None:
+                dst.scales = scales
+
+            offsets = gdal_metadata.get("offsets", None)
+            if offsets is not None:
+                dst.offsets = offsets
+
+            units = gdal_metadata.get("units", None)
+            if units is not None:
+                dst.units = units
 
-            dst.write(block, indexes=band, window=win)
+            # Add any tags that are present
+            tags = gdal_metadata.get("tags", None)
+            if tags is not None:
+                dst.update_tags(**tags)
 
     # Deal efficiently with "no overviews needed case"
     if len(overview_levels) == 0:
@@ -227,6 +280,7 @@ def write_cog(
     overview_levels: Optional[List[int]] = None,
     use_windowed_writes: bool = False,
     intermediate_compression: Union[bool, str, Dict[str, Any]] = False,
+    tags: Optional[Dict[str, Any]] = None,
     **extra_rio_opts,
 ) -> Union[Path, bytes]:
     """
@@ -244,7 +298,10 @@ def write_cog(
     :param nodata: Set ``nodata`` flag to this value if supplied, by default ``nodata`` is
                    read from the attributes of the input array (``geo_im.attrs['nodata']``).
     :param use_windowed_writes: Write image block by block (might need this for large images)
-    :param intermediate_compression: Configure compression settings for first pass write, default is no compression
+    :param intermediate_compression: Configure compression settings for first pass write
+                    , default is no compression
+    :param tags: Dictionary of tags to write into the output file. These are written as
+                    GDAL Metadata items in the GeoTIFF file.
     :param extra_rio_opts: Any other option is passed to ``rasterio.open``
 
     :returns: Path to which output was written
@@ -273,6 +330,7 @@ def write_cog(
             ovr_blocksize=ovr_blocksize,
             use_windowed_writes=use_windowed_writes,
             intermediate_compression=intermediate_compression,
+            tags=tags,
             **extra_rio_opts,
         )
         assert result is not None
@@ -299,6 +357,7 @@ def write_cog(
         overview_levels=overview_levels,
         use_windowed_writes=use_windowed_writes,
         intermediate_compression=intermediate_compression,
+        gdal_metadata=_get_gdal_metadata(geo_im, tags),
         **extra_rio_opts,
     )
 
@@ -312,6 +371,7 @@ def to_cog(
     overview_levels: Optional[List[int]] = None,
     use_windowed_writes: bool = False,
     intermediate_compression: Union[bool, str, Dict[str, Any]] = False,
+    tags: Optional[Dict[str, Any]] = None,
     **extra_rio_opts,
 ) -> bytes:
     """
@@ -329,7 +389,10 @@ def to_cog(
     :param nodata: Set ``nodata`` flag to this value if supplied, by default ``nodata`` is
                    read from the attributes of the input array (``geo_im.attrs['nodata']``).
     :param use_windowed_writes: Write image block by block (might need this for large images)
-    :param intermediate_compression: Configure compression settings for first pass write, default is no compression
+    :param intermediate_compression: Configure compression settings for first pass write
+                    , default is no compression
+    :param tags: Dictionary of tags to write into the output file. These are written as
+                    GDAL Metadata items in the GeoTIFF file.
     :param extra_rio_opts: Any other option is passed to ``rasterio.open``
 
     :returns: In-memory GeoTiff file as bytes
@@ -345,6 +408,7 @@ def to_cog(
         overview_levels=overview_levels,
         use_windowed_writes=use_windowed_writes,
         intermediate_compression=intermediate_compression,
+        tags=tags,
         **extra_rio_opts,
     )
 
@@ -381,6 +445,7 @@ def write_cog_layers(
     ovr_blocksize: Optional[int] = None,
     intermediate_compression: Union[bool, str, Dict[str, Any]] = False,
     use_windowed_writes: bool = False,
+    tags: Optional[Dict[str, Any]] = None,
     **extra_rio_opts,
 ) -> Union[Path, bytes, None]:
     """
@@ -417,6 +482,7 @@ def write_cog_layers(
         "blocksize": blocksize,
         "nodata": rio_opts.get("nodata", None),
         "use_windowed_writes": use_windowed_writes,
+        "gdal_metadata": _get_gdal_metadata(xx, tags),
         **_norm_compression_opts(intermediate_compression),
     }
 


Test patch:
diff --git a/tests/test_cog.py b/tests/test_cog.py
index 05dd9aa..db740d0 100644
--- a/tests/test_cog.py
+++ b/tests/test_cog.py
@@ -6,6 +6,8 @@ from typing import Optional, Tuple
 import numpy as np
 import pytest
 from dask import array as da
+from rasterio import MemoryFile
+from rasterio import open as rio_open
 
 from odc.geo.cog import CogMeta, cog_gbox, save_cog_with_dask
 from odc.geo.cog._shared import compute_cog_spec, num_overviews
@@ -55,6 +57,51 @@ def test_write_cog(gbox: GeoBox):
     assert len(img_bytes) == len(img_bytes2)
 
 
+@pytest.mark.parametrize(
+    ["scales", "offsets", "units"],
+    [
+        (0.1, 999, "fridges"),
+        (0.1, 999, None),
+    ],
+)
+def test_write_metadata(gbox: GeoBox, scales, offsets, units):
+    RANDOM = "random_value"
+    img = xr_zeros(gbox, dtype="uint16")
+    assert img.odc.geobox == gbox
+
+    img.attrs.update(scales=scales, offsets=offsets, units=units)
+
+    assert img.attrs["scales"] == scales
+
+    img_bytes = img.odc.to_cog(
+        blocksize=32,
+        tags=dict(random_tag=RANDOM),
+    )
+    with MemoryFile(img_bytes) as memfile:
+        with rio_open(memfile) as src:
+            # Make sure values are set and retrieved correctly
+            if isinstance(scales, (list, tuple)):
+                assert src.scales == tuple(scales)
+            else:
+                assert src.scales[0] == scales
+
+            if isinstance(offsets, (list, tuple)):
+                assert src.offsets == tuple(offsets)
+            else:
+                assert src.offsets[0] == offsets
+
+            if isinstance(units, (list, tuple)):
+                assert src.units == tuple(units)
+            else:
+                assert src.units[0] == units
+
+            # Check tags stick
+            assert src.tags()["random_tag"] == RANDOM
+
+            # Check the underlying data asn't altered
+            assert (src.read() == img.data).all()
+
+
 def test_write_cog_ovr(gbox: GeoBox):
     img = xr_zeros(gbox, dtype="uint16")
     assert img.odc.geobox == gbox
@@ -516,3 +563,4 @@ def test_stats_from_layer(array, nodata, minimum, maximum, mean, stddev, valid_p
     assert stats["mean"] == mean
     assert stats["stddev"] == stddev
     assert stats["valid_percent"] == valid_percent
+    assert stats["valid_percent"] == valid_percent