From e1d1c310dee402ce340d1acd451147e55ca1980f Mon Sep 17 00:00:00 2001 From: Lixin Xue Date: Wed, 29 Mar 2023 12:04:16 +0200 Subject: [PATCH] Add image mask for mask out dynamic parts in the images for feature extraction. Copy the code from #190 --- hloc/extract_features.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/hloc/extract_features.py b/hloc/extract_features.py index 4b9a617e..94092b85 100644 --- a/hloc/extract_features.py +++ b/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: