From fd29b5cd2a052deb85ab8d7ded4122de4f8e689b Mon Sep 17 00:00:00 2001 From: Thomas Tilak Date: Tue, 29 Sep 2026 09:01:36 +1100 Subject: [PATCH] Fix SAFE/S2GM band selection picking up auxiliary rasters (AOT/WVP/TCI). --- opensr_utils/pipeline.py | 10 ++++++++-- tests/test_pipeline_helpers.py | 30 ++++++++++++++++++++++++++++++ 2 files changed, 38 insertions(+), 2 deletions(-) diff --git a/opensr_utils/pipeline.py b/opensr_utils/pipeline.py index 22a8fb0..03a8f6f 100644 --- a/opensr_utils/pipeline.py +++ b/opensr_utils/pipeline.py @@ -353,7 +353,9 @@ def safe_band_name(path): used = set(lr_file_dict.values()) for path in jp2s: if path not in used: - lr_file_dict[safe_band_name(path)] = path + name = safe_band_name(path) + if name.startswith("B") and name[1:].isalnum(): + lr_file_dict[name] = path if not lr_file_dict: raise FileNotFoundError(f"No usable SAFE bands found in {self.root}.") @@ -407,7 +409,11 @@ def safe_band_name(path): lr_file_dict[label] = by_name[band_name] used = set(lr_file_dict.values()) for band_name, path in by_name.items(): - if path not in used: + if ( + path not in used + and band_name.startswith("B") + and band_name[1:].isalnum() + ): lr_file_dict[band_name] = path self.image_meta["lr_file_dict"] = lr_file_dict diff --git a/tests/test_pipeline_helpers.py b/tests/test_pipeline_helpers.py index 9114e9d..dd10906 100644 --- a/tests/test_pipeline_helpers.py +++ b/tests/test_pipeline_helpers.py @@ -438,3 +438,33 @@ def test_write_to_file_creates_placeholder_from_prediction_index( with rasterio.open(runner.placeholder_filepath) as src: assert (src.width, src.height, src.count) == (9, 6, 2) assert src.dtypes == ("uint16", "uint16") + + +def test_safe_band_selection_excludes_auxiliary_rasters(write_raster, tmp_path): + safe = tmp_path / "S2A_TEST.SAFE" + granule = safe / "GRANULE" / "L1C_TEST" / "IMG_DATA" / "R10m" + granule.mkdir(parents=True) + for band in ("B02", "B03", "B04", "B08", "AOT", "WVP", "TCI"): + write_raster( + granule / f"T_TEST_{band}_10m.jp2", + data=np.ones((1, 16, 16), dtype=np.uint16), + width=16, + height=16, + ) + (safe / "manifest.safe").write_text("") + + runner = pipeline.large_file_processing( + root=str(safe), + model=None, + window_size=(8, 8), + factor=2, + overlap=4, + eliminate_border_px=0, + device="cpu", + batch_size=1, + num_workers=0, + prefetch_factor=0, + ) + + assert runner.image_meta["bands"] == 4 + assert set(runner.image_meta["lr_file_dict"]) == {"R", "G", "B", "NIR"}