Skip to content

struan-robertson/DEM-Void-Infilling

 
 

Repository files navigation

Terrain Model Processing with Machine Learning

A fork of this PyTorch implementation of the network described in the paper Generative Image Inpainting with Contextual Attention. This fork allows for the training of the network on all DEM filetypes (GeoTiff, NASA PDS3/4) supported by GDAL. The checkpoint images from this training are saved using the Matplotlib terrain colormap to allow for better visualisation.

Requirements

Known working versions

  • Python3 3.10
  • PyTorch 1.13
  • Torchvision 0.14
  • Numpy 1.24
  • GDAL 3.6
  • Matplotlib 3.6

Train the model

Edit config in train.py

  • dataset should be the path to a folder containing one or more DEM files, not to the file itself.
  • Each DEM is tiled into tiles of size specified by the user and then loaded into memory.
python train.py

Checkpoint images and models will be saved to the out dir, which can be changed in config.

Test trained model

The org file test.org can be used to test a trained model using RMSE, MAE and SSIM. The file can be run using org-babel with the Emacs editor, requiring the emacs-jupyter plugin, however all of the code would run the exact same in a Jupyter notebook. Org files can be converted into Jupiter notebooks using Pandoc.

Use trained model

The org file infill.org can be used to load a pre-trained model and infill DEMs with it.

Training Results

Trained over 64 high-resolution NAC DTMs with another 32 saved for testing. The network was trained on a local GPU with a batch size of 12 to allow for more accurate initial training. For the remaining epochs the network was trained on an Nvidia A100 80GB with a batch size of 64 until converged.

The network was trained with masks in the shape of a square to simulate smaller naturally occurring no-data voids such as craters, and splits down the entire image to simulate a DEM joining task.

Training Checkpoints

Checkpoint images are saved in a grid where:

  • The top is the masked image
  • The middle is the infilling result
  • The bottom is the ground truth

Additionaly the top half of the image is elevation values and the bottom half is slope values.

Box Mask

examples/box_training.png

Split Mask

examples/split_training.png

Infill Example

Void

examples/400_void.png

Infilled

examples/400.png

About

Filling of lunar voids with machine learning

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages

  • Python 99.3%
  • Shell 0.7%