From dea04c1dd9587cb04190cdb10661eaac6ac57483 Mon Sep 17 00:00:00 2001 From: Simon Donike Date: Fri, 17 Jul 2026 15:21:54 +0200 Subject: [PATCH 01/10] gitignore --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index 8077c01..aaf7b79 100644 --- a/.gitignore +++ b/.gitignore @@ -17,6 +17,7 @@ gitlog.txt *.vscode .vscode __pycache__/ +playground # Test / coverage artifacts .coverage From 7b74ff23f7482cc1d0cb05de935487cf4caa8f8e Mon Sep 17 00:00:00 2001 From: Simon Donike Date: Fri, 17 Jul 2026 15:22:32 +0200 Subject: [PATCH 02/10] Respect explicit training output directories --- opensr_srgan/train.py | 43 +++++++++++++++---- .../test_train_checkpoint_loading.py | 17 ++++++++ 2 files changed, 51 insertions(+), 9 deletions(-) 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/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 = {} From cd226145851153bef402235f3c1afdc012c2b9d8 Mon Sep 17 00:00:00 2001 From: Simon Donike Date: Fri, 17 Jul 2026 16:17:32 +0200 Subject: [PATCH 03/10] Add Hydra training entry point and config presets --- opensr_srgan/configs/__init__.py | 1 + opensr_srgan/configs/hydra/__init__.py | 1 + opensr_srgan/configs/hydra/data/example.yaml | 9 ++ .../configs/hydra/data/lrhr_folder.yaml | 10 ++ opensr_srgan/configs/hydra/data/sen2naip.yaml | 9 ++ .../configs/hydra/experiment/10m.yaml | 8 ++ .../configs/hydra/experiment/20m.yaml | 16 +++ .../configs/hydra/experiment/example.yaml | 8 ++ opensr_srgan/configs/hydra/hydra/default.yaml | 14 +++ opensr_srgan/configs/hydra/logging/csv.yaml | 8 ++ .../configs/hydra/logging/wandb_10m.yaml | 8 ++ .../configs/hydra/logging/wandb_20m.yaml | 8 ++ .../configs/hydra/model/10m_esrgan.yaml | 25 ++++ .../configs/hydra/model/20m_rcab.yaml | 24 ++++ .../configs/hydra/model/example_srresnet.yaml | 24 ++++ opensr_srgan/configs/hydra/train.yaml | 11 ++ opensr_srgan/configs/hydra/training/10m.yaml | 50 ++++++++ opensr_srgan/configs/hydra/training/20m.yaml | 49 ++++++++ .../configs/hydra/training/cpu_debug.yaml | 49 ++++++++ .../configs/hydra/training/example.yaml | 49 ++++++++ opensr_srgan/train_hydra.py | 51 ++++++++ pyproject.toml | 2 + requirements.txt | 1 + tests/test_training/test_hydra_config.py | 111 ++++++++++++++++++ 24 files changed, 546 insertions(+) create mode 100644 opensr_srgan/configs/__init__.py create mode 100644 opensr_srgan/configs/hydra/__init__.py create mode 100644 opensr_srgan/configs/hydra/data/example.yaml create mode 100644 opensr_srgan/configs/hydra/data/lrhr_folder.yaml create mode 100644 opensr_srgan/configs/hydra/data/sen2naip.yaml create mode 100644 opensr_srgan/configs/hydra/experiment/10m.yaml create mode 100644 opensr_srgan/configs/hydra/experiment/20m.yaml create mode 100644 opensr_srgan/configs/hydra/experiment/example.yaml create mode 100644 opensr_srgan/configs/hydra/hydra/default.yaml create mode 100644 opensr_srgan/configs/hydra/logging/csv.yaml create mode 100644 opensr_srgan/configs/hydra/logging/wandb_10m.yaml create mode 100644 opensr_srgan/configs/hydra/logging/wandb_20m.yaml create mode 100644 opensr_srgan/configs/hydra/model/10m_esrgan.yaml create mode 100644 opensr_srgan/configs/hydra/model/20m_rcab.yaml create mode 100644 opensr_srgan/configs/hydra/model/example_srresnet.yaml create mode 100644 opensr_srgan/configs/hydra/train.yaml create mode 100644 opensr_srgan/configs/hydra/training/10m.yaml create mode 100644 opensr_srgan/configs/hydra/training/20m.yaml create mode 100644 opensr_srgan/configs/hydra/training/cpu_debug.yaml create mode 100644 opensr_srgan/configs/hydra/training/example.yaml create mode 100644 opensr_srgan/train_hydra.py create mode 100644 tests/test_training/test_hydra_config.py 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..13020f6 --- /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: discriminator/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..a48a887 --- /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: discriminator/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..f6b5361 --- /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: discriminator/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..f25849c --- /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: discriminator/adversarial_loss + patience_g: 10 + patience_d: 10 + factor_g: 0.5 + factor_d: 0.5 diff --git a/opensr_srgan/train_hydra.py b/opensr_srgan/train_hydra.py new file mode 100644 index 0000000..05aa075 --- /dev/null +++ b/opensr_srgan/train_hydra.py @@ -0,0 +1,51 @@ +"""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 hydra +import torch +from hydra.core.hydra_config import HydraConfig +from omegaconf import DictConfig, OmegaConf, open_dict + + +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 61da19e..358fe99 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,6 +19,7 @@ dependencies = [ "numpy", "kornia", "wandb", + "hydra-core>=1.3,<2", "omegaconf", "matplotlib", "rasterio", @@ -80,3 +81,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_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 From e45e8acebf4f822fae8d61adb7def2331071f316 Mon Sep 17 00:00:00 2001 From: Simon Donike Date: Fri, 17 Jul 2026 16:44:32 +0200 Subject: [PATCH 04/10] Document Hydra experiment workflow --- docs/api-reference/training.md | 14 +++++ docs/hydra.md | 107 +++++++++++++++++++++++++++++++++ mkdocs.yml | 9 ++- 3 files changed, 128 insertions(+), 2 deletions(-) create mode 100644 docs/hydra.md 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/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/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: From 8c70296b69f0d74e0502d617a73ff85b3d2ad6b0 Mon Sep 17 00:00:00 2001 From: Simon Donike Date: Fri, 17 Jul 2026 17:31:32 +0200 Subject: [PATCH 05/10] Add deployment regression tests --- .../test_deployment_cli_slurm_config_paths.py | 296 +++++++ .../test_deployment_core_utils.py | 72 ++ .../test_deployment_extra_paths2.py | 789 ++++++++++++++++++ .../test_deployment_inference.py | 285 +++++++ .../test_deployment_patching.py | 86 ++ .../test_deployment_staging_extra_paths.py | 158 ++++ 6 files changed, 1686 insertions(+) create mode 100644 tests/test_deployment/test_deployment_cli_slurm_config_paths.py create mode 100644 tests/test_deployment/test_deployment_core_utils.py create mode 100644 tests/test_deployment/test_deployment_extra_paths2.py create mode 100644 tests/test_deployment/test_deployment_inference.py create mode 100644 tests/test_deployment/test_deployment_patching.py create mode 100644 tests/test_deployment/test_deployment_staging_extra_paths.py 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..5c22e66 --- /dev/null +++ b/tests/test_deployment/test_deployment_cli_slurm_config_paths.py @@ -0,0 +1,296 @@ +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", + aoi_layer="named-layer", + 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"), + "--layer", + "named-layer", + "--script-path", + "/tmp/slurm.sh", + "--dry-run", + ], + ) + assert cli_module.main() == 0 + aoi_output = capsys.readouterr().out + assert '"aoi_layer": "named-layer"' 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..92f6540 --- /dev/null +++ b/tests/test_deployment/test_deployment_core_utils.py @@ -0,0 +1,72 @@ +from __future__ import annotations + +import hashlib +import json +import logging +from pathlib import Path + +import pytest + +from deployment.srgan_hpc.checkpoint import resolve_checkpoint_path, sha256sum +from deployment.srgan_hpc.logging_utils import configure_logging +from deployment.srgan_hpc.metadata import write_software_metadata + + +def test_resolve_checkpoint_path_accepts_none_and_existing_file( + tmp_path: Path, +) -> None: + checkpoint = tmp_path / "model.ckpt" + checkpoint.write_bytes(b"weights") + + assert resolve_checkpoint_path(None) is None + assert resolve_checkpoint_path(str(checkpoint)) == checkpoint.resolve() + + +def test_resolve_checkpoint_path_rejects_missing_file(tmp_path: Path) -> None: + with pytest.raises(FileNotFoundError, match="Checkpoint not found"): + resolve_checkpoint_path(str(tmp_path / "missing.ckpt")) + + +def test_sha256sum_streams_file_contents(tmp_path: Path) -> None: + payload = b"abc" + (b"0123456789" * 200_000) + checkpoint = tmp_path / "model.ckpt" + checkpoint.write_bytes(payload) + + assert sha256sum(checkpoint) == hashlib.sha256(payload).hexdigest() + + +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..c7bdb6b --- /dev/null +++ b/tests/test_deployment/test_deployment_extra_paths2.py @@ -0,0 +1,789 @@ +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, + aoi_layer="layer-name", + edge_size=64, + resolution_m=10.0, + overlap_meters=0.0, + ) + assert selection.aoi_layer == "layer-name" + 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, + aoi_layer=None, + 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"] From 83569b508e22677239582c78a90135f1e1a08d1a Mon Sep 17 00:00:00 2001 From: Simon Donike Date: Fri, 17 Jul 2026 18:22:36 +0200 Subject: [PATCH 06/10] docs for hydra + readme --- README.md | 22 +++++++++++++-------- docs/architecture.md | 8 ++++---- docs/configuration.md | 27 +++++++++++++++++++------ docs/data.md | 14 ++++++++++++- docs/getting-started.md | 15 +++++++------- docs/index.md | 29 +++++++++++---------------- docs/inference.md | 24 ++++++++++++++-------- docs/results.md | 4 ++-- docs/training.md | 44 ++++++++++++++++++++++------------------- 9 files changed, 113 insertions(+), 74 deletions(-) diff --git a/README.md b/README.md index 78ade9d..19834dd 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ | **PyPI** | **Versions** | **Docs & License** | **Tests** | **Reference** | |:---------:|:-------------:|:------------------:|:----------:|:--------------:| -| [![PyPI](https://img.shields.io/pypi/v/opensr-srgan)](https://pypi.org/project/opensr-srgan/) | ![PythonVersion](https://img.shields.io/badge/Python-v3.10%20v3.12-blue.svg)
![PLVersion](https://img.shields.io/badge/PytorchLightning-v2%2B-blue.svg) | [![Docs](https://img.shields.io/badge/docs-mkdocs%20material-brightgreen)](https://srgan.opensr.eu)
![License: Apache](https://img.shields.io/badge/license-Apache%20License%202.0-blue) | [![CI](https://github.com/ESAOpenSR/SRGAN/actions/workflows/ci.yml/badge.svg)](https://github.com/ESAOpenSR/SRGAN/actions/workflows/ci.yml)
[![codecov](https://codecov.io/github/ESAOpenSR/SRGAN/graph/badge.svg?token=LQ69MIMLVE)](https://codecov.io/github/ESAOpenSR/SRGAN) | [![arXiv](https://img.shields.io/badge/arXiv-2511.10461-b31b1b.svg)](https://arxiv.org/abs/2511.10461)
[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.20088994.svg)](https://doi.org/10.5281/zenodo.20088994) +| [![PyPI](https://img.shields.io/pypi/v/opensr-srgan)](https://pypi.org/project/opensr-srgan/) | ![PythonVersion](https://img.shields.io/badge/Python-v3.10%20v3.12-blue.svg)
![PLVersion](https://img.shields.io/badge/PyTorchLightning-v2%2B-blue.svg) | [![Docs](https://img.shields.io/badge/docs-mkdocs%20material-brightgreen)](https://srgan.opensr.eu)
![License: Apache](https://img.shields.io/badge/license-Apache%20License%202.0-blue) | [![CI](https://github.com/ESAOpenSR/SRGAN/actions/workflows/ci.yml/badge.svg)](https://github.com/ESAOpenSR/SRGAN/actions/workflows/ci.yml)
[![codecov](https://codecov.io/github/ESAOpenSR/SRGAN/graph/badge.svg?token=LQ69MIMLVE)](https://codecov.io/github/ESAOpenSR/SRGAN) | [![arXiv](https://img.shields.io/badge/arXiv-2511.10461-b31b1b.svg)](https://arxiv.org/abs/2511.10461)
[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.17590993.svg)](https://doi.org/10.5281/zenodo.17590993) @@ -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,7 +83,7 @@ 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! [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](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/). +* **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/getting-started/inference/). --- 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 fc627de..97f1616 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..13ef184 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -27,7 +27,7 @@ For the fastest start, open the interactive notebook in Google Colab and run thr ``` 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 +39,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 +50,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 +68,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 +87,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 +104,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="opensr_srgan/configs/config_training_example.yaml", checkpoint_uri="https://example.com/checkpoints/srgan.ckpt", map_location="cuda", # optional ) diff --git a/docs/index.md b/docs/index.md index 0e25b44..6160550 100644 --- a/docs/index.md +++ b/docs/index.md @@ -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 ![RCAB 8× example](assets/x8_6band_example_1.png) diff --git a/docs/training.md b/docs/training.md index 73e27ff..f110a30 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 | | --- | --- | --- | @@ -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. From 0372dc5a74168659db5ee966dbb1c4308a638d7a Mon Sep 17 00:00:00 2001 From: Simon Donike Date: Fri, 17 Jul 2026 18:48:25 +0200 Subject: [PATCH 07/10] update to include P3.13, P3.14 --- .github/dependabot.yml | 11 +++++++++++ .github/workflows/ci.yml | 22 ++++++++++++++-------- README.md | 2 +- docs/getting-started.md | 2 +- docs/index.md | 2 +- opensr_srgan/train_hydra.py | 25 +++++++++++++++++++++++++ pyproject.toml | 4 +++- 7 files changed, 56 insertions(+), 12 deletions(-) create mode 100644 .github/dependabot.yml 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..ffc2370 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.10", "3.11", "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/README.md b/README.md index 19834dd..fb35318 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ | **PyPI** | **Versions** | **Docs & License** | **Tests** | **Reference** | |:---------:|:-------------:|:------------------:|:----------:|:--------------:| -| [![PyPI](https://img.shields.io/pypi/v/opensr-srgan)](https://pypi.org/project/opensr-srgan/) | ![PythonVersion](https://img.shields.io/badge/Python-v3.10%20v3.12-blue.svg)
![PLVersion](https://img.shields.io/badge/PyTorchLightning-v2%2B-blue.svg) | [![Docs](https://img.shields.io/badge/docs-mkdocs%20material-brightgreen)](https://srgan.opensr.eu)
![License: Apache](https://img.shields.io/badge/license-Apache%20License%202.0-blue) | [![CI](https://github.com/ESAOpenSR/SRGAN/actions/workflows/ci.yml/badge.svg)](https://github.com/ESAOpenSR/SRGAN/actions/workflows/ci.yml)
[![codecov](https://codecov.io/github/ESAOpenSR/SRGAN/graph/badge.svg?token=LQ69MIMLVE)](https://codecov.io/github/ESAOpenSR/SRGAN) | [![arXiv](https://img.shields.io/badge/arXiv-2511.10461-b31b1b.svg)](https://arxiv.org/abs/2511.10461)
[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.17590993.svg)](https://doi.org/10.5281/zenodo.17590993) +| [![PyPI](https://img.shields.io/pypi/v/opensr-srgan)](https://pypi.org/project/opensr-srgan/) | ![PythonVersion](https://img.shields.io/badge/Python-v3.10--v3.14-blue.svg)
![PLVersion](https://img.shields.io/badge/PyTorchLightning-v2%2B-blue.svg) | [![Docs](https://img.shields.io/badge/docs-mkdocs%20material-brightgreen)](https://srgan.opensr.eu)
![License: Apache](https://img.shields.io/badge/license-Apache%20License%202.0-blue) | [![CI](https://github.com/ESAOpenSR/SRGAN/actions/workflows/ci.yml/badge.svg)](https://github.com/ESAOpenSR/SRGAN/actions/workflows/ci.yml)
[![codecov](https://codecov.io/github/ESAOpenSR/SRGAN/graph/badge.svg?token=LQ69MIMLVE)](https://codecov.io/github/ESAOpenSR/SRGAN) | [![arXiv](https://img.shields.io/badge/arXiv-2511.10461-b31b1b.svg)](https://arxiv.org/abs/2511.10461)
[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.17590993.svg)](https://doi.org/10.5281/zenodo.17590993) diff --git a/docs/getting-started.md b/docs/getting-started.md index 13ef184..22accae 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.10-3.14, PyTorch Lightning, and Weights & Biases for experiment tracking. ## Try it in Colab first diff --git a/docs/index.md b/docs/index.md index 6160550..611be3e 100644 --- a/docs/index.md +++ b/docs/index.md @@ -4,7 +4,7 @@ | **PyPI** | **Versions** | **Docs & License** | **Tests** | **Reference** | |:---------:|:-------------:|:------------------:|:----------:|:--------------:| -| [![PyPI](https://img.shields.io/pypi/v/opensr-srgan)](https://pypi.org/project/opensr-srgan/) | ![PythonVersion](https://img.shields.io/badge/Python-v3.10%20v3.12-blue.svg)
![PLVersion](https://img.shields.io/badge/PyTorchLightning-v2%2B-blue.svg) | [![Docs](https://img.shields.io/badge/docs-mkdocs%20material-brightgreen)](https://srgan.opensr.eu)
![License: Apache](https://img.shields.io/badge/license-Apache%20License%202.0-blue) | [![CI](https://github.com/ESAOpenSR/SRGAN/actions/workflows/ci.yml/badge.svg)](https://github.com/ESAOpenSR/SRGAN/actions/workflows/ci.yml)
[![codecov](https://codecov.io/github/ESAOpenSR/SRGAN/graph/badge.svg?token=LQ69MIMLVE)](https://codecov.io/github/ESAOpenSR/SRGAN) | [![arXiv](https://img.shields.io/badge/arXiv-2511.10461-b31b1b.svg)](https://arxiv.org/abs/2511.10461)
[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.17590993.svg)](https://doi.org/10.5281/zenodo.17590993) | +| [![PyPI](https://img.shields.io/pypi/v/opensr-srgan)](https://pypi.org/project/opensr-srgan/) | ![PythonVersion](https://img.shields.io/badge/Python-v3.10--v3.14-blue.svg)
![PLVersion](https://img.shields.io/badge/PyTorchLightning-v2%2B-blue.svg) | [![Docs](https://img.shields.io/badge/docs-mkdocs%20material-brightgreen)](https://srgan.opensr.eu)
![License: Apache](https://img.shields.io/badge/license-Apache%20License%202.0-blue) | [![CI](https://github.com/ESAOpenSR/SRGAN/actions/workflows/ci.yml/badge.svg)](https://github.com/ESAOpenSR/SRGAN/actions/workflows/ci.yml)
[![codecov](https://codecov.io/github/ESAOpenSR/SRGAN/graph/badge.svg?token=LQ69MIMLVE)](https://codecov.io/github/ESAOpenSR/SRGAN) | [![arXiv](https://img.shields.io/badge/arXiv-2511.10461-b31b1b.svg)](https://arxiv.org/abs/2511.10461)
[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.17590993.svg)](https://doi.org/10.5281/zenodo.17590993) | ![Super-resolved Sentinel-2 example](assets/6band_banner.png) diff --git a/opensr_srgan/train_hydra.py b/opensr_srgan/train_hydra.py index 05aa075..4706039 100644 --- a/opensr_srgan/train_hydra.py +++ b/opensr_srgan/train_hydra.py @@ -6,12 +6,37 @@ 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, "") diff --git a/pyproject.toml b/pyproject.toml index 358fe99..35fe5d5 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.10,<3.15" dependencies = [ "torch>=2.1", "pytorch-lightning>=2.1,<3.0", @@ -33,6 +33,8 @@ classifiers = [ "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", From a045783e400668c6f86e9f97b96e87768524c0a7 Mon Sep 17 00:00:00 2001 From: Simon Donike Date: Fri, 17 Jul 2026 19:00:02 +0200 Subject: [PATCH 08/10] fix test import --- pyproject.toml | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 35fe5d5..037bd50 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -69,7 +69,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", From d148512ecfcf19a9e1176d49d7a5aa110db8eefa Mon Sep 17 00:00:00 2001 From: Simon Donike Date: Thu, 23 Jul 2026 12:31:29 +0200 Subject: [PATCH 09/10] Align user documentation after main merge --- README.md | 4 ++-- docs/getting-started.md | 7 ++++--- docs/training.md | 2 +- opensr_srgan/_factory.py | 2 +- 4 files changed, 8 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index fb35318..bacd7a0 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ | **PyPI** | **Versions** | **Docs & License** | **Tests** | **Reference** | |:---------:|:-------------:|:------------------:|:----------:|:--------------:| -| [![PyPI](https://img.shields.io/pypi/v/opensr-srgan)](https://pypi.org/project/opensr-srgan/) | ![PythonVersion](https://img.shields.io/badge/Python-v3.10--v3.14-blue.svg)
![PLVersion](https://img.shields.io/badge/PyTorchLightning-v2%2B-blue.svg) | [![Docs](https://img.shields.io/badge/docs-mkdocs%20material-brightgreen)](https://srgan.opensr.eu)
![License: Apache](https://img.shields.io/badge/license-Apache%20License%202.0-blue) | [![CI](https://github.com/ESAOpenSR/SRGAN/actions/workflows/ci.yml/badge.svg)](https://github.com/ESAOpenSR/SRGAN/actions/workflows/ci.yml)
[![codecov](https://codecov.io/github/ESAOpenSR/SRGAN/graph/badge.svg?token=LQ69MIMLVE)](https://codecov.io/github/ESAOpenSR/SRGAN) | [![arXiv](https://img.shields.io/badge/arXiv-2511.10461-b31b1b.svg)](https://arxiv.org/abs/2511.10461)
[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.17590993.svg)](https://doi.org/10.5281/zenodo.17590993) +| [![PyPI](https://img.shields.io/pypi/v/opensr-srgan)](https://pypi.org/project/opensr-srgan/) | ![PythonVersion](https://img.shields.io/badge/Python-v3.10--v3.14-blue.svg)
![PLVersion](https://img.shields.io/badge/PyTorchLightning-v2%2B-blue.svg) | [![Docs](https://img.shields.io/badge/docs-mkdocs%20material-brightgreen)](https://srgan.opensr.eu)
![License: Apache](https://img.shields.io/badge/license-Apache%20License%202.0-blue) | [![CI](https://github.com/ESAOpenSR/SRGAN/actions/workflows/ci.yml/badge.svg)](https://github.com/ESAOpenSR/SRGAN/actions/workflows/ci.yml)
[![codecov](https://codecov.io/github/ESAOpenSR/SRGAN/graph/badge.svg?token=LQ69MIMLVE)](https://codecov.io/github/ESAOpenSR/SRGAN) | [![arXiv](https://img.shields.io/badge/arXiv-2511.10461-b31b1b.svg)](https://arxiv.org/abs/2511.10461)
[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.20088994.svg)](https://doi.org/10.5281/zenodo.20088994) @@ -84,7 +84,7 @@ 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! [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](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 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/getting-started/inference/). +* **Inference:** Ready-made presets and large-scene pipelines are described in the [inference section](https://srgan.opensr.eu/inference/). --- diff --git a/docs/getting-started.md b/docs/getting-started.md index 22accae..366f086 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -21,10 +21,11 @@ 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. * Keep `Logging.wandb.enabled: false` for local CSV logs when you do not want to use W&B. @@ -104,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_training_example.yaml", + config_path="path/to/your_config.yaml", checkpoint_uri="https://example.com/checkpoints/srgan.ckpt", map_location="cuda", # optional ) diff --git a/docs/training.md b/docs/training.md index 8926d5a..965b9a9 100644 --- a/docs/training.md +++ b/docs/training.md @@ -74,7 +74,7 @@ logs make it clear 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. | 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)) From a5acb5e9efc71d98a5005cd6fe935e9aeb48de59 Mon Sep 17 00:00:00 2001 From: Simon Donike Date: Thu, 23 Jul 2026 12:38:40 +0200 Subject: [PATCH 10/10] remove 3.10 and 3.11 support --- .github/workflows/ci.yml | 2 +- .github/workflows/docs.yml | 2 +- .github/workflows/publish.yml | 2 +- README.md | 2 +- docs/getting-started.md | 2 +- docs/index.md | 2 +- pyproject.toml | 4 +--- 7 files changed, 7 insertions(+), 9 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index ffc2370..7931643 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -34,7 +34,7 @@ jobs: strategy: fail-fast: false matrix: - python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"] + python-version: ["3.12", "3.13", "3.14"] steps: - uses: actions/checkout@v4 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 bacd7a0..0413913 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ | **PyPI** | **Versions** | **Docs & License** | **Tests** | **Reference** | |:---------:|:-------------:|:------------------:|:----------:|:--------------:| -| [![PyPI](https://img.shields.io/pypi/v/opensr-srgan)](https://pypi.org/project/opensr-srgan/) | ![PythonVersion](https://img.shields.io/badge/Python-v3.10--v3.14-blue.svg)
![PLVersion](https://img.shields.io/badge/PyTorchLightning-v2%2B-blue.svg) | [![Docs](https://img.shields.io/badge/docs-mkdocs%20material-brightgreen)](https://srgan.opensr.eu)
![License: Apache](https://img.shields.io/badge/license-Apache%20License%202.0-blue) | [![CI](https://github.com/ESAOpenSR/SRGAN/actions/workflows/ci.yml/badge.svg)](https://github.com/ESAOpenSR/SRGAN/actions/workflows/ci.yml)
[![codecov](https://codecov.io/github/ESAOpenSR/SRGAN/graph/badge.svg?token=LQ69MIMLVE)](https://codecov.io/github/ESAOpenSR/SRGAN) | [![arXiv](https://img.shields.io/badge/arXiv-2511.10461-b31b1b.svg)](https://arxiv.org/abs/2511.10461)
[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.20088994.svg)](https://doi.org/10.5281/zenodo.20088994) +| [![PyPI](https://img.shields.io/pypi/v/opensr-srgan)](https://pypi.org/project/opensr-srgan/) | ![PythonVersion](https://img.shields.io/badge/Python-v3.12--v3.14-blue.svg)
![PLVersion](https://img.shields.io/badge/PyTorchLightning-v2%2B-blue.svg) | [![Docs](https://img.shields.io/badge/docs-mkdocs%20material-brightgreen)](https://srgan.opensr.eu)
![License: Apache](https://img.shields.io/badge/license-Apache%20License%202.0-blue) | [![CI](https://github.com/ESAOpenSR/SRGAN/actions/workflows/ci.yml/badge.svg)](https://github.com/ESAOpenSR/SRGAN/actions/workflows/ci.yml)
[![codecov](https://codecov.io/github/ESAOpenSR/SRGAN/graph/badge.svg?token=LQ69MIMLVE)](https://codecov.io/github/ESAOpenSR/SRGAN) | [![arXiv](https://img.shields.io/badge/arXiv-2511.10461-b31b1b.svg)](https://arxiv.org/abs/2511.10461)
[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.20088994.svg)](https://doi.org/10.5281/zenodo.20088994) diff --git a/docs/getting-started.md b/docs/getting-started.md index 366f086..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.14, 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 diff --git a/docs/index.md b/docs/index.md index 3e115f5..f919621 100644 --- a/docs/index.md +++ b/docs/index.md @@ -4,7 +4,7 @@ | **PyPI** | **Versions** | **Docs & License** | **Tests** | **Reference** | |:---------:|:-------------:|:------------------:|:----------:|:--------------:| -| [![PyPI](https://img.shields.io/pypi/v/opensr-srgan)](https://pypi.org/project/opensr-srgan/) | ![PythonVersion](https://img.shields.io/badge/Python-v3.10--v3.14-blue.svg)
![PLVersion](https://img.shields.io/badge/PyTorchLightning-v2%2B-blue.svg) | [![Docs](https://img.shields.io/badge/docs-mkdocs%20material-brightgreen)](https://srgan.opensr.eu)
![License: Apache](https://img.shields.io/badge/license-Apache%20License%202.0-blue) | [![CI](https://github.com/ESAOpenSR/SRGAN/actions/workflows/ci.yml/badge.svg)](https://github.com/ESAOpenSR/SRGAN/actions/workflows/ci.yml)
[![codecov](https://codecov.io/github/ESAOpenSR/SRGAN/graph/badge.svg?token=LQ69MIMLVE)](https://codecov.io/github/ESAOpenSR/SRGAN) | [![arXiv](https://img.shields.io/badge/arXiv-2511.10461-b31b1b.svg)](https://arxiv.org/abs/2511.10461)
[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.20088994.svg)](https://doi.org/10.5281/zenodo.20088994) | +| [![PyPI](https://img.shields.io/pypi/v/opensr-srgan)](https://pypi.org/project/opensr-srgan/) | ![PythonVersion](https://img.shields.io/badge/Python-v3.12--v3.14-blue.svg)
![PLVersion](https://img.shields.io/badge/PyTorchLightning-v2%2B-blue.svg) | [![Docs](https://img.shields.io/badge/docs-mkdocs%20material-brightgreen)](https://srgan.opensr.eu)
![License: Apache](https://img.shields.io/badge/license-Apache%20License%202.0-blue) | [![CI](https://github.com/ESAOpenSR/SRGAN/actions/workflows/ci.yml/badge.svg)](https://github.com/ESAOpenSR/SRGAN/actions/workflows/ci.yml)
[![codecov](https://codecov.io/github/ESAOpenSR/SRGAN/graph/badge.svg?token=LQ69MIMLVE)](https://codecov.io/github/ESAOpenSR/SRGAN) | [![arXiv](https://img.shields.io/badge/arXiv-2511.10461-b31b1b.svg)](https://arxiv.org/abs/2511.10461)
[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.20088994.svg)](https://doi.org/10.5281/zenodo.20088994) | ![Super-resolved Sentinel-2 example](assets/6band_banner.png) diff --git a/pyproject.toml b/pyproject.toml index 43d97d8..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.15" +requires-python = ">=3.12,<3.15" dependencies = [ "torch>=2.1", "pytorch-lightning>=2.1,<3.0", @@ -30,8 +30,6 @@ 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",