gdal/apps/nearblack_lib_floodfill.cpp

623 lines
20 KiB
C++

/******************************************************************************
*
* Project: GDAL Utilities
* Purpose: Convert nearly black or nearly white border to exact black/white
* using the flood fill algorithm.
* Author: Even Rouault <even dot rouault at spatialys.com>
*
* ****************************************************************************
* Copyright (c) 2023, Even Rouault <even dot rouault at spatialys.com>
*
* SPDX-License-Identifier: MIT
****************************************************************************/
#include "gdal_priv.h"
#include "nearblack_lib.h"
#include <algorithm>
#include <memory>
#include <queue>
/************************************************************************/
/* GDALNearblackFloodFillAlg */
/************************************************************************/
// Implements the "final, combined-scan-and-fill span filler was then published
// in 1990" algorithm of https://en.wikipedia.org/wiki/Flood_fill#Span_filling
struct GDALNearblackFloodFillAlg
{
// Input arguments of the algorithm
const GDALNearblackOptions *m_psOptions = nullptr;
GDALDataset *m_poSrcDataset = nullptr;
GDALDataset *m_poDstDS = nullptr;
GDALRasterBand *m_poMaskBand = nullptr;
int m_nSrcBands = 0;
int m_nDstBands = 0;
bool m_bSetMask = false;
Colors m_oColors{};
GByte m_nReplacevalue = 0;
// As we (generally) do not modify the value of pixels that are "black"
// we need to keep track of the pixels we visited
// Cf https://en.wikipedia.org/wiki/Flood_fill#Disadvantages_2
// and https://en.wikipedia.org/wiki/Flood_fill#Adding_pattern_filling_support
// for the requirement to add that extra sentinel
std::unique_ptr<GDALDataset> m_poVisitedDS = nullptr;
// Active line for the m_abyLine, m_abyLineMustSet, m_abyMask buffers
int m_nLoadedLine = -1;
// Whether Set(..., m_nLoadedLine) has been called
bool m_bLineModified = true;
// Content of m_poSrcDataset/m_poDstDS for m_nLoadedLine
// Contains m_nDstBands * nXSize values in the order (R,G,B),(R,G,B),...
std::vector<GByte> m_abyLine{};
static constexpr GByte MUST_FILL_UNINIT = 0; // must be 0
static constexpr GByte MUST_FILL_FALSE = 1;
static constexpr GByte MUST_FILL_TRUE = 2;
// Content of m_poVisitedDS for m_nLoadedLine
std::vector<GByte> m_abyLineMustSet{};
// Only use if m_bSetMask
std::vector<GByte> m_abyMask{};
// Used for progress bar. Incremented the first time a line ifs loaded
int m_nCountLoadedOnce = 0;
// m_abLineLoadedOnce[line] is set to true after the first time the line
// of m_poSrcDataset is loaded by LoadLine(line)
std::vector<bool> m_abLineLoadedOnce{};
// m_abLineSavedOnce[line] is set to true after the first time the line
// of m_poDstDS is written by LoadLine()
std::vector<bool> m_abLineSavedOnce{};
#ifdef DEBUG
size_t m_nMaxQueueSize = 0;
#endif
// Entry point
bool Process();
private:
bool Fill(int iX, int iY);
bool LoadLine(int iY);
bool MustSet(int iX, int iY);
void Set(int iX, int iY);
};
/************************************************************************/
/* GDALNearblackFloodFillAlg::MustSet() */
/* */
/* Called Inside() in https://en.wikipedia.org/wiki/Flood_fill */
/************************************************************************/
// Returns true if the pixel (iX, iY) is "black" (or more generally transparent
// according to m_oColors)
bool GDALNearblackFloodFillAlg::MustSet(int iX, int iY)
{
CPLAssert(iX >= 0);
CPLAssert(iX < m_poSrcDataset->GetRasterXSize());
CPLAssert(iY >= 0);
CPLAssert(iY < m_poSrcDataset->GetRasterYSize());
CPLAssert(iY == m_nLoadedLine);
CPL_IGNORE_RET_VAL(iY);
if (m_abyLineMustSet[iX] != MUST_FILL_UNINIT)
{
return m_abyLineMustSet[iX] == MUST_FILL_TRUE;
}
/***** loop over the colors *****/
for (int iColor = 0; iColor < static_cast<int>(m_oColors.size()); iColor++)
{
const Color &oColor = m_oColors[iColor];
/***** loop over the bands *****/
bool bIsNonBlack = false;
for (int iBand = 0; iBand < m_nSrcBands; iBand++)
{
const int nPix = m_abyLine[iX * m_nDstBands + iBand];
if (oColor[iBand] - nPix > m_psOptions->nNearDist ||
nPix > m_psOptions->nNearDist + oColor[iBand])
{
bIsNonBlack = true;
break;
}
}
if (!bIsNonBlack)
{
m_abyLineMustSet[iX] = MUST_FILL_TRUE;
return true;
}
}
m_abyLineMustSet[iX] = MUST_FILL_FALSE;
return false;
}
/************************************************************************/
/* GDALNearblackFloodFillAlg::LoadLine() */
/************************************************************************/
// Load the new line iY, and saves if needed buffer of the previous loaded
// line (m_nLoadedLine).
// Returns true if no error
bool GDALNearblackFloodFillAlg::LoadLine(int iY)
{
if (iY != m_nLoadedLine)
{
#ifdef DEBUG
// CPLDebug("GDAL", "GDALNearblackFloodFillAlg::LoadLine(%d)", iY);
#endif
const int nXSize = m_poSrcDataset->GetRasterXSize();
if (m_nLoadedLine >= 0)
{
if (m_bLineModified || (m_poDstDS != m_poSrcDataset &&
!m_abLineSavedOnce[m_nLoadedLine]))
{
if (m_poDstDS->RasterIO(
GF_Write, 0, m_nLoadedLine, nXSize, 1, m_abyLine.data(),
nXSize, 1, GDT_Byte, m_nDstBands, nullptr, m_nDstBands,
static_cast<GSpacing>(nXSize) * m_nDstBands, 1,
nullptr) != CE_None)
{
return false;
}
}
if (m_bSetMask &&
(m_bLineModified || !m_abLineSavedOnce[m_nLoadedLine]))
{
if (m_poMaskBand->RasterIO(GF_Write, 0, m_nLoadedLine, nXSize,
1, m_abyMask.data(), nXSize, 1,
GDT_Byte, 0, 0, nullptr) != CE_None)
{
return false;
}
}
m_abLineSavedOnce[m_nLoadedLine] = true;
}
if (iY >= 0)
{
if (m_poDstDS != m_poSrcDataset && m_abLineSavedOnce[iY])
{
// If the output dataset is different from the source one,
// load from the output dataset if we have already written the
// line of interest
if (m_poDstDS->RasterIO(
GF_Read, 0, iY, nXSize, 1, m_abyLine.data(), nXSize, 1,
GDT_Byte, m_nDstBands, nullptr, m_nDstBands,
static_cast<GSpacing>(nXSize) * m_nDstBands, 1,
nullptr) != CE_None)
{
return false;
}
}
else
{
// Otherwise load from the source data
if (m_poSrcDataset->RasterIO(
GF_Read, 0, iY, nXSize, 1, m_abyLine.data(), nXSize, 1,
GDT_Byte,
// m_nSrcBands intended
m_nSrcBands,
// m_nDstBands intended
nullptr, m_nDstBands,
static_cast<GSpacing>(nXSize) * m_nDstBands, 1,
nullptr) != CE_None)
{
return false;
}
// Initialize the alpha component to 255 if it is the first time
// we load that line.
if (m_psOptions->bSetAlpha && !m_abLineLoadedOnce[iY])
{
for (int iCol = 0; iCol < nXSize; iCol++)
{
m_abyLine[iCol * m_nDstBands + m_nDstBands - 1] = 255;
}
}
}
if (m_bSetMask)
{
if (!m_abLineLoadedOnce[iY])
{
for (int iCol = 0; iCol < nXSize; iCol++)
{
m_abyMask[iCol] = 255;
}
}
else
{
if (m_poMaskBand->RasterIO(
GF_Read, 0, iY, nXSize, 1, m_abyMask.data(), nXSize,
1, GDT_Byte, 0, 0, nullptr) != CE_None)
{
return false;
}
}
}
if (!m_abLineLoadedOnce[iY])
{
m_nCountLoadedOnce++;
// Very rough progression report based on the first time
// we load a line...
// We arbitrarily consider that it's 90% of the processing time
const int nYSize = m_poSrcDataset->GetRasterYSize();
if (!(m_psOptions->pfnProgress(
0.9 *
(m_nCountLoadedOnce / static_cast<double>(nYSize)),
nullptr, m_psOptions->pProgressData)))
{
return false;
}
m_abLineLoadedOnce[iY] = true;
}
}
if (m_nLoadedLine >= 0)
{
if (m_poVisitedDS->GetRasterBand(1)->RasterIO(
GF_Write, 0, m_nLoadedLine, nXSize, 1,
m_abyLineMustSet.data(), nXSize, 1, GDT_Byte, 0, 0,
nullptr) != CE_None)
{
return false;
}
}
if (iY >= 0)
{
if (m_poVisitedDS->GetRasterBand(1)->RasterIO(
GF_Read, 0, iY, nXSize, 1, m_abyLineMustSet.data(), nXSize,
1, GDT_Byte, 0, 0, nullptr) != CE_None)
{
return false;
}
}
m_bLineModified = false;
m_nLoadedLine = iY;
}
return true;
}
/************************************************************************/
/* GDALNearblackFloodFillAlg::Set() */
/************************************************************************/
// Mark the pixel as transparent
void GDALNearblackFloodFillAlg::Set(int iX, int iY)
{
CPLAssert(iY == m_nLoadedLine);
CPL_IGNORE_RET_VAL(iY);
m_bLineModified = true;
m_abyLineMustSet[iX] = MUST_FILL_FALSE;
for (int iBand = 0; iBand < m_nSrcBands; iBand++)
m_abyLine[iX * m_nDstBands + iBand] = m_nReplacevalue;
/***** alpha *****/
if (m_nDstBands > m_nSrcBands)
m_abyLine[iX * m_nDstBands + m_nDstBands - 1] = 0;
if (m_bSetMask)
m_abyMask[iX] = 0;
}
/************************************************************************/
/* GDALNearblackFloodFillAlg::Fill() */
/************************************************************************/
/* Implements the "final, combined-scan-and-fill span filler was then published
* in 1990" algorithm of https://en.wikipedia.org/wiki/Flood_fill#Span_filling
* with the following enhancements:
* - extra bound checking to avoid calling MustSet() outside the raster
* - extra bound checking to avoid pushing spans outside the raster
*
* Returns true if no error.
*/
bool GDALNearblackFloodFillAlg::Fill(int iXInit, int iYInit)
{
const int nXSize = m_poSrcDataset->GetRasterXSize();
const int nYSize = m_poSrcDataset->GetRasterYSize();
struct Span
{
int x1;
int x2;
int y;
int dy;
Span(int x1In, int x2In, int yIn, int dyIn)
: x1(x1In), x2(x2In), y(yIn), dy(dyIn)
{
}
};
if (!LoadLine(iYInit))
return false;
if (!MustSet(iXInit, iYInit))
{
// nothing to do
return true;
}
std::queue<Span> queue;
queue.emplace(Span(iXInit, iXInit, iYInit, 1));
if (iYInit > 0)
{
queue.emplace(Span(iXInit, iXInit, iYInit - 1, -1));
}
while (!queue.empty())
{
#ifdef DEBUG
m_nMaxQueueSize = std::max(m_nMaxQueueSize, queue.size());
#endif
const Span s = queue.front();
queue.pop();
CPLAssert(s.x1 >= 0);
CPLAssert(s.x1 < nXSize);
CPLAssert(s.x2 >= 0);
CPLAssert(s.x2 < nXSize);
CPLAssert(s.x2 >= s.x1);
CPLAssert(s.y >= 0);
CPLAssert(s.y < nYSize);
int iX = s.x1;
const int iY = s.y;
if (!LoadLine(iY))
return false;
if (iX > 0 && MustSet(iX, iY))
{
while (MustSet(iX - 1, iY))
{
Set(iX - 1, iY);
iX--;
if (iX == 0)
break;
}
}
if (iX >= 0 && iX <= s.x1 - 1 && iY - s.dy >= 0 && iY - s.dy < nYSize)
{
queue.emplace(Span(iX, s.x1 - 1, iY - s.dy, -s.dy));
}
int iX1 = s.x1;
const int iX2 = s.x2;
while (iX1 <= iX2)
{
while (MustSet(iX1, iY))
{
Set(iX1, iY);
iX1++;
if (iX1 == nXSize)
break;
}
if (iX <= iX1 - 1 && iY + s.dy >= 0 && iY + s.dy < nYSize)
{
queue.emplace(Span(iX, iX1 - 1, iY + s.dy, s.dy));
}
if (iX1 - 1 > iX2 && iY - s.dy >= 0 && iY - s.dy < nYSize)
{
queue.emplace(Span(iX2 + 1, iX1 - 1, iY - s.dy, -s.dy));
}
iX1++;
while (iX1 < iX2 && !MustSet(iX1, iY))
iX1++;
iX = iX1;
}
}
return true;
}
/************************************************************************/
/* GDALNearblackFloodFillAlg::Process() */
/************************************************************************/
// Entry point.
// Returns true if no error.
bool GDALNearblackFloodFillAlg::Process()
{
const int nXSize = m_poSrcDataset->GetRasterXSize();
const int nYSize = m_poSrcDataset->GetRasterYSize();
/* -------------------------------------------------------------------- */
/* Allocate working buffers. */
/* -------------------------------------------------------------------- */
try
{
m_abyLine.resize(static_cast<size_t>(nXSize) * m_nDstBands);
m_abyLineMustSet.resize(nXSize);
if (m_bSetMask)
m_abyMask.resize(nXSize);
if (m_psOptions->nMaxNonBlack > 0)
{
m_abLineLoadedOnce.resize(nYSize, true);
m_abLineSavedOnce.resize(nYSize, true);
}
else
{
m_abLineLoadedOnce.resize(nYSize);
m_abLineSavedOnce.resize(nYSize);
}
}
catch (const std::exception &e)
{
CPLError(CE_Failure, CPLE_OutOfMemory,
"Cannot allocate working buffers: %s", e.what());
return false;
}
/* -------------------------------------------------------------------- */
/* Create a temporary dataset to save visited state */
/* -------------------------------------------------------------------- */
// For debugging / testing purposes only
const char *pszTmpDriver =
CPLGetConfigOption("GDAL_TEMP_DRIVER_NAME", nullptr);
if (!pszTmpDriver)
{
pszTmpDriver =
(nXSize < 100 * 1024 * 1024 / nYSize ||
(m_poDstDS->GetDriver() &&
strcmp(m_poDstDS->GetDriver()->GetDescription(), "MEM") == 0))
? "MEM"
: "GTiff";
}
GDALDriverH hDriver = GDALGetDriverByName(pszTmpDriver);
if (!hDriver)
{
CPLError(CE_Failure, CPLE_AppDefined,
"Cannot find driver %s for temporary file", pszTmpDriver);
return false;
}
std::string osVisitedDataset = m_poDstDS->GetDescription();
VSIStatBuf sStat;
if (strcmp(pszTmpDriver, "MEM") == 0 ||
STARTS_WITH(osVisitedDataset.c_str(), "/vsimem/") ||
// Regular VSIStat() (not VSIStatL()) intended to check this is
// a real file
VSIStat(osVisitedDataset.c_str(), &sStat) == 0)
{
osVisitedDataset += ".visited";
}
else
{
osVisitedDataset = CPLGenerateTempFilename(osVisitedDataset.c_str());
}
CPLStringList aosOptions;
if (strcmp(pszTmpDriver, "GTiff") == 0)
{
aosOptions.SetNameValue("SPARSE_OK", "YES");
aosOptions.SetNameValue("COMPRESS", "LZW");
osVisitedDataset += ".tif";
}
m_poVisitedDS.reset(GDALDataset::FromHandle(
GDALCreate(hDriver, osVisitedDataset.c_str(), nXSize, nYSize, 1,
GDT_Byte, aosOptions.List())));
if (!m_poVisitedDS)
return false;
if (strcmp(pszTmpDriver, "MEM") != 0)
{
VSIUnlink(osVisitedDataset.c_str());
}
m_poVisitedDS->MarkSuppressOnClose();
/* -------------------------------------------------------------------- */
/* Iterate over the border of the raster */
/* -------------------------------------------------------------------- */
// Fill from top line
for (int iX = 0; iX < nXSize; iX++)
{
if (!Fill(iX, 0))
return false;
}
// Fill from left and right side
for (int iY = 1; iY < nYSize - 1; iY++)
{
if (!Fill(0, iY))
return false;
if (!Fill(nXSize - 1, iY))
return false;
}
// Fill from bottom line
for (int iX = 0; iX < nXSize; iX++)
{
if (!Fill(iX, nYSize - 1))
return false;
}
if (!(m_psOptions->pfnProgress(1.0, nullptr, m_psOptions->pProgressData)))
{
return false;
}
#ifdef DEBUG
CPLDebug("GDAL", "flood fill max queue size = %u",
unsigned(m_nMaxQueueSize));
#endif
// Force update of last visited line
return LoadLine(-1);
}
/************************************************************************/
/* GDALNearblackFloodFill() */
/************************************************************************/
// Entry point.
// Returns true if no error.
bool GDALNearblackFloodFill(const GDALNearblackOptions *psOptions,
GDALDatasetH hSrcDataset, GDALDatasetH hDstDS,
GDALRasterBandH hMaskBand, int nSrcBands,
int nDstBands, bool bSetMask, const Colors &oColors)
{
GDALNearblackFloodFillAlg alg;
alg.m_psOptions = psOptions;
alg.m_poSrcDataset = GDALDataset::FromHandle(hSrcDataset);
alg.m_poDstDS = GDALDataset::FromHandle(hDstDS);
alg.m_poMaskBand = GDALRasterBand::FromHandle(hMaskBand);
alg.m_nSrcBands = nSrcBands;
alg.m_nDstBands = nDstBands;
alg.m_bSetMask = bSetMask;
alg.m_oColors = oColors;
alg.m_nReplacevalue = psOptions->bNearWhite ? 255 : 0;
if (psOptions->nMaxNonBlack > 0)
{
// First pass: use the TwoPasses algorithm to deal with nMaxNonBlack
GDALNearblackOptions sOptionsTmp(*psOptions);
sOptionsTmp.pProgressData = GDALCreateScaledProgress(
0, 0.5, psOptions->pfnProgress, psOptions->pProgressData);
sOptionsTmp.pfnProgress = GDALScaledProgress;
bool bRet = GDALNearblackTwoPassesAlgorithm(
&sOptionsTmp, hSrcDataset, hDstDS, hMaskBand, nSrcBands, nDstBands,
bSetMask, oColors);
GDALDestroyScaledProgress(sOptionsTmp.pProgressData);
if (!bRet)
return false;
// Second pass: use flood fill
sOptionsTmp.pProgressData = GDALCreateScaledProgress(
0.5, 1, psOptions->pfnProgress, psOptions->pProgressData);
sOptionsTmp.pfnProgress = GDALScaledProgress;
alg.m_psOptions = &sOptionsTmp;
bRet = alg.Process();
GDALDestroyScaledProgress(sOptionsTmp.pProgressData);
return bRet;
}
else
{
return alg.Process();
}
}