Skip to content

Commit

Permalink
VRT Processed Dataset: Add BandScaleOffset algorithm
Browse files Browse the repository at this point in the history
  • Loading branch information
@dbaston
dbaston committed Jan 3, 2025
1 parent 2c3a671 commit 151ad77
Show file tree
Hide file tree
Showing 4 changed files with 290 additions and 2 deletions.
82 changes: 82 additions & 0 deletions autotest/gdrivers/vrtprocesseddataset.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -589,6 +589,88 @@ def test_vrtprocesseddataset_lut_errors(tmp_vsimem):
)


###############################################################################
# Test nominal case of BandScaleOffset algorithm


@pytest.mark.parametrize(
"dtype",
(
gdal.GDT_Byte,
gdal.GDT_Int8,
gdal.GDT_UInt16,
gdal.GDT_Int16,
gdal.GDT_UInt32,
gdal.GDT_Int32,
gdal.GDT_UInt64,
gdal.GDT_Int64,
gdal.GDT_Float32,
gdal.GDT_Float64,
),
ids=gdal.GetDataTypeName,
)
def test_vrtprocesseddataset_bandscaleoffset_nominal(tmp_vsimem, dtype):

src_filename = tmp_vsimem / "src.tif"

nx = 2
ny = 3
nz = 2

if type == gdal.GDT_Float32:
nodata = float("nan")
else:
nodata = 7

data = np.arange(nx * ny * nz).reshape(nz, ny, nx)
data[1, 2, 1] = nodata

offsets = [i + 2 for i in range(nz)]
scales = [(i + 1) / 4 for i in range(nz)]

with gdal.GetDriverByName("GTiff").Create(
src_filename, nx, ny, nz, eType=dtype
) as src_ds:
src_ds.WriteArray(data)
for i in range(src_ds.RasterCount):
bnd = src_ds.GetRasterBand(i + 1)
bnd.SetOffset(offsets[i])
bnd.SetScale(scales[i])
bnd.SetNoDataValue(nodata)

ds = gdal.Open(
f"""
<VRTDataset subclass='VRTProcessedDataset'>
<Input>
<SourceFilename>{src_filename}</SourceFilename>
</Input>
<ProcessingSteps>
<Step>
<Algorithm>BandScaleOffset</Algorithm>
</Step>
</ProcessingSteps>
<OutputBands count="FROM_SOURCE" dataType="Float64"/>
</VRTDataset>"""
)

for i in range(ds.RasterCount):
bnd = ds.GetRasterBand(i + 1)
assert bnd.GetScale() in (None, 1)
assert bnd.GetOffset() in (None, 0)

result = np.ma.stack(
[ds.GetRasterBand(i + 1).ReadAsMaskedArray() for i in range(ds.RasterCount)]
)

expected = np.ma.masked_array(
np.stack([data[i, :, :] * scales[i] + offsets[i] for i in range(nz)]),
data == nodata,
)

np.testing.assert_array_equal(result.mask, expected.mask)
np.testing.assert_array_equal(result, expected)


###############################################################################
# Test nominal case of LocalScaleOffset algorithm

Expand Down
3 changes: 3 additions & 0 deletions doc/source/drivers/raster/vrt_processed_dataset.rst
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,9 @@ to apply chained processing steps that may apply to several bands at the same ti
The following built-in algorithms are introduced, and may typically be applied
in the following order:

- BandScaleOffset (since GDAL 3.11): apply per-band scale and offset values from
:cpp:func:`GDALRasterBand::GetScale` and :cpp:func:`GDALRasterBand::GetOffset`.

- LocalScaleOffset: apply per-pixel gain and offset coming (typically subsampled)
from auxiliary datasets. Can be used for dehazing processing.

Expand Down
9 changes: 7 additions & 2 deletions frmts/vrt/vrtprocesseddataset.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -574,6 +574,9 @@ CPLErr VRTProcessedDataset::Init(const CPLXMLNode *psTree,
auto poBand =
new VRTProcessedRasterBand(this, i + 1, eOutputBandType);
poBand->CopyCommonInfoFrom(poSrcBand);
poBand->SetNoDataValue(adfNoData[i]);
poBand->SetScale(1);
poBand->SetOffset(0);
SetBand(i + 1, poBand);
}
}
Expand Down Expand Up @@ -713,7 +716,8 @@ bool VRTProcessedDataset::ParseStep(const CPLXMLNode *psStep, bool bIsFinalStep,
for (int i = 1; i <= m_poSrcDS->GetRasterCount(); ++i)
{
int bHasVal = false;
const double dfVal = GetRasterBand(i)->GetOffset(&bHasVal);
const double dfVal =
m_poSrcDS->GetRasterBand(i)->GetOffset(&bHasVal);
oStep.aosArguments.AddNameValue(
CPLSPrintf("offset_%d", i),
CPLSPrintf("%.17g", bHasVal ? dfVal : 0.0));
Expand All @@ -726,7 +730,8 @@ bool VRTProcessedDataset::ParseStep(const CPLXMLNode *psStep, bool bIsFinalStep,
for (int i = 1; i <= m_poSrcDS->GetRasterCount(); ++i)
{
int bHasVal = false;
const double dfVal = GetRasterBand(i)->GetScale(&bHasVal);
const double dfVal =
m_poSrcDS->GetRasterBand(i)->GetScale(&bHasVal);
oStep.aosArguments.AddNameValue(
CPLSPrintf("scale_%d", i),
CPLSPrintf("%.17g", bHasVal ? dfVal : 1.0));
Expand Down
198 changes: 198 additions & 0 deletions frmts/vrt/vrtprocesseddatasetfunctions.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@

#include "cpl_minixml.h"
#include "cpl_string.h"
#include "gdal_typetraits.h"
#include "vrtdataset.h"

#include <algorithm>
Expand Down Expand Up @@ -604,6 +605,193 @@ LUTProcess(const char * /*pszFuncName*/, void * /*pUserData*/,
return CE_None;
}

/************************************************************************/
/* BandScaleOffsetData */
/************************************************************************/

constexpr std::array<GDALDataType, 10> BandScaleOffsetSupportedTypes{
GDT_Byte, GDT_Int8, GDT_UInt16, GDT_Int16, GDT_UInt32,
GDT_Int32, GDT_UInt64, GDT_Int64, GDT_Float32, GDT_Float64};

namespace
{
struct BandScaleOffsetData
{
BandScaleOffsetData(int nInBands, const double *padfNoData,
CSLConstList papszFunctionArgs)
: m_adfInNoData(padfNoData, padfNoData + nInBands),
m_adfOffsets(nInBands), m_adfScales(nInBands)
{
for (const auto &[pszKey, pszValue] :
cpl::IterateNameValue(papszFunctionArgs))
{
if (EQUALN(pszKey, "scale_", 6))
{
int i = std::atoi(pszKey + 6);
m_adfScales[i - 1] = CPLAtof(pszValue);
}
else if (EQUALN(pszKey, "offset_", 7))
{
int i = std::atoi(pszKey + 7);
m_adfOffsets[i - 1] = CPLAtof(pszValue);
}
}
}

template <typename T>
void applyScaleOffset(size_t nElts, const void *pInBuffer,
void *pOutBuffer) const
{
constexpr double nan = std::numeric_limits<double>::quiet_NaN();

const T *CPL_RESTRICT paSrc = static_cast<const T *>(pInBuffer);
double *CPL_RESTRICT padfDst = static_cast<double *>(pOutBuffer);

for (std::size_t iElt = 0; iElt < nElts; ++iElt)
{
for (size_t iBand = 0; iBand < m_adfOffsets.size(); iBand++)
{

double dfSrc = static_cast<double>(*paSrc++);
*padfDst++ =
(dfSrc == m_adfInNoData[iBand])
? nan
: (dfSrc * m_adfScales[iBand] + m_adfOffsets[iBand]);
}
}
}

std::vector<double> m_adfInNoData;
std::vector<double> m_adfOffsets;
std::vector<double> m_adfScales;
};
} // namespace

static CPLErr
BandScaleOffsetInit(const char * /*pszFuncName*/, void * /*pUserData*/,
CSLConstList papszFunctionArgs, int nInBands,
GDALDataType /* eInDT */, double *padfInNoData,
int *pnOutBands, GDALDataType *peOutDT,
double **ppadfOutNoData, const char * /* pszVRTPath */,
VRTPDWorkingDataPtr *ppWorkingData)
{
const bool bIsFinalStep = *pnOutBands != 0;
*peOutDT = GDT_Float64;
*ppWorkingData = nullptr;

if (!bIsFinalStep)
{
*pnOutBands = nInBands;
}

auto data = std::make_unique<BandScaleOffsetData>(nInBands, padfInNoData,
papszFunctionArgs);

// Input NoData value is meaningless after scaling, so we set the output NoData value to NaN.
for (int i = 0; i < *pnOutBands; i++)
{
(*ppadfOutNoData)[i] = std::numeric_limits<double>::quiet_NaN();
}

*ppWorkingData = data.release();

return CE_None;
}

static void BandScaleOffsetFree(const char * /*pszFuncName*/,
void * /*pUserData*/,
VRTPDWorkingDataPtr pWorkingData)
{
BandScaleOffsetData *data =
static_cast<BandScaleOffsetData *>(pWorkingData);
delete data;
}

static CPLErr BandScaleOffsetProcess(
const char * /*pszFuncName*/, void * /*pUserData*/,
VRTPDWorkingDataPtr pWorkingData, CSLConstList /* papszFunctionArgs*/,
int nBufXSize, int nBufYSize, const void *pInBuffer,
size_t /* nInBufferSize */, GDALDataType eInDT, int nInBands,
const double *CPL_RESTRICT /* padfInNoData */, void *pOutBuffer,
size_t /* nOutBufferSize */, GDALDataType eOutDT, int nOutBands,
const double *CPL_RESTRICT /* padfOutNoData */, double /* dfSrcXOff */,
double /* dfSrcYOff */, double /* dfSrcXSize */, double /* dfSrcYSize */,
const double /* adfSrcGT */[], const char * /* pszVRTPath */,
CSLConstList /*papszExtra*/)
{
CPL_IGNORE_RET_VAL(nInBands);
CPL_IGNORE_RET_VAL(nOutBands);
CPLAssert(nInBands == nOutBands);
CPL_IGNORE_RET_VAL(eOutDT);
CPLAssert(eOutDT = GDT_Float64);

const size_t nElts = static_cast<size_t>(nBufXSize) * nBufYSize;

BandScaleOffsetData *data =
static_cast<BandScaleOffsetData *>(pWorkingData);

switch (eInDT)
{
case GDT_Byte:
data->template applyScaleOffset<
typename gdal::GDALDataTypeTraits<GDT_Byte>::type>(
nElts, pInBuffer, pOutBuffer);
break;
case GDT_Int8:
data->template applyScaleOffset<
typename gdal::GDALDataTypeTraits<GDT_Int8>::type>(
nElts, pInBuffer, pOutBuffer);
break;
case GDT_UInt16:
data->template applyScaleOffset<
typename gdal::GDALDataTypeTraits<GDT_UInt16>::type>(
nElts, pInBuffer, pOutBuffer);
break;
case GDT_Int16:
data->template applyScaleOffset<
typename gdal::GDALDataTypeTraits<GDT_Int16>::type>(
nElts, pInBuffer, pOutBuffer);
break;
case GDT_UInt32:
data->template applyScaleOffset<
typename gdal::GDALDataTypeTraits<GDT_UInt32>::type>(
nElts, pInBuffer, pOutBuffer);
break;
case GDT_Int32:
data->template applyScaleOffset<
typename gdal::GDALDataTypeTraits<GDT_Int32>::type>(
nElts, pInBuffer, pOutBuffer);
break;
case GDT_UInt64:
data->template applyScaleOffset<
typename gdal::GDALDataTypeTraits<GDT_UInt64>::type>(
nElts, pInBuffer, pOutBuffer);
break;
case GDT_Int64:
data->template applyScaleOffset<
typename gdal::GDALDataTypeTraits<GDT_Int64>::type>(
nElts, pInBuffer, pOutBuffer);
break;
case GDT_Float32:
data->template applyScaleOffset<
typename gdal::GDALDataTypeTraits<GDT_Float32>::type>(
nElts, pInBuffer, pOutBuffer);
break;
case GDT_Float64:
data->template applyScaleOffset<
typename gdal::GDALDataTypeTraits<GDT_Float64>::type>(
nElts, pInBuffer, pOutBuffer);
break;
default:
CPLError(CE_Failure, CPLE_AppDefined,
"Data type is not handled by BandScaleOffsetData");
return CE_Failure;
break;
}

return CE_None;
}

/************************************************************************/
/* LocalScaleOffsetData */
/************************************************************************/
Expand Down Expand Up @@ -1514,6 +1702,16 @@ void GDALVRTRegisterDefaultProcessedDatasetFuncs()
GDT_Float64, nullptr, 0, nullptr, 0, LUTInit, LUTFree, LUTProcess,
nullptr);

GDALVRTRegisterProcessedDatasetFunc(
"BandScaleOffset", nullptr,
"<ProcessedDatasetFunctionArgumentsList>"
" <Argument name='scale_{band}' type='builtin' />"
" <Argument name='offset_{band}' type='builtin' />"
"</ProcessedDatasetFunctionArgumentsList>",
GDT_Float64, BandScaleOffsetSupportedTypes.data(),
BandScaleOffsetSupportedTypes.size(), nullptr, 0, BandScaleOffsetInit,
BandScaleOffsetFree, BandScaleOffsetProcess, nullptr);

GDALVRTRegisterProcessedDatasetFunc(
"LocalScaleOffset", nullptr,
"<ProcessedDatasetFunctionArgumentsList>"
Expand Down

0 comments on commit 151ad77

Please sign in to comment.