Callback function to log Masked Autoencoder reconstructions to WandB

To visually inspect how the Masked Autoencoder is performing over the training run in terms of reconstructing to original image.

Implemented as a Lightning Callback function which runs at the end of a validation loop's mini-batch (on_validation_batch_end). A sample of 6 image pairs (original + reconstructed, so total of 12) are uploaded online.

Example usage with LightningCLI:

Samples of histogram equalized RGB images, and the reconstructed outputs (only random noise as this is early in the training).

python trainer.py fit --trainer.max_epochs=2 \
                      --trainer.precision=bf16-mixed \
                      --data.data_path=data/32VLM \
                      --data.num_workers=4 \
                      --trainer.logger=WandbLogger \
                      --trainer.logger.project=clay \
                      --trainer.logger.save_dir=checkpoints \
                      --trainer.callbacks+=LogMAEReconstruction \

TODO:

[x] Add wandb dependency
[x] Initial implementation to upload RGB images
[x] Do proper histogram equalization of the images
[x] Add a unit test

TODO in the future:

[ ] Upload SAR and DEM images

References:

Clay-foundation / model

Callback function to log Masked Autoencoder reconstructions to WandB #88