Repository: opendatacube/odc-geo (python)

Issue:
Feature request: have save_cog_with_dask write band descriptions
Using `write_cog` we get band descriptions, but these don't currently show up if we use `save_cog_with_dask`. Could this be added? For example if I open a COG created with `write_cog` in a text editor I see a block like:
```
<GDALMetadata>
  <Item name="OVR_RESAMPLING_ALG">NEAREST</Item>
  <Item name="DESCRIPTION" sample="0" role="description">green</Item>
  <Item name="DESCRIPTION" sample="1" role="description">red</Item>
  <Item name="DESCRIPTION" sample="2" role="description">nir</Item>
  <Item name="DESCRIPTION" sample="3" role="description">swir1</Item>
  <Item name="DESCRIPTION" sample="4" role="description">swir2</Item>
  <Item name="DESCRIPTION" sample="5" role="description">swir3</Item>
  <Item name="DESCRIPTION" sample="6" role="description">swir4</Item>
  <Item name="DESCRIPTION" sample="7" role="description">swir5</Item>
  <Item name="DESCRIPTION" sample="8" role="description">swir6</Item>
  <Item name="DESCRIPTION" sample="9" role="description">tir1</Item>
  <Item name="DESCRIPTION" sample="10" role="description">tir2</Item>
  <Item name="DESCRIPTION" sample="11" role="description">tir3</Item>
  <Item name="DESCRIPTION" sample="12" role="description">tir4</Item>
  <Item name="DESCRIPTION" sample="13" role="description">tir5</Item>
</GDALMetadata>
```
Which means in QGIS I see band labels like "Band 04: swir1" whereas for a COG made by `save_cog_with_dask` I just see "Band 04"

cc @mplough-kobold @Kirill888 

Pull request description:
Write band descriptions in save_cog_with_dask
The xarray and rioxarray packages use a dataset's `long_name` attribute to store band names.  This PR saves those band names in COGs written with `save_cog_with_dask`.  The names are written to be as compatible as possible with GDAL's band name writes.

This mechanism may also be used to write other GDAL metadata.

Resolves #138.

Patch:
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 2720365..bd87e5a 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -218,7 +218,7 @@ jobs:
         if: |
           github.repository == 'opendatacube/odc-geo'
 
-        uses: codecov/codecov-action@v3
+        uses: codecov/codecov-action@v4
         with:
           fail_ci_if_error: false
           verbose: false
diff --git a/odc/geo/cog/_tifffile.py b/odc/geo/cog/_tifffile.py
index e69f93d..6d95bd9 100644
--- a/odc/geo/cog/_tifffile.py
+++ b/odc/geo/cog/_tifffile.py
@@ -11,6 +11,7 @@ import itertools
 from functools import partial
 from io import BytesIO
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
+from xml.sax.saxutils import escape as xml_escape
 
 import numpy as np
 import xarray as xr
@@ -39,10 +40,11 @@ if TYPE_CHECKING:
 
 
 def _render_gdal_metadata(
-    band_stats: list[dict[str, float]] | dict[str, float],
+    band_stats: list[dict[str, float]] | dict[str, float] | None,
     precision: int = 10,
     pad: int = 0,
     eol: str = "",
+    gdal_metadata_extra: Optional[List[str]] = None,
 ) -> str:
     def _item(sample: int, stats: dict[str, float]) -> str:
         return eol.join(
@@ -52,10 +54,17 @@ def _render_gdal_metadata(
             ]
         )
 
+    if band_stats is None:
+        band_stats = []
     if isinstance(band_stats, dict):
         band_stats = [band_stats]
 
-    body = eol.join([_item(sample, stats) for sample, stats in enumerate(band_stats)])
+    gdal_metadata_extra = [] if gdal_metadata_extra is None else gdal_metadata_extra
+
+    body = eol.join(
+        [_item(sample, stats) for sample, stats in enumerate(band_stats)]
+        + gdal_metadata_extra
+    )
     return eol.join(["<GDALMetadata>", body, "</GDALMetadata>"])
 
 
@@ -469,6 +478,7 @@ def _patch_hdr(
     meta: CogMeta,
     hdr0: bytes,
     stats: Optional[list[dict[str, float]]] = None,
+    gdal_metadata_extra: Optional[List[str]] = None,
 ) -> bytes:
     # pylint: disable=import-outside-toplevel,import-error
     from tifffile import TiffFile, TiffPage
@@ -479,10 +489,12 @@ def _patch_hdr(
     _bio = BytesIO(hdr0)
     with TiffFile(_bio, mode="r+", name=":mem:") as tr:
         assert len(tile_info) == len(tr.pages)
-        if stats is not None:
+        if stats is not None or gdal_metadata_extra:
             md_tag = tr.pages.first.tags.get(42112, None)
             assert md_tag is not None
-            gdal_metadata = _render_gdal_metadata(stats, precision=6)
+            gdal_metadata = _render_gdal_metadata(
+                stats, precision=6, gdal_metadata_extra=gdal_metadata_extra
+            )
             md_tag.overwrite(gdal_metadata)
 
         hdr_sz = len(_bio.getbuffer())
@@ -573,6 +585,40 @@ def _norm_compression_tifffile(
     return (predictor, compression, compressionargs)
 
 
+def _gdal_sample_description(sample: int, description: str) -> str:
+    """Make XML line of GDAL metadata.
+
+    :param band: Sample / band number in data array.
+    :param description: Band name in data array.
+
+    :return: GDAL XML metadata line to place in TIFF file.
+    """
+    # GDAL does double escaping; see frmts/gtiff/geotiff.cpp.
+    # We also double escape to maximize compatibility with tools expecting GDAL-generated metadata.
+    double_escaped_description = xml_escape(xml_escape(description))
+    return f'<Item name="DESCRIPTION" sample="{sample}" role="description">{double_escaped_description}</Item>'
+
+
+def _band_names(xx: xr.DataArray) -> List[str]:
+    if "band" in xx.coords and xx.coords["band"].dtype.type is np.str_:
+        return list(xx["band"].values)
+    if "long_name" in xx.attrs:
+        long_name = xx.attrs["long_name"]
+        return [long_name] if isinstance(long_name, str) else long_name
+    return []
+
+
+def _gdal_sample_descriptions(descriptions: List[str]) -> List[str]:
+    """Convert band names to GDAL sample descriptions.
+
+    :return: List of GDAL XML metadata lines to place in TIFF file.
+    """
+    return [
+        _gdal_sample_description(sample, description)
+        for sample, description in enumerate(descriptions)
+    ]
+
+
 def save_cog_with_dask(
     xx: xr.DataArray,
     dst: str = "",
@@ -635,7 +681,10 @@ def save_cog_with_dask(
     if isinstance(blocksize, Unset):
         blocksize = [data_chunks, int(max(*data_chunks) // 2)]
 
-    gdal_metadata = None if stats is False else ""
+    band_names = _band_names(xx)
+    sample_descriptions_metadata = _gdal_sample_descriptions(band_names)
+    no_metadata = (stats is False) and not band_names
+    gdal_metadata = None if no_metadata else ""
 
     meta, hdr0 = _make_empty_cog(
         xx.shape,
@@ -652,6 +701,11 @@ def save_cog_with_dask(
     )
     hdr0 = bytes(hdr0)
 
+    if band_names and len(band_names) != meta.nsamples:
+        raise ValueError(
+            f"Found {len(band_names)} band names ({band_names}) but there are {meta.nsamples} bands."
+        )
+
     layers = _pyramids_from_cog_metadata(xx, meta, resampling=overview_resampling)
 
     if stats is True:
@@ -695,7 +749,12 @@ def save_cog_with_dask(
             tiles_write_order,
             write,
             mk_header=_patch_hdr,
-            user_kw={"meta": meta, "hdr0": hdr0, "stats": _stats},
+            user_kw={
+                "meta": meta,
+                "hdr0": hdr0,
+                "stats": _stats,
+                "gdal_metadata_extra": sample_descriptions_metadata,
+            },
             **upload_params,
         )
 
@@ -710,7 +769,12 @@ def save_cog_with_dask(
     return s3_sink.upload(
         tiles_write_order,
         mk_header=_patch_hdr,
-        user_kw={"meta": meta, "hdr0": hdr0, "stats": _stats},
+        user_kw={
+            "meta": meta,
+            "hdr0": hdr0,
+            "stats": _stats,
+            "gdal_metadata_extra": sample_descriptions_metadata,
+        },
         client=client,
         **upload_params,
     )


Test patch:
diff --git a/tests/test_cog.py b/tests/test_cog.py
index fefd59c..ee1561e 100644
--- a/tests/test_cog.py
+++ b/tests/test_cog.py
@@ -9,6 +9,9 @@ from odc.geo.cog import CogMeta, cog_gbox, save_cog_with_dask
 from odc.geo.cog._shared import compute_cog_spec, num_overviews
 from odc.geo.cog._tifffile import (
     GEOTIFF_TAGS,
+    _band_names,
+    _gdal_sample_description,
+    _gdal_sample_descriptions,
     _make_empty_cog,
     _norm_compression_tifffile,
     geotiff_metadata,
@@ -331,6 +334,77 @@ def test_geotiff_metadata(gbox: GeoBox, nodata, gdal_metadata: Optional[str]):
         assert md["GTModelTypeGeoKey"] == 32767
 
 
+@pytest.mark.parametrize(
+    ("sample", "description", "expected"),
+    [
+        (
+            0,
+            "easy",
+            '<Item name="DESCRIPTION" sample="0" role="description">easy</Item>',
+        ),
+        (
+            1,
+            "<",
+            '<Item name="DESCRIPTION" sample="1" role="description">&amp;lt;</Item>',
+        ),
+        (
+            2,
+            ">",
+            '<Item name="DESCRIPTION" sample="2" role="description">&amp;gt;</Item>',
+        ),
+        (
+            3,
+            "&",
+            '<Item name="DESCRIPTION" sample="3" role="description">&amp;amp;</Item>',
+        ),
+        (
+            4,
+            "& <a>",
+            '<Item name="DESCRIPTION" sample="4" role="description">&amp;amp; &amp;lt;a&amp;gt;</Item>',
+        ),
+    ],
+)
+def test_gdal_sample_description(sample: int, description: str, expected: str):
+    assert _gdal_sample_description(sample, description) == expected
+
+
+def test_gdal_sample_descriptions():
+    assert _gdal_sample_descriptions(["red", "green", "blue"]) == [
+        '<Item name="DESCRIPTION" sample="0" role="description">red</Item>',
+        '<Item name="DESCRIPTION" sample="1" role="description">green</Item>',
+        '<Item name="DESCRIPTION" sample="2" role="description">blue</Item>',
+    ]
+
+
+def test_band_names(gbox: GeoBox):
+    gbox = gbox.zoom_to(1024)
+    dtype = "float32"
+    n = 512
+    nodata = -9999
+
+    img = xr_zeros(gbox, dtype, chunks=(n, n), nodata=nodata)
+    img.attrs["long_name"] = []
+
+    assert _band_names(img) == []
+
+    img.attrs["long_name"] = "first band"
+    assert _band_names(img) == ["first band"]
+
+    img.attrs["long_name"] = ["first band"]
+    assert _band_names(img) == ["first band"]
+
+    img = img.odc.colorize()
+    assert img.ndim == 3
+
+    # Obtain descriptions from strings in "band" coordinate
+    assert _band_names(img) == ["r", "g", "b", "a"]
+
+    # Obtain descriptions from "long_names" attribute
+    img["band"] = [0, 1, 2, 3]
+    img.attrs["long_name"] = ["red", "green", "blue", "alpha"]
+    assert _band_names(img) == ["red", "green", "blue", "alpha"]
+
+
 @pytest.mark.parametrize("dtype", ["int16", "float32"])
 def test_cog_with_dask_smoke_test(gbox: GeoBox, tmp_path: Path, dtype):
     gbox = gbox.zoom_to(1024)
@@ -363,3 +437,16 @@ def test_cog_with_dask_smoke_test(gbox: GeoBox, tmp_path: Path, dtype):
     fut = save_cog_with_dask(img, fname, compression="deflate", level=2)
     rr = fut.compute()
     assert str(rr) == fname
+
+    # Band names
+    img.attrs["long_name"] = ["2000", "2001"]
+    fname = str(tmp_path / "cog-bandnames.tif")
+    fut = save_cog_with_dask(img, fname, compression="deflate", level=2)
+    rr = fut.compute()
+    assert str(rr) == fname
+
+    # Incorrect number of band names
+    img.attrs["long_name"] = ["red", "green", "blue"]
+    fname = str(tmp_path / "cog-bandnames-incorrect.tif")
+    with pytest.raises(ValueError):
+        save_cog_with_dask(img, fname, compression="deflate", level=2)