frasertheking / blindzone_inpainting

Public data repository for the AIES spaceborne blind zone inpainting model

BlindPaint: a 3Net+ for Inpainting Radar Blind Zones, maintained by Fraser King

---

Development of a full-scale connected U-Net for reflectivity inpainting in spaceborne radar blind zones

This project is currently being written into a journal article for Artificial Intelligence for the Earth Systems.

Snowfall is a critical contributor to the global water-energy budget, with important connections to water resource management, flood mitigation, and ecosystem sustainability. However, traditional spaceborne remote monitoring of snowfall faces challenges due to a near-surface radar blind zone, which masks a portion of the atmosphere. Here, a deep learning model was developed to fill in missing data across these regions using surface radar and atmospheric climate variables. The model accurately predicts reflectivity, with significant improvements over conventional methods. This innovative approach enhances our understanding of reflectivity patterns and atmospheric interactions, bolstering advances in remote snowfall prediction.

This repository contains the model structure, training and testing code for the project. We also include a handful of example cases and preprocessed data on Zenodo if you wish to test the model on your own hardware.

Model Structure

Through an adaptation of the architecture outlined in Huang et al., 2020 for image segmentation, we developed a full-scale connected U-Net with deep supervision for block inpainting missing regions in vertically pointing radar. The encoder-decoder architecture allows the model to learn deep latent features in aloft cloud and surrounding climate state variables to intelligently reconstruct near-surface reflectivity values. This project also strongly builds on the findings of Geiss and Hardin, 2021. Our 3Net+ model architecture diagram is below:

Performance

Current pixel-scale and scene-scale accuracy levels are substantially improved over traditional linear inpainting techniques, with the largest improvements noted in multi-layer clouds, near surface reflectivity gradients, shallow snowfall and virga cases, and the prediction of mixed-phase clouds during data sparse conditions. An example performance diagram showing the detection capabilities of the model is shown below for a handful of cases:

Computational Costs

The UNet models we examine here have around 7 million trainable parameters. We used a combination of Azure cluster GPUs (V100s) and a high performance desktop (RTX 4090) for most training, which would take about half a day on average. Some additional performance details are included below:

Installation

To perform your own training tests, you'll want to clone this repo, and create a Conda environment to install the necessary packages using the below command. Note, you will need the GPU enabled version of Tensorflow.

conda env create -f environment.yml
source activate tflow

Examples

If you want to see how data is loaded, models are defined and training is performed, check out our example Jupyter Notebook (example.ipynb) which covers most of the basics to get you up and running with a small example dataset.

Structure

The repository is broken up into multiple modules, responsible for different parts of the preprocessing, training, testing, evaluating and plotting processes used throughout the project. For instance, loss functions are kept in the bp_loss.py file, while the general configuration settings are in the bp_configs.py file. Each of these are then imported when necessary into the larger run scripts to produce less coupled and more cohesive code.

Configs

As previously mentioned, most of the code editing you'll need to do to configure the project for your own needs can be found in bbp_configs.py. This file allows you to change the number of training channels on the fly, tracking information such as the run identifier, inpainting sizes, MC dropout values and other post-processing constants. I've included some examples below.

CHANNELS = 5
USE_DOP_SPW = False
RUN_CASE='128_' + str(CHANNELS) + 'chan_era5_nsa_oli_10km_dsv_long_hybrid'
RUN_CASE_DATA='128_' + str(CHANNELS) + 'chan_era5_nsa_oli_10km_dsv_long_hybrid/'
RUN_NAME='test-run-001'
DOWNFILL_SIZES = [16]
N_MC_TESTS = 50 
STD_CUTOFF = 1
PERC_CLOUD = 0.05
CLOUD_MASK = -0.5
PRECIPITATION_DBZ = -20
MAX_DBZ = 20

This is also where you can point to the downloaded test code on your system.

Contributing

Pull requests are welcome. For major changes, please open an issue first to discuss what you would like to change. Note that, as a living project, code is not as clean as it could (should) be, and unit tests need to be produced in future iterations to maintain stability.

Authors & Contact

• Fraser King, University of Michigan, [email protected]
• Claire Pettersen, University of Michigan
• Chris Fletcher, University of Waterloo
• Andrew Geiss, Pacific Northwest National Laboratory

Funding

This project was funded by NASA New (Early Career) Investigator Program (NIP) grant at the University of Michigan.