The watershed transform technique of https://arxiv.org/abs/1611.08303 is implemented to predict the number of gravitational lensing regions on a synthetic star image dataset. Each image has either 0, 1, or 2 gravitational lens fields. Here, watershed levels correspond to binning distances from the centers of the lenses, and running a classifier (trained using multiclass cross entropy loss) on those distances (we chose 16 bins non-uniformly spaced bins). A sample training image (which has two lenses) and its watershed level predictions are shown below:
Training/validation/test images are zipped. Labels (indicating how many lenses are in the image) are in training-labels.csv and validation-labels.csv, while coordinates of the centers of the lenses are in training-lenses.csv and validataion-lenses.csv.
To start training, unzip all image files in training.zip and validation.zip and run
python main.py
As a baseline, a resnet34 ternary classifier can also be trained using
python main.py --watershed False
To facilitate multiple training runs and training monitoring, detailed tensorboard summaries are saved for each run. Hyperparameters (such as batch size and learning rate) can easily be modified via the args available in cmd_args.py.
On an initial run with no hyperparameter fine tuning, we got 92% train / 88% validation for the watershed transform model
and a 99% train / 83% validation for the resnet34 classifier. Here, we chose clustered the union of level sets 0-9 (out of 16) to determine which regions are to count as lenses (the threshold 9 had the best performance on the train/validataion data). You can validate your trained models modifying the weights paths in validate.py and running the latter. Test images are stored in test.zip and sample test predictions of the aforementioned trained watershed model can be found in test_predictions.zip.
The conclusion is that the watershed transform learns a better set of features than a simple ternary classification and thus generalizes better. Moreover, the predictions of the watershed transform are more informative - instead of a single integer prediction, we get the fine-grained information of where the network thinks lens clusters are.