Strong Gravitational Lensing Parameter Estimation with Vision Transformer

This repo contains all our code for using the Vision Transformer models on the imaging multi-regression task of parameter and uncertainty estimation for strong lensing systems. The paper is published in the ECCV 2022 Workshops and this is the arXiv link.

Authors: Kuan-Wei Huang, Geoff Chih-Fan Chen, Po-Wen Chang, Sheng-Chieh Lin, ChiaJung Hsu, Vishal Thengane, and Joshua Yao-Yu Lin

Tag v3.0.0 is the code version when the paper was submitted.

Data generation / preparation using Lenstronomy

• This notebook is used to generate the images (data) and paramters (targets) as the dataset for the strong lensing systems.
• This notebook is used to process the dataset: data split and target normalization.

Source code for training models

• train_model.py is the main code to train models (ViT and ResNet).
• The src folder contains scripts for helper functions.

Train models

• This notebook is used to train a ViT model.
• The training_eccv folder contains the notebooks used to train models for our ECCV paper.
• The training_icml folder contains the notebooks used to train models for our ICML paper.

Model Prediction and visulization

• predict.py is the source code to make prediction using a trained model.
• This notebook uses predict.py to make predictions for our ECCV paper.
• visualization.py contains objects and functions for visulization.
• This notebook uses visualization.py to make figures for our ECCV paper.