| """ | |
| 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, | |
| } | |