SOLAR: Second-Order Loss and Attention for Image Retrieval

This repository contains the PyTorch implementation of our paper:

"SOLAR: Second-Order Loss and Attention for Image Retrieval"
Tony Ng, Vassileios Balntas, Yurun Tian, Krystian Mikolajczyk. ECCV 2020.
[arXiv] [short video] [long video] [ECCV Daily feature article] [OpenCV blog]

Before going further, please check out Filip Radenovic's great repository on image retrieval. Our solar-global module is heavily built upon it. If you use this code in your research, please also cite their work! [link to license]

Features

• Complete test scripts for large-scale image retrieval with solar-global
• Inference code for extracting local descriptors with solar-local
• Second-order attention map visualisation for large images
• Image matching visualisation
• Training code for image retrieval

Requirements

• Python 3
• PyTorch tested on 1.3.0 - 1.5.1, torchvision 0.5+
• opencv-python (cv2) tested on 3.3.0.10
• TensorBoard tested on 2.0.0+
• numpy
• PIL
• tqdm

Download model weights and descriptors

Begin with downloading our best models (both global and local) described in the paper, as well as the pre-computed descriptors of the 1M distractors set.

sh download.sh

The global model is saved at data/networks/resnet101-solar-best.pth and the local model at solar_local/weights/local-solar-345-liberty.pth. The descriptors of the 1M distractors are saved in the main directory (the file is quite big ~8GB, so it might take a while to download).

Testing our global descriptor

Here you can try out our pretrained model resnet101-solar-best.pth on the Revisiting Oxford and Paris dataset

Testing on R-Oxford, R-Paris

Once you've successfully downloaded the global model weights, run

python3 -m solar_global.examples.test

This script automatically downloads roxford5k,rparis6k into data/test/ and evaluates SOLAR on them. After a while, you should be able to get results as below

>> roxford5k: mAP E: 85.88, M: 69.9, H: 47.91
>> roxford5k: mP@k[1, 5, 10] E: [94.12 92.45 88.8 ], M: [94.29 90.86 86.71], H: [88.57 74.29 63.  ]

>> rparis6k: mAP E: 92.95, M: 81.57, H: 64.45
>> rparis6k: mP@k[1, 5, 10] E: [100.   96.57 95.43], M: [100.   98.   97.14], H: [97.14 94.57 93.  ]

Retrieval rankings are visualised in specs/ using

tensorboard --logdir specs/ --samples_per_plugin images=1000

You can view them on your browser at localhost:6006 in the IMAGES tab. Here's an example

You can also switch to the PROJECTOR tab and play around with TensorBoard's embedding visualisation tool. Here's an example of the 6322 database images in R-Paris, visualised with t-SNE

Testing with the extra 1-million distractors

If you decide to extract the descriptors on your own, run

(Note: this step takes a lot of time and storage, and we only provide it for verification. You can skip to the next command if you've already downloded the pre-computed descriptors from the previous step!)

python3 -m solar_global.examples.extract_1m

This script would download and extract the 1M distractors set and save them into data/test/revisitop1m/. This dataset is quite large (400GB+), so depending on your network & GPU, the whole process of downloading + extracting descriptors can take from a couple of days to a week. In our setting (~100MBps, V100), the download + extraction takes ~10 hours and the descriptors ~30 hours to be computed.

Now, make sure that resnet101-solar-best.pth_vecs_revisitop1m.pt is in the main directory, whether from the extraction step above or from the download ealier. Then you can run

python3 -m solar_global.examples.test_1m

and get results as below

>> roxford5k: mAP E: 72.04, M: 53.49, H: 29.89
>> roxford5k: mP@k[1, 5, 10] E: [88.24 81.99 76.96], M: [88.57 82.29 76.71], H: [74.29 58.29 48.86]

>> rparis6k: mAP E: 83.35, M: 59.19, H: 33.41
>> rparis6k: mP@k[1, 5, 10] E: [98.57 95.14 93.57], M: [98.57 96.29 94.86], H: [92.86 89.14 81.57]

Visualising second-order attention maps

Using our interactive visualisation tool

We provide a small demo for you to click around an image and interactively visualise the second-order attention (SOA) maps at different locations you select. (c.f. Section 4.3 in the paper for an in-depth analysis)

First, run

python3 -m demo.interactive_soa

This gorgeous image of the Eiffel Tower should pop up in a new window

Try drawing a (light green) rectangle centred at the location you would like to visualise the SOA map

A new window titled Second order attention with the SOA from the closest location in the feature map overlaid on the image, and a white dot indicating where you've selected should appear as below

Now, try drawing a rectangle in the sky, you should see the SOA more spread-out and silhouetting the main landmarks like this

You can keep clicking around the image to visualise more SOAs. Remember, the white dot in the SOA map tells you where the currently displayed attention map is selected from!

You can also try out different images by parsing the programme with

python3 -m demo.interactive_soa --image PATH/TO/YOUR/IMAGE

Jupyter-Notebook

Coming Soon!

Testing our local descriptor

Simple inference

We provide a bare-bones inference code for the local counterpart of SOLAR (Section 5.3 in the paper), so you can plug it into whatever applications you have for local descriptors.

To check that it works, run

python3 -m solar_local.example

If successful, it should display the following message

SOLAR_LOCAL - SOSNet w/ SOA layers:
SOA_3:
Num channels:    in   out   mid
                 64    64    16
SOA_4:
Num channels:    in   out   mid
                 64    64    16
SOA_5:
Num channels:    in   out   mid
                128   128    64
Descriptors shape torch.Size([512, 128])

Jupyter-Notebook

Follow our demo notebook to see a comparison between solar_local and the baseline SOSNet on an image-matching toy example.

Training SOLAR

Pre-processing the training set

As the GL18 dataset consists of only URLs, many of which have already expired, this part of the code lets you download the images we had at the time of training our models. However, this also means that extra storage space would be required for extracting tarballs, so please expect to have ~700GB upwards available. Otherwise, you could still download using GL18's downloader and save the images at data/train/gl18/jpg.

To download the images and pre-process them for training, simply run

sh gl18_preprocessing.sh

This would take sometime but you should then see around 1-million images in data/train/gl18/jpg and the pickle file data/train/gl18/db_gl18.pkl required for training.

If you downloaded the images from the URLs directly, please also make sure you download train.csv, boxes_split1.csv and boxes_split2.csv and save them into data/train/gl18. Then you can run

cd data/train/gl18 && python3 create_db_pickle.py

You should then see data/train/gl18/db_gl18.pkl successfully created.

Training

Once you've downloaded and pre-processed GL18, you can start the training with the settings described in the paper by running

python3 -m solar_global.examples.train specs/gl18 --training-dataset 'gl18' --test-datasets 'roxford5k,rparis6k' --arch 'resnet101' --pool 'gem' --p 3 --loss 'triplet' --pretrained-type 'gl18' --loss-margin 1.25 --optimizer 'adam' --lr 1e-6 -ld 1e-2 --neg-num 5 --query-size 2000 --pool-size 20000 --batch-size 8 --image-size 1024 --update-every 1 --whitening --soa --soa-layers '45' --sos --lambda 10 --no-val --print-freq 10 --flatten-desc

You can monitor the training losses and image pairs with tensorboard

tensorboard --logdir specs/

Citation

If you use this repository in your work, please cite our paper:

@inproceedings{ng2020solar,
    author    = {Ng, Tony and Balntas, Vassileios and Tian, Yurun and Mikolajczyk, Krystian},
    title     = {{SOLAR}: Second-Order Loss and Attention for Image Retrieval},
    booktitle = {ECCV},
    year      = {2020}
}