diff --git a/.github/dependabot.yml b/.github/dependabot.yml
new file mode 100644
index 0000000..696b5d7
--- /dev/null
+++ b/.github/dependabot.yml
@@ -0,0 +1,11 @@
+version: 2
+updates:
+ - package-ecosystem: "github-actions"
+ directory: "/"
+ schedule:
+ interval: "weekly"
+
+ - package-ecosystem: "pip"
+ directory: "/"
+ schedule:
+ interval: "weekly"
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index f41335e..7931643 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -8,8 +8,12 @@ on:
# 1) manual run from the Actions tab
workflow_dispatch:
- # 2) tagged releases (v*)
- # 3) any branch push that changes code/tests/dependency manifests
+ # 2) weekly run to catch new dependency releases against supported Python versions
+ schedule:
+ - cron: "17 3 * * 1"
+
+ # 3) tagged releases (v*)
+ # 4) any branch push that changes code/tests/dependency manifests
push:
branches:
- '**' # allow all branches (needed so branches still trigger)
@@ -30,7 +34,7 @@ jobs:
strategy:
fail-fast: false
matrix:
- python-version: ["3.10", "3.12"]
+ python-version: ["3.12", "3.13", "3.14"]
steps:
- uses: actions/checkout@v4
@@ -42,11 +46,13 @@ jobs:
- name: Install dependencies
run: |
python -m pip install --upgrade pip
- pip install -e .
- pip install pytest pytest-cov
- pip install --index-url https://download.pytorch.org/whl/cpu torch
- pip install -r requirements.txt
- pip install numpy
+ python -m pip install --index-url https://download.pytorch.org/whl/cpu torch torchvision
+ python -m pip install -e ".[tests]"
+
+ - name: Show dependency versions
+ run: |
+ python -m pip list
+ python -m pip check
- name: Run tests with coverage
run: |
diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
index f391c02..ea47c75 100644
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@@ -30,7 +30,7 @@ jobs:
- uses: actions/setup-python@v5
with:
- python-version: "3.11"
+ python-version: "3.12"
- name: Install doc dependencies
run: |
diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml
index d76709e..02512db 100644
--- a/.github/workflows/publish.yml
+++ b/.github/workflows/publish.yml
@@ -14,7 +14,7 @@ jobs:
- uses: actions/checkout@v4
- uses: actions/setup-python@v5
with:
- python-version: "3.10"
+ python-version: "3.12"
- name: Build
run: |
python -m pip install --upgrade build
diff --git a/README.md b/README.md
index 78ade9d..0413913 100644
--- a/README.md
+++ b/README.md
@@ -2,7 +2,7 @@
| **PyPI** | **Versions** | **Docs & License** | **Tests** | **Reference** |
|:---------:|:-------------:|:------------------:|:----------:|:--------------:|
-| [](https://pypi.org/project/opensr-srgan/) | 
 | [](https://srgan.opensr.eu)
 | [](https://github.com/ESAOpenSR/SRGAN/actions/workflows/ci.yml)
[](https://codecov.io/github/ESAOpenSR/SRGAN) | [](https://arxiv.org/abs/2511.10461)
[](https://doi.org/10.5281/zenodo.20088994)
+| [](https://pypi.org/project/opensr-srgan/) | 
 | [](https://srgan.opensr.eu)
 | [](https://github.com/ESAOpenSR/SRGAN/actions/workflows/ci.yml)
[](https://codecov.io/github/ESAOpenSR/SRGAN) | [](https://arxiv.org/abs/2511.10461)
[](https://doi.org/10.5281/zenodo.20088994)
@@ -26,19 +26,25 @@ Full docs live at **[srgan.opensr.eu](https://srgan.opensr.eu/)**. They cover us
* **Remote-sensing aware losses:** combine spectral, perceptual, and adversarial objectives with tunable weights.
* **Stable training loop:** generator pretraining, adversarial ramp-ups, EMA, and multi-GPU Lightning support out of the box.
* **PyPI distribution:** `pip install opensr-srgan` for ready-to-use presets or custom configs.
-* **Extensive Logging:** Logging all important information automatically to `WandB` for optimal insights.
+* **Experiment tracking:** Log metrics and validation imagery to Weights & Biases, or use local CSV logs for lightweight runs.
---
## 🏗️ Configuration Examples
-All key knobs are exposed via YAML in the `opensr_srgan/configs` folder:
+All key knobs are exposed via YAML in the `opensr_srgan/configs` folder. You can keep using the legacy single-file workflow, or use the Hydra entry point for composable experiment presets and command-line overrides:
-* **Model**: `in_channels`, `n_channels`, `n_blocks`, `scale`, ESRGAN knobs (`growth_channels`, `res_scale`, `out_channels`), `block_type ∈ {SRResNet, res, rcab, rrdb, lka}`
+```bash
+python -m opensr_srgan.train --config opensr_srgan/configs/config_training_example.yaml
+python -m opensr_srgan.train_hydra experiment=10m Training.max_epochs=5 Logging.wandb.enabled=false
+srgan-train experiment=20m Training.gpus=[0,1]
+```
+
+* **Model**: `Model.in_bands`, `Generator.n_channels`, `Generator.n_blocks`, `Generator.scaling_factor`, ESRGAN knobs (`growth_channels`, `res_scale`, `out_channels`, `use_icnr`), and `Generator.block_type ∈ {standard, res, rcab, rrdb, lka}`
* **Losses**: `l1_weight`, `sam_weight`, `perceptual_weight`, `tv_weight`, `adv_loss_beta`
-* **Training**: `pretrain_g_only`, `g_pretrain_steps` (`-1` keeps generator-only pretraining active indefinitely), `adv_loss_ramp_steps`, `label_smoothing`, generator LR warmup (`Schedulers.g_warmup_steps`, `Schedulers.g_warmup_type`), discriminator cadence controls
+* **Training**: `pretrain_g_only`, `g_pretrain_steps` (`-1` keeps generator-only pretraining active indefinitely), `adv_loss_ramp_steps`, `label_smoothing`, generator LR warmup (`Schedulers.g_warmup_steps`, `Schedulers.g_warmup_type`), EMA, and manual gradient clipping
* **Adversarial mode**: `Training.Losses.adv_loss_type` (`bce`/`wasserstein`) and optional `Training.Losses.relativistic_average_d` for BCE-based relativistic-average GAN updates
-* **Data**: band order, normalization stats, crop sizes, augmentations
+* **Data**: `Data.dataset_type`, dataloader sizes/workers, dataset paths/manifests, and normalization policy
---
@@ -60,7 +66,7 @@ The schedule and ramp make training **easier, safer, and more reproducible**.
| Component | Options | Config keys |
|-----------|---------|-------------|
| **Generators** | `SRResNet`, `res`, `rcab`, `rrdb`, `lka`, `esrgan`, `stochastic_gan` | `Generator.model_type`, depth via `Generator.n_blocks`, width via `Generator.n_channels`, kernels/scale plus ESRGAN-specific `growth_channels`, `res_scale`, `out_channels`. |
-| **Discriminators** | `standard` `SRGAN`, `CNN`, `patchgan`, `esrgan` | `Discriminator.model_type`, granularity with `Discriminator.n_blocks`, spectral norm toggle via `Discriminator.use_spectral_norm`, ESRGAN-specific `base_channels`, `linear_size`. |
+| **Discriminators** | `standard` SRGAN, `patchgan`, `esrgan` | `Discriminator.model_type`, granularity with `Discriminator.n_blocks`, spectral norm toggle via `Discriminator.use_spectral_norm`, ESRGAN-specific `base_channels`, `linear_size`. |
| **Content losses** | L1, Spectral Angle Mapper, VGG19/LPIPS perceptual metrics, Total Variation | Weighted by `Training.Losses.*` (e.g. `l1_weight`, `sam_weight`, `perceptual_weight`, `perceptual_metric`, `tv_weight`). |
| **Adversarial loss** | BCE‑with‑logits or Wasserstein critic | Controlled by `Training.Losses.adv_loss_type`, warmup via `Training.pretrain_g_only`, ramped by `adv_loss_ramp_steps`, capped at `adv_loss_beta`, optional label smoothing. For BCE, enable `Training.Losses.relativistic_average_d` for RaGAN-style relative logits. |
@@ -77,8 +83,8 @@ Follow the [installation instructions](https://srgan.opensr.eu/getting-started/)
* To test the package immediately, launch the Google Colab right now and follow along the introduction! [](https://colab.research.google.com/drive/16W0FWr6py1J8P4po7JbNDMaepHUM97yL?usp=sharing)
* **Datasets:** Grab the bundled example dataset or learn how to register your own sources in the [data guide](https://srgan.opensr.eu/data/).
-* **Training:** Launch training with `python -m opensr_srgan.train --config opensr_srgan/configs/config_10m.yaml` or import `train` from the package as described in the [training walkthrough](https://srgan.opensr.eu/training/).
-* **Inference:** Ready-made presets and large-scene pipelines are described in the [inference section](https://srgan.opensr.eu/getting-started/inference/).
+* **Training:** Launch a smoke test with `python -m opensr_srgan.train --config opensr_srgan/configs/config_training_example.yaml`, or use Hydra with `python -m opensr_srgan.train_hydra experiment=example`. Use the `10m` and `20m` presets once the matching datasets are configured.
+* **Inference:** Ready-made presets and large-scene pipelines are described in the [inference section](https://srgan.opensr.eu/inference/).
---
diff --git a/docs/api-reference/training.md b/docs/api-reference/training.md
index 1990724..bd90d3b 100644
--- a/docs/api-reference/training.md
+++ b/docs/api-reference/training.md
@@ -17,6 +17,20 @@ OmegaConf-based experiment files.
- "!^_"
members_order: source
+## Hydra entry point
+
+::: opensr_srgan.train_hydra
+ handler: python
+ options:
+ show_root_heading: false
+ show_root_full_path: false
+ show_source: true
+ members:
+ - main
+ filters:
+ - "!^_"
+ members_order: source
+
## Trainer configuration helpers
::: opensr_srgan.utils.build_trainer_kwargs
diff --git a/docs/architecture.md b/docs/architecture.md
index 0ef3a1d..fec281c 100644
--- a/docs/architecture.md
+++ b/docs/architecture.md
@@ -23,7 +23,7 @@ Because every generator variant (residual, RCAB, RRDB, large-kernel attention, E
`Generator.model_type`/`block_type` and `Discriminator.model_type`. Unsupported combinations fail fast with clear error
messages.
* **Loss construction.** `GeneratorContentLoss` (from `opensr_srgan.model.loss`) provides L1, spectral angle mapper (SAM), perceptual, and
- total-variation terms. Adversarial supervision uses `torch.nn.BCEWithLogitsLoss` with optional label smoothing.
+ total-variation terms. Adversarial supervision uses BCE-with-logits or Wasserstein objectives, with optional label smoothing, relativistic-average BCE, and R1 penalty depending on the config.
* **Optimiser scheduling.** `configure_optimizers()` returns paired Adam optimisers (generator + discriminator) with
`ReduceLROnPlateau` schedulers that monitor a configurable validation metric.
* **Training orchestration.** `setup_lightning()` binds `training_step_PL2()` and enables manual optimisation
@@ -56,7 +56,7 @@ The generator zoo lives under `opensr_srgan/model/generators/` and can be select
* **Stochastic GAN generator (`cgan_generator.py`).** Extends the flexible generator with conditioning inputs and latent noise,
enabling experiments where auxiliary metadata influences the super-resolution output.
* **ESRGAN generator (`esrgan.py`).** Implements the RRDBNet trunk introduced with ESRGAN, exposing `n_blocks`, `growth_channels`,
- and `res_scale` so you can dial in deeper receptive fields and sharper textures. The implementation supports original features like Relativistic Average GAN (RaGAN), and the codebase allows a two-step training phase (content-oriented pretraining of the generator followed by adversarial training with the discriminator), as originally proposed by the ESRGAN authors.
+ and `res_scale` so you can dial in deeper receptive fields and sharper textures. The shared training loop supports optional relativistic-average BCE and a two-step training phase (content-oriented pretraining of the generator followed by adversarial training with the discriminator).
* **Advanced variants (`SRGAN_advanced.py`).** Provides additional block implementations and compatibility aliases exposed in
`__init__.py` for backwards compatibility.
@@ -89,9 +89,9 @@ The same module exposes `return_metrics()` so validation can log PSNR/SSIM-style
## Data flow and normalisation
The Lightning module expects batches of `(lr_imgs, hr_imgs)` tensors supplied by the `LightningDataModule` returned from
-`opensr_srgan/data/dataset_selector.py`. `predict_step()` and the validation hooks rely on two utilities from `opensr_srgan.utils.spectral_helpers`:
+`opensr_srgan/data/dataset_selector.py`. `predict_step()` uses the configured `opensr_srgan.data.utils.Normalizer` plus histogram matching from `opensr_srgan.utils.radiometrics`:
-* `normalise_10k`: Converts Sentinel-2 style reflectance values between `[0, 10000]` and `[0, 1]`.
+* `Normalizer`: Applies the configured `Data.normalization` strategy, such as `normalise_10k`, `zero_one`, or a custom callable.
* `histogram`: Matches the SR histogram to the LR reference to minimise domain gaps during inference.
These helpers allow the generator to operate in a normalised space while still reporting outputs in physical units when needed.
diff --git a/docs/configuration.md b/docs/configuration.md
index 90dc18b..67f071d 100644
--- a/docs/configuration.md
+++ b/docs/configuration.md
@@ -1,6 +1,6 @@
# Configuration
-ESA OpenSR relies on YAML files to control every aspect of the training pipeline. This page documents the available keys and how they influence the underlying code. Use `opensr_srgan/configs/config_20m.yaml` and `opensr_srgan/configs/config_10m.yaml` as starting points.
+ESA OpenSR relies on YAML files to control every aspect of the training pipeline. This page documents the available keys and how they influence the underlying code. Use `opensr_srgan/configs/config_training_example.yaml` for a smoke test, `config_10m.yaml` for SEN2NAIP-style runs, or the Hydra presets described below for composable experiments.
## File structure
@@ -23,9 +23,9 @@ Each section maps directly to parameters consumed inside `opensr_srgan/model/SRG
| Key | Default | Description |
| --- | --- | --- |
-| `train_batch_size` | 12 | Mini-batch size for the training dataloader. Falls back to `batch_size` if set. |
+| `train_batch_size` | 8 | Mini-batch size for the training dataloader. Falls back to `batch_size` if set. |
| `val_batch_size` | 8 | Batch size for validation. |
-| `num_workers` | 6 | Number of worker processes for both dataloaders. |
+| `num_workers` | 4 | Number of worker processes for both dataloaders. |
| `prefetch_factor` | 2 | Additional batches prefetched by each worker. Ignored when `num_workers == 0`. |
| `dataset_type` | `ExampleDataset` | Dataset selector consumed by `opensr_srgan.data.dataset_selector.select_dataset`. |
| `normalization` | `'sen2_stretch'` | Normalisation strategy applied to input tensors. Accepts a string alias or a mapping (see below). |
@@ -211,7 +211,7 @@ Both optimisers share the same configuration keys because they use `torch.optim.
| Key | Default | Description |
| --- | --- | --- |
-| `metric` | `val_metrics/l1` | Validation metric monitored for plateau detection. |
+| `metric` | `val_g_loss` / `val_d_loss` fallback | Shared fallback monitor when `metric_g` or `metric_d` is omitted. Shipped configs set explicit monitor keys. |
| `metric_g` | — | Optional override for the generator scheduler monitor. |
| `metric_d` | — | Optional override for the discriminator scheduler monitor. |
| `patience_g` | `100` | Epochs with no improvement before reducing the generator LR. |
@@ -220,7 +220,6 @@ Both optimisers share the same configuration keys because they use `torch.optim.
| `factor_d` | `0.5` | Multiplicative factor applied to the discriminator LR upon plateau. |
| `cooldown` | `0` | Number of epochs to wait after an LR drop before resuming plateau checks. |
| `min_lr` | `1e-7` | Minimum learning rate allowed for both schedulers. |
-| `verbose` | `True` | Enables scheduler logging messages. |
| `g_warmup_steps` | `2000` | Number of optimiser steps used for generator LR warmup. Set to `0` to disable. |
| `g_warmup_type` | `cosine` | Warmup curve for the generator LR (`cosine` or `linear`). |
@@ -235,7 +234,23 @@ different validation metrics.
| Key | Default | Description |
| --- | --- | --- |
-| `num_val_images` | `5` | Number of validation batches visualised and logged to Weights & Biases each epoch. |
+| `num_val_images` | `5` | Number of validation batches visualised and logged to Weights & Biases when W&B is enabled. |
+| `wandb.enabled` | `False` in example configs | Enables `WandbLogger`; when false, training uses `CSVLogger`. |
+| `wandb.entity` | `opensr` | W&B entity/team name. |
+| `wandb.project` | `SRGAN_10m` | Project name used by W&B and by the default checkpoint directory. |
+| `output_dir` | unset | Optional explicit run directory. Hydra sets this to its runtime output directory; legacy YAML runs fall back to `logs//`. |
+
+## Hydra experiment configs
+
+The legacy single-file configs remain supported, but new experiments can also be launched through Hydra:
+
+```bash
+python -m opensr_srgan.train_hydra experiment=example
+python -m opensr_srgan.train_hydra experiment=10m Training.max_epochs=5 Logging.wandb.enabled=false
+srgan-train experiment=20m Training.gpus=[0,1]
+```
+
+Hydra configs live under `opensr_srgan/configs/hydra`. The root `train.yaml` selects simple defaults, while `experiment/example.yaml`, `experiment/10m.yaml`, and `experiment/20m.yaml` reproduce the shipped flat YAML presets through grouped `data`, `model`, `training`, and `logging` configs. Choose `experiment=...` for a full known setup; use direct overrides for a temporary change without editing files. See [Hydra Experiments](hydra.md) for a complete walkthrough.
## Tips for managing configurations
diff --git a/docs/data.md b/docs/data.md
index 0e498e9..11f6e5a 100644
--- a/docs/data.md
+++ b/docs/data.md
@@ -21,7 +21,7 @@ from opensr_srgan.data.example_data.download_example_dataset import get_example_
get_example_dataset()
```
-The helper pulls `example_dataset.zip` from the [`simon-donike/SR-GAN`](https://huggingface.co/simon-donike/SR-GAN) repository, extracts it, and removes the temporary archive once the copy completes.
+The helper pulls `example_dataset.zip` from the [`simon-donike/SR-GAN`](https://huggingface.co/simon-donike/SR-GAN) repository through the Hugging Face cache and extracts it into the requested output directory.
### Directory layout
@@ -35,6 +35,18 @@ Data:
The training loop automatically instantiates `opensr_srgan.data.example_data.example_dataset.ExampleDataset` for both the
training and validation dataloaders.
+## Built-in dataset selectors
+
+`Data.dataset_type` is resolved by `opensr_srgan.data.dataset_selector.select_dataset`. The currently supported selectors are:
+
+| `Data.dataset_type` | Required fields | Notes |
+| --- | --- | --- |
+| `ExampleDataset` | none beyond the downloaded `example_dataset/` folder | Quick smoke-test dataset generated from HR `.npz` chips. |
+| `SEN2NAIP` / `sen2naip` | `Data.sen2naip_taco_file`; optional `Data.sen2naip_val_fraction` | Requires `tacoreader` and a Taco manifest. Used by the 10 m preset once the manifest path is configured. |
+| `LRHRFolderDataset` | `Data.root_dir` or `Data.dataset_root` | Expects paired LR/HR folders; see the dataset class tests for the accepted layout. |
+
+Older labels such as `S2_6b`, `SPOT6`, or `cv` are not registered selector keys in the current code. If you start from an older config that uses one of those names, update `Data.dataset_type` or add a matching branch to the selector.
+
## Adding new datasets
When you are ready to move beyond the bundled sample data you can register a custom dataset. The repository uses a single factory function `opensr_srgan.data.dataset_selector.select_dataset` to keep the training script agnostic of individual
diff --git a/docs/getting-started.md b/docs/getting-started.md
index 96eeecb..00f1989 100644
--- a/docs/getting-started.md
+++ b/docs/getting-started.md
@@ -1,6 +1,6 @@
# Getting started
-This guide walks through installing dependencies, configuring datasets, and launching your first ESA OpenSR experiment. The stack supports Python 3.10-3.12, PyTorch Lightning, and Weights & Biases for experiment tracking.
+This guide walks through installing dependencies, configuring datasets, and launching your first ESA OpenSR experiment. The stack supports Python 3.12-3.14, PyTorch Lightning, and Weights & Biases for experiment tracking.
## Try it in Colab first
@@ -21,13 +21,14 @@ For the fastest start, open the interactive notebook in Google Colab and run thr
python -m venv .venv
source .venv/bin/activate
```
-2. **Install Python dependencies.**
+2. **Install the project and Python dependencies.**
```bash
- pip install -r requirements.txt
+ python -m pip install -e .
```
+ This editable install also exposes the `srgan-train` and `srgan-hpc` commands. Contributors can install the test and documentation tools with `python -m pip install -e ".[tests,docs]"`; cluster users can add the HPC dependencies with `python -m pip install -e ".[hpc]"`.
3. **Authenticate logging backends (optional but recommended).**
* Run `wandb login` to capture metrics and images in your W&B workspace.
- * Start `tensorboard --logdir logs/` if you prefer local dashboards.
+ * Keep `Logging.wandb.enabled: false` for local CSV logs when you do not want to use W&B.
## 2. Gather training data
@@ -39,7 +40,7 @@ from opensr_srgan.data.example_data.download_example_dataset import get_example_
get_example_dataset() # downloads into ./example_dataset/
```
-The script downloads `example_dataset.zip` from the Hugging Face Hub, extracts it to `example_dataset/`, and removes the archive once extraction finishes. The configuration only needs to specify the dataset type:
+The script downloads `example_dataset.zip` through the Hugging Face Hub cache and extracts it to `example_dataset/`. The configuration only needs to specify the dataset type:
```yaml
Data:
@@ -50,10 +51,10 @@ When you are ready to integrate your own collections, follow the guidance in [Da
## 3. Configure the experiment
-Use of the provided YAML presets or copy and edit one:
+Use a provided YAML preset or copy and edit one. For the bundled example dataset, start from the example config:
```bash
-cp opensr_srgan/configs/config_10m.yaml opensr_srgan/configs/my_experiment.yaml
+cp opensr_srgan/configs/config_training_example.yaml opensr_srgan/configs/my_experiment.yaml
```
Update at least the following fields:
@@ -68,10 +69,11 @@ See [Configuration](configuration.md) for a full breakdown of available options.
## 4. Launch training
-Run the training script with your customised config:
+Run the training script with your customised config, or use the Hydra example preset:
```bash
python -m opensr_srgan.train --config opensr_srgan/configs/my_experiment.yaml
+python -m opensr_srgan.train_hydra experiment=example
```
Prefer to stay inside Python? Import the helper exposed by the package:
@@ -86,7 +88,7 @@ Both entry points will:
1. Instantiate the `SRGAN_model` Lightning module from the YAML file.
2. Build the appropriate dataset pair and wrap it in a `LightningDataModule`.
-3. Configure Weights & Biases and TensorBoard loggers alongside checkpointing and learning-rate monitoring callbacks.
+3. Configure a Weights & Biases or CSV logger, checkpointing, early stopping, and per-step learning-rate logging.
4. Start alternating generator/discriminator optimisation according to your warm-start schedule.
Training resumes automatically if `Model.continue_training` points to a Lightning checkpoint. If you interrupt training, always use the `Model.continue_training` flag to pass the generated checkpoint, since that restores all optimizers, schedulers, EMA etc. Do not set `Model.load_checkpoint` and `Model.continue_training` at the same time.
@@ -103,7 +105,7 @@ Training resumes automatically if `Model.continue_training` points to a Lightnin
# Option A – bring your own config + checkpoint (local path or URL)
custom_model = load_from_config(
- config_path="opensr_srgan/configs/config_10m.yaml",
+ config_path="path/to/your_config.yaml",
checkpoint_uri="https://example.com/checkpoints/srgan.ckpt",
map_location="cuda", # optional
)
diff --git a/docs/hydra.md b/docs/hydra.md
new file mode 100644
index 0000000..7ce0a05
--- /dev/null
+++ b/docs/hydra.md
@@ -0,0 +1,107 @@
+# Hydra experiment configs
+
+Hydra is the recommended entry point when you want to run repeatable experiments without copying and editing a full YAML file for every small change. The old single-file config workflow still works; Hydra is an additional layer that composes the same `Data`, `Model`, `Training`, `Generator`, `Discriminator`, `Optimizers`, `Schedulers`, and `Logging` sections that the SRGAN code already understands.
+
+## Mental model
+
+The legacy command loads one complete YAML file:
+
+```bash
+python -m opensr_srgan.train --config opensr_srgan/configs/config_training_example.yaml
+```
+
+The Hydra command starts from a small default config, selects reusable config groups, applies an optional experiment preset, and then applies command-line overrides:
+
+```bash
+python -m opensr_srgan.train_hydra experiment=10m Training.max_epochs=5 Logging.wandb.enabled=false
+```
+
+After composition, `opensr_srgan.train_hydra` forwards the resolved OmegaConf object to the same `opensr_srgan.train.train()` function used by the legacy workflow. This keeps the model, data, checkpoint, and logger code paths shared.
+
+## Config layout
+
+Hydra configs live in `opensr_srgan/configs/hydra/`:
+
+| Path | Purpose |
+| --- | --- |
+| `train.yaml` | Default Hydra entry config. It selects the bundled example dataset, small CPU-safe training settings, CSV logging, and `experiment: null`. |
+| `data/` | Dataset and dataloader presets such as `example`, `sen2naip`, and `lrhr_folder`. |
+| `model/` | Combined `Model`, `Generator`, and `Discriminator` presets. |
+| `training/` | Combined `Training`, `Optimizers`, and `Schedulers` presets. |
+| `logging/` | CSV and W&B logging presets. |
+| `experiment/` | Full experiment presets that reproduce the shipped legacy YAMLs. |
+| `hydra/default.yaml` | Hydra's own run-directory and working-directory settings. |
+
+The important convention is that each group writes into the existing top-level SRGAN schema. For example, `model/10m_esrgan.yaml` contributes `Model`, `Generator`, and `Discriminator`, not a new lowercase `model` object.
+
+## Choosing a config
+
+Use `experiment=...` when you want a known full setup:
+
+```bash
+python -m opensr_srgan.train_hydra experiment=example
+python -m opensr_srgan.train_hydra experiment=10m
+python -m opensr_srgan.train_hydra experiment=20m
+```
+
+Use direct overrides when you only want to change one or two values for a run:
+
+```bash
+python -m opensr_srgan.train_hydra experiment=10m Training.max_epochs=20
+python -m opensr_srgan.train_hydra experiment=10m Generator.n_blocks=8 Optimizers.optim_g_lr=5e-5
+python -m opensr_srgan.train_hydra experiment=20m Training.gpus=[0,1] Logging.wandb.enabled=false
+```
+
+Use group overrides when you want to mix reusable pieces:
+
+```bash
+python -m opensr_srgan.train_hydra data=example model=example_srresnet training=cpu_debug logging=csv
+python -m opensr_srgan.train_hydra data=sen2naip model=10m_esrgan training=10m logging=wandb_10m
+```
+
+If a field does not already exist in the composed config, Hydra requires a `+` prefix. For normal SRGAN fields this is rarely needed because the presets define the expected sections up front.
+
+## Output directories
+
+Hydra creates a run directory under:
+
+```text
+logs//
+```
+
+The Hydra entry point copies that directory into `Logging.output_dir` before calling the existing trainer. The trainer then writes checkpoints and the resolved `config.yaml` there.
+
+Hydra is configured with:
+
+```yaml
+hydra:
+ job:
+ chdir: false
+```
+
+This means relative paths keep behaving as if you launched the command from the repository root. Dataset paths such as `example_dataset/` therefore do not silently start resolving inside the run directory.
+
+## Inspecting configs
+
+Use Hydra's built-in config printing when you want to check exactly what will be passed to training:
+
+```bash
+python -m opensr_srgan.train_hydra --cfg job experiment=10m
+python -m opensr_srgan.train_hydra --cfg hydra experiment=10m
+```
+
+The first command prints the resolved SRGAN job config. The second prints Hydra's runtime config, including output-directory settings.
+
+## When to use legacy YAML
+
+Use the legacy YAML entry point when you already have a complete hand-written config file or when a notebook/script expects a path. For the bundled example dataset, use:
+
+```bash
+python -m opensr_srgan.train --config opensr_srgan/configs/config_training_example.yaml
+```
+
+Use Hydra when you want quick overrides, reusable experiment pieces, or clean per-run config snapshots without creating many near-duplicate files. Use `experiment=10m` or `experiment=20m` after configuring the matching real datasets. Both paths call the same trainer once the config is loaded.
+
+## References
+
+Hydra composes these configs through a Defaults List. See the Hydra Defaults List documentation for the underlying mechanism: . The SRGAN entry point also uses an explicit `version_base` and `config_path`, following Hydra's upgrade guidance: .
diff --git a/docs/index.md b/docs/index.md
index 80f10ff..f919621 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -4,7 +4,7 @@
| **PyPI** | **Versions** | **Docs & License** | **Tests** | **Reference** |
|:---------:|:-------------:|:------------------:|:----------:|:--------------:|
-| [](https://pypi.org/project/opensr-srgan/) | 
 | [](https://srgan.opensr.eu)
 | [](https://github.com/ESAOpenSR/SRGAN/actions/workflows/ci.yml)
[](https://codecov.io/github/ESAOpenSR/SRGAN) | [](https://arxiv.org/abs/2511.10461)
[](https://doi.org/10.5281/zenodo.20088994) |
+| [](https://pypi.org/project/opensr-srgan/) | 
 | [](https://srgan.opensr.eu)
 | [](https://github.com/ESAOpenSR/SRGAN/actions/workflows/ci.yml)
[](https://codecov.io/github/ESAOpenSR/SRGAN) | [](https://arxiv.org/abs/2511.10461)
[](https://doi.org/10.5281/zenodo.20088994) |
@@ -18,25 +18,20 @@ Whether you are reproducing published results, exploring new remote-sensing moda
## Why this repository?
-* **One configuration, many models.** Swap between RCAN-style residual channel attention, RRDB, SRResNet, SwinIR-inspired
- backbones, and matching discriminators ranging from classic SRGAN to PatchGAN by editing a single config block.
-* **Loss combinations that just work.** Mix pixel, perceptual, style, and adversarial objectives with sensible defaults for
- weights, schedules, and warm-up durations.
-* **Battle-tested training loop.** PyTorch Lightning handles mixed precision, gradient accumulation, multi-GPU training, and
- restartable checkpoints while the repo layers in GAN-specific tweaks such as adversarial weight ramping and learning-rate
- restarts.
+* **One configuration, many models.** Swap between SRResNet, RCAB, RRDB, large-kernel attention, ESRGAN, and stochastic generators, plus standard SRGAN, PatchGAN, and ESRGAN discriminators, by editing config values.
+* **Loss combinations that just work.** Mix pixel, spectral-angle, perceptual, total-variation, and adversarial objectives with sensible defaults for weights, schedules, and warm-up durations.
+* **Battle-tested training loop.** PyTorch Lightning handles device placement, multi-GPU DDP, and restartable checkpoints while the repo layers in GAN-specific manual optimisation, adversarial weight ramping, gradient clipping, and learning-rate scheduling.
* **Lightning 2+ training path.** Training uses manual optimization (`automatic_optimization=False`) and a single GAN training-step implementation.
-* **Remote-sensing aware defaults.** Normalisation, histogram matching, spectral-band handling, and Sentinel-2 SAFE ingestion are
- ready-made for 10 m and 20 m bands and easily extendable to other sensors.
+* **Remote-sensing aware defaults.** Normalisation, histogram matching, spectral-band handling, and Sentinel-2-oriented presets are ready-made for RGB-NIR and SWIR-style workflows and are extensible to other sensors.
## What you get out of the box
| Capability | Highlights |
| --- | --- |
-| **Generators & discriminators** | RCAB, RRDB, residual-in-residual, large-kernel attention, PatchGAN, UNet-based discriminators, and more. |
-| **Losses** | Weighted combinations of L1/L2, perceptual (VGG/LPIPS), style, and relativistic adversarial losses. |
-| **Training utilities** | Generator warm-up phases, on-plateau learning-rate schedules, adversarial-weight ramping, EMA tracking, and mixed-precision support. |
-| **Experiment management** | Configurable logging (Weights & Biases, TensorBoard), checkpointing, and experiment reproducibility hooks. |
+| **Generators & discriminators** | SRResNet, RCAB, RRDB, large-kernel attention, ESRGAN, stochastic generators, standard SRGAN, PatchGAN, and ESRGAN discriminators. |
+| **Losses** | Weighted combinations of L1, spectral-angle, perceptual (VGG/LPIPS), total-variation, BCE/Wasserstein, and optional relativistic adversarial losses. |
+| **Training utilities** | Generator warm-up phases, on-plateau learning-rate schedules, adversarial-weight ramping, gradient clipping, and EMA tracking. |
+| **Experiment management** | Weights & Biases or CSV logging, checkpointing, saved resolved configs, and Hydra experiment presets. |
| **Datasets** | Bundled example dataset for quick smoke tests plus a selector designed for custom collections. |
| **Deployment** | PyPI package (`opensr-srgan`) with helpers to load Lightning modules from configs or download pre-trained presets from the Hugging Face Hub. |
@@ -57,16 +52,14 @@ Whether you are reproducing published results, exploring new remote-sensing moda
| `opensr_srgan/data/` | Dataset wrappers and helper utilities for custom and preconfigured SR datasets. |
| `opensr_srgan/configs/` | Ready-to-run YAML presets covering common scale factors, band selections, and architecture pairings. |
| `opensr_srgan/utils/` | Logging helpers, spectral normalisation utilities, and model summary functions used across the stack. |
-| `opensr_srgan/train.py` | Command-line entry point that wires configuration, data module, loggers, and the Lightning trainer together. |
+| `opensr_srgan/train.py` / `opensr_srgan/train_hydra.py` | Legacy YAML and Hydra entry points that wire configuration, data module, loggers, and the Lightning trainer together. |
## Typical workflow
-1. **Pick a configuration.** Start from a preset in `opensr_srgan/configs/` and adapt dataset paths, scale, generator, discriminator, and loss
- options to match your experiment.
+1. **Pick a configuration.** Start from `opensr_srgan/configs/config_training_example.yaml` for a smoke test, or use a Hydra preset such as `experiment=example`, `experiment=10m`, or `experiment=20m` and adapt dataset paths, scale, generator, discriminator, and loss options to match your experiment.
2. **Prepare datasets.** Download the bundled example dataset or register your own source with the dataset selector (see
[Data](data.md)).
-3. **Launch training.** Run `python -m opensr_srgan.train --config ` or import `train` from the package to instantiate the
- Lightning module, configure optimisers and callbacks, and start adversarial training (see [Training](training.md)).
+3. **Launch training.** Run `python -m opensr_srgan.train --config `, `python -m opensr_srgan.train_hydra experiment=example`, or import `train` from the package to instantiate the Lightning module, configure optimisers and callbacks, and start adversarial training (see [Training](training.md)).
4. **Monitor progress.** Use the included Weights & Biases logging to track perceptual losses, adversarial
metrics, and validation imagery.
5. **Deploy or evaluate.** The Lightning module exposes `predict_step` for batched inference, automatically normalising inputs and
diff --git a/docs/inference.md b/docs/inference.md
index 7096d5a..0ac9901 100644
--- a/docs/inference.md
+++ b/docs/inference.md
@@ -6,13 +6,14 @@ This walkthrough covers the fastest path from zero to an end-to-end super-resolu
```bash
pip install opensr-srgan
+pip install opensr-utils # optional, for full-tile inference
```
* `opensr-srgan` exposes helpers that reconstruct Lightning checkpoints from YAML configs or download ready-to-run presets.
-* The optional `huggingface` extra adds `huggingface-hub`, which `load_inference_model` uses internally when fetching preset weights from the Hub.
-* `opensr-utils` provides the tiling/mosaicking pipeline that can super-resolve whole Sentinel-2 SAFE folders, GeoTIFFs, or other large rasters.
+* `huggingface-hub` is installed with `opensr-srgan` and is used by `load_inference_model` when fetching preset weights from the Hub.
+* Install `opensr-utils` separately when you want the large-raster tiling/mosaicking pipeline for Sentinel-2 SAFE folders, GeoTIFFs, or other large rasters.
-## 2.1 Instantiate a Preset
+## 2.1 Instantiate a preset
```python
from opensr_srgan import load_inference_model
@@ -20,24 +21,31 @@ from opensr_srgan import load_inference_model
model = load_inference_model("RGB-NIR", map_location="cuda")
```
-`load_inference_model` retrieves the configuration and checkpoint that correspond to the selected preset (here the four-band RGB-NIR model), restores the Lightning module, and switches it to evaluation mode so that it is ready for inference.【F:opensr_srgan/_factory.py†L107-L150】 If you run on CPU, change `map_location="cuda"` to `map_location="cpu"`.
+`load_inference_model` retrieves the configuration and checkpoint that correspond to the selected preset (here the four-band RGB-NIR model), restores the Lightning module, and switches it to evaluation mode so that it is ready for inference. If you run on CPU, change `map_location="cuda"` to `map_location="cpu"`.
+
+## 2.2 Instantiate your own model
-## 2.2 Instanciate your own Model
```python
from opensr_srgan import load_from_config
model = load_from_config(config_path="YOUR_CONFIG_PATH",
checkpoint_uri="YOUR_CKPT_PATH")
```
+
Using the path to your trained model as well as the config file that was used to train your model, you can load your model for inference.
-## 3.1 Run SR on your Images
-After the model has been created, please use the pytorch-lightning native `predict_step` function. Sticking to the pytorch-lightning workflow includes the selected normalization procedures that the model was trained on, and also seamlessly enables multi-GPU processing and other tweaks. Be aware that this only SRs raw tensors, so the patching and stitching needs to be done manually.
+## 3.1 Run SR on tensors
+
+After the model has been created, use `predict_step` for tensor inference. This applies the normalization procedure configured for the model and returns a CPU tensor. It only processes raw tensors, so patching and stitching must be handled separately for large rasters.
+
```python
sr = model.predict_step(lr)
```
+
## 3.2 Super-resolve a full tile with OpenSR-Utils
-You can build on top of out utility package in order to do all of the stitching, georeferencing and patching automatically (Note: this currently only works for RGB-NIR Sentinel-2 images)
+
+You can build on top of the `opensr-utils` package to handle stitching, georeferencing, and patching automatically. This currently targets RGB-NIR Sentinel-2 inference.
+
```python
import opensr_utils
diff --git a/docs/results.md b/docs/results.md
index e7591e9..fa51d96 100644
--- a/docs/results.md
+++ b/docs/results.md
@@ -1,6 +1,6 @@
# Results
-This page shows some visual results for different configuations and data types.
+This page shows some visual results for different configurations and data types.
## Example 1: RCAB generator + standard discriminator (8× Sentinel-2 20 m → 2.5 m)
@@ -10,7 +10,7 @@ This page shows some visual results for different configuations and data types.
- Upscaling factor: 8×
**Dataset**
-- Sentinel-2 Level-2A tiles using the six 20 m bands (B5, B6, B7, B8A, B11, B12) at 20m
+- Sentinel-2 Level-2A tiles using the six 20 m bands (B5, B6, B7, B8A, B11, B12) at 20 m
- Low-resolution inputs generated via bicubic downsampling to 160 m ground sampling distance
diff --git a/docs/training.md b/docs/training.md
index ae8a561..965b9a9 100644
--- a/docs/training.md
+++ b/docs/training.md
@@ -17,7 +17,7 @@ In order to train, you need a dataset. `Data.dataset_type` decides which dataset
You can launch training from the CLI or by importing the helper inside Python.
```bash
-python opensr_srgan.train --config path/to/config.yaml
+python -m opensr_srgan.train --config path/to/config.yaml
```
```python
@@ -26,11 +26,21 @@ from opensr_srgan import train
train("path/to/config.yaml")
```
-Both entry points accept the same configuration file. The CLI exposes a single optional argument:
+For experiment management, the Hydra entry point composes grouped configs from `opensr_srgan/configs/hydra` and forwards the resolved config to the same `train()` function:
-* `--config / -c`: Path to a YAML file describing the experiment. Defaults to `opensr_srgan/configs/config_20m.yaml`.
+```bash
+python -m opensr_srgan.train_hydra experiment=example
+python -m opensr_srgan.train_hydra experiment=10m Training.max_epochs=5 Logging.wandb.enabled=false
+srgan-train experiment=20m Training.gpus=[0,1]
+```
+
+Use `experiment=...` when you want a complete preset (`example`, `10m`, or `20m`). Use direct overrides such as `Training.max_epochs=5` or `Generator.n_blocks=8` for quick one-off changes. See [Hydra Experiments](hydra.md) for the config layout, override patterns, and output-directory behavior.
+
+Both legacy entry points accept the same configuration file. The legacy CLI exposes a single optional argument:
+
+* `--config / -c`: Path to a YAML file describing the experiment. Defaults to `opensr_srgan/configs/config_10m.yaml`. For a local smoke test with the bundled dataset, pass `opensr_srgan/configs/config_training_example.yaml`.
-GPU assignment is handled directly in the configuration. Set `Training.gpus` to a list of device indices (for example `[0, 1, 2, 3]`) to enable multi-GPU training; a single value such as `[0]` keeps the run on one card. When more than one device is listed the trainer automatically activates PyTorch Lightning's Distributed Data Parallel (DDP) backend for significantly faster epochs.
+GPU assignment is handled directly in the configuration. Set `Training.gpus` to a list of device indices (for example `[0, 1, 2, 3]`) to enable multi-GPU training; a single value such as `[0]` keeps the run on one card. When more than one GPU is listed the trainer activates PyTorch Lightning DDP using `ddp_find_unused_parameters_true` by default, which fits the manual GAN update pattern.
## Initialisation steps - Overview
The code performs the following, regardless of whether the script is launched from the CLI or via import.
@@ -39,7 +49,7 @@ The code performs the following, regardless of whether the script is launched fr
3. **Load configuration.** `OmegaConf.load()` parses the YAML file into an object used throughout the run.
4. **Construct the model.**
* If `Model.load_checkpoint` is set, the script calls `model.load_weights_from_checkpoint()` to import learned weights only while
- respecting the new configuration values. If `Model.continue_training` is passed with a path to a pretrained checkpoint, all scheduler states, epochs and step numbers, EMA weights, etc are loaded in order to seamlessly continue training from a previous run.
+ respecting the new configuration values. If `Model.continue_training` is passed with a path to a Lightning checkpoint, optimizer/scheduler state, epoch counters, global step, and EMA state are restored by Lightning before continuing the run.
* `Model.load_checkpoint` and `Model.continue_training` are mutually exclusive. Use only one, depending on whether you want weight initialization or full training-state resume.
* Otherwise, it initialises a fresh `SRGAN_model`, which immediately builds the generator/discriminator and prints a
parameter summary.
@@ -49,16 +59,14 @@ The code performs the following, regardless of whether the script is launched fr
## Logging setup
* **Weights & Biases.** `WandbLogger` records scalar metrics, adversarial diagnostics, and validation image panels.
-* **TensorBoard.** `TensorBoardLogger` writes the same scalar metrics locally under `logs//`.
-* **Manual SummaryWriter.** A temporary TensorBoard writer (`logs/tmp`) remains available for quick custom logging if needed.
+* **CSV logs.** When `Logging.wandb.enabled: false`, `CSVLogger` writes lightweight local logs under `logs/` or under `Logging.output_dir` for Hydra runs.
-To disable W&B logging, either remove the logger from the list or unset your API key before launching the script.
+To disable W&B logging, set `Logging.wandb.enabled: false` in YAML or override it from Hydra with `Logging.wandb.enabled=false`.
## Metrics
-The Lightning module pushes the same scalar streams to both TensorBoard and W&B so you can monitor convergence from either
-interface. Generator-only pretraining, adversarial training, and the EMA helper each contribute their own indicators, so the
-dashboard quickly reveals which subsystem is active at any given step.
+The Lightning module logs scalar metrics through the active Lightning logger, either W&B or CSV. Generator-only pretraining, adversarial training, and the EMA helper each contribute their own indicators, so the
+logs make it clear which subsystem is active at any given step.
| Metric | Description | Expected behaviour |
| --- | --- | --- |
@@ -66,7 +74,7 @@ dashboard quickly reveals which subsystem is active at any given step.
| `discriminator/adversarial_loss` | Binary cross-entropy loss of the discriminator on real vs. fake batches. | Drops below ~0.7 as the discriminator learns; continues trending down when D keeps up. |
| `discriminator/D(y)_prob` | Mean discriminator confidence that HR inputs are real. | Rises toward 0.8–1.0 during stable training. |
| `discriminator/D(G(x))_prob` | Mean discriminator confidence that SR predictions are real. | Starts low (~0.0–0.2) and climbs toward 0.5 as the generator improves. |
-| `train_metrics/l1` | Mean absolute error between SR and HR tensors. In generator-only pretraining this is the hardwired optimization target. | Decreases toward 0 as reconstructions sharpen. |
+| `train_metrics/l1` | Mean absolute error between SR and HR tensors and one component of the configured content objective. | Decreases toward 0 as reconstructions sharpen. |
| `train_metrics/sam` | Spectral angle mapper (radians) averaged over pixels. | Falls toward 0; values <0.1 indicate strong spectral fidelity. |
| `train_metrics/perceptual` | Perceptual distance (VGG or LPIPS) on selected RGB bands. | Decreases as textures align; exact range depends on the chosen metric. |
| `train_metrics/tv` | Total variation penalty capturing SR smoothness. | Remains small; near-zero means little high-frequency noise. |
@@ -99,23 +107,19 @@ The following callbacks are registered with the Lightning trainer:
| Callback | Purpose |
| --- | --- |
-| `ModelCheckpoint` | Saves the top two checkpoints according to `Schedulers.metric` and always keeps the last epoch. |
-| `LearningRateMonitor` | Logs learning rates for both optimisers every epoch. |
+| `ModelCheckpoint` | Saves the top two checkpoints according to `Schedulers.metric_g` and always keeps the last epoch. |
| `EarlyStopping` | Monitors the same metric as the schedulers with a patience of 250 epochs and finite-check enabled. |
-Checkpoint directories are nested under the TensorBoard log folder using the W&B project name and a timestamp, making it easy to
-correlate files across tooling.
+Checkpoint directories default to `logs//`. Hydra runs set `Logging.output_dir`, so checkpoints and the resolved `config.yaml` land in Hydra's run directory.
## Trainer configuration
The script builds a `Trainer` with the following notable arguments:
-* `accelerator='cuda'` with `devices=config.Training.gpus`. When more than one device index is provided the script selects the
- `ddp` strategy automatically, so scaling across multiple GPUs is as simple as enumerating them in the config.
-* `check_val_every_n_epoch=1` to evaluate after every epoch.
+* `accelerator='gpu'` with `devices=config.Training.gpus` for CUDA/GPU runs, or `accelerator='cpu'` with `devices=1` for CPU runs. When more than one GPU is requested, the script selects `ddp_find_unused_parameters_true` by default because manual GAN optimisation can leave one branch unused on a given step.
* `limit_val_batches=250` as a safeguard against excessive validation time on large datasets.
-* `logger=[wandb_logger]` to register external logging backends (add `tb_logger` if you prefer TensorBoard-driven monitoring).
-* `callbacks=[checkpoint_callback, early_stop_callback, lr_monitor]` to activate the components described above.
+* `logger=[wandb_logger]` to register the active W&B or CSV logger.
+* `callbacks=[checkpoint_callback, early_stop_callback]` to activate checkpointing and finite-check early stopping. Learning rates are logged by the Lightning module at the end of each training batch.
Finally, `trainer.fit(model, datamodule=pl_datamodule)` launches the optimisation loop and `wandb.finish()` ensures clean shutdown
of the W&B session.
diff --git a/mkdocs.yml b/mkdocs.yml
index 7eb3215..1d6450c 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -1,7 +1,7 @@
site_name: "SRGAN Docs"
site_description: "OpenSR-SRGAN: Flexible Super-Resolution GANs for satellite imagery"
site_url: https://srgan.opensr.eu/
-copyright: "2024-2025, Simon Donike et al."
+copyright: "2024-2026, Simon Donike et al."
repo_url: https://github.com/ESAOpenSR/SRGAN
repo_name: ESAOpenSR/SRGAN
edit_uri: edit/main/docs/
@@ -42,7 +42,11 @@ plugins:
handlers:
python:
options:
- docstring_style: auto
+ docstring_style: google
+ docstring_options:
+ warnings: false
+ warn_missing_types: false
+ warn_unknown_params: false
show_source: true
markdown_extensions:
@@ -92,6 +96,7 @@ nav:
- Inference Guide: inference.md
- HPC Inference: hpc-inference.md
- Configuration: configuration.md
+ - Hydra Experiments: hydra.md
- Data Preparation: data.md
- Training:
diff --git a/opensr_srgan/_factory.py b/opensr_srgan/_factory.py
index e9ec5dc..600a058 100644
--- a/opensr_srgan/_factory.py
+++ b/opensr_srgan/_factory.py
@@ -8,7 +8,7 @@
Typical usage
-------------
->>> from opensr_srgan.model.loading import load_inference_model
+>>> from opensr_srgan import load_inference_model
>>> model = load_inference_model("RGB-NIR")
>>> model.eval()
>>> sr = model(torch.randn(1, 4, 64, 64))
diff --git a/opensr_srgan/configs/__init__.py b/opensr_srgan/configs/__init__.py
new file mode 100644
index 0000000..300ac00
--- /dev/null
+++ b/opensr_srgan/configs/__init__.py
@@ -0,0 +1 @@
+"""Configuration resources for opensr_srgan."""
diff --git a/opensr_srgan/configs/hydra/__init__.py b/opensr_srgan/configs/hydra/__init__.py
new file mode 100644
index 0000000..300ac00
--- /dev/null
+++ b/opensr_srgan/configs/hydra/__init__.py
@@ -0,0 +1 @@
+"""Configuration resources for opensr_srgan."""
diff --git a/opensr_srgan/configs/hydra/data/example.yaml b/opensr_srgan/configs/hydra/data/example.yaml
new file mode 100644
index 0000000..fcf1337
--- /dev/null
+++ b/opensr_srgan/configs/hydra/data/example.yaml
@@ -0,0 +1,9 @@
+# @package _global_
+
+Data:
+ train_batch_size: 2
+ val_batch_size: 2
+ num_workers: 1
+ prefetch_factor: 2
+ dataset_type: ExampleDataset
+ normalization: normalise_10k
diff --git a/opensr_srgan/configs/hydra/data/lrhr_folder.yaml b/opensr_srgan/configs/hydra/data/lrhr_folder.yaml
new file mode 100644
index 0000000..6bf334c
--- /dev/null
+++ b/opensr_srgan/configs/hydra/data/lrhr_folder.yaml
@@ -0,0 +1,10 @@
+# @package _global_
+
+Data:
+ train_batch_size: 2
+ val_batch_size: 2
+ num_workers: 1
+ prefetch_factor: 2
+ dataset_type: LRHRFolderDataset
+ root_dir: null
+ normalization: normalise_10k
diff --git a/opensr_srgan/configs/hydra/data/sen2naip.yaml b/opensr_srgan/configs/hydra/data/sen2naip.yaml
new file mode 100644
index 0000000..c11639a
--- /dev/null
+++ b/opensr_srgan/configs/hydra/data/sen2naip.yaml
@@ -0,0 +1,9 @@
+# @package _global_
+
+Data:
+ train_batch_size: 6
+ val_batch_size: 4
+ num_workers: 6
+ prefetch_factor: 2
+ dataset_type: SEN2NAIP
+ normalization: normalise_10k
diff --git a/opensr_srgan/configs/hydra/experiment/10m.yaml b/opensr_srgan/configs/hydra/experiment/10m.yaml
new file mode 100644
index 0000000..4a84019
--- /dev/null
+++ b/opensr_srgan/configs/hydra/experiment/10m.yaml
@@ -0,0 +1,8 @@
+# @package _global_
+
+# Reproduces opensr_srgan/configs/config_10m.yaml.
+defaults:
+ - override /data: sen2naip
+ - override /model: 10m_esrgan
+ - override /training: 10m
+ - override /logging: wandb_10m
diff --git a/opensr_srgan/configs/hydra/experiment/20m.yaml b/opensr_srgan/configs/hydra/experiment/20m.yaml
new file mode 100644
index 0000000..ff17ad2
--- /dev/null
+++ b/opensr_srgan/configs/hydra/experiment/20m.yaml
@@ -0,0 +1,16 @@
+# @package _global_
+
+# Reproduces opensr_srgan/configs/config_20m.yaml.
+defaults:
+ - override /data: sen2naip
+ - override /model: 20m_rcab
+ - override /training: 20m
+ - override /logging: wandb_20m
+
+Data:
+ train_batch_size: 24
+ val_batch_size: 8
+ num_workers: 8
+ prefetch_factor: 4
+ dataset_type: S2_6b
+ normalization: normalise_10k
diff --git a/opensr_srgan/configs/hydra/experiment/example.yaml b/opensr_srgan/configs/hydra/experiment/example.yaml
new file mode 100644
index 0000000..20559e7
--- /dev/null
+++ b/opensr_srgan/configs/hydra/experiment/example.yaml
@@ -0,0 +1,8 @@
+# @package _global_
+
+# Reproduces opensr_srgan/configs/config_training_example.yaml.
+defaults:
+ - override /data: example
+ - override /model: example_srresnet
+ - override /training: example
+ - override /logging: csv
diff --git a/opensr_srgan/configs/hydra/hydra/default.yaml b/opensr_srgan/configs/hydra/hydra/default.yaml
new file mode 100644
index 0000000..c3186d3
--- /dev/null
+++ b/opensr_srgan/configs/hydra/hydra/default.yaml
@@ -0,0 +1,14 @@
+# @package hydra
+
+# Keep Hydra runs reproducible without changing the process working directory.
+defaults:
+ - override hydra_logging: default
+ - override job_logging: default
+
+run:
+ dir: logs/${Logging.wandb.project}/${now:%Y-%m-%d_%H-%M-%S}
+sweep:
+ dir: logs/${Logging.wandb.project}/multiruns/${now:%Y-%m-%d_%H-%M-%S}
+ subdir: ${hydra.job.num}
+job:
+ chdir: false
diff --git a/opensr_srgan/configs/hydra/logging/csv.yaml b/opensr_srgan/configs/hydra/logging/csv.yaml
new file mode 100644
index 0000000..c72c363
--- /dev/null
+++ b/opensr_srgan/configs/hydra/logging/csv.yaml
@@ -0,0 +1,8 @@
+# @package _global_
+
+Logging:
+ num_val_images: 5
+ wandb:
+ enabled: false
+ entity: opensr
+ project: SRGAN_10m
diff --git a/opensr_srgan/configs/hydra/logging/wandb_10m.yaml b/opensr_srgan/configs/hydra/logging/wandb_10m.yaml
new file mode 100644
index 0000000..9ce4584
--- /dev/null
+++ b/opensr_srgan/configs/hydra/logging/wandb_10m.yaml
@@ -0,0 +1,8 @@
+# @package _global_
+
+Logging:
+ num_val_images: 5
+ wandb:
+ enabled: true
+ entity: opensr
+ project: SRGAN_10m
diff --git a/opensr_srgan/configs/hydra/logging/wandb_20m.yaml b/opensr_srgan/configs/hydra/logging/wandb_20m.yaml
new file mode 100644
index 0000000..5e03eab
--- /dev/null
+++ b/opensr_srgan/configs/hydra/logging/wandb_20m.yaml
@@ -0,0 +1,8 @@
+# @package _global_
+
+Logging:
+ num_val_images: 5
+ wandb:
+ enabled: true
+ entity: opensr
+ project: SRGAN_20m
diff --git a/opensr_srgan/configs/hydra/model/10m_esrgan.yaml b/opensr_srgan/configs/hydra/model/10m_esrgan.yaml
new file mode 100644
index 0000000..b42155f
--- /dev/null
+++ b/opensr_srgan/configs/hydra/model/10m_esrgan.yaml
@@ -0,0 +1,25 @@
+# @package _global_
+
+Model:
+ in_bands: 4
+ continue_training: false
+ load_checkpoint: false
+
+Generator:
+ model_type: esrgan
+ block_type: rrdb
+ large_kernel_size: 9
+ small_kernel_size: 3
+ n_channels: 64
+ n_blocks: 16
+ scaling_factor: 4
+ growth_channels: 32
+ res_scale: 0.2
+ use_icnr: true
+
+Discriminator:
+ model_type: esrgan
+ n_blocks: 8
+ use_spectral_norm: false
+ base_channels: 64
+ linear_size: 1024
diff --git a/opensr_srgan/configs/hydra/model/20m_rcab.yaml b/opensr_srgan/configs/hydra/model/20m_rcab.yaml
new file mode 100644
index 0000000..70800c5
--- /dev/null
+++ b/opensr_srgan/configs/hydra/model/20m_rcab.yaml
@@ -0,0 +1,24 @@
+# @package _global_
+
+Model:
+ in_bands: 6
+ continue_training: false
+ load_checkpoint: false
+
+Generator:
+ model_type: SRResNet
+ block_type: rcab
+ large_kernel_size: 9
+ small_kernel_size: 3
+ n_channels: 96
+ n_blocks: 32
+ scaling_factor: 8
+ growth_channels: 32
+ res_scale: 0.2
+
+Discriminator:
+ model_type: standard
+ n_blocks: 8
+ use_spectral_norm: false
+ base_channels: 64
+ linear_size: 1024
diff --git a/opensr_srgan/configs/hydra/model/example_srresnet.yaml b/opensr_srgan/configs/hydra/model/example_srresnet.yaml
new file mode 100644
index 0000000..6e8bc01
--- /dev/null
+++ b/opensr_srgan/configs/hydra/model/example_srresnet.yaml
@@ -0,0 +1,24 @@
+# @package _global_
+
+Model:
+ in_bands: 4
+ continue_training: false
+ load_checkpoint: false
+
+Generator:
+ model_type: SRResNet
+ block_type: rrdb
+ large_kernel_size: 9
+ small_kernel_size: 3
+ n_channels: 32
+ n_blocks: 4
+ scaling_factor: 4
+ growth_channels: 32
+ res_scale: 0.2
+
+Discriminator:
+ model_type: standard
+ n_blocks: 2
+ use_spectral_norm: false
+ base_channels: 32
+ linear_size: 1024
diff --git a/opensr_srgan/configs/hydra/train.yaml b/opensr_srgan/configs/hydra/train.yaml
new file mode 100644
index 0000000..c8533f8
--- /dev/null
+++ b/opensr_srgan/configs/hydra/train.yaml
@@ -0,0 +1,11 @@
+# @package _global_
+
+# Composable Hydra entry point for SRGAN training.
+defaults:
+ - _self_
+ - data: example
+ - model: example_srresnet
+ - training: cpu_debug
+ - logging: csv
+ - hydra: default
+ - experiment: null
diff --git a/opensr_srgan/configs/hydra/training/10m.yaml b/opensr_srgan/configs/hydra/training/10m.yaml
new file mode 100644
index 0000000..25f5853
--- /dev/null
+++ b/opensr_srgan/configs/hydra/training/10m.yaml
@@ -0,0 +1,50 @@
+# @package _global_
+
+Training:
+ device: cuda
+ gpus: [0, 1]
+ max_epochs: 9999
+ val_check_interval: 1.0
+ limit_val_batches: 100
+ pretrain_g_only: true
+ g_pretrain_steps: 5000
+ adv_loss_ramp_steps: 500
+ label_smoothing: true
+ EMA:
+ enabled: false
+ decay: 0.999
+ update_after_step: 0
+ use_num_updates: true
+ Losses:
+ adv_loss_type: bce
+ relativistic_average_d: false
+ adv_loss_beta: 0.001
+ adv_loss_schedule: cosine
+ r1_gamma: 0.0
+ l1_weight: 1.0
+ sam_weight: 0.00
+ perceptual_weight: 0.05
+ perceptual_metric: vgg
+ fixed_idx: [0, 1, 2]
+ tv_weight: 0.0
+ max_val: 1.0
+ ssim_win: 11
+
+Optimizers:
+ optim_g_lr: 1e-5
+ optim_d_lr: 5e-5
+ gradient_clip_val: 0
+ betas: [0.0, 0.99]
+ eps: 1.0e-7
+ weight_decay_g: 0.0
+ weight_decay_d: 0.0
+
+Schedulers:
+ g_warmup_steps: 0
+ g_warmup_type: linear
+ metric_g: val_metrics/l1
+ metric_d: validation/DISC_adversarial_loss
+ patience_g: 10
+ patience_d: 10
+ factor_g: 0.5
+ factor_d: 0.5
diff --git a/opensr_srgan/configs/hydra/training/20m.yaml b/opensr_srgan/configs/hydra/training/20m.yaml
new file mode 100644
index 0000000..f4b056f
--- /dev/null
+++ b/opensr_srgan/configs/hydra/training/20m.yaml
@@ -0,0 +1,49 @@
+# @package _global_
+
+Training:
+ device: cuda
+ gpus: [2, 3]
+ max_epochs: 9999
+ val_check_interval: 1.0
+ limit_val_batches: 250
+ pretrain_g_only: true
+ g_pretrain_steps: 15000
+ adv_loss_ramp_steps: 2500
+ label_smoothing: true
+ EMA:
+ enabled: false
+ decay: 0.999
+ update_after_step: 0
+ use_num_updates: true
+ Losses:
+ adv_loss_type: bce
+ relativistic_average_d: false
+ adv_loss_beta: 1e-3
+ adv_loss_schedule: cosine
+ r1_gamma: 0.0
+ l1_weight: 1.0
+ sam_weight: 0.05
+ perceptual_weight: 0.1
+ perceptual_metric: vgg
+ tv_weight: 0.0
+ max_val: 1.0
+ ssim_win: 11
+
+Optimizers:
+ optim_g_lr: 1e-4
+ optim_d_lr: 1e-6
+ gradient_clip_val: 1.0
+ betas: [0.0, 0.99]
+ eps: 1.0e-7
+ weight_decay_g: 0.0
+ weight_decay_d: 0.0
+
+Schedulers:
+ g_warmup_steps: 2500
+ g_warmup_type: cosine
+ metric_g: val_metrics/l1
+ metric_d: validation/DISC_adversarial_loss
+ patience_g: 50
+ patience_d: 50
+ factor_g: 0.5
+ factor_d: 0.5
diff --git a/opensr_srgan/configs/hydra/training/cpu_debug.yaml b/opensr_srgan/configs/hydra/training/cpu_debug.yaml
new file mode 100644
index 0000000..8fe2c1a
--- /dev/null
+++ b/opensr_srgan/configs/hydra/training/cpu_debug.yaml
@@ -0,0 +1,49 @@
+# @package _global_
+
+Training:
+ device: cpu
+ gpus: []
+ max_epochs: 1
+ val_check_interval: 1.0
+ limit_val_batches: 1
+ pretrain_g_only: true
+ g_pretrain_steps: 10
+ adv_loss_ramp_steps: 10
+ label_smoothing: true
+ EMA:
+ enabled: false
+ decay: 0.999
+ update_after_step: 0
+ use_num_updates: true
+ Losses:
+ adv_loss_type: bce
+ relativistic_average_d: false
+ adv_loss_beta: 0.001
+ adv_loss_schedule: cosine
+ r1_gamma: 0.0
+ l1_weight: 1.0
+ sam_weight: 0.05
+ perceptual_weight: 0.2
+ perceptual_metric: vgg
+ tv_weight: 0.0
+ max_val: 1.0
+ ssim_win: 11
+
+Optimizers:
+ optim_g_lr: 1e-4
+ optim_d_lr: 1e-6
+ gradient_clip_val: 1.0
+ betas: [0.0, 0.99]
+ eps: 1.0e-7
+ weight_decay_g: 0.0
+ weight_decay_d: 0.0
+
+Schedulers:
+ g_warmup_steps: 10
+ g_warmup_type: cosine
+ metric_g: val_metrics/l1
+ metric_d: validation/DISC_adversarial_loss
+ patience_g: 10
+ patience_d: 10
+ factor_g: 0.5
+ factor_d: 0.5
diff --git a/opensr_srgan/configs/hydra/training/example.yaml b/opensr_srgan/configs/hydra/training/example.yaml
new file mode 100644
index 0000000..f50e152
--- /dev/null
+++ b/opensr_srgan/configs/hydra/training/example.yaml
@@ -0,0 +1,49 @@
+# @package _global_
+
+Training:
+ device: cuda
+ gpus: [0]
+ max_epochs: 10
+ val_check_interval: 0.25
+ limit_val_batches: 250
+ pretrain_g_only: true
+ g_pretrain_steps: 1000
+ adv_loss_ramp_steps: 500
+ label_smoothing: true
+ EMA:
+ enabled: false
+ decay: 0.999
+ update_after_step: 0
+ use_num_updates: true
+ Losses:
+ adv_loss_type: bce
+ relativistic_average_d: false
+ adv_loss_beta: 0.001
+ adv_loss_schedule: cosine
+ r1_gamma: 0.0
+ l1_weight: 1.0
+ sam_weight: 0.05
+ perceptual_weight: 0.2
+ perceptual_metric: vgg
+ tv_weight: 0.0
+ max_val: 1.0
+ ssim_win: 11
+
+Optimizers:
+ optim_g_lr: 1e-4
+ optim_d_lr: 1e-6
+ gradient_clip_val: 1.0
+ betas: [0.0, 0.99]
+ eps: 1.0e-7
+ weight_decay_g: 0.0
+ weight_decay_d: 0.0
+
+Schedulers:
+ g_warmup_steps: 10
+ g_warmup_type: cosine
+ metric_g: val_metrics/l1
+ metric_d: validation/DISC_adversarial_loss
+ patience_g: 10
+ patience_d: 10
+ factor_g: 0.5
+ factor_d: 0.5
diff --git a/opensr_srgan/train.py b/opensr_srgan/train.py
index 6a75b02..d324404 100644
--- a/opensr_srgan/train.py
+++ b/opensr_srgan/train.py
@@ -16,6 +16,26 @@
import pytorch_lightning as pl
+def _is_path_set(value) -> bool:
+ """Return True when an optional path-like config value should be used."""
+
+ return value not in (False, None, "")
+
+
+def _resolve_experiment_dir(config) -> str:
+ """Resolve the directory used for checkpoints and the saved config."""
+
+ configured_output_dir = getattr(config.Logging, "output_dir", None)
+ if _is_path_set(configured_output_dir):
+ return os.path.normpath(os.path.expanduser(str(configured_output_dir)))
+
+ return os.path.join(
+ os.path.normpath("logs/"),
+ config.Logging.wandb.project,
+ datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S"),
+ )
+
+
def train(config):
"""
Train SRGAN from a configuration.
@@ -91,31 +111,36 @@ def _checkpoint_is_set(value) -> bool:
""" Configure Trainer """
#############################################################################################################
+ output_dir_is_set = _is_path_set(getattr(config.Logging, "output_dir", None))
+ experiment_dir = _resolve_experiment_dir(config)
+
# Configure Logger
if config.Logging.wandb.enabled:
# set up logging
from pytorch_lightning.loggers import WandbLogger
wandb_project = config.Logging.wandb.project # whatever you want
- wandb_logger = WandbLogger(
- project=wandb_project, entity=config.Logging.wandb.entity, log_model=False
- )
+ wandb_kwargs = {
+ "project": wandb_project,
+ "entity": config.Logging.wandb.entity,
+ "log_model": False,
+ }
+ if output_dir_is_set:
+ wandb_kwargs["save_dir"] = experiment_dir
+ wandb_logger = WandbLogger(**wandb_kwargs)
else:
print("Not using Weights & Biases logging, reduced CSV logs written locally.")
from pytorch_lightning.loggers import CSVLogger
+ csv_save_dir = experiment_dir if output_dir_is_set else "logs/"
wandb_logger = CSVLogger(
- save_dir="logs/",
+ save_dir=csv_save_dir,
)
# Configure Saving Checkpoints
from pytorch_lightning.callbacks import ModelCheckpoint
- dir_save_checkpoints = os.path.join(
- os.path.normpath("logs/"),
- config.Logging.wandb.project,
- datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S"),
- )
+ dir_save_checkpoints = experiment_dir
from opensr_srgan.utils.gpu_rank import (
_is_global_zero,
) # make dir only on main process
diff --git a/opensr_srgan/train_hydra.py b/opensr_srgan/train_hydra.py
new file mode 100644
index 0000000..4706039
--- /dev/null
+++ b/opensr_srgan/train_hydra.py
@@ -0,0 +1,76 @@
+"""Hydra entry point for SRGAN training.
+
+This module composes grouped Hydra configs and forwards the resolved OmegaConf
+object to the existing :func:`opensr_srgan.train.train` function.
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+
+import hydra
+import torch
+from hydra.core.hydra_config import HydraConfig
+from omegaconf import DictConfig, OmegaConf, open_dict
+
+
+def _patch_python314_argparse_help_validation() -> None:
+ """Allow Hydra lazy completion help objects on Python 3.14+."""
+
+ if sys.version_info < (3, 14):
+ return
+
+ original_check_help = argparse.ArgumentParser._check_help
+ if getattr(original_check_help, "_srgan_hydra_compat", False):
+ return
+
+ def _check_help(self, action):
+ if action.help is not None and not isinstance(action.help, str):
+ return
+ return original_check_help(self, action)
+
+ _check_help._srgan_hydra_compat = True
+ argparse.ArgumentParser._check_help = _check_help
+
+
+_patch_python314_argparse_help_validation()
+
+
+def _has_value(value) -> bool:
+ return value not in (False, None, "")
+
+
+def _attach_hydra_output_dir(cfg: DictConfig) -> DictConfig:
+ """Populate ``Logging.output_dir`` from Hydra runtime when it is unset."""
+
+ output_dir = OmegaConf.select(cfg, "Logging.output_dir", default=None)
+ if _has_value(output_dir):
+ return cfg
+
+ try:
+ hydra_output_dir = HydraConfig.get().runtime.output_dir
+ except ValueError:
+ return cfg
+
+ with open_dict(cfg):
+ if "Logging" not in cfg:
+ cfg.Logging = {}
+ cfg.Logging.output_dir = hydra_output_dir
+ return cfg
+
+
+@hydra.main(version_base="1.3", config_path="configs/hydra", config_name="train")
+def main(cfg: DictConfig) -> None:
+ """Compose Hydra config and launch the existing SRGAN trainer."""
+
+ torch.set_float32_matmul_precision("medium")
+ cfg = _attach_hydra_output_dir(cfg)
+
+ from opensr_srgan.train import train
+
+ train(cfg)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/pyproject.toml b/pyproject.toml
index 20bdde8..3de70ea 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -11,7 +11,7 @@ authors = [
{ name = "Simon Donike, ESAOpenSR Team" }
]
license = { file = "LICENSE" }
-requires-python = ">=3.10,<3.13"
+requires-python = ">=3.12,<3.15"
dependencies = [
"torch>=2.1",
"pytorch-lightning>=2.1,<3.0",
@@ -19,6 +19,7 @@ dependencies = [
"numpy",
"kornia",
"wandb",
+ "hydra-core>=1.3,<2",
"omegaconf",
"matplotlib",
"rasterio",
@@ -29,9 +30,9 @@ dependencies = [
keywords = ["super-resolution", "gan", "remote-sensing", "pytorch", "lightning"]
classifiers = [
"Programming Language :: Python",
- "Programming Language :: Python :: 3.10",
- "Programming Language :: Python :: 3.11",
"Programming Language :: Python :: 3.12",
+ "Programming Language :: Python :: 3.13",
+ "Programming Language :: Python :: 3.14",
"License :: OSI Approved :: Apache Software License",
"Intended Audience :: Science/Research",
"Topic :: Scientific/Engineering :: Artificial Intelligence",
@@ -66,7 +67,14 @@ omit = [
]
[project.optional-dependencies]
-tests = ["pytest", "pytest-cov","numpy"]
+tests = [
+ "pytest",
+ "pytest-cov",
+ "numpy",
+ "pyproj>=3.4",
+ "pyshp>=2.3.1",
+ "shapely>=2.0.2",
+]
docs = ["mkdocs-material", "mkdocstrings[python]", "pymdown-extensions"]
hpc = [
"PyYAML>=6.0",
@@ -80,3 +88,4 @@ hpc = [
[project.scripts]
srgan-hpc = "deployment.srgan_hpc.cli:main"
+srgan-train = "opensr_srgan.train_hydra:main"
diff --git a/requirements.txt b/requirements.txt
index 1417e71..406a9c1 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -5,6 +5,7 @@ torchvision
numpy
kornia
wandb
+hydra-core>=1.3,<2
omegaconf
matplotlib
rasterio
diff --git a/tests/test_deployment/test_deployment_cli_slurm_config_paths.py b/tests/test_deployment/test_deployment_cli_slurm_config_paths.py
new file mode 100644
index 0000000..298a3f4
--- /dev/null
+++ b/tests/test_deployment/test_deployment_cli_slurm_config_paths.py
@@ -0,0 +1,293 @@
+from __future__ import annotations
+
+import json
+import subprocess
+import sys
+import types
+from pathlib import Path
+
+import pytest
+
+import deployment.srgan_hpc.cli as cli_module
+from deployment.srgan_hpc.aoi import AoiSelection
+from deployment.srgan_hpc.config import (
+ EnvironmentConfig,
+ RuntimeConfig,
+ SlurmConfig,
+ get_product_config,
+ load_runtime_config,
+)
+from deployment.srgan_hpc.manifests import read_yaml
+from deployment.srgan_hpc.patching import Patch
+from deployment.srgan_hpc.slurm import (
+ SlurmJobSpec,
+ build_sbatch_command,
+ parse_job_id,
+ submit_job,
+)
+
+
+def test_cli_grid_aoi_and_collect_commands(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+ capsys: pytest.CaptureFixture[str],
+) -> None:
+ config_path = tmp_path / "runtime.yaml"
+ config_path.write_text(
+ f"output_root: {tmp_path / 'runs'}\nproject_name: cli\nmode: rgbnir\n",
+ encoding="utf-8",
+ )
+
+ monkeypatch.setattr(
+ sys,
+ "argv",
+ [
+ "srgan-hpc",
+ "submit",
+ "grid",
+ "--config",
+ str(config_path),
+ "--output-root",
+ str(tmp_path / "override-runs"),
+ "--project-name",
+ "cli-grid",
+ "--start-date",
+ "2025-01-01",
+ "--end-date",
+ "2025-01-02",
+ "--lat1",
+ "45.0",
+ "--lon1",
+ "9.0",
+ "--lat2",
+ "45.001",
+ "--lon2",
+ "9.001",
+ "--script-path",
+ "/tmp/slurm.sh",
+ "--dry-run",
+ ],
+ )
+ assert cli_module.main() == 0
+ grid_output = capsys.readouterr().out
+ assert '"patches": 1' in grid_output
+ assert "override-runs" in grid_output
+
+ patch = Patch(
+ patch_id="patch_000001",
+ latitude=45.0,
+ longitude=9.0,
+ edge_size=512,
+ row_index=0,
+ row_count=1,
+ column_index=0,
+ column_count=1,
+ )
+ monkeypatch.setattr(
+ "deployment.srgan_hpc.aoi.select_aoi_patches",
+ lambda **_kwargs: AoiSelection(
+ aoi_path=tmp_path / "area.shp",
+ geometry=types.SimpleNamespace(),
+ patches=[patch],
+ ),
+ )
+ monkeypatch.setattr(
+ sys,
+ "argv",
+ [
+ "srgan-hpc",
+ "submit",
+ "aoi",
+ "--config",
+ str(config_path),
+ "--start-date",
+ "2025-01-01",
+ "--end-date",
+ "2025-01-02",
+ "--aoi-path",
+ str(tmp_path / "area.shp"),
+ "--script-path",
+ "/tmp/slurm.sh",
+ "--dry-run",
+ ],
+ )
+ assert cli_module.main() == 0
+ aoi_output = capsys.readouterr().out
+ assert '"patches": 1' in aoi_output
+
+ run_dir = tmp_path / "collect-run"
+ source = run_dir / "patches" / "patch_000001" / "outputs" / "fused_sr.tif"
+ source.parent.mkdir(parents=True)
+ source.write_bytes(b"output")
+ dest = tmp_path / "collected"
+ monkeypatch.setattr(
+ sys,
+ "argv",
+ ["srgan-hpc", "collect", "--run-dir", str(run_dir), "--dest", str(dest)],
+ )
+ assert cli_module.main() == 0
+ collect_payload = json.loads(capsys.readouterr().out)
+ assert collect_payload["copied"] == 1
+ assert Path(collect_payload["destination"]) == dest.resolve()
+
+
+def test_slurm_command_includes_optional_resources() -> None:
+ base = {
+ "job_name": "srgan",
+ "script_path": Path("/tmp/slurm.sh"),
+ "manifest_path": Path("/tmp/manifest.yaml"),
+ "output_path": Path("/tmp/out.log"),
+ "error_path": Path("/tmp/err.log"),
+ "environment": EnvironmentConfig(
+ python_executable="/work/envs/srgan/bin/python",
+ modules=["cuda", "gdal"],
+ conda_env="srgan",
+ ),
+ }
+
+ command = build_sbatch_command(
+ SlurmJobSpec(
+ **base,
+ slurm=SlurmConfig(
+ partition="gpu",
+ gres="gpu:a100:1",
+ account="proj",
+ qos="normal",
+ extra_args=["--exclusive"],
+ ),
+ array="0-3",
+ )
+ )
+ assert "--partition=gpu" in command
+ assert "--gres=gpu:a100:1" in command
+ assert "--account=proj" in command
+ assert "--qos=normal" in command
+ assert "--array=0-3" in command
+ assert "--exclusive" in command
+ assert any("SRGAN_HPC_MODULES=cuda,gdal" in part for part in command)
+ assert any("SRGAN_HPC_CONDA_ENV=srgan" in part for part in command)
+
+ gpu_command = build_sbatch_command(
+ SlurmJobSpec(
+ **base,
+ slurm=SlurmConfig(gpu_type="A100", gpus=2),
+ )
+ )
+ assert "--gpus=A100:2" in gpu_command
+
+ plain_gpu_command = build_sbatch_command(
+ SlurmJobSpec(
+ **base,
+ slurm=SlurmConfig(gpus=1),
+ )
+ )
+ assert "--gpus=1" in plain_gpu_command
+
+
+def test_submit_job_success_and_error_paths(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ spec = SlurmJobSpec(
+ job_name="srgan",
+ script_path=Path("/tmp/slurm.sh"),
+ manifest_path=Path("/tmp/manifest.yaml"),
+ output_path=Path("/tmp/out.log"),
+ error_path=Path("/tmp/err.log"),
+ slurm=SlurmConfig(gpus=0),
+ environment=EnvironmentConfig(python_executable="python"),
+ )
+ captured: dict[str, object] = {}
+
+ def fake_run(cmd, check, capture_output, text):
+ captured["cmd"] = cmd
+ captured["check"] = check
+ captured["capture_output"] = capture_output
+ captured["text"] = text
+ return types.SimpleNamespace(
+ stdout="Submitted batch job 12345\n",
+ stderr="",
+ )
+
+ monkeypatch.setattr("deployment.srgan_hpc.slurm.subprocess.run", fake_run)
+
+ payload = submit_job(spec, tmp_path / "success")
+
+ assert payload["job_id"] == "12345"
+ assert captured["cmd"][0] == "sbatch"
+ assert (
+ json.loads(
+ (tmp_path / "success" / "slurm_job_ids.json").read_text(encoding="utf-8")
+ )["job_id"]
+ == "12345"
+ )
+
+ def fake_error(cmd, check, capture_output, text):
+ raise subprocess.CalledProcessError(
+ returncode=7,
+ cmd=cmd,
+ output="stdout text",
+ stderr="stderr text",
+ )
+
+ monkeypatch.setattr("deployment.srgan_hpc.slurm.subprocess.run", fake_error)
+
+ with pytest.raises(RuntimeError, match="sbatch submission failed"):
+ submit_job(spec, tmp_path / "error")
+
+ error_payload = json.loads(
+ (tmp_path / "error" / "slurm_job_ids.json").read_text(encoding="utf-8")
+ )
+ assert error_payload["mode"] == "error"
+ assert error_payload["returncode"] == "7"
+ assert error_payload["stderr"] == "stderr text"
+
+
+def test_parse_job_id_rejects_empty_output() -> None:
+ with pytest.raises(ValueError, match="Could not parse"):
+ parse_job_id("")
+
+
+def test_runtime_config_extra_loading_paths(tmp_path: Path) -> None:
+ with pytest.raises(ValueError, match="Unknown product"):
+ get_product_config(RuntimeConfig(), "unknown")
+
+ non_mapping = tmp_path / "bad.yaml"
+ non_mapping.write_text("- item\n", encoding="utf-8")
+ with pytest.raises(ValueError, match="Expected mapping"):
+ load_runtime_config(non_mapping)
+
+ model_config = tmp_path / "model.yaml"
+ model_checkpoint = tmp_path / "model.ckpt"
+ aoi_path = tmp_path / "area.shp"
+ model_config.write_text("Model: {}\n", encoding="utf-8")
+ model_checkpoint.write_bytes(b"checkpoint")
+ aoi_path.touch()
+ config_path = tmp_path / "runtime.yaml"
+ config_path.write_text(
+ """
+mode: rgbnir
+output_root: runs
+aoi:
+ path: area.shp
+rgbnir:
+ config_path: model.yaml
+ checkpoint_path: model.ckpt
+ model:
+ preset: null
+ cache_dir: cache
+""",
+ encoding="utf-8",
+ )
+
+ config = load_runtime_config(
+ config_path,
+ overrides={"staging": {"edge_size": 128}},
+ )
+
+ assert config.output_root == (tmp_path / "runs").resolve()
+ assert config.aoi.path == str(aoi_path.resolve())
+ assert config.rgbnir.model.config_path == str(model_config.resolve())
+ assert config.rgbnir.model.checkpoint_path == str(model_checkpoint.resolve())
+ assert config.rgbnir.model.cache_dir == str((tmp_path / "cache").resolve())
+ assert config.staging.edge_size == 128
diff --git a/tests/test_deployment/test_deployment_core_utils.py b/tests/test_deployment/test_deployment_core_utils.py
new file mode 100644
index 0000000..542a903
--- /dev/null
+++ b/tests/test_deployment/test_deployment_core_utils.py
@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+
+from deployment.srgan_hpc.logging_utils import configure_logging
+from deployment.srgan_hpc.metadata import write_software_metadata
+
+
+def test_configure_logging_sets_stream_and_file_handlers(tmp_path: Path) -> None:
+ log_path = tmp_path / "logs" / "run.log"
+
+ logger = configure_logging(log_path=log_path, verbose=True)
+ logger.info("hello logging")
+ for handler in logger.handlers:
+ handler.flush()
+
+ assert logger.level == logging.DEBUG
+ assert len(logger.handlers) == 2
+ assert "hello logging" in log_path.read_text(encoding="utf-8")
+
+
+def test_configure_logging_replaces_existing_handlers() -> None:
+ logger = configure_logging(verbose=False)
+ first_handlers = list(logger.handlers)
+
+ logger = configure_logging(verbose=False)
+
+ assert logger.level == logging.INFO
+ assert len(logger.handlers) == 1
+ assert logger.handlers != first_handlers
+
+
+def test_write_software_metadata_includes_environment_and_extra_fields(
+ tmp_path: Path,
+) -> None:
+ metadata_path = tmp_path / "metadata" / "software.json"
+
+ write_software_metadata(metadata_path, extra={"run_id": "srgan_001"})
+
+ payload = json.loads(metadata_path.read_text(encoding="utf-8"))
+ assert payload["run_id"] == "srgan_001"
+ assert "python_version" in payload
+ assert "platform" in payload
diff --git a/tests/test_deployment/test_deployment_extra_paths2.py b/tests/test_deployment/test_deployment_extra_paths2.py
new file mode 100644
index 0000000..e457f88
--- /dev/null
+++ b/tests/test_deployment/test_deployment_extra_paths2.py
@@ -0,0 +1,786 @@
+from __future__ import annotations
+
+import json
+import os
+import sys
+import types
+from pathlib import Path
+
+import numpy as np
+import pytest
+import rasterio
+from rasterio.transform import from_origin
+
+import deployment.srgan_hpc.aoi as aoi_module
+import deployment.srgan_hpc.cli as cli_module
+import deployment.srgan_hpc.delivery as delivery_module
+import deployment.srgan_hpc.run_task as run_task_module
+import deployment.srgan_hpc.staging as staging_module
+import deployment.srgan_hpc.submit as submit_module
+from deployment.srgan_hpc.aoi import (
+ load_aoi_geometry,
+ patch_footprint,
+ resolve_aoi_source_path,
+ select_aoi_patches,
+)
+from deployment.srgan_hpc.config import (
+ RuntimeConfig,
+ StagingConfig,
+ runtime_config_to_dict,
+ validate_runtime_config,
+)
+from deployment.srgan_hpc.delivery import deliver_bbox_outputs
+from deployment.srgan_hpc.manifests import read_yaml, write_yaml
+from deployment.srgan_hpc.patching import Patch
+from deployment.srgan_hpc.raster import (
+ _as_scalar,
+ compute_centroid_lat_lon,
+ ensure_proj_env,
+ guess_utm_epsg,
+ parse_epsg,
+ raster_validity_stats,
+ scale_to_uint16,
+)
+from deployment.srgan_hpc.run_task import run_task
+from deployment.srgan_hpc.staging import (
+ SkipTileError,
+ _select_or_mosaic_time_items,
+ ensure_cube_has_valid_data,
+ stage_cutout,
+)
+from deployment.srgan_hpc.submit import submit_grid_run, submit_patch_run
+
+
+def _write_tif(
+ path: Path,
+ data: np.ndarray,
+ transform,
+ *,
+ crs: str | None = "EPSG:32632",
+ nodata: int | None = 0,
+) -> None:
+ with rasterio.open(
+ path,
+ "w",
+ driver="GTiff",
+ width=data.shape[-1],
+ height=data.shape[-2],
+ count=data.shape[0],
+ dtype=data.dtype,
+ crs=crs,
+ transform=transform,
+ nodata=nodata,
+ ) as dst:
+ dst.write(data)
+
+
+def _write_polygon_shapefile(path: Path) -> None:
+ import shapefile
+ from pyproj import CRS
+
+ writer = shapefile.Writer(str(path), shapeType=shapefile.POLYGON)
+ try:
+ writer.field("name", "C")
+ writer.poly([[(0.0, 0.0), (0.02, 0.0), (0.02, 0.02), (0.0, 0.02), (0.0, 0.0)]])
+ writer.record("area")
+ finally:
+ writer.close()
+ path.with_suffix(".prj").write_text(CRS.from_epsg(4326).to_wkt(), encoding="utf-8")
+
+
+def _write_point_shapefile(path: Path) -> None:
+ import shapefile
+ from pyproj import CRS
+
+ writer = shapefile.Writer(str(path), shapeType=shapefile.POINT)
+ try:
+ writer.field("name", "C")
+ writer.point(0.0, 0.0)
+ writer.record("point")
+ finally:
+ writer.close()
+ path.with_suffix(".prj").write_text(CRS.from_epsg(4326).to_wkt(), encoding="utf-8")
+
+
+def test_aoi_source_and_geometry_extra_paths(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ with pytest.raises(FileNotFoundError, match="AOI path not found"):
+ resolve_aoi_source_path(tmp_path / "missing.shp")
+
+ with pytest.raises(ValueError, match="No .shp file"):
+ resolve_aoi_source_path(tmp_path)
+
+ text_path = tmp_path / "area.geojson"
+ text_path.write_text("{}", encoding="utf-8")
+ with pytest.raises(ValueError, match="must be a .shp file"):
+ resolve_aoi_source_path(text_path)
+
+ shp_path = tmp_path / "area.shp"
+ _write_polygon_shapefile(shp_path)
+ resolved_path, geometry = load_aoi_geometry(shp_path)
+ assert resolved_path == shp_path.resolve()
+ assert geometry.bounds == pytest.approx((0.0, 0.0, 0.02, 0.02))
+
+ missing_prj = tmp_path / "missing_prj.shp"
+ missing_prj.touch()
+ with pytest.raises(ValueError, match="missing .prj"):
+ load_aoi_geometry(missing_prj)
+
+ empty_prj = tmp_path / "empty_prj.shp"
+ empty_prj.touch()
+ empty_prj.with_suffix(".prj").write_text("", encoding="utf-8")
+ with pytest.raises(ValueError, match="empty .prj"):
+ load_aoi_geometry(empty_prj)
+
+ point_path = tmp_path / "point.shp"
+ _write_point_shapefile(point_path)
+ with pytest.raises(ValueError, match="polygon geometries"):
+ load_aoi_geometry(point_path)
+
+ selection = select_aoi_patches(
+ aoi_path=shp_path,
+ edge_size=64,
+ resolution_m=10.0,
+ overlap_meters=0.0,
+ )
+ assert selection.patches
+
+ from shapely.geometry import box
+
+ monkeypatch.setattr(
+ aoi_module, "load_aoi_geometry", lambda _path: (shp_path, box(0, 0, 1, 1))
+ )
+ monkeypatch.setattr(aoi_module, "build_patches", lambda *args, **kwargs: [])
+ with pytest.raises(ValueError, match="No SR cutouts intersect"):
+ select_aoi_patches(
+ aoi_path=shp_path,
+ edge_size=64,
+ resolution_m=10.0,
+ overlap_meters=0.0,
+ )
+
+
+def test_patch_footprint_builds_bounds_around_center() -> None:
+ patch = Patch(
+ patch_id="patch_000001",
+ latitude=0.0,
+ longitude=0.0,
+ edge_size=100,
+ row_index=0,
+ row_count=1,
+ column_index=0,
+ column_count=1,
+ )
+
+ footprint = patch_footprint(patch, resolution_m=10.0)
+
+ assert footprint.bounds[0] < 0.0
+ assert footprint.bounds[1] < 0.0
+ assert footprint.bounds[2] > 0.0
+ assert footprint.bounds[3] > 0.0
+
+
+def test_raster_small_helpers_and_validity_stats(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ monkeypatch.setenv("PROJ_LIB", "/existing/proj")
+ ensure_proj_env()
+ assert os.environ["PROJ_LIB"] == "/existing/proj"
+
+ monkeypatch.delenv("PROJ_LIB", raising=False)
+ monkeypatch.delenv("PROJ_DATA", raising=False)
+ ensure_proj_env()
+ assert os.environ["PROJ_LIB"]
+ assert os.environ["PROJ_DATA"]
+
+ assert guess_utm_epsg(45.0, 9.0) == 32632
+ assert guess_utm_epsg(-45.0, 9.0) == 32732
+ assert parse_epsg("EPSG:32633", 45.0, 9.0) == 32633
+ assert parse_epsg("unknown", 45.0, 9.0) == 32632
+ assert _as_scalar(type("Computed", (), {"compute": lambda self: "2.5"})()) == 2.5
+ monkeypatch.delenv("PROJ_LIB", raising=False)
+ monkeypatch.delenv("PROJ_DATA", raising=False)
+
+ scaled_unit = scale_to_uint16(
+ np.array([[[0.0, 0.5, 1.0, np.nan]]], dtype="float32")
+ )
+ scaled_reflectance = scale_to_uint16(
+ np.array([[[-1.0, 50.0, 10001.0]]], dtype="float32")
+ )
+ scaled_integer = scale_to_uint16(np.array([[[1, 2, 3]]], dtype="int16"))
+ assert scaled_unit.tolist() == [[[0, 5000, 10000, 0]]]
+ assert scaled_reflectance.tolist() == [[[0, 50, 10000]]]
+ assert scaled_integer.dtype == np.dtype("uint16")
+
+ source = tmp_path / "source.tif"
+ _write_tif(
+ source,
+ np.ones((1, 10, 10), dtype="uint16"),
+ from_origin(9.0, 46.0, 0.1, 0.1),
+ crs="EPSG:4326",
+ )
+ lat, lon = compute_centroid_lat_lon(source)
+ assert lat == pytest.approx(45.5)
+ assert lon == pytest.approx(9.5)
+
+ no_crs = tmp_path / "no_crs.tif"
+ _write_tif(
+ no_crs,
+ np.ones((1, 2, 2), dtype="uint16"),
+ from_origin(0.0, 0.0, 1.0, 1.0),
+ crs=None,
+ nodata=None,
+ )
+ with pytest.raises(ValueError, match="lacks a CRS"):
+ compute_centroid_lat_lon(no_crs)
+
+ validity = tmp_path / "validity.tif"
+ _write_tif(
+ validity,
+ np.array([[[0, 1], [2, 0]]], dtype="uint16"),
+ from_origin(0.0, 2.0, 1.0, 1.0),
+ )
+ assert raster_validity_stats(validity) == {
+ "total_pixels": 4,
+ "valid_pixels": 2,
+ "nonzero_pixels": 2,
+ }
+
+
+def _write_patch_output(run_dir: Path, patch_id: str, output_name: str) -> Path:
+ output = run_dir / "patches" / patch_id / "outputs" / output_name
+ output.parent.mkdir(parents=True, exist_ok=True)
+ output.write_bytes(b"patch")
+ return output
+
+
+def test_deliver_bbox_outputs_direct_and_nested_runs(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ run_dir = tmp_path / "run_001"
+ write_yaml(run_dir / "run_manifest.yaml", {"start_date": "2025-01-01"})
+ input_path = _write_patch_output(run_dir, "patch_000001", "custom_sr.tif")
+ captured: dict[str, object] = {}
+
+ def fake_merge_and_clip_bbox(*, input_paths, output_path, bbox, nodata=0):
+ captured["input_paths"] = input_paths
+ captured["bbox"] = bbox
+ output_path.write_text("merged", encoding="utf-8")
+ return output_path
+
+ monkeypatch.setattr(
+ delivery_module, "merge_and_clip_bbox", fake_merge_and_clip_bbox
+ )
+ destination, delivered = deliver_bbox_outputs(
+ run_root=run_dir,
+ bbox=(9.0, 45.0, 10.0, 46.0),
+ output_name="custom_sr.tif",
+ )
+ assert destination == run_dir / "delivery_clipped"
+ assert captured["input_paths"] == [input_path]
+ assert delivered[0]["date"] == "2025-01-01"
+ manifest = json.loads(
+ (destination / "delivery_manifest.json").read_text(encoding="utf-8")
+ )
+ assert manifest["bbox"] == [9.0, 45.0, 10.0, 46.0]
+
+ run_a = tmp_path / "run_a"
+ run_b = tmp_path / "run_b"
+ write_yaml(run_a / "run_manifest.yaml", {})
+ write_yaml(run_b / "run_manifest.yaml", {"start_date": "2025-02-01"})
+ _write_patch_output(run_a, "patch_000001", "fused_sr.tif")
+ _write_patch_output(run_b, "patch_000001", "fused_sr.tif")
+ _write_patch_output(run_dir, "patch_000001", "fused_sr.tif")
+ nested_destination, nested = deliver_bbox_outputs(
+ run_root=tmp_path,
+ bbox=(0.0, 0.0, 1.0, 1.0),
+ )
+ assert nested_destination == tmp_path / "delivery_clipped"
+ assert [Path(item["run_dir"]).name for item in nested] == [
+ "run_001",
+ "run_a",
+ "run_b",
+ ]
+ assert nested[0]["date"] == "2025-01-01"
+ assert nested[1]["date"] == "run_a"
+ assert nested[2]["date"] == "2025-02-01"
+
+
+def test_deliver_bbox_outputs_rejects_runs_without_outputs(tmp_path: Path) -> None:
+ run_dir = tmp_path / "run_001"
+ write_yaml(run_dir / "run_manifest.yaml", {})
+
+ with pytest.raises(FileNotFoundError, match="No fused_sr.tif files"):
+ deliver_bbox_outputs(run_root=run_dir, bbox=(0.0, 0.0, 1.0, 1.0))
+
+
+class FakeCoord:
+ def __init__(self, values: list[str]) -> None:
+ self.values = values
+
+
+class FakeRio:
+ def __init__(self, cube: "FakeCube") -> None:
+ self.cube = cube
+
+ def write_crs(self, epsg_code: int, inplace: bool = False):
+ self.cube.written_crs = epsg_code
+ return self.cube
+
+ def write_nodata(self, nodata: int, encoded: bool = True, inplace: bool = False):
+ self.cube.written_nodata = nodata
+ return self.cube
+
+ def to_raster(self, output_path: Path, **kwargs) -> None:
+ self.cube.raster_kwargs = kwargs
+ output_path.write_bytes(b"raster")
+
+
+class FakeCube:
+ def __init__(
+ self,
+ data: np.ndarray,
+ dims: tuple[str, ...] = ("band", "y", "x"),
+ coords: dict[str, object] | None = None,
+ attrs: dict[str, object] | None = None,
+ ) -> None:
+ self.data = np.asarray(data)
+ self.dims = dims
+ self.coords = coords or {}
+ self.attrs = attrs or {}
+ self.sizes = {dim: size for dim, size in zip(dims, self.data.shape)}
+ self.rio = FakeRio(self)
+
+ def transpose(self, *dims: str):
+ axes = [self.dims.index(dim) for dim in dims]
+ return FakeCube(
+ self.data.transpose(axes), tuple(dims), self.coords, dict(self.attrs)
+ )
+
+ def isel(self, **kwargs):
+ time_index = kwargs["time"]
+ axis = self.dims.index("time")
+ dims = tuple(dim for dim in self.dims if dim != "time")
+ return FakeCube(
+ np.take(self.data, time_index, axis=axis),
+ dims,
+ self.coords,
+ dict(self.attrs),
+ )
+
+ def copy(self, data=None):
+ return FakeCube(
+ self.data.copy() if data is None else data,
+ self.dims,
+ self.coords,
+ dict(self.attrs),
+ )
+
+
+def test_staging_validation_and_selection_extra_paths(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+ caplog: pytest.LogCaptureFixture,
+) -> None:
+ with pytest.raises(SkipTileError, match="empty_cube"):
+ ensure_cube_has_valid_data(types.SimpleNamespace(data=np.empty((1, 0, 0))))
+ with pytest.raises(SkipTileError, match="all_nan_cube"):
+ ensure_cube_has_valid_data(
+ types.SimpleNamespace(data=np.full((1, 2, 2), np.nan))
+ )
+ with pytest.raises(SkipTileError, match="all_zero_cube"):
+ ensure_cube_has_valid_data(types.SimpleNamespace(data=np.zeros((1, 2, 2))))
+
+ no_time_cube = FakeCube(np.ones((1, 4, 4), dtype="uint16"))
+ selected, diagnostics = _select_or_mosaic_time_items(no_time_cube, StagingConfig())
+ assert selected is no_time_cube
+ assert diagnostics["candidate_count"] == 1
+
+ timed_cube = FakeCube(
+ np.stack([np.zeros((1, 4, 4)), np.ones((1, 4, 4))]),
+ dims=("time", "band", "y", "x"),
+ coords={"time": FakeCoord(["first", "second"])},
+ )
+ selected, diagnostics = _select_or_mosaic_time_items(
+ timed_cube,
+ StagingConfig(item_strategy="fixed_index", image_index=1),
+ )
+ assert np.all(selected.data == 1)
+ assert diagnostics["selected_labels"] == ["second"]
+
+ with pytest.raises(IndexError, match="image_index=3"):
+ _select_or_mosaic_time_items(
+ timed_cube,
+ StagingConfig(item_strategy="fixed_index", image_index=3),
+ )
+
+ sparse = np.zeros((1, 1, 4, 4), dtype="uint16")
+ sparse[0, :, 1:3, 1:3] = 5
+ with caplog.at_level("WARNING", logger="srgan-hpc"):
+ _select_or_mosaic_time_items(
+ FakeCube(sparse, dims=("time", "band", "y", "x")),
+ StagingConfig(
+ min_center_nonzero_fraction=0.0, min_full_nonzero_fraction=0.8
+ ),
+ )
+ assert "Low full cutout coverage" in caplog.text
+
+ monkeypatch.setattr(staging_module, "ensure_proj_env", lambda: None)
+ monkeypatch.setattr(
+ staging_module,
+ "create_cube_with_retry",
+ lambda **_kwargs: (
+ FakeCube(np.ones((1, 2, 2), dtype="float32"), attrs={"epsg": "EPSG:32632"}),
+ [{"id": "S2_ITEM"}],
+ ),
+ )
+ output_path = tmp_path / "inputs" / "rgbnir.tif"
+ metadata_path = tmp_path / "metadata" / "rgbnir.json"
+ result = stage_cutout(
+ latitude=45.0,
+ longitude=9.0,
+ start_date="2025-01-01",
+ end_date="2025-01-02",
+ config=StagingConfig(rate_limit_retry_delays_seconds=[]),
+ bands=["B04"],
+ edge_size=64,
+ resolution=10,
+ output_path=output_path,
+ metadata_path=metadata_path,
+ patch_id="patch_000001",
+ product_name="rgbnir",
+ )
+ assert result == output_path.resolve()
+ payload = json.loads(metadata_path.read_text(encoding="utf-8"))
+ assert payload["auto_selected_items"] == [{"id": "S2_ITEM"}]
+ assert payload["validity_stats"]["nonzero_pixels"] == 4
+
+
+def _write_task_manifest(
+ tmp_path: Path,
+ *,
+ mode: str,
+ products: list[str],
+) -> Path:
+ config = RuntimeConfig(output_root=tmp_path / "runs", mode=mode)
+ manifest_path = tmp_path / "patches" / "patch_000001" / "manifest.yaml"
+ input_paths = {product: f"inputs/{product}.tif" for product in products}
+ write_yaml(
+ manifest_path,
+ {
+ "patch_id": "patch_000001",
+ "products": products,
+ "paths": {
+ "inputs": input_paths,
+ "output_dir": "outputs",
+ "metadata_dir": "metadata",
+ },
+ "config": runtime_config_to_dict(config),
+ },
+ )
+ for relative_path in input_paths.values():
+ input_path = manifest_path.parent / relative_path
+ input_path.parent.mkdir(parents=True, exist_ok=True)
+ input_path.write_bytes(b"input")
+ return manifest_path
+
+
+def test_run_task_skip_single_and_fused_paths(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ manifest_path = _write_task_manifest(
+ tmp_path / "skip", mode="rgbnir", products=["rgbnir"]
+ )
+ monkeypatch.setattr(
+ run_task_module,
+ "raster_validity_stats",
+ lambda _path: {"total_pixels": 4, "valid_pixels": 0, "nonzero_pixels": 0},
+ )
+ assert run_task(manifest_path) is None
+ skip_payload = json.loads(
+ (manifest_path.parent / "metadata" / "rgbnir_skip.json").read_text(
+ encoding="utf-8"
+ )
+ )
+ assert skip_payload["reason"] == "empty_input_raster"
+
+ manifest_path = _write_task_manifest(
+ tmp_path / "single", mode="rgbnir", products=["rgbnir"]
+ )
+ monkeypatch.setattr(
+ run_task_module,
+ "raster_validity_stats",
+ lambda _path: {"total_pixels": 4, "valid_pixels": 4, "nonzero_pixels": 4},
+ )
+
+ def fake_run_inference(*, output_dir, product_name, **_kwargs):
+ output = output_dir / f"{product_name}_sr.tif"
+ output.write_bytes(b"sr")
+ return output
+
+ monkeypatch.setattr(run_task_module, "run_inference", fake_run_inference)
+ result = run_task(manifest_path)
+ assert result == manifest_path.parent / "outputs" / "rgbnir_sr.tif"
+ result_payload = json.loads(
+ (manifest_path.parent / "metadata" / "result.json").read_text(encoding="utf-8")
+ )
+ assert result_payload["status"] == "completed"
+
+ fused_manifest = _write_task_manifest(
+ tmp_path / "fused", mode="fused", products=["rgbnir", "swir"]
+ )
+ captured: dict[str, object] = {}
+
+ def fake_stack_geotiffs(**kwargs):
+ captured.update(kwargs)
+ kwargs["output_path"].write_bytes(b"fused")
+ return kwargs["output_path"]
+
+ monkeypatch.setattr(run_task_module, "stack_geotiffs", fake_stack_geotiffs)
+ fused_result = run_task(fused_manifest)
+ assert fused_result == fused_manifest.parent / "outputs" / "fused_sr.tif"
+ assert captured["band_names"] == [
+ "B04",
+ "B03",
+ "B02",
+ "B08",
+ "B05",
+ "B06",
+ "B07",
+ "B8A",
+ "B11",
+ "B12",
+ ]
+
+ missing_manifest = _write_task_manifest(
+ tmp_path / "missing", mode="fused", products=["rgbnir"]
+ )
+ with pytest.raises(RuntimeError, match="missing: swir"):
+ run_task(missing_manifest)
+
+
+def test_submit_skip_paths_write_manifests(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ def fake_stage_cutout(**_kwargs):
+ raise SkipTileError("all_zero_cube", details={"total_pixels": 4})
+
+ monkeypatch.setattr(submit_module, "stage_cutout", fake_stage_cutout)
+ config = RuntimeConfig(output_root=tmp_path / "runs", mode="rgbnir")
+ patch = Patch(
+ patch_id="patch_000001",
+ latitude=45.0,
+ longitude=9.0,
+ edge_size=512,
+ row_index=0,
+ row_count=1,
+ column_index=0,
+ column_count=1,
+ )
+
+ _, run_dir, submission = submit_patch_run(
+ config=config,
+ patch=patch,
+ start_date="2025-01-01",
+ end_date="2025-01-02",
+ script_path=Path("/tmp/slurm.sh"),
+ )
+ assert submission["mode"] == "skipped"
+ assert read_yaml(run_dir / "run_manifest.yaml")["tasks"] == []
+ assert (
+ read_yaml(run_dir / "patches" / "patch_000001" / "manifest.yaml")["status"]
+ == "skipped"
+ )
+
+ _, grid_dir, grid_submission = submit_grid_run(
+ config=config,
+ patches=[patch],
+ start_date="2025-01-01",
+ end_date="2025-01-02",
+ script_path=Path("/tmp/slurm.sh"),
+ )
+ assert grid_submission["reason"] == "no_submittable_patches"
+ assert read_yaml(grid_dir / "run_manifest.yaml")["skipped_count"] == 1
+
+
+@pytest.mark.parametrize(
+ ("mutate", "message"),
+ [
+ (lambda config, tmp: setattr(config.staging, "edge_size", 0), "edge_size"),
+ (
+ lambda config, tmp: setattr(config.staging, "overlap_meters", -1),
+ "overlap_meters",
+ ),
+ (
+ lambda config, tmp: setattr(
+ config.staging, "rate_limit_retry_delays_seconds", [0]
+ ),
+ "rate_limit",
+ ),
+ (
+ lambda config, tmp: setattr(config.staging, "item_strategy", "bad"),
+ "item_strategy",
+ ),
+ (
+ lambda config, tmp: setattr(
+ config.staging, "min_center_nonzero_fraction", 1.5
+ ),
+ "min_center",
+ ),
+ (
+ lambda config, tmp: setattr(
+ config.staging, "min_full_nonzero_fraction", -0.1
+ ),
+ "min_full",
+ ),
+ (
+ lambda config, tmp: setattr(config.staging, "auto_select_item_limit", 0),
+ "auto_select_item_limit",
+ ),
+ (
+ lambda config, tmp: setattr(config.staging, "search_max_items", 0),
+ "search_max_items",
+ ),
+ (
+ lambda config, tmp: setattr(config.staging, "search_limit", 0),
+ "search_limit",
+ ),
+ (
+ lambda config, tmp: setattr(config.inference, "window_size", (128,)),
+ "window_size",
+ ),
+ (lambda config, tmp: setattr(config.inference, "batch_size", 0), "batch_size"),
+ (lambda config, tmp: setattr(config.slurm, "gpus", -1), "gpus"),
+ (lambda config, tmp: setattr(config.slurm, "mem_gb", 0), "mem_gb"),
+ (
+ lambda config, tmp: setattr(config.slurm, "cpus_per_task", 0),
+ "cpus_per_task",
+ ),
+ (
+ lambda config, tmp: setattr(config.aoi, "path", str(tmp / "missing.shp")),
+ "AOI path",
+ ),
+ (lambda config, tmp: setattr(config.rgbnir, "resolution", 0), "resolution"),
+ (lambda config, tmp: setattr(config.rgbnir, "factor", 0), "factor"),
+ (lambda config, tmp: setattr(config.rgbnir, "bands", []), "bands"),
+ (
+ lambda config, tmp: (
+ setattr(config.rgbnir.model, "preset", None),
+ setattr(config.rgbnir.model, "config_path", None),
+ ),
+ "model",
+ ),
+ (lambda config, tmp: setattr(config.staging, "edge_size", 3), "whole-number"),
+ ],
+)
+def test_validate_runtime_config_rejects_invalid_values(
+ tmp_path: Path,
+ mutate,
+ message: str,
+) -> None:
+ config = RuntimeConfig(mode="rgbnir")
+ if message == "whole-number":
+ config.mode = "fused"
+ mutate(config, tmp_path)
+
+ with pytest.raises((ValueError, FileNotFoundError), match=message):
+ validate_runtime_config(config)
+
+
+def test_cli_main_dispatches_common_commands(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+ capsys: pytest.CaptureFixture[str],
+) -> None:
+ config_path = tmp_path / "runtime.yaml"
+ config_path.write_text(
+ f"output_root: {tmp_path / 'runs'}\nproject_name: cli\nmode: rgbnir\n",
+ encoding="utf-8",
+ )
+
+ monkeypatch.setattr(
+ sys, "argv", ["srgan-hpc", "validate-config", "--config", str(config_path)]
+ )
+ assert cli_module.main() == 0
+ assert "Configuration valid" in capsys.readouterr().out
+
+ monkeypatch.setattr(
+ sys,
+ "argv",
+ [
+ "srgan-hpc",
+ "submit",
+ "patch",
+ "--config",
+ str(config_path),
+ "--start-date",
+ "2025-01-01",
+ "--end-date",
+ "2025-01-02",
+ "--lat",
+ "45.0",
+ "--lon",
+ "9.0",
+ "--script-path",
+ "/tmp/slurm.sh",
+ "--dry-run",
+ ],
+ )
+ assert cli_module.main() == 0
+ assert '"run_id"' in capsys.readouterr().out
+
+ run_dir = tmp_path / "status_run"
+ (run_dir / "patches" / "patch_000001").mkdir(parents=True)
+ monkeypatch.setattr(sys, "argv", ["srgan-hpc", "status", "--run-dir", str(run_dir)])
+ assert cli_module.main() == 0
+ status_payload = json.loads(capsys.readouterr().out)
+ assert status_payload["patch_count"] == 1
+
+ called: dict[str, object] = {}
+ monkeypatch.setattr(
+ "deployment.srgan_hpc.run_task.run_task",
+ lambda manifest, task_index=None: called.update(
+ {"manifest": manifest, "task_index": task_index}
+ )
+ or None,
+ )
+ monkeypatch.setenv("SLURM_ARRAY_TASK_ID", "4")
+ manifest = tmp_path / "manifest.yaml"
+ manifest.write_text("tasks: []\n", encoding="utf-8")
+ monkeypatch.setattr(
+ sys, "argv", ["srgan-hpc", "run", "task", "--manifest", str(manifest)]
+ )
+ assert cli_module.main() == 0
+ assert called["task_index"] == 4
+ assert capsys.readouterr().out.strip() == "skipped"
+
+ monkeypatch.setattr(
+ "deployment.srgan_hpc.delivery.deliver_bbox_outputs",
+ lambda **_kwargs: (tmp_path / "delivery", [{"output": "out.tif"}]),
+ )
+ monkeypatch.setattr(
+ sys,
+ "argv",
+ [
+ "srgan-hpc",
+ "deliver-bbox",
+ "--run-root",
+ str(tmp_path),
+ "--west",
+ "0",
+ "--south",
+ "0",
+ "--east",
+ "1",
+ "--north",
+ "1",
+ ],
+ )
+ assert cli_module.main() == 0
+ assert "out.tif" in capsys.readouterr().out
diff --git a/tests/test_deployment/test_deployment_inference.py b/tests/test_deployment/test_deployment_inference.py
new file mode 100644
index 0000000..f4161e0
--- /dev/null
+++ b/tests/test_deployment/test_deployment_inference.py
@@ -0,0 +1,285 @@
+from __future__ import annotations
+
+import sys
+import types
+from pathlib import Path
+
+import pytest
+
+import deployment.srgan_hpc.inference as inference_module
+from deployment.srgan_hpc.config import (
+ InferenceConfig,
+ ModelSourceConfig,
+ ProductConfig,
+)
+from deployment.srgan_hpc.inference import (
+ _build_runner,
+ load_srgan_model,
+ run_inference,
+)
+
+
+class FakeModel:
+ def __init__(self) -> None:
+ self.device: str | None = None
+
+ def to(self, device: str):
+ self.device = device
+ return self
+
+
+def test_load_srgan_model_uses_explicit_config(monkeypatch: pytest.MonkeyPatch) -> None:
+ captured: dict[str, object] = {}
+
+ def fake_load_from_config(*args, **kwargs):
+ captured["args"] = args
+ captured["kwargs"] = kwargs
+ return FakeModel()
+
+ monkeypatch.setattr(inference_module.torch.cuda, "is_available", lambda: False)
+ monkeypatch.setattr("opensr_srgan.load_from_config", fake_load_from_config)
+
+ model, device = load_srgan_model(
+ ProductConfig(
+ bands=["B04"],
+ resolution=10,
+ factor=4,
+ model=ModelSourceConfig(
+ preset=None,
+ config_path="/tmp/config.yaml",
+ checkpoint_path="/tmp/model.ckpt",
+ ),
+ )
+ )
+
+ assert device == "cpu"
+ assert model.device == "cpu"
+ assert captured["args"][:2] == ("/tmp/config.yaml", "/tmp/model.ckpt")
+ assert captured["kwargs"]["mode"] == "eval"
+
+
+def test_load_srgan_model_uses_preset(monkeypatch: pytest.MonkeyPatch) -> None:
+ captured: dict[str, object] = {}
+
+ def fake_load_inference_model(*args, **kwargs):
+ captured["args"] = args
+ captured["kwargs"] = kwargs
+ return FakeModel()
+
+ monkeypatch.setattr(inference_module.torch.cuda, "is_available", lambda: False)
+ monkeypatch.setattr("opensr_srgan.load_inference_model", fake_load_inference_model)
+
+ model, device = load_srgan_model(
+ ProductConfig(
+ bands=["B04"],
+ resolution=10,
+ factor=4,
+ model=ModelSourceConfig(preset="RGB-NIR", cache_dir="/tmp/cache"),
+ )
+ )
+
+ assert device == "cpu"
+ assert model.device == "cpu"
+ assert captured["args"] == ("RGB-NIR",)
+ assert captured["kwargs"]["cache_dir"] == "/tmp/cache"
+
+
+def test_load_srgan_model_requires_source() -> None:
+ with pytest.raises(ValueError, match="preset or config_path"):
+ load_srgan_model(
+ ProductConfig(
+ bands=["B04"],
+ resolution=10,
+ factor=4,
+ model=ModelSourceConfig(preset=None),
+ )
+ )
+
+
+def test_build_runner_creates_datamodule(monkeypatch: pytest.MonkeyPatch) -> None:
+ created: dict[str, object] = {}
+
+ class BaseLargeFileProcessing:
+ def __init__(self, *args, **kwargs) -> None:
+ self.input_type = "geotiff"
+ self.root = kwargs["root"]
+ self.image_meta = {
+ "image_windows": ["window"],
+ "lr_file_dict": {"B04": "input.tif"},
+ }
+ self.messages: list[str] = []
+
+ def _log(self, message: str) -> None:
+ self.messages.append(message)
+
+ class FakePredictionDataModule:
+ def __init__(self, **kwargs) -> None:
+ created.update(kwargs)
+ self.dataset = [object(), object()]
+ self.did_setup = False
+
+ def setup(self) -> None:
+ self.did_setup = True
+
+ monkeypatch.setitem(
+ sys.modules,
+ "opensr_utils.data_utils.datamodule",
+ types.SimpleNamespace(PredictionDataModule=FakePredictionDataModule),
+ )
+ runner_cls = _build_runner(
+ types.SimpleNamespace(large_file_processing=BaseLargeFileProcessing),
+ InferenceConfig(batch_size=7),
+ )
+
+ runner = runner_cls(root="/tmp/input.tif", batch_size=7)
+ runner.create_datamodule()
+
+ assert created["batch_size"] == 7
+ assert created["num_workers"] == 4
+ assert runner.datamodule.did_setup is True
+ assert "batch_size=7" in runner.messages[0]
+
+
+def test_run_inference_compresses_runner_output(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ input_tif = tmp_path / "input.tif"
+ input_tif.write_bytes(b"input")
+ output_dir = tmp_path / "outputs"
+ output_dir.mkdir()
+ sr_path = output_dir / "raw_sr.tif"
+ sr_path.write_bytes(b"raw")
+ captured: dict[str, object] = {}
+
+ class FakeLargeFileProcessing:
+ final_path = sr_path
+
+ def __init__(self, **kwargs) -> None:
+ captured["runner_kwargs"] = kwargs
+ self.final_sr_path = self.final_path
+
+ def start_super_resolution(self) -> None:
+ captured["started"] = True
+
+ def fake_compress(src_path: Path, dest_path: Path, band_names: list[str]):
+ captured["compress"] = (src_path, dest_path, band_names)
+ dest_path.write_bytes(src_path.read_bytes())
+ return dest_path
+
+ monkeypatch.setitem(
+ sys.modules,
+ "opensr_utils",
+ types.SimpleNamespace(large_file_processing=FakeLargeFileProcessing),
+ )
+ monkeypatch.setattr(
+ inference_module, "load_srgan_model", lambda product: (FakeModel(), "cpu")
+ )
+ monkeypatch.setattr(inference_module, "compress_geotiff", fake_compress)
+
+ result = run_inference(
+ input_tif=input_tif,
+ output_dir=output_dir,
+ product_name="rgbnir",
+ product=ProductConfig(
+ bands=["B04", "B03"],
+ resolution=10,
+ factor=4,
+ model=ModelSourceConfig(preset="RGB-NIR"),
+ ),
+ inference=InferenceConfig(window_size=(64, 96), batch_size=3, overlap=8),
+ )
+
+ assert captured["started"] is True
+ assert captured["runner_kwargs"]["root"] == str(input_tif)
+ assert captured["runner_kwargs"]["window_size"] == (64, 96)
+ assert "batch_size" not in captured["runner_kwargs"]
+ assert result == output_dir / "rgbnir_sr.tif"
+ assert captured["compress"] == (sr_path, result, ["B04", "B03"])
+ assert not sr_path.exists()
+
+
+def test_run_inference_uses_default_sr_path_when_runner_does_not_set_one(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ input_tif = tmp_path / "input.tif"
+ input_tif.write_bytes(b"input")
+ output_dir = tmp_path / "outputs"
+ output_dir.mkdir()
+ default_sr = output_dir / "sr.tif"
+ default_sr.write_bytes(b"raw")
+
+ class FakeLargeFileProcessing:
+ def __init__(self, **_kwargs) -> None:
+ pass
+
+ def start_super_resolution(self) -> None:
+ pass
+
+ monkeypatch.setitem(
+ sys.modules,
+ "opensr_utils",
+ types.SimpleNamespace(large_file_processing=FakeLargeFileProcessing),
+ )
+ monkeypatch.setattr(
+ inference_module, "load_srgan_model", lambda product: (FakeModel(), "cpu")
+ )
+ monkeypatch.setattr(
+ inference_module,
+ "compress_geotiff",
+ lambda src_path, dest_path, band_names: dest_path.write_bytes(
+ src_path.read_bytes()
+ ),
+ )
+
+ result = run_inference(
+ input_tif=input_tif,
+ output_dir=output_dir,
+ product_name="swir",
+ product=ProductConfig(
+ bands=["B11"],
+ resolution=20,
+ factor=8,
+ model=ModelSourceConfig(preset="SWIR"),
+ ),
+ inference=InferenceConfig(),
+ )
+
+ assert result == output_dir / "swir_sr.tif"
+ assert not default_sr.exists()
+
+
+def test_run_inference_rejects_missing_runner_output(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ class FakeLargeFileProcessing:
+ def __init__(self, **_kwargs) -> None:
+ self.final_sr_path = tmp_path / "missing.tif"
+
+ def start_super_resolution(self) -> None:
+ pass
+
+ monkeypatch.setitem(
+ sys.modules,
+ "opensr_utils",
+ types.SimpleNamespace(large_file_processing=FakeLargeFileProcessing),
+ )
+ monkeypatch.setattr(
+ inference_module, "load_srgan_model", lambda product: (FakeModel(), "cpu")
+ )
+
+ with pytest.raises(FileNotFoundError, match="Expected SR output"):
+ run_inference(
+ input_tif=tmp_path / "input.tif",
+ output_dir=tmp_path,
+ product_name="rgbnir",
+ product=ProductConfig(
+ bands=["B04"],
+ resolution=10,
+ factor=4,
+ model=ModelSourceConfig(preset="RGB-NIR"),
+ ),
+ inference=InferenceConfig(),
+ )
diff --git a/tests/test_deployment/test_deployment_patching.py b/tests/test_deployment/test_deployment_patching.py
new file mode 100644
index 0000000..0aface7
--- /dev/null
+++ b/tests/test_deployment/test_deployment_patching.py
@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+import pytest
+
+from deployment.srgan_hpc.patching import (
+ build_patches,
+ clamp_center,
+ compute_centers,
+ meters_to_lat_deg,
+ meters_to_lon_deg,
+)
+
+
+def test_patch_degree_helpers_and_clamping() -> None:
+ assert meters_to_lat_deg(111_320.0) == pytest.approx(1.0)
+ assert meters_to_lon_deg(111_320.0, 0.0) == pytest.approx(1.0)
+ assert clamp_center(0.0, -1.0, 1.0, 0.25) == 0.0
+ assert clamp_center(-2.0, -1.0, 1.0, 0.25) == -0.75
+ assert clamp_center(2.0, -1.0, 1.0, 0.25) == 0.75
+ assert clamp_center(5.0, 0.0, 1.0, 2.0) == 0.5
+
+
+def test_meters_to_lon_deg_rejects_polar_latitudes() -> None:
+ with pytest.raises(ValueError, match="Cannot compute longitude"):
+ meters_to_lon_deg(10.0, 90.000001)
+
+
+def test_compute_centers_handles_small_span_and_multiple_steps() -> None:
+ assert compute_centers(0.0, 0.1, patch_deg=1.0, step_deg=0.5) == [0.05]
+
+ centers = compute_centers(0.0, 3.0, patch_deg=1.0, step_deg=0.75)
+
+ assert centers == pytest.approx([0.5, 1.25, 2.0, 2.5])
+
+
+def test_compute_centers_rejects_non_positive_sizes() -> None:
+ with pytest.raises(ValueError, match="must be positive"):
+ compute_centers(0.0, 1.0, patch_deg=0.0, step_deg=1.0)
+
+ with pytest.raises(ValueError, match="must be positive"):
+ compute_centers(0.0, 1.0, patch_deg=1.0, step_deg=0.0)
+
+
+def test_build_patches_assigns_grid_metadata() -> None:
+ patches = build_patches(
+ lat1=0.0,
+ lon1=0.0,
+ lat2=0.05,
+ lon2=0.05,
+ edge_size=100,
+ resolution_m=10.0,
+ overlap_meters=0.0,
+ )
+
+ assert len(patches) > 1
+ assert patches[0].patch_id == "patch_000001"
+ assert patches[0].row_index == 0
+ assert patches[0].column_index == 0
+ assert patches[-1].patch_id == f"patch_{len(patches):06d}"
+ assert patches[-1].row_count >= 1
+ assert patches[-1].column_count >= 1
+
+
+@pytest.mark.parametrize(
+ ("kwargs", "message"),
+ [
+ ({"lon1": -179.0, "lon2": 179.5}, "antimeridian"),
+ ({"edge_size": 0}, "edge_size"),
+ ({"resolution_m": 0.0}, "resolution_m"),
+ ({"overlap_meters": 1000.0}, "overlap_meters"),
+ ],
+)
+def test_build_patches_validates_inputs(kwargs: dict[str, float], message: str) -> None:
+ params = {
+ "lat1": 0.0,
+ "lon1": 0.0,
+ "lat2": 0.01,
+ "lon2": 0.01,
+ "edge_size": 100,
+ "resolution_m": 10.0,
+ "overlap_meters": 0.0,
+ }
+ params.update(kwargs)
+
+ with pytest.raises(ValueError, match=message):
+ build_patches(**params)
diff --git a/tests/test_deployment/test_deployment_staging_extra_paths.py b/tests/test_deployment/test_deployment_staging_extra_paths.py
new file mode 100644
index 0000000..7f67397
--- /dev/null
+++ b/tests/test_deployment/test_deployment_staging_extra_paths.py
@@ -0,0 +1,158 @@
+from __future__ import annotations
+
+import sys
+import types
+
+import numpy as np
+import pytest
+
+from deployment.srgan_hpc.config import StagingConfig
+from deployment.srgan_hpc.staging import (
+ SkipTileError,
+ _auto_select_item_ids,
+ _format_cloud_cover_range,
+ _search_kwargs_from_config,
+ _select_or_mosaic_time_items,
+ candidate_coverage_report,
+ create_cube_with_retry,
+ is_rate_limit_error,
+ mosaic_candidates_by_valid_pixels,
+)
+from tests.test_deployment.test_deployment_extra_paths2 import FakeCube
+
+
+def test_staging_coverage_helpers_handle_empty_and_invalid_inputs() -> None:
+ assert candidate_coverage_report(np.empty((1, 0, 0))) == {
+ "full_nonzero_fraction": 0.0,
+ "center_nonzero_fraction": 0.0,
+ }
+
+ with pytest.raises(ValueError, match="Expected candidate data"):
+ mosaic_candidates_by_valid_pixels(np.zeros((1, 2, 2)), [0])
+
+ with pytest.raises(ValueError, match="must not be empty"):
+ mosaic_candidates_by_valid_pixels(np.zeros((1, 1, 2, 2)), [])
+
+
+def test_staging_small_private_helpers() -> None:
+ assert is_rate_limit_error(RuntimeError("too many requests"))
+ assert (
+ _format_cloud_cover_range([{"cloud_cover": None}, {"cloud_cover": "bad"}])
+ == "n/a"
+ )
+ assert _format_cloud_cover_range([{"cloud_cover": "1.25"}]) == "1.250%"
+ assert (
+ _format_cloud_cover_range([{"cloud_cover": 3}, {"cloud_cover": 1}])
+ == "1.000-3.000%"
+ )
+
+ assert _search_kwargs_from_config(
+ StagingConfig(
+ search_query={"eo:cloud_cover": {"lt": 10}},
+ search_max_items=3,
+ search_limit=2,
+ )
+ ) == {
+ "query": {"eo:cloud_cover": {"lt": 10}},
+ "max_items": 3,
+ "limit": 2,
+ }
+
+
+def test_auto_select_item_ids_skips_when_stac_search_returns_no_items(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ class EmptySearch:
+ def items(self) -> list[object]:
+ return []
+
+ class FakeCatalog:
+ def search(self, **_kwargs):
+ return EmptySearch()
+
+ monkeypatch.setitem(
+ sys.modules,
+ "pystac_client",
+ types.SimpleNamespace(
+ Client=types.SimpleNamespace(open=lambda _url: FakeCatalog())
+ ),
+ )
+
+ with pytest.raises(SkipTileError, match="no_stac_items") as exc_info:
+ _auto_select_item_ids(
+ latitude=45.0,
+ longitude=9.0,
+ start_date="2025-01-01",
+ end_date="2025-01-02",
+ config=StagingConfig(auto_select_item=True),
+ )
+
+ assert exc_info.value.details == {"latitude": 45_000_000, "longitude": 9_000_000}
+
+
+def test_create_cube_with_retry_returns_from_retry_loop(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ captured: dict[str, object] = {}
+ cube = object()
+
+ def fake_create(**kwargs):
+ captured.update(kwargs)
+ return cube
+
+ monkeypatch.setitem(sys.modules, "cubo", types.SimpleNamespace(create=fake_create))
+ monkeypatch.setitem(sys.modules, "rioxarray", types.SimpleNamespace())
+
+ result, reports = create_cube_with_retry(
+ latitude=45.0,
+ longitude=9.0,
+ start_date="2025-01-01",
+ end_date="2025-01-02",
+ config=StagingConfig(rate_limit_retry_delays_seconds=[1]),
+ bands=["B04"],
+ edge_size=64,
+ resolution=10,
+ )
+
+ assert result is cube
+ assert reports == []
+ assert captured["bands"] == ["B04"]
+ assert captured["edge_size"] == 64
+
+
+def test_select_or_mosaic_time_items_empty_and_low_center_paths() -> None:
+ with pytest.raises(SkipTileError, match="empty_cube"):
+ _select_or_mosaic_time_items(
+ FakeCube(np.empty((0, 1, 4, 4)), dims=("time", "band", "y", "x")),
+ StagingConfig(),
+ )
+
+ with pytest.raises(SkipTileError, match="low_center_coverage") as exc_info:
+ _select_or_mosaic_time_items(
+ FakeCube(np.zeros((1, 1, 4, 4)), dims=("time", "band", "y", "x")),
+ StagingConfig(min_center_nonzero_fraction=0.5),
+ )
+
+ assert exc_info.value.details == {
+ "candidate_count": 1,
+ "center_nonzero_percent": 0,
+ }
+
+
+def test_select_or_mosaic_time_items_falls_back_to_index_label() -> None:
+ class BrokenCoord:
+ @property
+ def values(self):
+ raise RuntimeError("cannot read labels")
+
+ data = np.ones((1, 1, 4, 4), dtype="uint16")
+ _, diagnostics = _select_or_mosaic_time_items(
+ FakeCube(
+ data,
+ dims=("time", "band", "y", "x"),
+ coords={"time": BrokenCoord()},
+ ),
+ StagingConfig(item_strategy="fixed_index", image_index=0),
+ )
+
+ assert diagnostics["selected_labels"] == ["0"]
diff --git a/tests/test_training/test_hydra_config.py b/tests/test_training/test_hydra_config.py
new file mode 100644
index 0000000..7b5f479
--- /dev/null
+++ b/tests/test_training/test_hydra_config.py
@@ -0,0 +1,111 @@
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import pytest
+from omegaconf import OmegaConf
+
+hydra = pytest.importorskip("hydra")
+from hydra.core.global_hydra import GlobalHydra
+
+CONFIG_DIR = Path(__file__).resolve().parents[2] / "opensr_srgan" / "configs" / "hydra"
+CONSUMED_SECTIONS = (
+ "Data",
+ "Model",
+ "Training",
+ "Generator",
+ "Discriminator",
+ "Optimizers",
+ "Schedulers",
+ "Logging",
+)
+
+
+@pytest.fixture(autouse=True)
+def clear_hydra_state():
+ GlobalHydra.instance().clear()
+ yield
+ GlobalHydra.instance().clear()
+
+
+def _compose(overrides: list[str] | None = None, *, include_hydra: bool = False):
+ with hydra.initialize_config_dir(config_dir=str(CONFIG_DIR), version_base="1.3"):
+ return hydra.compose(
+ config_name="train",
+ overrides=overrides or [],
+ return_hydra_config=include_hydra,
+ )
+
+
+def _consumed_sections(cfg):
+ return {
+ section: OmegaConf.to_container(cfg[section], resolve=True)
+ for section in CONSUMED_SECTIONS
+ }
+
+
+def _legacy_sections(path: str):
+ return _consumed_sections(OmegaConf.load(path))
+
+
+def test_default_hydra_config_is_cpu_safe_and_uses_example_dataset():
+ cfg = _compose()
+
+ assert cfg.Data.dataset_type == "ExampleDataset"
+ assert cfg.Training.device == "cpu"
+ assert cfg.Training.gpus == []
+ assert cfg.Logging.wandb.enabled is False
+
+
+def test_hydra_config_sets_run_directory_without_changing_cwd():
+ cfg = _compose(include_hydra=True)
+
+ assert cfg.hydra.job.chdir is False
+ run_cfg = OmegaConf.to_container(cfg.hydra.run, resolve=False)
+ assert run_cfg["dir"] == ("logs/${Logging.wandb.project}/${now:%Y-%m-%d_%H-%M-%S}")
+
+
+@pytest.mark.parametrize(
+ ("experiment", "legacy_path"),
+ [
+ ("example", "opensr_srgan/configs/config_training_example.yaml"),
+ ("10m", "opensr_srgan/configs/config_10m.yaml"),
+ ("20m", "opensr_srgan/configs/config_20m.yaml"),
+ ],
+)
+def test_hydra_experiments_match_legacy_yaml_sections(experiment, legacy_path):
+ cfg = _compose([f"experiment={experiment}"])
+
+ assert _consumed_sections(cfg) == _legacy_sections(legacy_path)
+
+
+def test_hydra_entrypoint_passes_composed_overrides_to_train(monkeypatch, tmp_path):
+ import opensr_srgan.train as train_module
+ from opensr_srgan import train_hydra
+
+ captured = {}
+
+ def fake_train(cfg):
+ captured["cfg"] = cfg
+
+ monkeypatch.setattr(train_module, "train", fake_train)
+ monkeypatch.setattr(
+ sys,
+ "argv",
+ [
+ "srgan-train",
+ "experiment=example",
+ "Training.max_epochs=2",
+ "Logging.wandb.enabled=false",
+ f"hydra.run.dir={tmp_path.as_posix()}",
+ ],
+ )
+
+ train_hydra.main()
+
+ cfg = captured["cfg"]
+ assert cfg.Data.dataset_type == "ExampleDataset"
+ assert cfg.Training.max_epochs == 2
+ assert cfg.Logging.wandb.enabled is False
+ assert Path(cfg.Logging.output_dir) == tmp_path
diff --git a/tests/test_training/test_train_checkpoint_loading.py b/tests/test_training/test_train_checkpoint_loading.py
index 6bb9a4c..004f9eb 100644
--- a/tests/test_training/test_train_checkpoint_loading.py
+++ b/tests/test_training/test_train_checkpoint_loading.py
@@ -287,6 +287,23 @@ def test_train_uses_wandb_logger_and_saves_config_when_global_zero(
), "expected config.yaml to be written in logs/unit-tests//"
+def test_train_respects_logging_output_dir(train_module, monkeypatch, tmp_path):
+ train_mod, state = train_module
+
+ gpu_rank_module = sys.modules["opensr_srgan.utils.gpu_rank"]
+ monkeypatch.setattr(gpu_rank_module, "_is_global_zero", lambda: True)
+
+ output_dir = tmp_path / "hydra-run"
+ config = _make_config(load_checkpoint=False, continue_training=False)
+ config.Logging.output_dir = str(output_dir)
+
+ train_mod.train(config)
+
+ assert (output_dir / "config.yaml").is_file()
+ assert state["builder_checkpoint_callback"].kwargs["dirpath"] == str(output_dir)
+ assert state["builder_logger"].kwargs["save_dir"] == str(output_dir)
+
+
def test_train_module_main_guard(monkeypatch):
state = {}