Merge remote-tracking branch 'origin/main'

Author: LucaRom

.github/workflows/build-on-tag.yml

@@ -1,5 +1,10 @@
+
+name: Build Docker Image on Tag (CUDA 12)
+
 on:
-  create
+  push:
+    tags:
+      - '*'
 
 env:
   REGISTRY: docker.io
@@ -28,16 +33,16 @@ jobs:
         images:  ${{ env.REGISTRY }}/${{ secrets.DOCKER_REPO }}
 
     - name: Build the docker image
-      run: docker build . --file Dockerfile --tag gdl-cuda11:${{ fromJSON(steps.meta.outputs.json).labels['org.opencontainers.image.version'] }} --build-arg GIT_TAG=${{ fromJSON(steps.meta.outputs.json).labels['org.opencontainers.image.version'] }}
+      run: docker build . --file Dockerfile --tag gdl-cuda12:${{ fromJSON(steps.meta.outputs.json).labels['org.opencontainers.image.version'] }} --build-arg GIT_TAG=${{ fromJSON(steps.meta.outputs.json).labels['org.opencontainers.image.version'] }}
 
     - name: Tag the docker image
-      run: docker tag gdl-cuda11:${{ fromJSON(steps.meta.outputs.json).labels['org.opencontainers.image.version'] }} ${{ secrets.DOCKER_REPO }}/gdl-cuda11:${{ fromJSON(steps.meta.outputs.json).labels['org.opencontainers.image.version'] }}
+      run: docker tag gdl-cuda12:${{ fromJSON(steps.meta.outputs.json).labels['org.opencontainers.image.version'] }} ${{ secrets.DOCKER_REPO }}/gdl-cuda12:${{ fromJSON(steps.meta.outputs.json).labels['org.opencontainers.image.version'] }}
 
     - name: Push the docker image
-      run: docker push ${{ secrets.DOCKER_REPO }}/gdl-cuda11:${{ fromJSON(steps.meta.outputs.json).labels['org.opencontainers.image.version'] }}
+      run: docker push ${{ secrets.DOCKER_REPO }}/gdl-cuda12:${{ fromJSON(steps.meta.outputs.json).labels['org.opencontainers.image.version'] }}
 
     - name: Tag the docker image to latest
-      run: docker tag gdl-cuda11:${{ fromJSON(steps.meta.outputs.json).labels['org.opencontainers.image.version'] }} ${{ secrets.DOCKER_REPO }}/gdl-cuda11:latest
+      run: docker tag gdl-cuda12:${{ fromJSON(steps.meta.outputs.json).labels['org.opencontainers.image.version'] }} ${{ secrets.DOCKER_REPO }}/gdl-cuda12:latest
 
     - name: Push the docker image (latest tag)
-      run: docker push ${{ secrets.DOCKER_REPO }}/gdl-cuda11:latest
+      run: docker push ${{ secrets.DOCKER_REPO }}/gdl-cuda12:latest

.gitignore

@@ -2,6 +2,9 @@
 *.idea**
 *.vscode**
 
+# Distribution / packaging
+*.egg-info/
+
 # Specific folders name
 waterloo_subset_512/
 mlruns/

Dockerfile

@@ -0,0 +1,30 @@
+# syntax=docker/dockerfile:1
+FROM nvidia/cuda:12.4.1-cudnn-runtime-ubuntu22.04
+
+RUN apt-get update && apt-get install -y curl bzip2 && \
+    curl -Ls https://micro.mamba.pm/api/micromamba/linux-64/latest | \
+    tar -xvj -C /usr/local/bin --strip-components=1 bin/micromamba && \
+    rm -rf /var/lib/apt/lists/*
+
+ENV MAMBA_DOCKERFILE_ACTIVATE=1 \
+    CONDA_ENV_NAME=geo-dl \
+    MAMBA_ROOT_PREFIX=/opt/conda \
+    PATH="/opt/conda/envs/geo-dl/bin:$PATH"
+
+WORKDIR /tmp
+COPY requirements.txt pyproject.toml ./
+
+RUN micromamba create -y -n $CONDA_ENV_NAME -c conda-forge python=3.10 pip && \
+    micromamba run -n $CONDA_ENV_NAME pip install --no-cache-dir -r requirements.txt && \
+    find $MAMBA_ROOT_PREFIX/envs/$CONDA_ENV_NAME -name "*.pyc" -delete 2>/dev/null || true && \
+    find $MAMBA_ROOT_PREFIX/envs/$CONDA_ENV_NAME -name "__pycache__" -type d -exec rm -rf {} + 2>/dev/null || true && \
+    micromamba clean -a -y
+
+RUN useradd -m -u 1000 gdl_user && mkdir -p /app && chown -R gdl_user /app
+USER gdl_user
+
+WORKDIR /app
+COPY --chown=gdl_user:gdl_user . /app
+
+ENTRYPOINT ["python"]
+CMD ["-m", "geo_deep_learning.train"]

README.md

@@ -4,17 +4,18 @@ A PyTorch Lightning-based framework for geospatial deep learning with multi-sens
 
 ## Overview
 
-Geo Deep Learning (GDL) is a modular framework designed for semantic segmentation of geospatial imagery using state-of-the-art deep learning models. Built on PyTorch Lightning, it provides efficient training pipelines for multi-sensor data with WebDataset support.
+Geo Deep Learning (GDL) is a modular framework designed to support a wide range of geospatial deep learning tasks such as semantic segmentation, object detection, and regression.
+Built on PyTorch Lightning, it provides efficient training pipelines for multi-sensor data.
 
 ## Features
 
-- **Multi-sensor Support**: Handle multiple Earth observation sensors simultaneously
-- **Modular Architecture**: Encoder-neck-decoder pattern with interchangeable components
-- **WebDataset Integration**: Efficient large-scale data loading and processing
-- **Multiple Model Types**: UNet++, SegFormer, DOFA (Dynamic-one-for-all Architecture)
-- **Distributed Training**: Multi-GPU training with DDP strategy
-- **MLflow Logging**: Comprehensive experiment tracking and model versioning
-- **Flexible Data Pipeline**: Support for CSV and WebDataset formats
+- **Multi-sensor Support**: Handle multiple Earth observation sensors simultaneously.
+- **Modular Architecture**: Encoder-neck-decoder pattern with interchangeable components.
+- **WebDataset Integration**: Efficient large-scale data loading and processing.
+- **Multiple Model Types**: UNet++, SegFormer, DOFA (Dynamic-one-for-all Architecture).
+- **Distributed Training**: Multi-GPU training with supported strategies.
+- **MLflow Logging**: Comprehensive experiment tracking and model versioning.
+- **Flexible Data Pipeline**: Support for CSV and WebDataset formats.
 
 ## Architecture
 
@@ -31,23 +32,47 @@ Geo Deep Learning (GDL) is a modular framework designed for semantic segmentatio
 └── samplers/              # Custom data sampling strategies
 ```
 
+## Requirements
+- Install [uv](https://docs.astral.sh/uv/) package manager for your OS.
+
 ## Quick Start
 
+1. **Clone the repository:**
 ```bash
-git clone <repository-url>
+git clone https://github.com/NRCan/geo-deep-learning.git
 cd geo-deep-learning
 ```
+2. **Install dependencies:**
 
-### Training
+For **GPU training** with CUDA 12.8:
+```bash
+uv sync --extra cu128
+```
 
+For **CPU-only** training:
 ```bash
-# Single GPU training
-python geo_deep_learning/train.py fit --config configs/dofa_config_RGB.yaml
+uv sync --extra cpu
 ```
+This creates a virtual environment in `.venv/` and installs all dependencies.
+
+3. **Activate the environment:**
+```bash
+# Linux/macOS
+source .venv/bin/activate
+
+# Windows
+.venv\Scripts\activate
+```
+
+Or use `uv run` to execute commands without manual activation:
+```bash
+uv run python geo_deep_learning/train.py fit --config configs/dofa_config_RGB.yaml
+```
+**Note:** *If you prefer to use conda or another environment manager, you can generate a `requirements.txt` file from the dependencies listed in `pyproject.toml` for manual installation.*
 
 ### Configuration
 
-Models are configured via YAML files in `configs/`:
+Models are configured via YAML files in the `configs/` directory:
 
 ```yaml
 model:
@@ -65,54 +90,53 @@ data:
     sensor_configs_path: "path/to/sensor_configs.yaml"
     batch_size: 16
     patch_size: [512, 512]
+
+trainer:
+  max_epochs: 100
+  precision: 16-mixed
+  accelerator: gpu
+  devices: 1
 ```
 
 ## Supported Models
 
-### DOFA (Domain-Oriented Foundation Architecture)
-- **DOFA Base**: 768-dim embeddings, suitable for most tasks
-- **DOFA Large**: 1024-dim embeddings, higher capacity
-- Multi-scale feature extraction with UperNet decoder
-- Support for wavelength-specific processing
-
 ### UNet++
-- Classic U-Net architecture with dense skip connections
-- Multiple encoder backbones (ResNet, EfficientNet, etc.)
-- Optimized for medical and satellite imagery
+- Classic U-Net architecture with dense skip connections.
+- Multiple encoder backbones (ResNet, EfficientNet, etc.).
+- Available through segmentation-models-pytorch.
 
 ### SegFormer
-- Transformer-based architecture for semantic segmentation
-- Hierarchical feature representation
-- Efficient attention mechanisms
+- Transformer-based architecture for semantic segmentation.
+- Hierarchical feature representation (MixTransformer encoder).
+- Multiple model sizes (B0-B5).
+
+### DOFA (Dynamic One-For-All foundation model)
+- **DOFA Base**: 768-dim embeddings, suitable for most tasks.
+- **DOFA Large**: 1024-dim embeddings, higher capacity.
+- Multi-scale feature extraction with UperNet decoder.
+- Support for wavelength-specific processing.
+
 
 ## Data Pipeline
 
 ### Multi-Sensor DataModule
-- **Sensor Mixing**: Combine data from multiple sensors during training
-- **WebDataset Format**: Efficient sharded data storage and loading
-- **Patch-based Processing**: Configurable patch sizes (default: 512x512)
-- **Data Augmentation**: Built-in augmentation pipeline
+- **Sensor Mixing**: Combine data from multiple sensors during training.
+- **WebDataset Format**: Efficient sharded data storage and loading.
 
 ### Supported Data Formats
-- **WebDataset**: Sharded tar files with metadata
-- **CSV**: Traditional CSV with file paths and labels
-- **Multi-sensor**: YAML configuration for sensor-specific settings
+- **WebDataset**: Sharded tar files with metadata.
+- **CSV**: Traditional CSV with file paths and labels.
+- **Multi-sensor**: YAML configuration for sensor-specific settings.
 
 ## Training Features
-
-- **Mixed Precision**: 16-bit mixed precision training
-- **Gradient Clipping**: Configurable gradient clipping
-- **Early Stopping**: Automatic training termination
-- **Model Checkpointing**: Best model saving based on validation metrics
-- **Visualization**: Built-in prediction visualization callbacks
-
-## Distributed Training
-
-The framework supports multi-GPU training with:
-- DDP (Distributed Data Parallel) strategy
-- Automatic mixed precision
-- Synchronized batch normalization
-- Efficient NCCL communication
+- **Large-scale training**: Distributed training strategies enabled with pytorch lightning.
+- **Mixed Precision Training**: 16-bit mixed precision for faster training.
+- **Gradient Clipping**: Configurable gradient clipping for stability.
+- **Early Stopping**: Automatic training termination based on validation metrics.
+- **Model Checkpointing**: Saves best models based on validation performance.
+- **MLflow Integration**: Experiment tracking, metrics logging, and model registry.
+- **Visualization Callbacks**: Built-in prediction visualization during training.
+- **Learning Rate Scheduling**: Cosine annealing, step decay, and more.
 
 ## Development
 

geo_deep_learning/config/__init__.py

@@ -0,0 +1 @@
+"""Logging configuration."""

configs/log_config.yaml geo_deep_learning/config/log_config.yamlconfigs/log_config.yaml renamed to geo_deep_learning/config/log_config.yaml

@@ -7,11 +7,12 @@
 
 import torch
 import torch.nn.functional as fn
-from models.segmentation.base import EncoderMixin
 from timm.layers import DropPath, to_2tuple, trunc_normal_
 from torch import Tensor, nn
 from torch.utils import model_zoo
 
+from geo_deep_learning.models.segmentation.base import EncoderMixin
+
 
 class Mlp(nn.Module):
     """MLP module."""

configs/logging_config.py …o_deep_learning/config/logging_config.pyconfigs/logging_config.py renamed to geo_deep_learning/config/logging_config.py configs/logging_config.py renamed to geo_deep_learning/config/logging_config.py

@@ -2,8 +2,12 @@
 
 import torch
 import torch.nn.functional as fn
-from models.decoders.segformer_mlp import Decoder
-from models.encoders.mix_transformer import DynamicMixTransformer, get_encoder
+
+from geo_deep_learning.models.decoders.segformer_mlp import Decoder
+from geo_deep_learning.models.encoders.mix_transformer import (
+    DynamicMixTransformer,
+    get_encoder,
+)
 
 from .base import BaseSegmentationModel
 

geo_deep_learning/models/encoders/mix_transformer.py

@@ -16,10 +16,10 @@
 from torchmetrics.segmentation import MeanIoU
 from torchmetrics.wrappers import ClasswiseWrapper
 
+from geo_deep_learning.models.segmentation.dofa import DOFASegmentationModel
+from geo_deep_learning.tools.visualization import visualize_prediction
 from geo_deep_learning.utils.models import load_weights_from_checkpoint
 from geo_deep_learning.utils.tensors import denormalization
-from models.segmentation.dofa import DOFASegmentationModel
-from tools.visualization import visualize_prediction
 
 # Ignore warning about default grid_sample and affine_grid behavior triggered by kornia
 warnings.filterwarnings(