SuperPoint Rpautrat (MIT license) (#1603)

* initial python implementation of superpoint rpautrat * working python implementation of superpoint * small tweaks to try and work around the GIL issue * fix missing os import * begging cpp impl of superpoint python model * finishing cpp superpoint impl using the same SPDetector interface * finalizing superpoint cpp impl, working with cpu but needs to be cleaned * fixing feature matching by reworking nms and filtering logic * speeding up nms with batched operations and cleaning up * updating conversion script * adding args for image dimensions and cuda usage to model tracer * wiring up UI to parse superpoint params * adding label to superpoint rpautrat ui * reverting some unintended ui changes * typo * oneline revert * removing nms and threshold filtering from cpp, this is handled internally by the superpoint model * using python interface to run the superpoint _to_torchscript.py script at runtime * generate and load model file on the first incoming frame * reverting unintented change * ui changes appear mysteriously again, reverting * remove topk from model to prevent issue when number of detected keypoints is lower then the k value (scripting fails in this case) * cleaning up for review * change dest for model file and remove debug logs * bump patch and add version comment * rm unucessary comment * use resources to load file to ensure it works when rtabmap isn't built from source * execute with pybind runpy instead of system call * only build superpoint rpautrat if we have torch and python support * changing param to accept a path to the weights .pth file directly * parse the default working directory to save the model file in * rm extraneous change * rm setters and re-initialize the detector whenever params change. We cannot support changing params after the model is constructed * update UI text * only enable superpoint rpautrat when built with python and torch * more build information regarding superpoint rpautrat * generate temporary python script in the working dir * rm unecessary changes to rtabmap_superpoint.py * fix comment and only add repo root to sys.path * introduce a new parameter for the superpoint python model definition * execute script from string instead of writing to a file * wrap parse params in a single compiler directive * remove descriptor spatial matching logic and rely on upstream RTAB-Map processes to take the top-K desc and kpts * only add python script to resources if built with superpoint rpautrat support. supress warning with type casting * isolate pybind11 setup so it can't affect any other modules or potential regnerations of the model file * update about dialog to show superpoint rpautrat * remove debug logs --------- Co-authored-by: Felix Toft <[email protected]>
2026-10-07 10:37:47 +08:00 · 2025-11-07 17:22:35 -08:00
parent f44a4fc478
commit 5e4fd171e2
16 changed files with 2083 additions and 1178 deletions
@@ -0,0 +1,71 @@
+#! /usr/bin/env python3
+#
+# Drop this file in the root folder of SuperPoint git: https://github.com/rpautrat/SuperPoint
+# To use with rtabmap:
+#   --Vis/FeatureType 15 --Kp/DetectorStrategy 15 --PyDetector/Path "~/SuperPoint/rtabmap_superpoint_rpautrat.py"
+#
+import numpy as np
+import os
+import torch
+from superpoint_pytorch import SuperPoint
+
+superpoint = []
+device = 'cpu'
+
+def init(cuda):
+    global superpoint, device
+
+    superpoint = SuperPoint().eval()
+
+    # set up device, gpu or cpu depending on the availability and the user's choice
+    device = 'cuda' if torch.cuda.is_available() and cuda else 'cpu'
+    
+    # Load weights directly to target device
+    # Get the directory where this script is located
+    script_dir = os.path.dirname(os.path.abspath(__file__))
+    weights_path = os.path.join(script_dir, 'weights', 'superpoint_v6_from_tf.pth')
+    
+    # Load model weights with proper error handling
+    try:
+        state_dict = torch.load(weights_path, map_location=device, weights_only=True)
+        superpoint.load_state_dict(state_dict)
+    except Exception as e:
+        print(f"Error loading weights: {e}")
+        raise
+    
+    # Move the model to the target device
+    superpoint.to(device)
+    
+    # Ensure model is in eval mode for inference
+    superpoint.eval()
+
+def detect(imageBuffer):
+    global superpoint, device
+    
+    image = np.asarray(imageBuffer)
+    image = (image.astype('float32') / 255.)
+    
+    try:
+        image_with_dims = image[None, None]  # Add batch and channel dims
+        image_tensor = torch.from_numpy(image_with_dims).float()
+        image_tensor = image_tensor.to(device)
+    except Exception as e:
+        print(f"Error creating tensor: {e}")
+        raise
+    # Result: (1, 1, H, W) - PyTorch tensor on correct device (CPU or GPU).
+    
+    with torch.no_grad():
+        pred = superpoint({'image': image_tensor})
+    
+    # Extract keypoints and descriptors
+    keypoints = pred['keypoints'][0].cpu().numpy()  # Shape: (N, 2)
+    keypoints_response = pred['keypoint_scores'][0].cpu().numpy()
+    keypoints_with_response = np.column_stack([keypoints, keypoints_response]).astype(np.float32)
+    # Result: (N, 3) with [x, y, response]
+
+    descriptors = pred['descriptors'][0].cpu().numpy()  
+    # Result: (N, descriptor_dim)    
+
+    desc = np.float32(descriptors).copy()
+    pts = np.float32(keypoints_with_response).copy()
+    return pts, desc