Add image mask for mask out dynamic parts in the images for feature extraction. Copy the code from cvg#190

Author: lxxue

hloc/extract_features.py

@@ -231,7 +231,8 @@ def main(conf: Dict,
          as_half: bool = True,
          image_list: Optional[Union[Path, List[str]]] = None,
          feature_path: Optional[Path] = None,
-         overwrite: bool = False) -> Path:
+         overwrite: bool = False,
+         mask_dir: Optional[Path] = None) -> Path:
     logger.info('Extracting local features with configuration:'
                 f'\n{pprint.pformat(conf)}')
 
@@ -256,6 +257,14 @@ def main(conf: Dict,
         name = dataset.names[idx]
         pred = model({'image': data['image'].to(device, non_blocking=True)})
         pred = {k: v[0].cpu().numpy() for k, v in pred.items()}
+        if mask_dir is not None:
+            mask_name = str(mask_dir / name) + '.png'
+            # print(mask_name)
+            mask = cv2.imread(mask_name)[:, :, 0]
+            valid_keypoint = mask[pred['keypoints'][:, 1].astype('int'), pred['keypoints'][:, 0].astype('int')]
+            pred['keypoints'] = pred['keypoints'][valid_keypoint > 0]
+            pred['descriptors'] = pred['descriptors'][:, valid_keypoint > 0]
+            pred['scores'] = pred['scores'][valid_keypoint > 0]
 
         pred['image_size'] = original_size = data['original_size'][0].numpy()
         if 'keypoints' in pred: