update readme

Signed-off-by: Isabelle Wittmann <isabelle.wittmann1@ibm.com>

README.md

@@ -5,44 +5,85 @@ homepage: https://github.com/IBM/TerraCodec
 ---
 
 # TerraCodec
+
 **Neural Compression for Earth Observation**
 
-TerraCodec (TEC) is a family of pretrained neural compression codecs for **multispectral Sentinel-2 satellite imagery**. The models compress optical Earth observation data using learned latent representations and entropy coding.
+![License](https://img.shields.io/badge/License-Apache%202.0-blue)
+![arXiv](https://img.shields.io/badge/arXiv-2510.12670-b31b1b)
+![GitHub](https://img.shields.io/badge/GitHub-IBM%2FTerraCodec-black?logo=github)
+![PyPI](https://img.shields.io/badge/PyPI-terracodec-blue?logo=pypi)
+
+TerraCodec (TEC) is a family of pretrained neural compression codecs for **multispectral Sentinel-2 satellite imagery**.  
+The models compress optical Earth observation data using learned latent representations and entropy coding.
 
-Compared to classical codecs such as JPEG2000 or WebP, TerraCodec achieves **3–10× higher compression at comparable reconstruction quality** on multispectral satellite imagery. Temporal models further improve compression by exploiting redundancy across seasonal image sequences.
+Compared to classical codecs such as JPEG2000 or WebP, TerraCodec achieves **3–10× higher compression at comparable reconstruction quality** on multispectral satellite imagery. Temporal models further improve compression by exploiting redundancy across seasonal image sequences of satellite imagery.
 
-📄 Paper: https://arxiv.org/abs/2510.12670  
-💻 GitHub: https://github.com/IBM/TerraCodec  
+![Reconstructions](assets/reconstructions.png)
 
 ---
 
-# Models
+# Model Family
 
 | Model | Available Checkpoints | Description |
 |---|---|---|
 | `terracodec_v1_fp_s2l2a` | λ = 0.5, 2, 10, 40, 200 | Factorized-prior image codec. Smallest model and strong baseline for multispectral image compression. |
-| `terracodec_v1_elic_s2l2a` | λ = 0.5, 2, 10, 40, 200 | Enhanced entropy model with spatial and channel context, providing improved rate–distortion performance for image compression. |
-| `terracodec_v1_tt_s2l2a` | λ = 0.4, 1, 5, 20, 100, 200, 700 | Temporal Transformer codec modeling redundancy across seasonal multispectral image sequences. |
-| `flextec_v1_s2l2a` | **Single checkpoint** (quality = 1–16) | Flexible-rate temporal codec. One model supports multiple compression levels via token-based quality settings at inference time. |
+| `terracodec_v1_elic_s2l2a` | λ = 0.5, 2, 10, 40, 200 | Enhanced entropy model with spatial and channel context for improved rate–distortion performance. |
+| `terracodec_v1_tt_s2l2a` | λ = 0.4, 1, 5, 20, 100, 200, 700 | Temporal Transformer codec modeling redundancy across seasonal image sequences. |
+| `flextec_v1_s2l2a` | **Single checkpoint** (quality = 1–16) | Flexible-rate temporal codec. One model supports multiple compression levels via token-based quality settings. |
+
+Lower λ / quality → **higher compression**  
+Higher λ / quality → **higher reconstruction quality**
+
+---
+
+# Model Architecture
 
-Lower λ/ quality → **higher compression**  
-Higher λ/ quality → **higher reconstruction quality**
+This repository contains the **TEC-FP (Factorized Prior)** variant of TerraCodec.
 
-See the paper and GitHub for details.
+[![TEC-ELIC Architecture](assets/TEC_image_architecture.png)](assets/TEC_image_architecture.png)
+
+TEC-FP is a convolutional encoder–decoder neural compression model with a fully factorized entropy model for the latent representation. Each quantized latent variable is modeled independently without spatial or channel context.
+This design enables efficient parallel entropy coding. TEC-FP is the smallest and fastest image codec in the TerraCodec family and is optimized for 12-band Sentinel-2 imagery.
+
+See the paper for additional architectural and training details.
 
 ---
 
-# Installation
+# Input Format
+
+| Codec type      | Expected shape      | Example                |
+|-----------------|---------------------|-------------------------|
+| Image codecs    | `[B, C, H, W]`      | `[1, 12, 256, 256]`     |
+| Temporal codecs | `[B, T, C, H, W]`   | `[1, 4, 12, 256, 256]`  |
 
-```bash
-pip install terracodec
-```
+- Inputs use **12 Sentinel‑2 L2A spectral bands**.  
+- Recommended spatial size: **256×256**. TEC‑FP supports arbitrary spatial sizes; other models expect 256×256.  
+- Temporal codecs were pretrained on four seasonal frames, but can process any number of timesteps during inference (higher T increases compute).  
 
 ---
 
-# QuickStart
+# Normalization
 
-```bash
+Models were trained on **[SSL4EO-S12 v1.1](https://huggingface.co/datasets/embed2scale/SSL4EO-S12-v1.1)**.
+
+Inputs should be standardized per spectral band using dataset statistics. For S2L2A:
+
+```python
+mean = torch.tensor([793.243, 924.863, 1184.553, 1340.936, 1671.402, 2240.082, 2468.412, 2563.243, 2627.704, 2711.071, 2416.714, 1849.625])
+
+std = torch.tensor([1160.144, 1201.092, 1219.943, 1397.225, 1400.035, 1373.136, 1429.170, 1485.025, 1447.836, 1652.703, 1471.002, 1365.307])
+```
+
+# Usage
+
+Install TerraCodec:
+
+```
+pip install terracodec
+```
+
+Load pretrained models:
+```python
 from terracodec import terracodec_v1_fp_s2l2a
 
 model = terracodec_v1_fp_s2l2a(
@@ -50,53 +91,28 @@ model = terracodec_v1_fp_s2l2a(
     compression=10
 )
 
-# Fast Reconstruction
+# Fast reconstruction (no bitstream)
 reconstruction = model(inputs)
 
-# True Compression
+# True compression
 compressed = model.compress(inputs)
 reconstruction = model.decompress(**compressed)
 ```
 
----
-
-# Input Format
-
-
-| Codec type | Shape | Example |
-|---|---|---|
-| Image codecs | `[B, C, H, W]` | `[1, 12, 256, 256]` |
-| Temporal codecs | `[B, T, C, H, W]` | `[1, 4, 12, 256, 256]` |
-
-- **12 spectral bands** (Sentinel-2 L2A)
-- **Spatial size:** 256×256 recommended. TEC-FP accepts arbitrary sizes; all other models expect 256×256.
-- **Temporal models:** Models are pretrained on four seasonal frames but can process an arbitrary number of input timesteps at inference time. Using more frames increases the computational cost and therefore the required inference time.
-
-### Normalization
+# Feedback
+If you have questions, encounter issues or want to discuss improvements:
+- open an issue or discussion on GitHub
+- or contribute directly to the repository
 
-Models were trained on [SSL4EO-S12 v1.1](https://huggingface.co/datasets/embed2scale/SSL4EO-S12-v1.1). Inputs should be standardized per spectral band using dataset statistics.
-For S2L2A:
+GitHub repository: https://github.com/IBM/TerraCodec
 
-```python
-mean = torch.tensor([793.243, 924.863, 1184.553, 1340.936, 1671.402, 2240.082, 2468.412, 2563.243, 2627.704, 2711.071, 2416.714, 1849.625])
-std = torch.tensor([1160.144, 1201.092, 1219.943, 1397.225, 1400.035, 1373.136, 1429.170, 1485.025, 1447.836, 1652.703, 1471.002, 1365.307])
+# Citation
+If you use TerraCodec in your research, please cite:
 ```
-
----
-
-## Citation
-
-```bibtex
 @article{terracodec2025,
   title   = {TerraCodec: Neural Codecs for Earth Observation},
   author  = {Costa Watanabe, Julen and Wittmann, Isabelle and Blumenstiel, Benedikt},
   journal = {arXiv preprint arXiv:2510.12670},
   year    = {2025}
 }
-```
-
----
-
-## License
-
-Apache 2.0.
+```