Line data Source code
1 : /******************************************************************************
2 : *
3 : * Project: GDAL
4 : * Purpose: "select" step of "raster pipeline"
5 : * Author: Even Rouault <even dot rouault at spatialys.com>
6 : *
7 : ******************************************************************************
8 : * Copyright (c) 2025, Even Rouault <even dot rouault at spatialys.com>
9 : *
10 : * SPDX-License-Identifier: MIT
11 : ****************************************************************************/
12 :
13 : #include "gdalalg_raster_select.h"
14 :
15 : #include "gdal_priv.h"
16 : #include "gdal_utils.h"
17 :
18 : #include <map>
19 : #include <set>
20 :
21 : //! @cond Doxygen_Suppress
22 :
23 : #ifndef _
24 : #define _(x) (x)
25 : #endif
26 :
27 : /************************************************************************/
28 : /* GDALRasterSelectAlgorithm::GDALRasterSelectAlgorithm() */
29 : /************************************************************************/
30 :
31 95 : GDALRasterSelectAlgorithm::GDALRasterSelectAlgorithm(bool standaloneStep)
32 : : GDALRasterPipelineStepAlgorithm(NAME, DESCRIPTION, HELP_URL,
33 95 : standaloneStep)
34 : {
35 : {
36 : auto &arg = AddArg("band", 'b',
37 : _("Band(s) (1-based index, 'mask', 'mask:<band>' or "
38 : "color interpretation such as 'red')"),
39 190 : &m_bands)
40 95 : .SetPositional()
41 95 : .SetRequired()
42 95 : .SetMinCount(1);
43 : arg.SetAutoCompleteFunction(
44 8 : [this](const std::string &)
45 : {
46 2 : std::vector<std::string> ret;
47 2 : std::unique_ptr<GDALDataset> poSrcDSTmp;
48 2 : GDALDataset *poSrcDS = m_inputDataset.empty()
49 2 : ? nullptr
50 1 : : m_inputDataset[0].GetDatasetRef();
51 2 : if (!poSrcDS && !m_inputDataset.empty())
52 : {
53 2 : CPLErrorStateBackuper oBackuper(CPLQuietErrorHandler);
54 1 : poSrcDSTmp.reset(GDALDataset::Open(
55 1 : m_inputDataset[0].GetName().c_str(), GDAL_OF_RASTER));
56 1 : poSrcDS = poSrcDSTmp.get();
57 : }
58 2 : if (poSrcDS)
59 : {
60 2 : std::set<GDALColorInterp> oSetColorInterp;
61 2 : for (int i = 1; i <= poSrcDS->GetRasterCount(); ++i)
62 : {
63 1 : ret.push_back(std::to_string(i));
64 1 : oSetColorInterp.insert(poSrcDS->GetRasterBand(i)
65 1 : ->GetColorInterpretation());
66 : }
67 1 : ret.push_back("mask");
68 2 : for (const auto eColorInterp : oSetColorInterp)
69 : {
70 2 : ret.push_back(CPLString(GDALGetColorInterpretationName(
71 : eColorInterp))
72 1 : .tolower());
73 : }
74 : }
75 4 : return ret;
76 95 : });
77 : arg.AddValidationAction(
78 48 : [&arg]()
79 : {
80 24 : int nColorInterpretations = 0;
81 : const auto paeColorInterp =
82 24 : GDALGetColorInterpretationList(&nColorInterpretations);
83 48 : std::set<std::string> oSetValidColorInterp;
84 840 : for (int i = 0; i < nColorInterpretations; ++i)
85 : oSetValidColorInterp.insert(
86 1632 : CPLString(
87 816 : GDALGetColorInterpretationName(paeColorInterp[i]))
88 816 : .tolower());
89 :
90 24 : const auto &val = arg.Get<std::vector<std::string>>();
91 60 : for (const auto &v : val)
92 : {
93 153 : if (!STARTS_WITH(v.c_str(), "mask") &&
94 36 : !(CPLGetValueType(v.c_str()) == CPL_VALUE_INTEGER &&
95 75 : atoi(v.c_str()) >= 1) &&
96 13 : !cpl::contains(oSetValidColorInterp,
97 52 : CPLString(v).tolower()))
98 : {
99 3 : CPLError(CE_Failure, CPLE_AppDefined,
100 : "Invalid band specification.");
101 3 : return false;
102 : }
103 : }
104 21 : return true;
105 95 : });
106 : }
107 :
108 95 : AddArg("exclude", 0, _("Exclude specified bands"), &m_exclude);
109 :
110 : {
111 : auto &arg = AddArg(
112 : "mask", 0,
113 : _("Mask band (1-based index, 'mask', 'mask:<band>' or 'none')"),
114 95 : &m_mask);
115 : arg.AddValidationAction(
116 3 : [&arg]()
117 : {
118 3 : const auto &v = arg.Get<std::string>();
119 3 : if (!STARTS_WITH(v.c_str(), "mask") &&
120 6 : !EQUAL(v.c_str(), "none") &&
121 3 : !(CPLGetValueType(v.c_str()) == CPL_VALUE_INTEGER &&
122 2 : atoi(v.c_str()) >= 1))
123 : {
124 1 : CPLError(CE_Failure, CPLE_AppDefined,
125 : "Invalid mask band specification.");
126 1 : return false;
127 : }
128 2 : return true;
129 95 : });
130 : }
131 95 : }
132 :
133 : /************************************************************************/
134 : /* GDALRasterSelectAlgorithm::RunStep() */
135 : /************************************************************************/
136 :
137 9 : bool GDALRasterSelectAlgorithm::RunStep(GDALPipelineStepRunContext &)
138 : {
139 9 : const auto poSrcDS = m_inputDataset[0].GetDatasetRef();
140 9 : CPLAssert(poSrcDS);
141 9 : CPLAssert(m_outputDataset.GetName().empty());
142 9 : CPLAssert(!m_outputDataset.GetDatasetRef());
143 :
144 18 : std::map<GDALColorInterp, std::vector<int>> oMapColorInterpToBands;
145 32 : for (int i = 1; i <= poSrcDS->GetRasterCount(); ++i)
146 : {
147 23 : oMapColorInterpToBands[poSrcDS->GetRasterBand(i)
148 46 : ->GetColorInterpretation()]
149 23 : .push_back(i);
150 : }
151 :
152 18 : CPLStringList aosOptions;
153 9 : aosOptions.AddString("-of");
154 9 : aosOptions.AddString("VRT");
155 9 : if (m_exclude)
156 : {
157 3 : if (m_bands.size() >= static_cast<size_t>(poSrcDS->GetRasterCount()))
158 : {
159 1 : ReportError(CE_Failure, CPLE_AppDefined,
160 : "Cannot exclude all input bands");
161 1 : return false;
162 : }
163 :
164 4 : std::set<int> excludedBandsFromColor;
165 4 : for (const std::string &v : m_bands)
166 : {
167 : const auto eColorInterp =
168 2 : GDALGetColorInterpretationByName(v.c_str());
169 2 : if (v == "undefined" || eColorInterp != GCI_Undefined)
170 : {
171 1 : const auto iter = oMapColorInterpToBands.find(eColorInterp);
172 1 : if (iter != oMapColorInterpToBands.end())
173 : {
174 2 : for (const int iBand : iter->second)
175 : {
176 1 : excludedBandsFromColor.insert(iBand);
177 : }
178 : }
179 : // We don't emit a warning if there are no bands matching
180 : // the color interpretation, because a potential use case
181 : // could be to run on a set of input files that might have or
182 : // might not have an alpha band, and remove it.
183 : }
184 : }
185 :
186 8 : for (int i = 1; i <= poSrcDS->GetRasterCount(); ++i)
187 : {
188 12 : const std::string iStr = std::to_string(i);
189 6 : if (std::find(m_bands.begin(), m_bands.end(), iStr) ==
190 17 : m_bands.end() &&
191 5 : !cpl::contains(excludedBandsFromColor, i))
192 : {
193 4 : aosOptions.AddString("-b");
194 4 : aosOptions.AddString(iStr);
195 : }
196 : }
197 : }
198 : else
199 : {
200 14 : for (const std::string &v : m_bands)
201 : {
202 : const auto eColorInterp =
203 10 : GDALGetColorInterpretationByName(v.c_str());
204 10 : if (v == "undefined" || eColorInterp != GCI_Undefined)
205 : {
206 4 : const auto iter = oMapColorInterpToBands.find(eColorInterp);
207 4 : if (iter == oMapColorInterpToBands.end())
208 : {
209 2 : ReportError(CE_Failure, CPLE_AppDefined,
210 : "No band has color interpretation %s",
211 : v.c_str());
212 2 : return false;
213 : }
214 4 : for (const int iBand : iter->second)
215 : {
216 2 : aosOptions.AddString("-b");
217 2 : aosOptions.AddString(std::to_string(iBand));
218 : }
219 : }
220 : else
221 : {
222 6 : aosOptions.AddString("-b");
223 6 : aosOptions.AddString(CPLString(v).replaceAll(':', ',').c_str());
224 : }
225 : }
226 : }
227 6 : if (!m_mask.empty())
228 : {
229 1 : aosOptions.AddString("-mask");
230 1 : aosOptions.AddString(CPLString(m_mask).replaceAll(':', ',').c_str());
231 : }
232 :
233 : GDALTranslateOptions *psOptions =
234 6 : GDALTranslateOptionsNew(aosOptions.List(), nullptr);
235 :
236 : auto poOutDS = std::unique_ptr<GDALDataset>(GDALDataset::FromHandle(
237 6 : GDALTranslate("", GDALDataset::ToHandle(poSrcDS), psOptions, nullptr)));
238 6 : GDALTranslateOptionsFree(psOptions);
239 6 : const bool bRet = poOutDS != nullptr;
240 6 : if (poOutDS)
241 : {
242 6 : m_outputDataset.Set(std::move(poOutDS));
243 : }
244 :
245 6 : return bRet;
246 : }
247 :
248 : GDALRasterSelectAlgorithmStandalone::~GDALRasterSelectAlgorithmStandalone() =
249 : default;
250 :
251 : //! @endcond
|
