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

This is the official repo for "DINO-Mix: Enhancing Visual Place Recognition with Foundational Vision Model and Feature Mixing"

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 link of this paper:[Paper][ArXiv]

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

Pretrained models of DINOv2

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:

• MixVPR
• GSV-Cities
• DINOv2