NASA-IBM Lunar-FM

Finetuning and inference release of the NASA-IBM Lunar Foundation Model foundation model. Two Python packages:

• ni_lfm/ — the model package (backbone, tokenizers, data utilities). Vendored; not edited in day-to-day work.
• terratorch_integration/ — TerraTorch-compatible datamodules, tasks, backbone wrappers, and runnable configs for lunar downstream tasks (crater detection, IMP segmentation, ice prospectivity, etc.). This is the working surface.

Pretraining code is not included.

Install

pyenv install -s 3.12.2
pyenv virtualenv 3.12.2 ni_lfm && pyenv activate ni_lfm
pip install -e .

(or conda create -n ni_lfm python=3.12 if you prefer conda.)

Weights and data

Model weights and config can be downloaded from HuggingFace using Python. See examples below:

from huggingface_hub import snapshot_download

# Only download model weights and config
snapshot_download(repo_id="nasa-ibm-ai4science/NASA-IBM-Lunar-Foundation-Model", allow_patterns="backbone/*", local_dir="./")

# Download entire model HuggingFace directory
snapshot_download(repo_id="nasa-ibm-ai4science/NASA-IBM-Lunar-Foundation-Model", local_dir="./")

# Download ice-prospectivity data
snapshot_download(repo_id="nasa-ibm-ai4science/Sombench-Ice-Prospectivity-Regression", local_dir="./")

Configs use two relative roots, data/ and backbone/, so no absolute paths are baked into any YAML. Point them at the shared release bundle with two symlinks:

B=<path_to_your_dir_containing_data_and_weights>
ln -sfn "$B/downstream_dataset"   data
ln -sfn "$B/checkpoints/backbone" backbone

That gives every config the paths it expects:

backbone/checkpoint.pt                 # base backbone checkpoint
backbone/config.yaml                   # pretraining config + per-modality info (required)
data/prospectivity_dataset/            # ice_prosp/
data/imp_dataset/                      # imp/
data/nac_craters_dataset/              # nac_craters/  (COCO: images/*.npy + annotations_min5px.json)
data/wac_craters_dataset/              # wac_craters/  (images_tiff/, metadata.parquet, train|val|test.json)

backbone_cfg is required for ni_lfm_v1_* backbones — the wrapper raises ValueError if missing.

To run against a different copy, you can either change the config path or re-point the symlinks. Single-value overrides also work, e.g. --model.init_args.model_args.backbone_checkpoint_path /other/checkpoint.pt.

Fine-tuning

Every YAML under terratorch_integration/configs/ is a runnable terratorch fit target:

PYTHONPATH=. terratorch fit -c terratorch_integration/configs/nac_craters/ni_lfm_ps8.yaml

Common overrides:

# Point at a specific data root without editing the yaml
PYTHONPATH=. terratorch fit -c <config>.yaml \
  --data.data_dir /path/to/data \
  --data.metadata_file /path/to/metadata.parquet \
  --data.annotations_file /path/to/annotations.json

Note: terratorch fit writes config.yaml/config_deploy.yaml to CWD by default — this is Lightning CLI's dumped merged config, not a project file. Delete after each run or configure save_config_kwargs to suppress.

Test / evaluate

PYTHONPATH=. terratorch test --config <config_from_finetuning>.yaml --ckpt_path <finetuned_model>.ckpt

Cluster (PBS / SLURM)

Example batch wrappers for cluster submission live at examples/pbs/run_finetuning.pbs (PBS) and examples/slurm/run_finetuning.sbatch (SLURM). Edit CFG_PATH, the scheduler directives (#PBS -W group_list / #SBATCH --account, etc.), and the conda env activation to match your site.

Repo layout

ni_lfm/
├── ni_lfm/                       # model package (backbone, tokenizers, data utils)
├── terratorch_integration/       # TerraTorch datamodules + tasks + configs
│   ├── README.md                 # package-level docs (backbones, tasks, determinism)
│   ├── configs/                  # runnable `terratorch fit` configs, grouped by task
│   │   ├── nac_craters/          # NAC crater detection
│   │   ├── wac_craters/          # WAC crater detection
│   │   │   ├── full_data/        #   100% of the train split
│   │   │   └── half_data/        #   50% of the train split (val/test still full)
│   │   ├── imp/                  # Irregular Mare Patch (IMP) segmentation
│   │   └── ice_prosp/            # Ice prospectivity
│   │       └── ablation/         #   modality-subset ablations (m2–m7)
│   ├── data_adapter.py           # LunarCraterDataModule, LunarNACDTMDataModule, LunarWACCraterDataModule
│   ├── data_utils.py             # D4DetectionTransform and related augmentations
│   ├── lunar_backbone.py         # TerraTorch backbone wrapper
│   ├── lunar_object_detection_task.py
│   ├── lunar_segmentation_task.py
│   ├── lunar_classification_task.py
│   ├── lunar_regression_task.py
│   ├── lunar_llrd_mixin.py       # layer-wise LR decay + split-group optimiser mixin
│   ├── lunar_register.py         # registers backbone variants with TerraTorch
│   ├── determinism.py            # deterministic drop-ins + Albumentations seeding callbacks
│   ├── necks.py                  # LearnedTokenProjection, SimpleFeaturePyramid, MultilayerSimpleFeaturePyramid
│   └── decoders.py               # SumFuseDeepGNDecoder
├── examples/
│   ├── pbs/                      # PBS batch scripts
│   ├── slurm/                    # SLURM batch scripts
├── README.md                     # this file
├── LICENSE                       # Apache-2.0
├── pyproject.toml
└── requirements.txt

Full documentation of TerraTorch is at https://torchgeo.org/terratorch/quick_start/.

License

Apache 2.0 — see LICENSE.