mirror of
https://github.com/introlab/rtabmap.git
synced 2026-10-08 11:07:46 +08:00
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]>
This commit is contained in:
Felix Toft
@@ -0,0 +1,230 @@
|
||||
/**
|
||||
* SuperPoint implementation based on the PyTorch version by Rémi Pautrat, Paul-Edouard Sarlin
|
||||
* Adapted for RTAB-Map integration
|
||||
*/
|
||||
|
||||
#include "SuperpointRpautrat.h"
|
||||
#include <rtabmap/utilite/ULogger.h>
|
||||
#include <rtabmap/utilite/UDirectory.h>
|
||||
#include <rtabmap/utilite/UFile.h>
|
||||
#include <rtabmap/utilite/UConversion.h>
|
||||
#include <pybind11/embed.h>
|
||||
#include <torch/torch.h>
|
||||
#include <torch/script.h>
|
||||
#include <opencv2/opencv.hpp>
|
||||
|
||||
#include <fstream>
|
||||
#include <sstream>
|
||||
#include "superpoint_to_torchscript_py.h"
|
||||
|
||||
namespace rtabmap
|
||||
{
|
||||
|
||||
// Run the python script to export the SuperPoint model file with the desired parameters
|
||||
static std::string exportSuperPointTorchScript(
|
||||
const std::string & superpointWeightsPath,
|
||||
const std::string & superpointModelPath,
|
||||
const std::string & outputDir,
|
||||
const int & width,
|
||||
const int & height,
|
||||
const float & threshold,
|
||||
const int & nms_radius,
|
||||
const bool & cuda)
|
||||
{
|
||||
// Validate output directory is explicitly set and exists
|
||||
if(outputDir.empty())
|
||||
{
|
||||
UERROR("Output directory is not set.");
|
||||
return std::string("");
|
||||
}
|
||||
if(!UDirectory::exists(outputDir))
|
||||
{
|
||||
UERROR("Output directory does not exist: %s", outputDir.c_str());
|
||||
return std::string("");
|
||||
}
|
||||
|
||||
// Resolve paths (no dependency on source tree)
|
||||
const std::string weightsPath = superpointWeightsPath;
|
||||
const std::string modelPath = superpointModelPath;
|
||||
const std::string output = std::string(outputDir + "/superpoint_v6_from_tf.pt");
|
||||
|
||||
// Sanity checks
|
||||
if(!UFile::exists(weightsPath)) {
|
||||
UERROR("Weights not found: %s", weightsPath.c_str());
|
||||
return "";
|
||||
}
|
||||
if(!UFile::exists(modelPath)) {
|
||||
UERROR("Model not found: %s", modelPath.c_str());
|
||||
return "";
|
||||
}
|
||||
|
||||
// Execute the script inside the embedded Python interpreter
|
||||
try
|
||||
{
|
||||
pybind11::gil_scoped_acquire acquire;
|
||||
pybind11::dict scope;
|
||||
scope["__builtins__"] = pybind11::module_::import("builtins");
|
||||
|
||||
// set sys.path to the location of the model definition so it can be imported
|
||||
std::string model_dir = UDirectory::getDir(modelPath);
|
||||
auto sys = pybind11::module_::import("sys");
|
||||
pybind11::list sys_path = sys.attr("path");
|
||||
sys_path.attr("insert")(0, model_dir);
|
||||
|
||||
try {
|
||||
// execute the script to generate the model
|
||||
pybind11::exec(uHex2Str(SUPERPOINT_TO_TORCHSCRIPT_PY), scope, scope);
|
||||
pybind11::function generate_model = scope["generate_model"].cast<pybind11::function>();
|
||||
pybind11::object result = generate_model(weightsPath, output, cuda, nms_radius, threshold, width, height);
|
||||
sys_path.attr("remove")(model_dir);
|
||||
}
|
||||
catch(...) {
|
||||
// Ensure sys.path cleanup on any exception
|
||||
sys_path.attr("remove")(model_dir);
|
||||
throw;
|
||||
}
|
||||
|
||||
}
|
||||
// pybind11 throws std::exception for RuntimeError
|
||||
catch (const std::exception &e)
|
||||
{
|
||||
UERROR("Python export failed: %s", e.what());
|
||||
return "";
|
||||
}
|
||||
|
||||
return output;
|
||||
}
|
||||
|
||||
SPDetectorRpautrat::SPDetectorRpautrat(std::string superpointWeightsPath, std::string superpointModelPath, std::string outputDir, float threshold, bool nms, int minDistance, bool cuda) :
|
||||
device_(torch::kCPU),
|
||||
superpointWeightsPath_(superpointWeightsPath),
|
||||
superpointModelPath_(superpointModelPath),
|
||||
outputDir_(outputDir),
|
||||
threshold_(threshold),
|
||||
nms_(nms),
|
||||
minDistance_(minDistance),
|
||||
detected_(false)
|
||||
{
|
||||
if(cuda && !torch::cuda::is_available())
|
||||
{
|
||||
UWARN("Cuda option is enabled but torch doesn't have cuda support on this platform, using CPU instead.");
|
||||
}
|
||||
cuda_ = cuda && torch::cuda::is_available();
|
||||
|
||||
if(!UFile::exists(superpointWeightsPath_)) {
|
||||
UERROR("Superpoint weights not found: %s", superpointWeightsPath_.c_str());
|
||||
}
|
||||
|
||||
// Update device based on cuda availability
|
||||
device_ = torch::Device(cuda_ ? torch::kCUDA : torch::kCPU);
|
||||
}
|
||||
|
||||
SPDetectorRpautrat::~SPDetectorRpautrat()
|
||||
{
|
||||
}
|
||||
|
||||
cv::Mat SPDetectorRpautrat::compute(const std::vector<cv::KeyPoint> &keypoints)
|
||||
{
|
||||
if(!detected_)
|
||||
{
|
||||
UERROR("SPDetector has been reset before extracting the descriptors! detect() should be called before compute().");
|
||||
return cv::Mat();
|
||||
}
|
||||
if(keypoints.empty())
|
||||
{
|
||||
return cv::Mat();
|
||||
}
|
||||
|
||||
// These should have the same size
|
||||
UASSERT(static_cast<size_t>(desc_.size(0)) == keypoints.size());
|
||||
|
||||
// Move to CPU and return descriptors computed in the forward pass
|
||||
torch::Tensor desc_cpu = desc_.to(torch::kCPU);
|
||||
cv::Mat desc_mat(cv::Size(desc_cpu.size(1), desc_cpu.size(0)), CV_32FC1, desc_cpu.data_ptr<float>());
|
||||
return desc_mat.clone();
|
||||
}
|
||||
|
||||
std::vector<cv::KeyPoint> SPDetectorRpautrat::detect(const cv::Mat &img, const cv::Mat & mask)
|
||||
{
|
||||
// On first frame, run a trace of the model with the desired parameters and load the model file
|
||||
if(!detected_)
|
||||
{
|
||||
// effectively disable nms if it is not enabled by setting radius to 0
|
||||
int nms_radius = nms_ ? minDistance_ : 0;
|
||||
|
||||
std::string modelPath = exportSuperPointTorchScript(
|
||||
superpointWeightsPath_,
|
||||
superpointModelPath_,
|
||||
outputDir_,
|
||||
img.cols,
|
||||
img.rows,
|
||||
threshold_,
|
||||
nms_radius,
|
||||
cuda_
|
||||
);
|
||||
|
||||
UDEBUG("Initializing SuperPoint Rpautrat detector with model: %s", modelPath.c_str());
|
||||
UDEBUG("modelPath=%s thr=%f nms=%d minDistance=%d cuda=%d", modelPath.c_str(), threshold_, nms_?1:0, minDistance_, cuda_?1:0);
|
||||
if(modelPath.empty())
|
||||
{
|
||||
UERROR("Model's path is empty! The model was not exported correctly.");
|
||||
return std::vector<cv::KeyPoint>();
|
||||
}
|
||||
if(!UFile::exists(modelPath))
|
||||
{
|
||||
UERROR("Model's path \"%s\" doesn't exist!", modelPath.c_str());
|
||||
return std::vector<cv::KeyPoint>();
|
||||
}
|
||||
|
||||
// Load TorchScript model
|
||||
model_ = torch::jit::load(modelPath);
|
||||
model_.eval(); // put in evaluation mode
|
||||
model_.to(device_);
|
||||
}
|
||||
|
||||
// format the input tensor for the model
|
||||
torch::NoGradGuard no_grad_guard;
|
||||
auto x = torch::from_blob(img.data, {1, 1, img.rows, img.cols}, torch::kByte);
|
||||
x = x.to(torch::kFloat) / 255;
|
||||
x = x.set_requires_grad(false).to(device_);
|
||||
|
||||
auto outputs = model_.forward({x}).toTuple();
|
||||
keypoints_tensor_ = outputs->elements()[0].toTensor(); // [N, 2] keypoint coordinates
|
||||
auto scores_tensor = outputs->elements()[1].toTensor(); // [N] keypoint scores
|
||||
desc_ = outputs->elements()[2].toTensor(); // [N, 256] descriptors
|
||||
|
||||
// Convert to CPU for processing
|
||||
auto keypoints_cpu = keypoints_tensor_.to(torch::kCPU);
|
||||
auto scores_cpu = scores_tensor.to(torch::kCPU);
|
||||
|
||||
std::vector<cv::KeyPoint> filtered_keypoints;
|
||||
std::vector<int64_t> keep_indices_vec;
|
||||
|
||||
// Apply mask filtering
|
||||
for(int i = 0; i < keypoints_cpu.size(0); i++) {
|
||||
float score = scores_cpu[i].item<float>();
|
||||
float x = keypoints_cpu[i][0].item<float>(); // x coordinate
|
||||
float y = keypoints_cpu[i][1].item<float>(); // y coordinate
|
||||
|
||||
// Check mask if provided
|
||||
if(mask.empty() || mask.at<unsigned char>((int)y, (int)x) != 0) {
|
||||
keep_indices_vec.push_back(i);
|
||||
filtered_keypoints.emplace_back(cv::KeyPoint(x, y, 8, -1, score));
|
||||
}
|
||||
}
|
||||
|
||||
// Update the stored tensors to maintain correspondence
|
||||
// This way if keypoints are re-ordered, we can still match kpts->descs in the compute step
|
||||
auto keep_indices = torch::from_blob(keep_indices_vec.data(), {(long int)keep_indices_vec.size()}, torch::kLong);
|
||||
keep_indices = keep_indices.to(keypoints_tensor_.device());
|
||||
auto filtered_keypoints_tensor = keypoints_tensor_.index_select(0, keep_indices);
|
||||
auto filtered_descriptors = desc_.index_select(0, keep_indices);
|
||||
|
||||
keypoints_tensor_ = filtered_keypoints_tensor;
|
||||
desc_ = filtered_descriptors;
|
||||
|
||||
detected_ = true;
|
||||
return filtered_keypoints;
|
||||
}
|
||||
|
||||
} // namespace rtabmap
|
||||
@@ -0,0 +1,52 @@
|
||||
/**
|
||||
* SuperPoint implementation based on the PyTorch version by Rémi Pautrat, Paul-Edouard Sarlin
|
||||
* Adapted for RTAB-Map integration
|
||||
*/
|
||||
|
||||
#ifndef SUPERPOINT_RPAUTRAT_H
|
||||
#define SUPERPOINT_RPAUTRAT_H
|
||||
|
||||
#include <torch/torch.h>
|
||||
#include <opencv2/opencv.hpp>
|
||||
#include <vector>
|
||||
#include <memory>
|
||||
|
||||
namespace rtabmap
|
||||
{
|
||||
|
||||
class SPDetectorRpautrat {
|
||||
public:
|
||||
SPDetectorRpautrat(
|
||||
std::string superpointWeightsPath,
|
||||
std::string superpointModelPath,
|
||||
std::string outputDir,
|
||||
float threshold = 0.005f,
|
||||
bool nms = true,
|
||||
int nmsRadius = 4,
|
||||
bool cuda = false
|
||||
);
|
||||
virtual ~SPDetectorRpautrat();
|
||||
std::vector<cv::KeyPoint> detect(const cv::Mat &img, const cv::Mat & mask = cv::Mat());
|
||||
cv::Mat compute(const std::vector<cv::KeyPoint> &keypoints);
|
||||
|
||||
private:
|
||||
torch::jit::script::Module model_;
|
||||
torch::Device device_;
|
||||
torch::Tensor desc_;
|
||||
torch::Tensor keypoints_tensor_;
|
||||
|
||||
std::string superpointWeightsPath_;
|
||||
std::string superpointModelPath_;
|
||||
std::string outputDir_;
|
||||
float threshold_;
|
||||
bool nms_;
|
||||
int minDistance_;
|
||||
bool cuda_;
|
||||
|
||||
bool detected_;
|
||||
};
|
||||
}
|
||||
|
||||
|
||||
|
||||
#endif // SUPERPOINT_RPAUTRAT_H
|
||||
@@ -0,0 +1,107 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Convert PyTorch weights to TorchScript format for C++ usage.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from superpoint_pytorch import SuperPoint
|
||||
|
||||
|
||||
def wrap_model(model: nn.Module):
|
||||
"""
|
||||
Simple wrapper to fix SuperPoint input format for TorchScript.
|
||||
Easier to call from C++ code since the input isn't a dictionary.
|
||||
"""
|
||||
class Wrapper(nn.Module):
|
||||
def __init__(self, net: nn.Module):
|
||||
super().__init__()
|
||||
self.net = net
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
# SuperPoint expects {"image": tensor} but TorchScript doesn't like dict indexing
|
||||
out = self.net.forward({"image": x})
|
||||
|
||||
# Return the format expected by C++ code: keypoints, scores, descriptors
|
||||
# For single batch item, take the first (and only) element
|
||||
keypoints = out["keypoints"][0] if out["keypoints"] else torch.empty(0, 2)
|
||||
scores = out["keypoint_scores"][0] if out["keypoint_scores"] else torch.empty(0)
|
||||
descriptors = out["descriptors"][0] if out["descriptors"] else torch.empty(0, 256)
|
||||
|
||||
return (keypoints, scores, descriptors)
|
||||
|
||||
return Wrapper(model)
|
||||
|
||||
def generate_model(
|
||||
weights_path: str,
|
||||
output_path: str,
|
||||
cuda: bool,
|
||||
nms_radius: int,
|
||||
threshold: float,
|
||||
width: int,
|
||||
height: int,
|
||||
):
|
||||
# Check if weights are already TorchScript
|
||||
try:
|
||||
scripted = torch.jit.load(weights_path, map_location="cpu")
|
||||
scripted.eval()
|
||||
torch.jit.save(scripted, output_path)
|
||||
print(f"Converted TorchScript file: {output_path}")
|
||||
return
|
||||
except:
|
||||
pass
|
||||
|
||||
device = "cuda" if cuda else "cpu"
|
||||
|
||||
# Load SuperPoint model and weights
|
||||
model = SuperPoint(
|
||||
nms_radius=nms_radius,
|
||||
threshold=threshold,
|
||||
).eval().to(device)
|
||||
|
||||
# Load weights without forcing CPU location to allow CUDA usage
|
||||
weights = torch.load(weights_path, map_location=None)
|
||||
if isinstance(weights, dict) and "state_dict" in weights:
|
||||
weights = weights["state_dict"]
|
||||
|
||||
model.load_state_dict(weights, strict=False)
|
||||
|
||||
wrapped = wrap_model(model)
|
||||
dummy = torch.randn(1, 1, height, width, device=device) # Dummy input, grayscale, using cuda.
|
||||
|
||||
# Convert to TorchScript using trace (SuperPoint has dynamic behavior that scripting can't handle)
|
||||
print("Using torch.jit.trace (SuperPoint has dynamic behavior)...")
|
||||
scripted = torch.jit.trace(wrapped, (dummy,), strict=False)
|
||||
print("Successfully traced SuperPoint model")
|
||||
|
||||
# Save output
|
||||
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
||||
torch.jit.save(scripted, output_path)
|
||||
print(f"Converted SuperPoint weights to TorchScript: {output_path}")
|
||||
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Convert SuperPoint weights to TorchScript")
|
||||
parser.add_argument("--weights", required=True, help="Path to weights file")
|
||||
parser.add_argument("--output", required=True, help="Output TorchScript file")
|
||||
parser.add_argument("--cuda", action="store_true", help="Use CUDA")
|
||||
parser.add_argument("--width", type=int, default=1920, help="Width of the input image")
|
||||
parser.add_argument("--height", type=int, default=288, help="Height of the input image")
|
||||
parser.add_argument("--nms_radius", type=int, default=4, help="NMS radius")
|
||||
parser.add_argument("--threshold", type=float, default=0.005, help="Confidence threshold")
|
||||
args = parser.parse_args()
|
||||
print(f"Generating model from weights: {args.weights} to output: {args.output}")
|
||||
|
||||
generate_model(
|
||||
weights_path=args.weights,
|
||||
output_path=args.output,
|
||||
cuda=args.cuda,
|
||||
nms_radius=args.nms_radius,
|
||||
threshold=args.threshold,
|
||||
width=args.width,
|
||||
height=args.height,
|
||||
)
|
||||
Reference in New Issue
Block a user