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 26 : static std::optional<std::vector<int>> ParseBandRange(const std::string &v,
28 : int nBands)
29 : {
30 52 : CPLStringList bandSel = cpl::tokenize_string(v, ":", CSLT_ALLOWEMPTYTOKENS);
31 26 : if (bandSel.Count() < 2 || bandSel.Count() > 3)
32 : {
33 1 : CPLError(CE_Failure, CPLE_IllegalArg, "Invalid value for --band: %s",
34 : v.c_str());
35 1 : return std::nullopt;
36 : }
37 25 : int nFirst = 1;
38 25 : const auto osvFirst = cpl::trim(bandSel[0]);
39 25 : if (!osvFirst.empty())
40 : {
41 20 : const auto maybeStart = cpl::strict_parse<int>(osvFirst);
42 20 : if (maybeStart.has_value())
43 : {
44 19 : nFirst = maybeStart.value();
45 19 : if (nFirst < 0)
46 : {
47 4 : nFirst += nBands + 1;
48 : }
49 19 : if (nFirst > nBands || nFirst <= 0)
50 : {
51 3 : CPLError(CE_Failure, CPLE_IllegalArg, "Invalid band: %s",
52 : bandSel[0]);
53 4 : return std::nullopt;
54 : }
55 : }
56 : else
57 : {
58 1 : CPLError(CE_Failure, CPLE_IllegalArg,
59 : "Failed to parse start value of --band range: %s",
60 : bandSel[0]);
61 1 : return std::nullopt;
62 : }
63 : }
64 21 : int nLast = nBands;
65 21 : const auto osvLast = cpl::trim(bandSel[1]);
66 21 : if (!osvLast.empty())
67 : {
68 14 : const auto maybeLast = cpl::strict_parse<int>(osvLast);
69 14 : if (maybeLast.has_value())
70 : {
71 13 : nLast = maybeLast.value();
72 13 : if (nLast < 0)
73 : {
74 2 : nLast += nBands + 1;
75 : }
76 13 : if (nLast > nBands || nLast <= 0)
77 : {
78 0 : CPLError(CE_Failure, CPLE_IllegalArg, "Invalid band: %s",
79 : bandSel[1]);
80 1 : return std::nullopt;
81 : }
82 : }
83 : else
84 : {
85 1 : CPLError(CE_Failure, CPLE_IllegalArg,
86 : "Failed to parse stop value of --band range: %s",
87 : bandSel[1]);
88 1 : return std::nullopt;
89 : }
90 : }
91 20 : int nStep = nFirst < nLast ? 1 : -1;
92 20 : if (bandSel.Count() == 3)
93 : {
94 13 : const auto maybeStep = cpl::strict_parse<int>(bandSel[2]);
95 13 : if (maybeStep.has_value())
96 : {
97 12 : nStep = maybeStep.value();
98 : }
99 : else
100 : {
101 1 : CPLError(CE_Failure, CPLE_IllegalArg,
102 : "Failed to parse step value of --band range: %s",
103 : bandSel[2]);
104 1 : return std::nullopt;
105 : }
106 : }
107 :
108 19 : if (nFirst < nLast && nStep <= 0)
109 : {
110 2 : CPLError(CE_Failure, CPLE_AppDefined, "Step value must be positive");
111 2 : return std::nullopt;
112 : }
113 17 : if (nFirst > nLast && nStep >= 0)
114 : {
115 2 : CPLError(CE_Failure, CPLE_AppDefined, "Step value must be negative");
116 2 : return std::nullopt;
117 : }
118 :
119 30 : std::vector<int> ret;
120 :
121 91 : for (int iBand = nFirst; nStep > 0 ? iBand <= nLast : iBand >= nLast;
122 76 : iBand += nStep)
123 : {
124 76 : if (iBand < 1 || iBand > nBands)
125 : {
126 0 : CPLError(CE_Failure, CPLE_IllegalArg, "Invalid band: %d", iBand);
127 0 : return std::nullopt;
128 : }
129 :
130 76 : ret.push_back(iBand);
131 : }
132 :
133 15 : return ret;
134 : }
135 :
136 : /************************************************************************/
137 : /* GDALRasterSelectAlgorithm::GDALRasterSelectAlgorithm() */
138 : /************************************************************************/
139 :
140 122 : GDALRasterSelectAlgorithm::GDALRasterSelectAlgorithm(bool standaloneStep)
141 : : GDALRasterPipelineStepAlgorithm(NAME, DESCRIPTION, HELP_URL,
142 122 : standaloneStep)
143 : {
144 : {
145 : auto &arg = AddArg("band", 'b',
146 : _("Band(s) (1-based index, 'mask', 'mask:<band>' or "
147 : "color interpretation such as 'red')"),
148 244 : &m_bands)
149 122 : .SetPositional()
150 122 : .SetRequired()
151 122 : .SetMinCount(1);
152 : arg.SetAutoCompleteFunction(
153 8 : [this](const std::string &)
154 : {
155 2 : std::vector<std::string> ret;
156 2 : std::unique_ptr<GDALDataset> poSrcDSTmp;
157 2 : GDALDataset *poSrcDS = m_inputDataset.empty()
158 2 : ? nullptr
159 1 : : m_inputDataset[0].GetDatasetRef();
160 2 : if (!poSrcDS && !m_inputDataset.empty())
161 : {
162 2 : CPLErrorStateBackuper oBackuper(CPLQuietErrorHandler);
163 1 : poSrcDSTmp.reset(GDALDataset::Open(
164 1 : m_inputDataset[0].GetName().c_str(), GDAL_OF_RASTER));
165 1 : poSrcDS = poSrcDSTmp.get();
166 : }
167 2 : if (poSrcDS)
168 : {
169 2 : std::set<GDALColorInterp> oSetColorInterp;
170 2 : for (int i = 1; i <= poSrcDS->GetRasterCount(); ++i)
171 : {
172 1 : ret.push_back(std::to_string(i));
173 1 : oSetColorInterp.insert(poSrcDS->GetRasterBand(i)
174 1 : ->GetColorInterpretation());
175 : }
176 1 : ret.push_back("mask");
177 2 : for (const auto eColorInterp : oSetColorInterp)
178 : {
179 2 : ret.push_back(CPLString(GDALGetColorInterpretationName(
180 : eColorInterp))
181 1 : .tolower());
182 : }
183 : }
184 4 : return ret;
185 122 : });
186 : arg.AddValidationAction(
187 150 : [&arg]()
188 : {
189 75 : int nColorInterpretations = 0;
190 : const auto paeColorInterp =
191 75 : GDALGetColorInterpretationList(&nColorInterpretations);
192 150 : std::set<std::string> oSetValidColorInterp;
193 2625 : for (int i = 0; i < nColorInterpretations; ++i)
194 : oSetValidColorInterp.insert(
195 5100 : CPLString(
196 2550 : GDALGetColorInterpretationName(paeColorInterp[i]))
197 2550 : .tolower());
198 :
199 75 : const auto &val = arg.Get<std::vector<std::string>>();
200 166 : for (const auto &v : val)
201 : {
202 91 : if (!STARTS_WITH(v.c_str(), "mask") &&
203 91 : v.find(":") == std::string::npos &&
204 198 : CPLGetValueType(v.c_str()) != CPL_VALUE_INTEGER &&
205 13 : !cpl::contains(oSetValidColorInterp,
206 107 : CPLString(v).tolower()))
207 : {
208 3 : CPLError(CE_Failure, CPLE_AppDefined,
209 : "Invalid band specification.");
210 3 : return false;
211 : }
212 : }
213 72 : return true;
214 122 : });
215 : }
216 :
217 122 : AddArg("exclude", 0, _("Exclude specified bands"), &m_exclude);
218 :
219 : {
220 : auto &arg = AddArg(
221 : "mask", 0,
222 : _("Mask band (1-based index, 'mask', 'mask:<band>' or 'none')"),
223 122 : &m_mask);
224 : arg.AddValidationAction(
225 3 : [&arg]()
226 : {
227 3 : const auto &v = arg.Get<std::string>();
228 3 : if (!STARTS_WITH(v.c_str(), "mask") &&
229 6 : !EQUAL(v.c_str(), "none") &&
230 3 : !(CPLGetValueType(v.c_str()) == CPL_VALUE_INTEGER &&
231 2 : atoi(v.c_str()) >= 1))
232 : {
233 1 : CPLError(CE_Failure, CPLE_AppDefined,
234 : "Invalid mask band specification.");
235 1 : return false;
236 : }
237 2 : return true;
238 122 : });
239 : }
240 122 : }
241 :
242 : /************************************************************************/
243 : /* GDALRasterSelectAlgorithm::RunStep() */
244 : /************************************************************************/
245 :
246 35 : bool GDALRasterSelectAlgorithm::RunStep(GDALPipelineStepRunContext &)
247 : {
248 35 : const auto poSrcDS = m_inputDataset[0].GetDatasetRef();
249 35 : CPLAssert(poSrcDS);
250 35 : CPLAssert(m_outputDataset.GetName().empty());
251 35 : CPLAssert(!m_outputDataset.GetDatasetRef());
252 :
253 70 : std::map<GDALColorInterp, std::vector<int>> oMapColorInterpToBands;
254 561 : for (int i = 1; i <= poSrcDS->GetRasterCount(); ++i)
255 : {
256 526 : oMapColorInterpToBands[poSrcDS->GetRasterBand(i)
257 1052 : ->GetColorInterpretation()]
258 526 : .push_back(i);
259 : }
260 :
261 70 : CPLStringList aosOptions;
262 35 : aosOptions.AddString("-of");
263 35 : aosOptions.AddString("VRT");
264 35 : if (m_exclude)
265 : {
266 3 : if (m_bands.size() >= static_cast<size_t>(poSrcDS->GetRasterCount()))
267 : {
268 1 : ReportError(CE_Failure, CPLE_AppDefined,
269 : "Cannot exclude all input bands");
270 1 : return false;
271 : }
272 :
273 4 : std::set<int> excludedBandsFromColor;
274 4 : for (const std::string &v : m_bands)
275 : {
276 : const auto eColorInterp =
277 2 : GDALGetColorInterpretationByName(v.c_str());
278 2 : if (v == "undefined" || eColorInterp != GCI_Undefined)
279 : {
280 1 : const auto iter = oMapColorInterpToBands.find(eColorInterp);
281 1 : if (iter != oMapColorInterpToBands.end())
282 : {
283 2 : for (const int iBand : iter->second)
284 : {
285 1 : excludedBandsFromColor.insert(iBand);
286 : }
287 : }
288 : // We don't emit a warning if there are no bands matching
289 : // the color interpretation, because a potential use case
290 : // could be to run on a set of input files that might have or
291 : // might not have an alpha band, and remove it.
292 : }
293 : }
294 :
295 8 : for (int i = 1; i <= poSrcDS->GetRasterCount(); ++i)
296 : {
297 12 : const std::string iStr = std::to_string(i);
298 6 : if (std::find(m_bands.begin(), m_bands.end(), iStr) ==
299 17 : m_bands.end() &&
300 5 : !cpl::contains(excludedBandsFromColor, i))
301 : {
302 4 : aosOptions.AddString("-b");
303 4 : aosOptions.AddString(iStr);
304 : }
305 : }
306 : }
307 : else
308 : {
309 57 : for (const std::string &v : m_bands)
310 : {
311 : const auto eColorInterp =
312 38 : GDALGetColorInterpretationByName(v.c_str());
313 38 : if (v == "undefined" || eColorInterp != GCI_Undefined)
314 : {
315 4 : const auto iter = oMapColorInterpToBands.find(eColorInterp);
316 4 : if (iter == oMapColorInterpToBands.end())
317 : {
318 2 : ReportError(CE_Failure, CPLE_AppDefined,
319 : "No band has color interpretation %s",
320 : v.c_str());
321 2 : return false;
322 : }
323 4 : for (const int iBand : iter->second)
324 : {
325 2 : aosOptions.AddString("-b");
326 2 : aosOptions.AddString(std::to_string(iBand));
327 : }
328 : }
329 34 : else if (v.find(':') != std::string::npos)
330 : {
331 : const auto &aiBands =
332 26 : ParseBandRange(v, poSrcDS->GetRasterCount());
333 26 : if (!aiBands.has_value())
334 : {
335 11 : return false;
336 : }
337 91 : for (int iBand : aiBands.value())
338 : {
339 76 : aosOptions.AddString("-b");
340 76 : aosOptions.AddString(std::to_string(iBand));
341 : }
342 : }
343 8 : else if (cpl::equals_ci(v, "mask"))
344 : {
345 1 : aosOptions.AddString("-b");
346 1 : aosOptions.AddString(v);
347 : }
348 : else
349 : {
350 7 : const auto maybeBand = cpl::strict_parse<int>(v);
351 7 : if (!maybeBand)
352 : {
353 0 : CPLError(CE_Failure, CPLE_IllegalArg, "Invalid band: %s",
354 : v.c_str());
355 0 : return false;
356 : }
357 7 : const int nBands = poSrcDS->GetRasterCount();
358 7 : int iBand = maybeBand.value();
359 7 : if (iBand < 0)
360 : {
361 1 : iBand += nBands + 1;
362 : }
363 :
364 7 : if (iBand > nBands || iBand < 1)
365 : {
366 0 : CPLError(CE_Failure, CPLE_IllegalArg, "Invalid band: %s",
367 : v.c_str());
368 0 : return false;
369 : }
370 :
371 7 : aosOptions.AddString("-b");
372 7 : aosOptions.AddString(std::to_string(iBand));
373 : }
374 : }
375 : }
376 21 : if (!m_mask.empty())
377 : {
378 1 : aosOptions.AddString("-mask");
379 1 : aosOptions.AddString(CPLString(m_mask).replaceAll(':', ',').c_str());
380 : }
381 :
382 : GDALTranslateOptions *psOptions =
383 21 : GDALTranslateOptionsNew(aosOptions.List(), nullptr);
384 :
385 : auto poOutDS = std::unique_ptr<GDALDataset>(GDALDataset::FromHandle(
386 21 : GDALTranslate("", GDALDataset::ToHandle(poSrcDS), psOptions, nullptr)));
387 21 : GDALTranslateOptionsFree(psOptions);
388 21 : const bool bRet = poOutDS != nullptr;
389 21 : if (poOutDS)
390 : {
391 21 : m_outputDataset.Set(std::move(poOutDS));
392 : }
393 :
394 21 : return bRet;
395 : }
396 :
397 : GDALRasterSelectAlgorithmStandalone::~GDALRasterSelectAlgorithmStandalone() =
398 : default;
399 :
400 : //! @endcond
|
