Skip to content

Commit

Permalink
Carry nodata in CogMeta, use in compressor
Browse files Browse the repository at this point in the history
If gdal nodata is configured then use it to pad out input blocks
when needed.
  • Loading branch information
Kirill888 committed Sep 29, 2023
1 parent 9cedb07 commit 5b4ff97
Show file tree
Hide file tree
Showing 2 changed files with 11 additions and 2 deletions.
4 changes: 3 additions & 1 deletion odc/geo/cog/_shared.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@

from ..geobox import GeoBox
from ..math import align_down_pow2, align_up
from ..types import Shape2d, SomeShape, shape_, wh_
from ..types import MaybeNodata, Shape2d, SomeShape, shape_, wh_

# pylint: disable=too-many-locals,too-many-branches,too-many-arguments,too-many-statements,too-many-instance-attributes

Expand Down Expand Up @@ -63,6 +63,7 @@ class CogMeta:
compressionargs: Dict[str, Any] = field(default_factory=dict, repr=False)
gbox: Optional[GeoBox] = None
overviews: Tuple["CogMeta", ...] = field(default=(), repr=False)
nodata: MaybeNodata = None

def _pix_shape(self, shape: Shape2d) -> Tuple[int, ...]:
if self.axis == "YX":
Expand Down Expand Up @@ -134,6 +135,7 @@ def __dask_tokenize__(self):
self.predictor,
self.gbox,
len(self.overviews),
self.nodata,
)


Expand Down
9 changes: 8 additions & 1 deletion odc/geo/cog/_tifffile.py
Original file line number Diff line number Diff line change
Expand Up @@ -131,6 +131,7 @@ def _sh(shape: Shape2d) -> Tuple[int, ...]:
predictor,
compressionargs=compressionargs,
gbox=gbox,
nodata=nodata,
)

if idx == 0:
Expand Down Expand Up @@ -165,12 +166,13 @@ def _cog_block_compressor(
encoder: Any = None,
predictor: Any = None,
axis: int = 1,
fill_value: Union[float, int] = 0,
**kw,
) -> bytes:
assert block.ndim == len(tile_shape)
if tile_shape != block.shape:
pad = tuple((0, want - have) for want, have in zip(tile_shape, block.shape))
block = np.pad(block, pad, "constant", constant_values=(0,))
block = np.pad(block, pad, "constant", constant_values=(fill_value,))

if predictor is not None:
block = predictor(block, axis=axis)
Expand All @@ -197,12 +199,17 @@ def _mk_tile_compressor(meta: CogMeta) -> Callable[[np.ndarray], bytes]:
if meta.predictor != 1:
predictor = TIFF.PREDICTORS[meta.predictor]

fill_value: Union[float, int] = 0
if meta.nodata is not None:
fill_value = float(meta.nodata) if isinstance(meta.nodata, str) else meta.nodata

return partial(
_cog_block_compressor,
tile_shape=tile_shape,
encoder=encoder,
predictor=predictor,
axis=axis,
fill_value=fill_value,
**meta.compressionargs,
)

Expand Down

0 comments on commit 5b4ff97

Please sign in to comment.