cosplace.py

'''
Code for loading models trained with CosPlace as a global features extractor
for geolocalization through image retrieval.
Multiple models are available with different backbones. Below is a summary of
models available (backbone : list of available output descriptors
dimensionality). For example you can use a model based on a ResNet50 with
descriptors dimensionality 1024.
    ResNet18:  [32, 64, 128, 256, 512]
    ResNet50:  [32, 64, 128, 256, 512, 1024, 2048]
    ResNet101: [32, 64, 128, 256, 512, 1024, 2048]
    ResNet152: [32, 64, 128, 256, 512, 1024, 2048]
    VGG16:     [    64, 128, 256, 512]

CosPlace paper: https://arxiv.org/abs/2204.02287
'''

import torch
import torchvision.transforms as tvf

from ..utils.base_model import BaseModel


class CosPlace(BaseModel):
    default_conf = {
        'backbone': 'ResNet50',
        'fc_output_dim' : 2048
    }
    required_inputs = ['image']
    def _init(self, conf):
        self.net = torch.hub.load(
            'gmberton/CosPlace',
            'get_trained_model',
            backbone=conf['backbone'],
            fc_output_dim=conf['fc_output_dim']
        ).eval()
        
        mean = [0.485, 0.456, 0.406]
        std = [0.229, 0.224, 0.225]
        self.norm_rgb = tvf.Normalize(mean=mean, std=std)

    def _forward(self, data):
        image = self.norm_rgb(data['image'])
        desc = self.net(image)
        return {
            'global_descriptor': desc,
        }