A PyTorch implementation of SuperGlue.
SuperGlue is a graph neural network that performs joint context aggregation, descriptors matching and filtering in a unified deep network architecture trained end-to-end. SuperGlue is made up of two major components: the Attentional Graph Neural Network and Optimal Matching Layer. The first component extracts matching descriptors from the input keypoints and descriptors using the graph neural network with attention mechanisms. The second one computes the score matrix with dustbins and finds optimal partial assignment using the Sinkhorn algorithm.
This repository provides two pretrained models:
- SuperGlue-DoGHardNet-200000steps.pt - pretrained model for keypoints detected with OpenCV SIFT (DoG detector) and HardNet descriptors
- SuperGlue-SIFT-50000steps.pt - model pretrained on OpenCV SIFT keypoints and descriptors
Note: Both of them are trained on less than 100k iterations which is significantly lower than the official pretrained models.
This code is tested with Python 3.7 in Linux.
You can install all required packages for both training and model inference from requirements.txt
pip install -r requirements.txtNavigate to the top level project directory and run
python demo.py <IMAGE0_PATH> <IMAGE1_PATH> <OUTPUT_IMAGE_PATH>For example, if you run demo script with these test images:
imgs/img0.png |
imgs/img1.png |
|---|---|
 |
 |
you should get the following result:
Also, you can run SMNN matcher on the same images by passing --do_smnn=True.
You'll get the following results, which are much worse than SuperGlue
You can launch training with default configuration with the following command from the project root directory:
CUDA_VISIBLE_DEVICES=0,1 PYTHONPATH="." python model_training/train.py model_training/config/train.yamlYou can change configuration by editing model_training/config/train.yaml configuration file.
For example, you can change dataset path by editing data.image_glob field in the configuration file.