diff --git a/src/cover_class/geotiff.py b/src/cover_class/geotiff.py new file mode 100644 index 0000000..4390e77 --- /dev/null +++ b/src/cover_class/geotiff.py @@ -0,0 +1,53 @@ +import numpy as np +from typing import List +from pathlib import Path + +BINARY_NO_DATA_VAL = 255 +PROBA_NO_DATA_VAL = -9999 + +numpy_to_gdal = { + np.dtype(np.float64): 7, + np.dtype(np.float32): 6, + np.dtype(np.int32): 5, + np.dtype(np.uint32): 4, + np.dtype(np.int16): 3, + np.dtype(np.uint16): 2, + np.dtype(np.uint8): 1, +} + +def make_geotiff(probabilities: np.ndarray, dataset_shape: tuple, outfp: str, proba: bool, thresholds: List[float]): + try: + from osgeo import gdal + except Exception as e: + raise ImportError("GDAL (osgeo) is required for make_geotiff(). Install it (e.g., conda-forge: gdal).") from e + + nan_mask = np.isnan(probabilities) + + mem_driver, tiff_driver = gdal.GetDriverByName('MEM'), gdal.GetDriverByName('GTiff') + opts = ['COMPRESS=LZW', 'COPY_SRC_OVERVIEWS=YES', 'TILED=YES', 'BLOCKXSIZE=256', 'BLOCKYSIZE=256'] + + if proba: + old_shape = probabilities.shape + probabilities[nan_mask] = PROBA_NO_DATA_VAL + probabilities = probabilities.reshape((dataset_shape[0], dataset_shape[1], 1)) + ds = mem_driver.Create('', probabilities.shape[1], probabilities.shape[0], probabilities.shape[2], numpy_to_gdal[probabilities.dtype]) + ds.GetRasterBand(1).WriteArray(probabilities[:,:,0]) + ds.GetRasterBand(1).SetNoDataValue(PROBA_NO_DATA_VAL) + + _op = Path(outfp) + _op_s = str(_op.with_name(f"{_op.stem}_prob{_op.suffix}")) + _ = tiff_driver.CreateCopy(_op_s, ds, options=opts) + + probabilities = probabilities.reshape(old_shape) + + binary_prob = (probabilities >= np.asarray(thresholds, dtype=probabilities.dtype)).astype(np.uint8) + binary_prob[nan_mask] = np.uint8(BINARY_NO_DATA_VAL) + + # Reshape into input shape + binary_prob = binary_prob.reshape((dataset_shape[0], dataset_shape[1], 1)) + + ds = mem_driver.Create('', binary_prob.shape[1], binary_prob.shape[0], binary_prob.shape[2], numpy_to_gdal[binary_prob.dtype]) + ds.GetRasterBand(1).WriteArray(binary_prob[:,:,0]) + ds.GetRasterBand(1).SetNoDataValue(BINARY_NO_DATA_VAL) + + _ = tiff_driver.CreateCopy(outfp, ds, options=opts) \ No newline at end of file