diff --git a/hloc/extract_features.py b/hloc/extract_features.py
index 1713f320..a2a82f85 100644
--- a/hloc/extract_features.py
+++ b/hloc/extract_features.py
@@ -158,6 +158,7 @@ def __init__(self, root, conf, paths=None):
         self.conf = conf = SimpleNamespace(**{**self.default_conf, **conf})
         self.root = root
 
+
         if paths is None:
             paths = []
             for g in conf.globs:
@@ -184,6 +185,7 @@ def __init__(self, root, conf, paths=None):
     def __getitem__(self, idx):
         name = self.names[idx]
         image = read_image(self.root / name, self.conf.grayscale)
+
         image = image.astype(np.float32)
         size = image.shape[:2][::-1]
 
@@ -217,10 +219,12 @@ 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)}')
-
+    
     loader = ImageDataset(image_dir, conf['preprocessing'], image_list)
     loader = torch.utils.data.DataLoader(loader, num_workers=1)
 
@@ -250,9 +254,18 @@ def main(conf: Dict,
             size = np.array(data['image'].shape[-2:][::-1])
             scales = (original_size / size).astype(np.float32)
             pred['keypoints'] = (pred['keypoints'] + .5) * scales[None] - .5
+            if mask_dir is not None:
+                mask = cv2.imread(str(mask_dir / name) + '.png')[:, :, 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]
             # add keypoint uncertainties scaled to the original resolution
             uncertainty = getattr(model, 'detection_noise', 1) * scales.mean()
 
+
+
         if as_half:
             for k in pred:
                 dt = pred[k].dtype