diff --git a/odc/geo/_cog.py b/odc/geo/_cog.py index 0c590435..089da2ec 100644 --- a/odc/geo/_cog.py +++ b/odc/geo/_cog.py @@ -8,6 +8,7 @@ import itertools import warnings from contextlib import contextmanager +from dataclasses import dataclass from io import BytesIO from pathlib import Path from typing import Any, Dict, Generator, Iterable, List, Literal, Optional, Tuple, Union @@ -25,11 +26,28 @@ from .types import MaybeNodata, Shape2d, SomeShape, Unset, shape_, wh_ from .warp import resampling_s2rio -# pylint: disable=too-many-locals,too-many-branches,too-many-arguments,too-many-statements +# pylint: disable=too-many-locals,too-many-branches,too-many-arguments,too-many-statements,too-many-instance-attributes AxisOrder = Union[Literal["YX"], Literal["YXS"], Literal["SYX"]] +@dataclass +class CogMeta: + """ + COG metadata. + """ + + axis: AxisOrder + shape: Shape2d + tile: Shape2d + pix_shape: Tuple[int, ...] + dtype: Any + compression: int + predictor: Any + gbox: Optional[GeoBox] = None + overviews: Optional[List["CogMeta"]] = None + + def _without(cfg: Dict[str, Any], *skip: str) -> Dict[str, Any]: skip = set(skip) return {k: v for k, v in cfg.items() if k not in skip} @@ -556,7 +574,7 @@ def make_empty_cog( predictor: Union[int, str, bool, Unset] = Unset(), blocksize: Union[int, List[Union[int, Tuple[int, int]]]] = 2048, **kw, -) -> memoryview: +) -> Tuple[CogMeta, memoryview]: # pylint: disable=import-outside-toplevel have.check_or_error("tifffile", "rasterio", "xarray") from tifffile import ( @@ -572,10 +590,10 @@ def make_empty_cog( compression = str(kw.pop("compress", "ADOBE_DEFLATE")) compression = compression.upper() compression = {"DEFLATE": "ADOBE_DEFLATE"}.get(compression, compression) - compression = COMPRESSION[compression] + _compression = int(COMPRESSION[compression]) if isinstance(predictor, Unset): - predictor = compression not in TIFF.IMAGE_COMPRESSIONS + predictor = _compression not in TIFF.IMAGE_COMPRESSIONS if isinstance(blocksize, int): blocksize = [blocksize] @@ -594,12 +612,6 @@ def make_empty_cog( else: nsamples = shape[0] - extratags: List[Tuple[int, int, int, Any]] = [] - if gbox is not None: - extratags, _ = geotiff_metadata( - gbox, nodata=nodata, gdal_metadata=gdal_metadata - ) - buf = BytesIO() opts_common = { @@ -607,7 +619,7 @@ def make_empty_cog( "photometric": photometric, "planarconfig": planarconfig, "predictor": predictor, - "compression": compression, + "compression": _compression, "software": False, **kw, } @@ -622,25 +634,45 @@ def _sh(shape: Shape2d) -> Tuple[int, ...]: tsz = _norm_blocksize(blocksize[-1]) im_shape, _, nlevels = _compute_cog_spec(im_shape, tsz) + extratags: List[Tuple[int, int, int, Any]] = [] + if gbox is not None: + gbox = gbox.expand(im_shape) + extratags, _ = geotiff_metadata( + gbox, nodata=nodata, gdal_metadata=gdal_metadata + ) + # TODO: support nodata/gdal_metadata without gbox? + _blocks = itertools.chain(iter(blocksize), itertools.repeat(blocksize[-1])) tw = TiffWriter(buf, bigtiff=True, shaped=False) + metas: List[CogMeta] = [] for tsz, idx in zip(_blocks, range(nlevels + 1)): + pix_shape = _sh(im_shape) + tile = _norm_blocksize(tsz) + meta = CogMeta( + ax, im_shape, shape_(tile), pix_shape, dtype, _compression, predictor + ) + if idx == 0: kw = {**opts_common, "extratags": extratags} + meta.gbox = gbox else: kw = {**opts_common, "subfiletype": FILETYPE.REDUCEDIMAGE} tw.write( itertools.repeat(b""), - shape=_sh(im_shape), - tile=_norm_blocksize(tsz), + shape=pix_shape, + tile=tile, **kw, ) + metas.append(meta) im_shape = im_shape.shrink2() + meta = metas[0] + meta.overviews = metas[1:] + tw.close() - return buf.getbuffer() + return meta, buf.getbuffer() diff --git a/tests/test_cog.py b/tests/test_cog.py index 5a1d88e7..5471e48b 100644 --- a/tests/test_cog.py +++ b/tests/test_cog.py @@ -150,7 +150,7 @@ def test_empty_cog(shape, blocksize, expect_ax, dtype, compression, expect_predi gbox = gbox.zoom_to(shape[:2]) assert gbox.shape == shape[:2] - mm = make_empty_cog( + meta, mm = make_empty_cog( shape, dtype, gbox=gbox, @@ -158,6 +158,12 @@ def test_empty_cog(shape, blocksize, expect_ax, dtype, compression, expect_predi compression=compression, ) assert isinstance(mm, memoryview) + assert meta.axis == expect_ax + assert meta.dtype == dtype + assert meta.shape[0] >= gbox.shape[0] + assert meta.shape[1] >= gbox.shape[1] + assert meta.gbox is not None + assert meta.shape == meta.gbox.shape f = tifffile.TiffFile(BytesIO(mm)) assert f.tiff.is_bigtiff