Skip to content

Commit

Permalink
Callback function to log Masked Autoencoder reconstructions to WandB (#…
Browse files Browse the repository at this point in the history
…88)

* ➕ Add wandb

A CLI and library for interacting with the Weights and Biases API!

* 🔊 Log Masked Autoencoder reconstructions to WandB

Created a custom callback function to log visualizations of the input and output images to the Masked Autoencoder. Only showing the RGB bands of Sentinel-2 for now. A sample of 6 image pairs (original + reconstructed, so 12 in total) is uploaded to Weights and Biases.

Example LightningCLI command: `python trainer.py fit --trainer.max_epochs=20 --data.data_path=data/32VLM --trainer.logger=WandbLogger --trainer.logger.project=clay --trainer.logger.save_dir=checkpoints --trainer.callbacks+=LogMAEReconstructedImage`.

* ➕ Add scikit-image

Image processing in Python!

* 📸 Apply histogram equalization to RGB images

Enhance low contrast images by applying a histogram equalization stretching algorithm on the RGB images, instead of dividing by a magic number like 6000.

* 🔧 Increase default sample size from 6 to 8

More samples to look at! Also only running einsum conversion on as many samples as needed rather than the whole batch, and handling cases where num_samples may be more than the batch_size.

* 🧑‍💻 Make wandb a somewhat optional dependency

Allows for `from src.callback_wandb import LogMAEReconstruction` to run, even without wandb being installed. Helpful if someone doesn't want to install wandb for whatever reason.

* ✅ Add unit test for LogMAEReconstruction

Testing that the LogMAEReconstruction callback works to save a set of images to WandB. Testing this in offline mode only, with checks that artifacts are saved locally, and that the wandb images have the correct caption and format.

* 🐛 Compare expected folders using set instead of list

Order of the folders could change, so using set instead of list.

* 🧪 Prevent WandB logger from saving logs to local drive for now

Setting WANDB_MODE="disabled", so no files are logged to disk, though the wandb.Image(s) are still created. See if this helps to resolve the exit code 255 issue on GitHub Actions.

* 📝 Fix a typo and improve docstring

Minor changes to the docstring of the on_validation_batch_end method, and a typo fix.
  • Loading branch information
weiji14 authored Dec 15, 2023
1 parent 9757bbf commit e259165
Show file tree
Hide file tree
Showing 7 changed files with 1,095 additions and 445 deletions.
Loading

0 comments on commit e259165

Please sign in to comment.