Skip to content

DINO-Mix: Enhancing Visual Place Recognition with Foundational Vision Model and Feature Mixing

License

Notifications You must be signed in to change notification settings

GaoShuang98/DINO-Mix

Repository files navigation

DINO-Mix: Enhancing Visual Place Recognition with Foundational Vision Model and Feature Mixing

Official Repo Paper Paper

Summary

Although most current VPR methods achieve favorable results under ideal conditions, their performance in complex environments, characterized by illumination, seasonal changes, and occlusions caused by moving objects, is generally unsatisfactory. Therefore, obtaining efficient and robust image feature descriptors even in complex environments is a pressing issue in VPR applications. In this study, we utilize the DINOv2 model as the backbone network for trimming and fine-tuning to extract robust image features. We propose a novel VPR architecture called DINO-Mix, which combines a foundational vision model with feature aggregation. This architecture relies on the powerful image feature extraction capabilities of foundational vision models. We employ an MLP-Mixer-based mix module to aggregate image features, resulting in globally robust and generalizable descriptors that enable high-precision VPR. We experimentally demonstrate that the proposed DINO-Mix architecture significantly outperforms current SOTA methods. In test sets having illumination, seasonal changes, and occlusions (Tokyo24/7, Nordland, SF-XL-Testv1), our proposed DINO-Mix architecture achieved Top-1 accuracy rates of 91.75%, 80.18%, and 82%, respectively. Compared with SOTA methods, our architecture exhibited an average accuracy improvement of 5.14%.

The framework of DINO-Mix as follows:

Trained models of DINO-Mix

All models have been trained on GSV-Cities dataset (https://github.com/amaralibey/gsv-cities).

Weights

Architecture Mix layer Output
dimension
Pitts30k-val Pitts30k-test Size Baidu Netdisk Password:DVPR
Google Drive
Top1 Top5 Top10 Top1 Top5 Top10
ViTg14-Mix 2 4096 92.34 98.59 99.20 87.71 94.25 96.11 4.2G LINK --
ViTl14-Mix 2 4096 93.86 99.13 99.68 91.27 96.43 97.62 1.1G LINK --
ViTb14-Mix 2 4096 94.37 98.86 99.41 92.03 95.89 97.17 334.4M LINK LINK(BEST)
ViTs14-Mix 2 4096 93.24 98.54 99.21 90.61 95.61 97.01 86.7M LINK --

Pretrained models of DINOv2

model # of
params
ImageNet
k-NN
ImageNet
linear
download
ViT-S/14 distilled 21 M 79.0% 81.1% backbone only
ViT-B/14 distilled 86 M 82.1% 84.5% backbone only
ViT-L/14 distilled 300 M 83.5% 86.3% backbone only
ViT-g/14 1,100 M 83.5% 86.5% backbone only

Code to load the ViTb14-mix model from torch_hub is as follows:

import torch
model = torch.hub.load('GaoShuang98/DINO-Mix', 'dino_mix', pretrained=True)

Code to load the pretrained weights is as follows:

from DINO_Mix import VPRModel
import torch

# Note that images must be resized to 224x224
model = VPRModel(
    backbone_arch='dinov2_vitb14',
    pretrained=True,
    layer1=7,
    use_cls=False,
    norm_descs=True,

    # ---- Aggregator
    agg_arch='DinoMixVPR',
    agg_config={'in_channels': 768,
                'in_h': 16,
                'in_w': 16,
                'out_channels': 1024,
                'mix_depth': 2,
                'mlp_ratio': 1,
                'out_rows': 4},
    )

checkpoint = torch.load(r"\DINO-Mix\dinov2_vitb14.ckpt", map_location='cuda')
if 'state_dict' in checkpoint:
    state_dict = checkpoint['state_dict']
else:
    state_dict = checkpoint
model_dict_weight = model.state_dict()
state_dict = {k: v for k, v in state_dict.items() if
              k in model_dict_weight}
model_dict_weight.update(state_dict)

# Find missing and unexpected weight parameters for pretrained models
missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
print("[missing_keys]:", *missing_keys, sep="\n")
print("[unexpected_keys]:", *unexpected_keys, sep="\n")
# Finally, load the content of the model pre-trained parameters
model.load_state_dict(model_dict_weight)

Bibtex

@article{huangDINOMixEnhancingVisual2024,
	title = {{DINO}-{Mix} enhancing visual place recognition with foundational vision model and feature mixing},
	volume = {14},
	issn = {2045-2322},
	url = {https://www.nature.com/articles/s41598-024-73853-3},
	doi = {10.1038/s41598-024-73853-3},
	language = {en},
	urldate = {2024-10-05},
	journal = {Scientific Reports},
	author = {Huang, Gaoshuang and Zhou, Yang and Hu, Xiaofei and Zhang, Chenglong and Zhao, Luying and Gan, Wenjian},
	month = sep,
	year = {2024},
	pages = {22100},
}

Acknowledgements

This code is based on the amazing work of: