Add CosPlace for global features extraction

README.md

@@ -95,7 +95,7 @@ We show in [`pipeline_SfM.ipynb`](https://nbviewer.jupyter.org/github/cvg/Hierar
 
 - Supported local feature extractors: [SuperPoint](https://arxiv.org/abs/1712.07629), [D2-Net](https://arxiv.org/abs/1905.03561), [SIFT](https://www.cs.ubc.ca/~lowe/papers/ijcv04.pdf), and [R2D2](https://arxiv.org/abs/1906.06195).
 - Supported feature matchers: [SuperGlue](https://arxiv.org/abs/1911.11763) and nearest neighbor search with ratio test, distance test, and/or mutual check.
-- Supported image retrieval: [NetVLAD](https://arxiv.org/abs/1511.07247), [AP-GeM/DIR](https://github.com/naver/deep-image-retrieval), and [OpenIBL](https://github.com/yxgeee/OpenIBL).
+- Supported image retrieval: [NetVLAD](https://arxiv.org/abs/1511.07247), [AP-GeM/DIR](https://github.com/naver/deep-image-retrieval), [OpenIBL](https://github.com/yxgeee/OpenIBL), and [CosPlace](https://github.com/gmberton/CosPlace).
 
 Using NetVLAD for retrieval, we obtain the following best results:
 

hloc/extract_features.py

@@ -133,6 +133,11 @@
         'output': 'global-feats-openibl',
         'model': {'name': 'openibl'},
         'preprocessing': {'resize_max': 1024},
+    },
+    'cosplace': {
+        'output': 'global-feats-cosplace',
+        'model': {'name': 'cosplace'},
+        'preprocessing': {'resize_max': 1024},
     }
 }
 

hloc/extractors/cosplace.py

@@ -0,0 +1,46 @@
+'''
+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,
+        }