Merge pull request #71 from simon-donike/dev

Merge Dev: PL 2 support (exc training)

Author: simon-donike

opensr_srgan/configs/config_10m.yaml

@@ -33,9 +33,10 @@ Model:
 # ============================================================================ #
 # 🏋️ TRAINING CONFIGURATION
 # ---------------------------------------------------------------------------- #
-Training:  
+Training:
   # --- Hardware Setup
-  gpus: [0,1,2,3]                    # Number of GPUs to use, individually in list form, e.g. [0] or [0,2]
+  device: "cuda"                # Runtime device backend: ['cuda', 'cpu']
+  gpus: [2,3]                    # Number of GPUs to use, individually in list form, e.g. [0] or [0,2]
   # --- General Training Setup
   max_epochs: 9999            # Maximum number of training epochs
   val_check_interval: 1.0     # Validate at x percent of epoch (float) or every N steps (int)
@@ -114,6 +115,6 @@ Schedulers:
 Logging:
   num_val_images: 5            # Number of validation images logged per epoch
   wandb:
-    enabled: False              # Toggle Weights & Biases logging on/off
+    enabled: True              # Toggle Weights & Biases logging on/off
     entity: "opensr"      # W&B entity or team name
     project: "SRGAN_10m"        # W&B project name

opensr_srgan/configs/config_20m.yaml

@@ -35,6 +35,7 @@ Model:
 # ---------------------------------------------------------------------------- #
 Training:
   # --- Hardware Setup
+  device: "cuda"                # Runtime device backend: ['cuda', 'cpu']
   gpus: [2,3]                    # Number of GPUs to use, individually in list form, e.g. [0] or [0,2]
   # --- General Training Setup
   max_epochs: 9999            # Maximum number of training epochs
@@ -113,6 +114,6 @@ Schedulers:
 Logging:
   num_val_images: 5            # Number of validation images logged per epoch
   wandb:
-    enabled: False              # Toggle Weights & Biases logging on/off
+    enabled: True              # Toggle Weights & Biases logging on/off
     entity: "opensr"      # W&B entity or team name
     project: "SRGAN_20m"        # W&B project name

opensr_srgan/configs/config_training_example.yaml

@@ -33,8 +33,9 @@ Model:
 # ============================================================================ #
 # 🏋️ TRAINING CONFIGURATION
 # ---------------------------------------------------------------------------- #
-Training:  
+Training:
   # --- Hardware Setup
+  device: "cpu"                # Runtime device backend: ['cuda', 'cpu']
   gpus: [0]                    # Number of GPUs to use, individually in list form, e.g. [0] or [0,2]
   # --- General Training Setup
   max_epochs: 5               # Maximum number of training epochs

opensr_srgan/data/__init__.py

@@ -1,5 +1,3 @@
-from pathlib import Path
-
 def select_dataset(config):
     """
     Build train/val datasets from `config` and wrap them into a LightningDataModule.
@@ -23,6 +21,10 @@ def select_dataset(config):
         A tiny DataModule that exposes train/val DataLoaders built from the selected datasets.
     """
     dataset_selection = config.Data.dataset_type
+    
+    # Please Note: The "S2_6b","S2_4b","SISR_WW" settings are leftover from previous versions
+    # I dont want to delete them in case they are needed again.
+    # Only the "ExampleDataset" is actively used in the current version.
 
     if dataset_selection == "S2_6b":
         # Import here to avoid import costs when other datasets are used elsewhere.
@@ -104,10 +106,10 @@ def select_dataset(config):
         path = "example_dataset/"
         ds_train = ExampleDataset(folder=path, phase="train")
         ds_val = ExampleDataset(folder=path, phase="val")
-        
     else:
         # Centralized error so unsupported keys fail loudly & clearly.
-        raise NotImplementedError(f"Dataset {dataset_selection} not implemented")
+        raise NotImplementedError(f"Dataset {dataset_selection} not implemented!"
+                                  f"Add your dataset in data/dataset_selector.py to train on that.")
 
     # Wrap the two datasets into a LightningDataModule with config-driven loader knobs.
     pl_datamodule = datamodule_from_datasets(config, ds_train, ds_val)

opensr_srgan/data/dataset_selector.py

@@ -1,24 +1,33 @@
 from huggingface_hub import hf_hub_download
-import zipfile
-import os
+import zipfile, os
 
 def get_example_dataset(out_dir: str = "example_dataset/"):
     """Download and extract the example dataset for SRGAN training."""
-    # make sure the target dir exists
     os.makedirs(out_dir, exist_ok=True)
 
-    # download the file from your repo
     repo_id = "simon-donike/SR-GAN"
     filename = "example_dataset.zip"
 
-    print("Downloading from Hugging Face Hub...")
+    print("📦 Downloading from Hugging Face Hub...")
     zip_path = hf_hub_download(repo_id=repo_id, filename=filename)
 
-    # unzip after download
     with zipfile.ZipFile(zip_path, "r") as z:
-        z.extractall(out_dir)
+        members = z.namelist()
 
+        # detect common top-level folder (e.g. "example_data/")
+        prefix = os.path.commonprefix(members)
+        if prefix and prefix.endswith("/"):
+            for member in members:
+                # strip the prefix
+                target = member[len(prefix):]
+                if not target:  # skip folder itself
+                    continue
+                target_path = os.path.join(out_dir, target)
+                os.makedirs(os.path.dirname(target_path), exist_ok=True)
+                with z.open(member) as src, open(target_path, "wb") as dst:
+                    dst.write(src.read())
+        else:
+            z.extractall(out_dir)
+
+    os.remove(zip_path)
     print(f"✅ Extracted dataset to: {os.path.abspath(out_dir)}")
-    
-    # delete the zip file to save space
-    os.remove(zip_path)

opensr_srgan/data/example_data/download_example_dataset.py

@@ -12,11 +12,11 @@
 from torch.optim.lr_scheduler import ReduceLROnPlateau
 
 # local imports
-from ..utils.logging_helpers import plot_tensors
-from ..utils.model_descriptions import print_model_summary
-from ..utils.radiometrics import histogram as histogram_match
-from ..utils.radiometrics import normalise_10k
-from .model_blocks import ExponentialMovingAverage
+from opensr_srgan.utils.logging_helpers import plot_tensors
+from opensr_srgan.utils.model_descriptions import print_model_summary
+from opensr_srgan.utils.radiometrics import histogram as histogram_match
+from opensr_srgan.utils.radiometrics import normalise_10k
+from opensr_srgan.model.model_blocks import ExponentialMovingAverage
 
 
 #############################################################################################################
@@ -44,15 +44,17 @@ def __init__(self, config="config.yaml", mode="train"):
         # SECTION: Load Configuration
         # Purpose: Load and parse model/training hyperparameters from YAML file.
         # ======================================================================
-        if isinstance(config, Path) or isinstance(config, str):
-            self.config = OmegaConf.load(config)  # load config file with OmegaConf
+        if isinstance(config, str) or isinstance(config, Path):
+            config = OmegaConf.load(config)     
         elif isinstance(config, dict):
-            self.config = OmegaConf.create(config)          # create config from dict
+            config = OmegaConf.create(config)
         elif OmegaConf.is_config(config):
-            self.config = config                           # already an OmegaConf object
+            pass
         else:
-            print("Invalid config type; must be file path or OmegaConf/dict.")
+            raise TypeError("Config must be a filepath (str or Path), dict, or OmegaConf object.")
         assert mode in {"train", "eval"}, "Mode must be 'train' or 'eval'"  # validate mode
+        
+        self.config = config
         self.mode = mode                                # store mode (train/eval)
 
         # --- Training settings ---
@@ -91,7 +93,7 @@ def __init__(self, config="config.yaml", mode="train"):
         # Purpose: Configure generator content loss and discriminator adversarial loss.
         # ======================================================================
         if self.mode == "train":
-            from .loss import GeneratorContentLoss
+            from opensr_srgan.model.loss import GeneratorContentLoss
             self.content_loss_criterion = GeneratorContentLoss(self.config)  # perceptual loss (VGG + pixel)
             self.adversarial_loss_criterion = torch.nn.BCEWithLogitsLoss()   # binary cross-entropy for D/G
 
@@ -109,7 +111,7 @@ def get_models(self, mode):
 
         if generator_type == 'SRResNet':
             # Standard SRResNet generator
-            from .generators.srresnet import Generator
+            from opensr_srgan.model.generators.srresnet import Generator
             self.generator = Generator(
                 in_channels=self.config.Model.in_bands,                # number of input channels
                 large_kernel_size=self.config.Generator.large_kernel_size,
@@ -120,7 +122,7 @@ def get_models(self, mode):
             )
         elif generator_type in ['res', 'rcab', 'rrdb', 'lka']:
             # Advanced generator variants (ResNet, RCAB, RRDB, etc.)
-            from .generators.flexible_generator import FlexibleGenerator
+            from opensr_srgan.model.generators.flexible_generator import FlexibleGenerator
             self.generator = FlexibleGenerator(
                 in_channels=self.config.Model.in_bands,
                 n_channels=self.config.Generator.n_channels,
@@ -131,7 +133,7 @@ def get_models(self, mode):
                 block_type=self.config.Generator.model_type
             )
         elif generator_type.lower() in ['conditional_cgan', 'cgan']:
-            from .generators import ConditionalGANGenerator
+            from opensr_srgan.model.generators import ConditionalGANGenerator
 
             self.generator = ConditionalGANGenerator(
                 in_channels=self.config.Model.in_bands,
@@ -156,7 +158,7 @@ def get_models(self, mode):
             n_blocks = getattr(self.config.Discriminator, 'n_blocks', None)
 
             if discriminator_type == 'standard':
-                from .discriminators.srgan_discriminator import Discriminator
+                from opensr_srgan.model.discriminators.srgan_discriminator import Discriminator
 
                 discriminator_kwargs = {
                     "in_channels": self.config.Model.in_bands,
@@ -166,7 +168,7 @@ def get_models(self, mode):
 
                 self.discriminator = Discriminator(**discriminator_kwargs)
             elif discriminator_type == 'patchgan':
-                from .discriminators.patchgan import PatchGANDiscriminator
+                from opensr_srgan.model.discriminators.patchgan import PatchGANDiscriminator
 
                 patchgan_layers = n_blocks if n_blocks is not None else 3
                 self.discriminator = PatchGANDiscriminator(
@@ -198,9 +200,9 @@ def predict_step(self, lr_imgs):
         lr_min, lr_max = lr_imgs.min().item(), lr_imgs.max().item()  # get value range
         if lr_max > 1.5:  # Sentinel-2 style raw reflectance → normalize
             lr_imgs = normalise_10k(lr_imgs, stage="norm")           # normalize to 0–1 range
-            normalized = True
+            needs_normalization = True
         else:
-            normalized = False                                       # already normalized
+            needs_normalization = False                                       # already normalized
 
         # --- Perform super-resolution (optionally using EMA weights) ---
         context = self.ema.average_parameters(self.generator) if self.ema is not None else nullcontext()
@@ -211,7 +213,7 @@ def predict_step(self, lr_imgs):
         sr_imgs = histogram_match(lr_imgs, sr_imgs)                  # match distributions
 
         # --- Denormalize only if normalization was applied ---
-        if normalized:
+        if needs_normalization:
             sr_imgs = normalise_10k(sr_imgs, stage="denorm")         # convert back to original scale
 
         # --- Move to CPU and return ---
@@ -300,6 +302,7 @@ def training_step(self,batch,batch_idx,optimizer_idx):
             # run discriminator and get loss between pred labels and true labels
             sr_discriminated = self.discriminator(sr_imgs)                             # D(SR): logits for generator outputs
             adversarial_loss = self.adversarial_loss_criterion(sr_discriminated, torch.ones_like(sr_discriminated)) # keep taargets 1.0 for G loss
+            self.log("generator/adversarial_loss",adversarial_loss,sync_dist=True)     # log unweighted adversarial loss
 
             """ 3. Weight the losses"""
             adv_weight = self._adv_loss_weight() # get adversarial weight based on current step
@@ -317,8 +320,9 @@ def optimizer_step(
         optimizer,
         optimizer_idx,
         optimizer_closure,
-        on_tpu=False,
+        on_tpu=False, # these arguments are needed in case we're running on PL>2.0
         using_lbfgs=False,
+        
     ):
         optimizer.step(closure=optimizer_closure)
         optimizer.zero_grad()
@@ -485,13 +489,13 @@ def configure_optimizers(self):
             optimizer_g, mode='min',
             factor=self.config.Schedulers.factor_g,
             patience=self.config.Schedulers.patience_g,
-            verbose=self.config.Schedulers.verbose
+            #verbose=self.config.Schedulers.verbose
         )
         scheduler_d = ReduceLROnPlateau(
             optimizer_d, mode='min',
             factor=self.config.Schedulers.factor_d,
             patience=self.config.Schedulers.patience_d,
-            verbose=self.config.Schedulers.verbose
+            #verbose=self.config.Schedulers.verbose
         )
 
         # optional generator warmup scheduler (step-based)
@@ -556,7 +560,7 @@ def on_fit_start(self):  # called once at the start of training
         # SECTION: Print Model Summary
         # Purpose: Output model architecture and parameter counts (only once).
         # ======================================================================
-        from ..utils.gpu_rank import _is_global_zero
+        from opensr_srgan.utils.gpu_rank import _is_global_zero
         if _is_global_zero():
             print_model_summary(self)  # print model summary to console
 
@@ -756,8 +760,7 @@ def load_from_checkpoint(self,ckpt_path):
 
 
 if __name__=="__main__":
-    config_path = Path(__file__).resolve().parents[1] / "configs" / "config_20m.yaml"
-    model = SRGAN_model(config_file_path=str(config_path))
-    model.forward(torch.randn(1,6,32,32))
-    
-    model.load_from_checkpoint("logs/SRGAN_6bands/2025-10-11_23-53-20/last.ckpt")
+    config_path = "opensr_srgan/configs/config_10m.yaml"
+    model = SRGAN_model(config=str(config_path))
+    model.forward(torch.randn(1,4,32,32))
+