Skip to content

Commit 7cc96d5

Browse files
authored
Merge pull request #87 from simon-donike/dev
Dev
2 parents 561629c + 8a7e37d commit 7cc96d5

20 files changed

Lines changed: 767 additions & 137 deletions

README.md

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ Full docs live at **[srgan.opensr.eu](https://srgan.opensr.eu/)**. They cover us
1818

1919
## 🧠 Highlights
2020

21-
* **Flexible models:** swap between SRResNet, RCAB, RRDB, and LKA-style generators with YAML-only changes.
21+
* **Flexible models:** swap between SRResNet, RCAB, RRDB, LKA, ESRGAN, and stochastic generators with YAML-only changes.
2222
* **Remote-sensing aware losses:** combine spectral, perceptual, and adversarial objectives with tunable weights.
2323
* **Stable training loop:** generator pretraining, adversarial ramp-ups, EMA, and multi-GPU Lightning support out of the box.
2424
* **PyPI distribution:** `pip install opensr-srgan` for ready-to-use presets or custom configs.
@@ -30,7 +30,7 @@ Full docs live at **[srgan.opensr.eu](https://srgan.opensr.eu/)**. They cover us
3030

3131
All key knobs are exposed via YAML in the `opensr_srgan/configs` folder:
3232

33-
* **Model**: `in_channels`, `n_channels`, `n_blocks`, `scale`, `block_type ∈ {SRResNet, res, rcab, rrdb, lka}`
33+
* **Model**: `in_channels`, `n_channels`, `n_blocks`, `scale`, ESRGAN knobs (`growth_channels`, `res_scale`, `out_channels`), `block_type ∈ {SRResNet, res, rcab, rrdb, lka}`
3434
* **Losses**: `l1_weight`, `sam_weight`, `perceptual_weight`, `tv_weight`, `adv_loss_beta`
3535
* **Training**: `pretrain_g_only`, `g_pretrain_steps`, `adv_loss_ramp_steps`, `label_smoothing`, generator LR warmup (`Schedulers.g_warmup_steps`, `Schedulers.g_warmup_type`), discriminator cadence controls
3636
* **Data**: band order, normalization stats, crop sizes, augmentations
@@ -52,8 +52,8 @@ The schedule and ramp make training **easier, safer, and more reproducible**.
5252

5353
| Component | Options | Config keys |
5454
|-----------|---------|-------------|
55-
| **Generators** | `SRResNet`, `res`, `rcab`, `rrdb`, `lka` | `Generator.model_type`, depth via `Generator.n_blocks`, width via `Generator.n_channels`, kernels and scale. |
56-
| **Discriminators** | `standard` SRGAN CNN, `patchgan` | `Discriminator.model_type`, granularity with `Discriminator.n_blocks`. |
55+
| **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`. |
56+
| **Discriminators** | `standard` SRGAN CNN, `patchgan`, `esrgan` | `Discriminator.model_type`, granularity with `Discriminator.n_blocks`, ESRGAN-specific `base_channels`, `linear_size`. |
5757
| **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`). |
5858
| **Adversarial loss** | BCE‑with‑logits on real/fake logits | Warmup via `Training.pretrain_g_only`, ramped by `adv_loss_ramp_steps`, capped at `adv_loss_beta`, optional label smoothing. |
5959

docs/architecture.md

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,9 @@ This document outlines how ESA OpenSR organises its super-resolution GAN, the ma
88

99
* **Configuration ingestion.** Uses OmegaConf to load hyperparameters, dataset choices, and logging options. Convenience helpers
1010
such as `_pretrain_check()` and `_compute_adv_loss_weight()` translate config values into runtime behaviour.
11-
* **Model factory.** `get_models()` builds the generator and discriminator at runtime based on `Generator.model_type` and
12-
`Discriminator.model_type`. Unsupported combinations fail fast with clear error messages.
11+
* **Model factory.** `get_models()` builds the generator and discriminator at runtime via the generator factory using
12+
`Generator.model_type`/`block_type` and `Discriminator.model_type`. Unsupported combinations fail fast with clear error
13+
messages.
1314
* **Loss construction.** `GeneratorContentLoss` (from `opensr_srgan.model.loss`) provides L1, spectral angle mapper (SAM), perceptual, and
1415
total-variation terms. Adversarial supervision uses `torch.nn.BCEWithLogitsLoss` with optional label smoothing.
1516
* **Optimiser scheduling.** `configure_optimizers()` returns paired Adam optimisers (generator + discriminator) with
@@ -40,21 +41,25 @@ The generator zoo lives under `opensr_srgan/model/generators/` and can be select
4041
* **Flexible residual families (`flexible_generator.py`).** Parameterised factory that instantiates residual, RCAB, RRDB, or
4142
large-kernel attention blocks while reusing the same interface. Channel counts, block depth, kernel sizes, and scaling factor
4243
are all read from the YAML file.
43-
* **Conditional GAN generator (`cgan_generator.py`).** Extends the flexible generator with conditioning inputs and latent noise,
44+
* **Stochastic GAN generator (`cgan_generator.py`).** Extends the flexible generator with conditioning inputs and latent noise,
4445
enabling experiments where auxiliary metadata influences the super-resolution output.
46+
* **ESRGAN generator (`esrgan.py`).** Implements the RRDBNet trunk introduced with ESRGAN, exposing `n_blocks`, `growth_channels`,
47+
and `res_scale` so you can dial in deeper receptive fields and sharper textures.
4548
* **Advanced variants (`SRGAN_advanced.py`).** Provides additional block implementations and compatibility aliases exposed in
4649
`__init__.py` for backwards compatibility.
4750

4851
Common traits across generators include configurable input channel counts (`Model.in_bands`), support for upscaling factors from 2× to 8×, and residual scaling to stabilise deeper networks.
4952

5053
## Discriminator options
5154

52-
`opensr_srgan/model/discriminators/` exposes two complementary discriminators:
55+
`opensr_srgan/model/discriminators/` exposes three complementary discriminators:
5356

5457
* **Standard SRGAN discriminator (`srgan_discriminator.py`).** Deep convolutional stack tailored for multispectral imagery. The
5558
number of convolutional blocks is configurable through `Discriminator.n_blocks`.
5659
* **PatchGAN discriminator (`patchgan.py`).** Operates on local patches, which can improve high-frequency fidelity when training
5760
with large images. The depth is controlled by `n_blocks` and defaults to three layers.
61+
* **ESRGAN discriminator (`esrgan.py`).** Deep VGG-style stack with configurable `base_channels` and `linear_size`; pairs well
62+
with RRDB generators when perceptual sharpness is the priority.
5863

5964
Both discriminators use LeakyReLU activations and strided convolutions to progressively downsample the input until a real/fake logit map is produced.
6065

docs/configuration.md

Lines changed: 23 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -79,19 +79,25 @@ stable validation imagery. The EMA is fully optional and controlled through the
7979

8080
| Key | Default | Description |
8181
| --- | --- | --- |
82-
| `model_type` | `cgan` | Generator architecture (`SRResNet`, `res`, `rcab`, `rrdb`, `lka`, `conditional_cgan`, `cgan`). |
82+
| `model_type` | `SRResNet` | Generator family (`SRResNet`, `stochastic_gan`, or `esrgan`). |
83+
| `block_type` | `standard` | SRResNet variant (`standard`, `res`, `rcab`, `rrdb`, `lka`). Ignored for `stochastic_gan`/`esrgan`. |
8384
| `large_kernel_size` | `9` | Kernel size for input/output convolution layers. |
8485
| `small_kernel_size` | `3` | Kernel size for residual/attention blocks. |
85-
| `n_channels` | `96` | Base number of feature channels. |
86-
| `n_blocks` | `32` | Number of residual/attention blocks. |
86+
| `n_channels` | `96` | Base number of feature channels (RRDB/ESRGAN trunk width). |
87+
| `n_blocks` | `32` | Number of residual/attention blocks (RRDB count when `model_type: esrgan`). |
8788
| `scaling_factor` | `8` | Super-resolution scale factor (2, 4, 8, ...). |
89+
| `growth_channels` | `32` | ESRGAN-only: growth channels inside each RRDB block. |
90+
| `res_scale` | `0.2` | Residual scaling used by stochastic/ESRGAN variants. |
91+
| `out_channels` | `Model.in_bands` | ESRGAN-only: override the number of output bands. |
8892

8993
## Discriminator
9094

9195
| Key | Default | Description |
9296
| --- | --- | --- |
93-
| `model_type` | `standard` | Discriminator architecture (`standard` SRGAN or `patchgan`). |
94-
| `n_blocks` | `8` | Number of convolutional blocks. PatchGAN defaults to 3 when unspecified. |
97+
| `model_type` | `standard` | Discriminator architecture (`standard`, `patchgan`, or `esrgan`). |
98+
| `n_blocks` | `8` | Number of convolutional blocks. PatchGAN defaults to 3 when unspecified (ignored by `esrgan`). |
99+
| `base_channels` | `64` | ESRGAN-only: base number of feature maps. |
100+
| `linear_size` | `1024` | ESRGAN-only: hidden dimension of the fully connected head. |
95101

96102
## Suggested settings
97103

@@ -109,24 +115,29 @@ performing sweeps:
109115

110116
| Generator type | Recommended `n_channels` | Recommended `n_blocks` | Typical `scaling_factor` | Notes |
111117
| --- | --- | --- | --- | --- |
112-
| `SRResNet` | 64 | 16 || Canonical baseline with batch-norm residual blocks; scale can be 2×/4×/8× as needed. |
113-
| `res` | 96 | 32 | 4×–8× | Lightweight residual blocks without batch norm; works well for high-scale (8×) Sentinel data. |
114-
| `rcab` | 96 | 32 | 4×–8× | Attention-enhanced residual blocks; keep depth high to exploit channel attention. |
115-
| `rrdb` | 96 | 32 | 4×–8× | Dense residual blocks expand receptive field; expect higher VRAM use at 32 blocks. |
116-
| `lka` | 96 | 24–32 | 4×–8× | Large-kernel attention blocks stabilise at moderate depth; drop to 24 blocks if memory bound. |
117-
| `conditional_cgan`/`cgan` | 96 | 16 || Latent-modulated residual stack; pair with noise_dim≈128 and res_scale≈0.2 defaults. |
118+
| `SRResNet` (`block_type: standard`) | 64 | 16 || Canonical baseline with batch-norm residual blocks; scale can be 2×/4×/8× as needed. |
119+
| `SRResNet` (`block_type: res`) | 96 | 32 | 4×–8× | Lightweight residual blocks without batch norm; works well for high-scale (8×) Sentinel data. |
120+
| `SRResNet` (`block_type: rcab`) | 96 | 32 | 4×–8× | Attention-enhanced residual blocks; keep depth high to exploit channel attention. |
121+
| `SRResNet` (`block_type: rrdb`) | 96 | 32 | 4×–8× | Dense residual blocks expand receptive field; expect higher VRAM use at 32 blocks. |
122+
| `SRResNet` (`block_type: lka`) | 96 | 24–32 | 4×–8× | Large-kernel attention blocks stabilise at moderate depth; drop to 24 blocks if memory bound. |
123+
| `stochastic_gan` | 96 | 16 || Latent-modulated residual stack; pair with `noise_dim ≈ 128` and `res_scale ≈ 0.2` defaults. |
124+
| `esrgan` | 64 | 23 || ESRGAN-style RRDB trunk; tune `growth_channels` (typically 32) and keep `res_scale ≈ 0.2` for stability. |
118125

119126
### Discriminator presets
120127

121128
Tune discriminator depth to match the generator capacity—too shallow and adversarial loss underfits, too deep and the training loop destabilises. These starting points mirror the architectures bundled with the repo:
122129

123130
| Discriminator type | Recommended depth parameter | Additional notes |
124131
| --- | --- | --- |
125-
| `standard` | `n_blocks = 8` | Mirrors the original SRGAN CNN with alternating stride-1/stride-2 blocks before the dense head. |
132+
| `standard` | `n_blocks = 8` | Mirrors the original SRGAN CNN with alternating stride-1/stride-2 blocks before the dense head. |
126133
| `patchgan` | `n_blocks = 3` | Maps to the 3-layer PatchGAN (a.k.a. `n_layers`); increase to 4–5 for larger crops or when the generator is particularly sharp. |
134+
| `esrgan` | `base_channels = 64`, `linear_size = 1024` | Deep VGG-style discriminator from ESRGAN; keep base width aligned with the generator feature count. |
127135

128136
When adjusting these presets, scale generator and discriminator together and monitor adversarial loss ramps defined in `Training.Losses` to keep training stable.
129137

138+
!!! note
139+
When you pick `model_type: esrgan` or `stochastic_gan`, SRResNet-only keys such as `block_type`, `large_kernel_size`, or `small_kernel_size` are automatically ignored. The model factory prints a console notice so you know which settings were overridden.
140+
130141
## Optimisers
131142

132143
The trainer instantiates independent Adam optimisers for the generator and discriminator and enables a Two-Time-Scale Update Rule (TTUR) setup by default. The discriminator learning rate automatically defaults to a slower schedule than the generator, which keeps adversarial updates balanced without extra configuration.

opensr_srgan/__init__.py

Lines changed: 27 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,30 @@
22

33
from __future__ import annotations
44

5-
from .model.SRGAN import SRGAN_model
6-
from .train import train
7-
from ._factory import load_from_config, load_inference_model
8-
9-
__all__ = [
10-
"SRGAN_model",
11-
"train",
12-
"load_from_config",
13-
"load_inference_model",
14-
]
5+
from typing import TYPE_CHECKING, Any
6+
7+
__all__ = ["SRGAN_model", "train", "load_from_config", "load_inference_model"]
8+
9+
if TYPE_CHECKING: # pragma: no cover - type checkers only
10+
from .model.SRGAN import SRGAN_model as SRGANModel
11+
from .train import train as _train
12+
13+
14+
def __getattr__(name: str) -> Any: # pragma: no cover - simple import proxy
15+
if name == "SRGAN_model":
16+
from .model.SRGAN import SRGAN_model as _cls
17+
18+
globals()[name] = _cls
19+
return _cls
20+
if name == "train":
21+
from .train import train as _train_fn
22+
23+
globals()[name] = _train_fn
24+
return _train_fn
25+
if name in {"load_from_config", "load_inference_model"}:
26+
from . import _factory
27+
28+
attr = getattr(_factory, name)
29+
globals()[name] = attr
30+
return attr
31+
raise AttributeError(f"module 'opensr_srgan' has no attribute '{name}'")

opensr_srgan/_factory.py

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -158,9 +158,16 @@ def load_inference_model(
158158

159159
if __name__ == "__main__":
160160
# simple test
161-
model = load_inference_model("RGB-NIR")
161+
# Create Model
162+
model = load_from_config("opensr_srgan/configs/config_playgound.yaml")
163+
#model = load_inference_model("RGB-NIR")
164+
165+
# Print Model Summary
166+
from opensr_srgan.utils.model_descriptions import print_model_summary
167+
print_model_summary(model)
168+
169+
# Simple test for funcionality
162170
import torch
163-
164171
lr = torch.randn(1, 4, 64, 64)
165-
sr = model.predict_step(lr)
166-
print(sr.shape)
172+
sr = model.forward(lr)
173+
print(lr.shape, "->", sr.shape)

opensr_srgan/configs/config_10m.yaml

Lines changed: 13 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -76,16 +76,22 @@ Training:
7676
# ---------------------------------------------------------------------------- #
7777
# See Docs for archtecture details and suggestions
7878
Generator:
79-
model_type: 'rrdb' # Block type: ['SRResNet', 'res', 'rcab', 'rrdb', 'lka', 'conditional_cgan'/'cgan']
80-
large_kernel_size: 9 # Kernel for head and tail conv layers
81-
small_kernel_size: 3 # Kernel for intermediate blocks
82-
n_channels: 64 # Number of feature channels (original 64)
83-
n_blocks: 16 # Number of residual/attention blocks (original 16)
79+
model_type: 'SRResNet' # Generator family: ['SRResNet', 'stochastic_gan', 'esrgan']
80+
block_type: 'rrdb' # SRResNet block variant: ['standard', 'res', 'rcab', 'rrdb', 'lka']
81+
large_kernel_size: 9 # Kernel for head and tail conv layers (SRResNet/stochastic)
82+
small_kernel_size: 3 # Kernel for intermediate blocks (SRResNet/stochastic)
83+
n_channels: 64 # Feature width (RRDB/ESRGAN uses this as trunk width)
84+
n_blocks: 16 # Residual/attention blocks (ESRGAN: number of RRDB blocks)
8485
scaling_factor: 4 # Upscaling factor (e.g., 2×, 4×, 8×)
86+
growth_channels: 32 # ESRGAN-specific RRDB growth channels (ignored otherwise)
87+
res_scale: 0.2 # Residual scaling used by stochastic/ESRGAN variants
88+
out_channels: 4 # Optional ESRGAN output band override (defaults to Model.in_bands)
8589

8690
Discriminator:
87-
model_type: 'standard' # Discriminator architecture selector ['standard', 'patchgan']
88-
n_blocks: 8 # Number of convolutional blocks / layers: ['standard': 8, 'patchgan': 3]
91+
model_type: 'standard' # Discriminator architecture selector ['standard', 'patchgan', 'esrgan']
92+
n_blocks: 8 # Convolutional depth for SRGAN/PatchGAN (ignored by ESRGAN)
93+
base_channels: 64 # ESRGAN discriminator base feature width (ignored otherwise)
94+
linear_size: 1024 # Hidden dim of ESRGAN discriminator head (ignored otherwise)
8995

9096
# ============================================================================ #
9197
# 🧮 OPTIMIZATION SETTINGS

opensr_srgan/configs/config_20m.yaml

Lines changed: 13 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -75,16 +75,22 @@ Training:
7575
# 🧩 ARCHITECTURAL PARAMETERS
7676
# ---------------------------------------------------------------------------- #
7777
Generator:
78-
model_type: 'rcab' # Block type: ['SRResNet', 'res', 'rcab', 'rrdb', 'lka', 'conditional_cgan'/'cgan']
79-
large_kernel_size: 9 # Kernel for head and tail conv layers
80-
small_kernel_size: 3 # Kernel for intermediate blocks
81-
n_channels: 96 # Number of feature channels (original 64)
82-
n_blocks: 32 # Number of residual/attention blocks (original 16)
78+
model_type: 'SRResNet' # Generator family: ['SRResNet', 'stochastic_gan', 'esrgan']
79+
block_type: 'rcab' # SRResNet block variant: ['standard', 'res', 'rcab', 'rrdb', 'lka']
80+
large_kernel_size: 9 # Kernel for head and tail conv layers (SRResNet/stochastic)
81+
small_kernel_size: 3 # Kernel for intermediate blocks (SRResNet/stochastic)
82+
n_channels: 96 # Feature width (RRDB/ESRGAN uses this as trunk width)
83+
n_blocks: 32 # Residual/attention blocks (ESRGAN: number of RRDB blocks)
8384
scaling_factor: 8 # Upscaling factor (e.g., 2×, 4×, 8×)
85+
growth_channels: 32 # ESRGAN-specific RRDB growth channels (ignored otherwise)
86+
res_scale: 0.2 # Residual scaling used by stochastic/ESRGAN variants
87+
out_channels: 6 # Optional ESRGAN output band override (defaults to Model.in_bands)
8488

8589
Discriminator:
86-
model_type: 'standard' # Discriminator architecture selector ['standard', 'patchgan']
87-
n_blocks: 8 # Number of convolutional blocks / layers: ['standard': 8, 'patchgan': 3]
90+
model_type: 'standard' # Discriminator architecture selector ['standard', 'patchgan', 'esrgan']
91+
n_blocks: 8 # Convolutional depth for SRGAN/PatchGAN (ignored by ESRGAN)
92+
base_channels: 64 # ESRGAN discriminator base feature width (ignored otherwise)
93+
linear_size: 1024 # Hidden dim of ESRGAN discriminator head (ignored otherwise)
8894

8995
# ============================================================================ #
9096
# 🧮 OPTIMIZATION SETTINGS

0 commit comments

Comments
 (0)