Add HardNet (#235)

cvg · Sep 20, 2022 · 25750cb · 25750cb
1 parent e566a8f
commit 25750cb
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/hloc/extractors/dog.py b/hloc/extractors/dog.py
@@ -36,6 +36,8 @@ class DoG(BaseModel):
     def _init(self, conf):
         if conf['descriptor'] == 'sosnet':
             self.describe = kornia.feature.SOSNet(pretrained=True)
+        elif conf['descriptor'] == 'hardnet':
+            self.describe = kornia.feature.HardNet(pretrained=True)
         elif conf['descriptor'] not in ['sift', 'rootsift']:
             raise ValueError(f'Unknown descriptor: {conf["descriptor"]}')
 
@@ -79,7 +81,7 @@ def _forward(self, data):
             if self.conf['descriptor'] == 'rootsift':
                 descriptors = sift_to_rootsift(descriptors)
             descriptors = torch.from_numpy(descriptors)
-        elif self.conf['descriptor'] == 'sosnet':
+        elif self.conf['descriptor'] in ('sosnet', 'hardnet'):
             center = keypoints[:, :2] + 0.5
             laf_scale = scales * self.conf['mr_size'] / 2
             laf_ori = -oris