Add image mask for mask out dynamic parts in the images for feature e…

…xtraction. Copy the code from cvg#190

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: