src/tracking/NeuralTracker.h - master - klischa/AstraScanner2

src/tracking/NeuralTracker.h
57 строк2 KB
Arena Agent
tracking: harden neural pose runtime validation
17 июн 2026, 19:47
17 июн 2026, 19:4747a6097
#ifndef NEURALTRACKER_H
#define NEURALTRACKER_H

#include <Eigen/Dense>
#include <pcl/point_cloud.h>
#include <pcl/point_types.h>
#include <memory>
#include <string>
#include <random>
#include <optional>
#include <QMutex>
#include <onnxruntime_cxx_api.h>
#include "NeuralTrackingGuards.h"

class NeuralTracker
{
public:
    NeuralTracker();
    ~NeuralTracker() = default;

    bool loadModel(const std::string& onnxPath, bool useCUDA = false);
    std::optional<Eigen::Matrix4f> tryEstimatePose(const pcl::PointCloud<pcl::PointXYZRGB>::Ptr& source,
                                                   const pcl::PointCloud<pcl::PointXYZRGB>::Ptr& target,
                                                   int maxPoints = 2048);
    // Backward-compatible wrapper: returns Identity only for callers that still
    // expect a matrix value. New code should prefer tryEstimatePose().
    Eigen::Matrix4f estimatePose(const pcl::PointCloud<pcl::PointXYZRGB>::Ptr& source,
                                 const pcl::PointCloud<pcl::PointXYZRGB>::Ptr& target,
                                 int maxPoints = 2048);
    bool isLoaded() const { return m_session != nullptr; }
    bool isUsingGPU() const { return m_usingGPU; }
    const std::string& lastRejectReason() const { return m_lastRejectReason; }

private:
    bool isCloudUsableForInference(const NeuralCloudStats& stats) const;
    std::vector<float> preprocessCloud(const pcl::PointCloud<pcl::PointXYZRGB>::Ptr& cloud,
                                       int maxPoints);
    std::vector<float> runInference(const std::vector<float>& srcData,
                                    const std::vector<float>& tgtData,
                                    int maxPoints);

    std::unique_ptr<Ort::Env> m_env;
    std::unique_ptr<Ort::Session> m_session;
    std::unique_ptr<Ort::SessionOptions> m_sessionOptions;
    std::unique_ptr<Ort::MemoryInfo> m_memoryInfo;
    bool m_usingGPU = false;
    std::string m_lastRejectReason;

    static constexpr const char* INPUT_SRC_NAME = "source";
    static constexpr const char* INPUT_TGT_NAME = "target";
    static constexpr const char* OUTPUT_NAME = "pose";

    std::mt19937 m_gen;
    mutable QMutex m_mutex;
};

#endif // NEURALTRACKER_H