diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 5548573..ff3f4c6 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -14,6 +14,8 @@ jobs: python-version: "3.12" - name: Install test runner run: python -m pip install --upgrade pip pytest + - name: Install package + run: python -m pip install --no-deps -e . - name: Syntax check run: python -m compileall opensr_utils setup.py demo.py tests - name: Shell syntax check diff --git a/README.md b/README.md index 41ccf0a..c7c107e 100644 --- a/README.md +++ b/README.md @@ -4,12 +4,6 @@ ![img1](resources/utils_poster.png) -*Info* -- It is functional for RGB+NIR images currently. -- Functionalities coming in the future: - - [x] Input many tiles at once - - [x] Pass more file types - - [ ] Run SR for a selectable amount of bands # Description This package performs super-resolution with any PyTorch or PyTorch lighning model for the Sentinel-2 10m bands (R-G-B-NIR). 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/setup.py b/setup.py index 1e2fa88..802807c 100644 --- a/setup.py +++ b/setup.py @@ -6,7 +6,7 @@ setup( name="opensr-utils", - version="1.2.0", + version="2.0.0", author="Simon Donike, Cesar Aybar, Luis Gomez Chova, Freddie Kalaitzis", author_email="accounts@donike.net", description="Utilities supporting the ESA opensr-model package for Super-Resolution of Sentinel-2 Imagery", 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"}