fix: rdd

imcui/hloc/extractors/rdd.py

@@ -9,6 +9,7 @@ sys.path.append(str(rdd_path))
 
 from RDD.RDD import build as build_rdd
 
+
 class Rdd(BaseModel):
     default_conf = {
         "keypoint_threshold": 0.1,
@@ -22,9 +23,7 @@ class Rdd(BaseModel):
         logger.info("Loading RDD model...")
         model_path = self._download_model(
             repo_id=MODEL_REPO_ID,
-            filename="{}/{}".format(
-                Path(__file__).stem, self.conf["model_name"]
-            ),
+            filename="{}/{}".format(Path(__file__).stem, self.conf["model_name"]),
         )
         config_path = rdd_path / "configs/default.yaml"
         with open(config_path, "r") as file:
@@ -38,6 +37,9 @@ class Rdd(BaseModel):
 
     def _forward(self, data):
         image = data["image"]
+        self.net.set_softdetect(
+            top_k=self.conf["max_keypoints"], scores_th=self.conf["keypoint_threshold"]
+        )
         pred = self.net.extract(image)[0]
         keypoints = pred["keypoints"]
         descriptors = pred["descriptors"]