Periodical synchronization between documentation from the working branch with main

docs/architecture.md

@@ -1,4 +1,4 @@
-# Overview (for developers)
+# Architecture Overview
 
 The main goal of the design is to extend TorchGeo's existing tasks to be able to handle Prithvi backbones with appropriate decoders and heads.
 At the same time, we wish to keep the existing TorchGeo functionality intact so it can be leveraged with pretrained models that are already included.
@@ -26,7 +26,7 @@ One of the most important design decisions was delegating the model construction
 - Prefers composition over inheritance
 - Allows new models to be easily added by introducing new factories
 
-Models are expected to be `torch.nn.Module`s and implement the [Model][terratorch.models.model.Model] interface, providing:
+Models are expected to be `torch.nn.Module` and implement the [Model][terratorch.models.model.Model] interface, providing:
 
 - `freeze_encoder()`
 - `freeze_decoder()`

docs/data.md

@@ -1,14 +1,28 @@
-# Data
-We rely on TorchGeo for the implementation of datasets and data modules.
+# Data Processing
 
-Check out [the TorchGeo tutorials on datasets](https://torchgeo.readthedocs.io/en/stable/tutorials/custom_raster_dataset.html) for more in depth information.
+In our workflow, we leverage TorchGeo to implement datasets and data modules, ensuring robust and flexible data handling. For a deeper dive into working with datasets using TorchGeo, please refer to the [TorchGeo tutorials on datasets](https://torchgeo.readthedocs.io/en/stable/tutorials/custom_raster_dataset.html).
 
-In general, it is reccomended you create a TorchGeo dataset specifically for your dataset. This gives you complete control and flexibility on how data is loaded, what transforms are done over it, and even how it is plotted if you log with tools like TensorBoard.
+In most cases, it’s best to create a custom TorchGeo dataset tailored to your specific data. Doing so gives you complete control over:
+- Data Loading: Customize how your data is read and organized.
+- Transforms: Decide which preprocessing or augmentation steps to apply.
+- Visualization: Define custom plotting methods (for example, when logging with TensorBoard).
 
-TorchGeo provides `GeoDataset` and `NonGeoDataset`.
+TorchGeo offers two primary classes to suit different data formats:
+- `NonGeoDataset`:
+  Use this if your dataset is already split into neatly tiled pieces ready for neural network consumption. Essentially, `NonGeoDataset` is a wrapper around a standard PyTorch dataset, making it straightforward to integrate into your pipeline.
+- `GeoDataset`:  
+  Opt for this class if your data comes in the form of large GeoTiff files from which you need to sample during training. `GeoDataset` automatically aligns your input data with corresponding labels and supports a range of geo-aware sampling techniques.
 
-- If your data is already nicely tiled and ready for consumption by a neural network, you can inherit from `NonGeoDataset`. This is essentially a wrapper of a regular torch dataset.
-- If your data consists of large GeoTiffs you would like to sample from during training, you can leverage the powerful `GeoDataset` from torch. This will automatically align your input data and labels and enable a variety of geo-aware samplers.
+
+In addition to these specialized TorchGeo datasets, TerraTorch offers generic datasets and data modules designed to work with directory-based data structures, similar to those used in MMLab libraries. These generic tools simplify data loading when your data is organized in conventional file directories:
+- The Generic Pixel-wise Dataset is ideal for tasks where each pixel represents a sample (e.g., segmentation or dense prediction problems).
+- The Generic Scalar Label Dataset is best suited for classification tasks where each sample is associated with a single label.
+
+TerraTorch also provides corresponding generic data modules that bundle the dataset with training, validation, and testing splits, integrating seamlessly with PyTorch Lightning. This arrangement makes it easy to manage data loading, batching, and preprocessing with minimal configuration.
+
+While generic datasets offer a quick start for common data structures, many projects require more tailored solutions. Custom datasets and data modules give you complete control over the entire data handling process—from fine-tuned data loading and specific transformations to enhanced visualization. By developing your own dataset and data module classes, you ensure that every step—from data ingestion to final model input—is optimized for your particular use case. TerraTorch’s examples provide an excellent starting point to build these custom components and integrate them seamlessly into your training pipeline.
+
+For additional examples on fine-tuning a TerraTorch model using these components, please refer to the [Prithvi EO Examples](https://github.com/NASA-IMPACT/Prithvi-EO-2.0) repository.
 
 ## Using Datasets already implemented in TorchGeo
 
@@ -73,8 +87,67 @@ For the `NonGeoDataset` case, we also provide "generic" datasets and datamodules
 
 ## Custom datasets and data modules
 
-Below is a documented example of how a custom dataset and data module class can be implemented.
+Our custom datasets and data modules are crafted to handle specific data, offering enhanced control and flexibility throughout the workflow. 
+In case you want to use TerraTorch on your specific data, we invite you to develop your own dataset and data module classes by following the examples below. 
+
+### Datasets
+#### :::terratorch.datasets.biomassters
+#### :::terratorch.datasets.burn_intensity
+#### :::terratorch.datasets.carbonflux
+#### :::terratorch.datasets.forestnet
+#### :::terratorch.datasets.fire_scars
+#### :::terratorch.datasets.landslide4sense
+#### :::terratorch.datasets.m_eurosat
+#### :::terratorch.datasets.m_bigearthnet
+#### :::terratorch.datasets.m_brick_kiln
+#### :::terratorch.datasets.m_forestnet
+#### :::terratorch.datasets.m_so2sat
+#### :::terratorch.datasets.m_pv4ger
+#### :::terratorch.datasets.m_cashew_plantation
+#### :::terratorch.datasets.m_nz_cattle
+#### :::terratorch.datasets.m_chesapeake_landcover
+#### :::terratorch.datasets.m_pv4ger_seg
+#### :::terratorch.datasets.m_SA_crop_type
+#### :::terratorch.datasets.m_neontree
+#### :::terratorch.datasets.multi_temporal_crop_classification
+#### :::terratorch.datasets.open_sentinel_map
+#### :::terratorch.datasets.openearthmap
+#### :::terratorch.datasets.pastis
+#### :::terratorch.datasets.sen1floods11
+#### :::terratorch.datasets.sen4agrinet
+#### :::terratorch.datasets.sen4map
+
+### Datamodules
+#### :::terratorch.datamodules.biomassters
+#### :::terratorch.datamodules.burn_intensity
+#### :::terratorch.datamodules.carbonflux
+#### :::terratorch.datamodules.forestnet
+#### :::terratorch.datamodules.fire_scars
+#### :::terratorch.datamodules.landslide4sense
+#### :::terratorch.datamodules.m_eurosat
+#### :::terratorch.datamodules.m_bigearthnet
+#### :::terratorch.datamodules.m_brick_kiln
+#### :::terratorch.datamodules.m_forestnet
+#### :::terratorch.datamodules.m_so2sat
+#### :::terratorch.datamodules.m_pv4ger
+#### :::terratorch.datamodules.m_cashew_plantation
+#### :::terratorch.datamodules.m_nz_cattle
+#### :::terratorch.datamodules.m_chesapeake_landcover
+#### :::terratorch.datamodules.m_pv4ger_seg
+#### :::terratorch.datamodules.m_SA_crop_type
+#### :::terratorch.datamodules.m_neontree
+#### :::terratorch.datamodules.multi_temporal_crop_classification
+#### :::terratorch.datamodules.open_sentinel_map
+#### :::terratorch.datamodules.openearthmap
+#### :::terratorch.datamodules.pastis
+#### :::terratorch.datamodules.sen1floods11
+#### :::terratorch.datamodules.sen4agrinet
+#### :::terratorch.datamodules.sen4map
+
+## Transforms
+The transforms module provides a set of specialized image transformations designed to manipulate spatial, temporal, and multimodal data efficiently. 
+These transformations allow for greater flexibility when working with multi-temporal, multi-channel, and multi-modal datasets, ensuring that data can be formatted appropriately for different model architectures.
+
+### :::terratorch.datasets.transforms
 
-### :::terratorch.datasets.fire_scars
 
-### :::terratorch.datamodules.fire_scars

docs/encoder_decoder_factory.md

@@ -1,8 +1,8 @@
 # EncoderDecoderFactory
+*Check the [Glossary](glossary.md) for more information about the terms used in this page.*
 
-## The Factory
-
-A special factory provided by terratorch is the [EncoderDecoderFactory][terratorch.models.encoder_decoder_factory.EncoderDecoderFactory].
+The [EncoderDecoderFactory][terratorch.models.encoder_decoder_factory.EncoderDecoderFactory] is the main class
+used to instantiate and compose models for general tasks. 
 
 This factory leverages the `BACKBONE_REGISTRY`, `DECODER_REGISTRY` and `NECK_REGISTRY` to compose models formed as encoder + decoder, with some optional glue in between provided by the necks.
 As most current models work this way, this is a particularly important factory, allowing for great flexibility in combining encoders and decoders from different sources.
@@ -25,14 +25,55 @@ The [EncoderDecoderFactory][terratorch.models.encoder_decoder_factory.EncoderDec
 
 ## Encoders
 
-To be a valid encoder, an object must be an `nn.Module` with an additional attribute `out_channels` which is a list of the channel dimension of the features it returns.
+To be a valid encoder, an object must be an `nn.Module` and contain an attribute `out_channels`, basically a list of the channel dimensions corresponding to
+the features it returns.
+The forward method of any encoder should return a list of `torch.Tensor`.
+
+```sh
+In [19]: backbone = BACKBONE_REGISTRY.build("prithvi_eo_v2_300", pretrained=True)
+
+In [20]: import numpy as np
+
+In [21]: import torch
+
+In [22]: input_image = torch.tensor(np.random.rand(1,6,224,224).astype("float32"))
+
+In [23]: output = backbone.forward(input_image)
 
-It's forward method should return a list of `torch.Tensor`.
+In [24]: [item.shape for item in output]
 
+Out[24]: 
+
+[torch.Size([1, 197, 1024]),
+ torch.Size([1, 197, 1024]),
+ torch.Size([1, 197, 1024]),
+ torch.Size([1, 197, 1024]),
+ torch.Size([1, 197, 1024]),
+ torch.Size([1, 197, 1024]),
+ torch.Size([1, 197, 1024]),
+ torch.Size([1, 197, 1024]),
+ torch.Size([1, 197, 1024]),
+ torch.Size([1, 197, 1024]),
+ torch.Size([1, 197, 1024]),
+ torch.Size([1, 197, 1024]),
+ torch.Size([1, 197, 1024]),
+ torch.Size([1, 197, 1024]),
+ torch.Size([1, 197, 1024]),
+ torch.Size([1, 197, 1024]),
+ torch.Size([1, 197, 1024]),
+ torch.Size([1, 197, 1024]),
+ torch.Size([1, 197, 1024]),
+ torch.Size([1, 197, 1024]),
+ torch.Size([1, 197, 1024]),
+ torch.Size([1, 197, 1024]),
+ torch.Size([1, 197, 1024]),
+ torch.Size([1, 197, 1024])]
+
+```
 
 ## Necks
 
-Necks are the glue between encoder and decoder. They can perform operations such as selecting elements from the output of the encoder ([SelectIndices][terratorch.models.necks.SelectIndices]), reshaping the outputs of ViTs so they are compatible with CNNs ([ReshapeTokensToImage][terratorch.models.necks.ReshapeTokensToImage]), amongst others.
+Necks are the connectors between encoder and decoder. They can perform operations such as selecting elements from the output of the encoder ([SelectIndices][terratorch.models.necks.SelectIndices]), reshaping the outputs of ViTs so they are compatible with CNNs ([ReshapeTokensToImage][terratorch.models.necks.ReshapeTokensToImage]), amongst others.
 
 Necks are `nn.Modules`, with an additional method `process_channel_list` which informs the [EncoderDecoderFactory][terratorch.models.encoder_decoder_factory.EncoderDecoderFactory] about how it will alter the channel list provided by `encoder.out_channels`.
 
@@ -53,12 +94,9 @@ Necks are `nn.Modules`, with an additional method `process_channel_list` which i
 ### :::terratorch.models.necks.LearnedInterpolateToPyramidal
 
 
-
-
-
 ## Decoders
 
-To be a valid decoder, an object must be an `nn.Module` with an additional attribute `out_channels` which is an `int` with the channel dimension of the output.
+To be a valid decoder, an object must be an `nn.Module` with an attribute `out_channels`, an `int` representing the channel dimension of the output.
 
 The first argument to its constructor will be a list of channel dimensions it should expect as input.
 
@@ -69,7 +107,7 @@ It's forward method should accept a list of embeddings.
 Most decoders require a final head to be added for a specific task (e.g. semantic segmentation vs pixel wise regression).
 
 Those registries producing decoders that dont require a head must expose the attribute `includes_head=True` so that a head is not added.
-Decoders passed as `nn.Modules` which dont require a head must expose the same attribute themselves.
+Decoders passed as `nn.Modules` which do not require a head must expose the same attribute themselves.
 
 ## :::terratorch.models.heads.classification_head.ClassificationHead
 

docs/examples.md

@@ -1,4 +1,34 @@
-# Examples
+# Performing an inference task with TerraTorch
+
+## Step 1: Download the test case from HuggingFace
+We will use the burn scars identification test case, in which we are interested in estimating the area
+affected by wildfires using a finetuned model (Prithvi-EO backbone + CNN decoder). To download the complete
+example, do:
+```sh
+git clone https://huggingface.co/ibm-nasa-geospatial/Prithvi-EO-2.0-300M-BurnScars/
+```
+## Step 2: Run the default inference case
+The example you download already contains some sample images to be used as input, so you just need to go to
+the local repository and create a directory to save the outputs:
+```sh
+cd Prithvi-EO-2.0-300M-BurnScars
+mkdir outputs
+```
+and to execute a command line like:
+```sh
+terratorch predict -c burn_scars_config.yaml --predict_output_dir outputs/ --data.init_args.predict_data_root examples/ --ckpt_path Prithvi_EO_V2_300M_BurnScars.pt
+```
+You will see the outputs being saved in the `outputs` directory. 
+
+### Input image (RGB components)
+
+![](figs/input.png)
+
+### Predicted mask
+
+![](figs/mask.png)}
+
+# More examples
 
 For some examples of training using the existing tasks, check out the following pages on our github repo:
 
@@ -16,4 +46,10 @@ Under `examples/confs`
 
 * Burn Scar Segmentation: `burn_scars.yaml`
 
-* Scene Classification: `eurosat.yaml`
+* Scene Classification: `eurosat.yaml`
+
+External examples available in [Prithvi-EO-2.0](https://github.com/NASA-IMPACT/Prithvi-EO-2.0)
+
+* [Carbon Flux](https://github.com/NASA-IMPACT/Prithvi-EO-2.0/tree/main/examples/carbon_flux)
+* [Landslide](https://github.com/NASA-IMPACT/Prithvi-EO-2.0/blob/main/examples/example_landslide4sense.ipynb)
+* [Multitemporal Crop](https://github.com/NASA-IMPACT/Prithvi-EO-2.0/blob/main/examples/example_multitemporalcrop.ipynb)

docs/glossary.md

@@ -0,0 +1,24 @@
+# Glossary of terms used in this Documentation and in the Geospatial AI area
+
+## Encoder
+The neural network used to map between the inputs and the intermdiary stage (usually referred as embedding
+or sometimes as latent space) of the forward step. The encoder is also frequently called backbone and, for 
+finetuning tasks, it is usually the part of the model which is not updated/trained. 
+
+## Decoder
+The neural network employed to map between the intermediary stage (embedding/latent space) and the target
+output. For finetuning tasks, the decoder is the most essential part, since it is trained to map the embedding
+produced by a previoulsy trained encoder to a new task. 
+
+## Head
+A network, usually very small when compared to the encoder and decoder, which is used as final step to adapt
+the decoder output to a specific task, for example, by applying a determined activation to it. 
+
+## Neck
+Necks are operations placed between the encoder and the decoder stages aimed at adjusting possible
+discrepancies, as incompatible shapes, or applying some specific transform, as a normalization required for the task being executed. 
+
+## Factory
+A Factory is a class which organizes the instantiation of a complete model, as a backbone-neck-decoder-head
+architecture. A class is intended to receive lists and dictionaries containing the required arguments used to
+build the model and returns a new instance already ready to be used. 

docs/index.md

@@ -1,26 +1,29 @@
 # Welcome to TerraTorch
+<img src="https://github.com/user-attachments/assets/f8c9586f-6220-4a53-9669-2aee3300b492" alt="TerraTorch"  width="400"/>
 
 ## Overview
 
-The purpose of this library is twofold:
-
-1. To integrate prithvi backbones into the TorchGeo framework.
-2. To provide generic LightningDataModules that can be built at runtime.
-3. To build a flexible fine-tuning framework based on TorchGeo which can be interacted with at different abstraction levels.
+The purpose of this package is to build a flexible fine-tuning framework for Geospatial Foundation Models (GFMs) based on TorchGeo and Lightning
+which can be employed at different abstraction levels. It currently supports models from the
+[Prithvi](https://huggingface.co/ibm-nasa-geospatial)
+and [Granite](https://huggingface.co/ibm-granite/granite-geospatial-land-surface-temperature) series, and also have been tested with others models available on HuggingFace. 
 
 This library provides:
 
 - All the functionality in TorchGeo.
-- Easy access to prithvi, timm and smp backbones.
+- Easy access to Prithvi, timm and smp backbones.
 - Flexible trainers for Image Segmentation, Pixel Wise Regression and Classification (more in progress).
-- Launching of fine-tuning tasks through powerful configuration files.
+    - Launching of fine-tuning tasks through powerful configuration files.
+
+    A good starting place is familiarization with [PyTorch Lightning](https://lightning.ai/docs/pytorch/stable/), which this project is built on. 
+    [TorchGeo](https://torchgeo.readthedocs.io/en/stable/) is also an important complementary reference. 
 
-A good starting place is familiarization with [PyTorch Lightning](https://lightning.ai/docs/pytorch/stable/), which this project is built on, and to a certain extent [TorchGeo](https://torchgeo.readthedocs.io/en/stable/)
+    Check out the [architecture overview](architecture.md) for a general description about how TerraTorch is
+    organized. 
 
 ## Quick start
 
 To get started, check out the [quick start guide](quick_start.md)
 
-## For developers
-
-Check out the [architecture overview](architecture.md)
+## License
+TerraTorch is distributed under the terms of License Apache 2.0, see [here](license.md) for more details. 

docs/models.md

@@ -1,19 +1,18 @@
 # Models
 
-To interface with terratorch tasks correctly, models must conform to the [Model][terratorch.models.model.Model] ABC:
+To interface with terratorch tasks correctly, models must inherit from the [Model][terratorch.models.model.Model] parent class:
 
 ::: terratorch.models.model.Model
 
-and have a forward method which returns a [ModelOutput][terratorch.models.model.ModelOutput]:
+and have a forward method which returns an object [ModelOutput][terratorch.models.model.ModelOutput]:
 
 :::terratorch.models.model.ModelOutput
 
 
 ## Model Factories
 
 In order to be used by tasks, models must have a Model Factory which builds them.
-
-Factories must conform to the [ModelFactory][terratorch.models.model.ModelFactory] ABC:
+Factories must conform to the [ModelFactory][terratorch.models.model.ModelFactory] parent class. 
 
 ::: terratorch.models.model.ModelFactory