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

src/tracking/NeuralTracker.cpp
236 строк9 KB
k k
refactor: remove MainWindow dead code + fix VoxelGrid warnings
14 июл 2026, 23:40
14 июл 2026, 23:401459aca
#include "NeuralTracker.h"
#include <QDebug>
#include <Eigen/Geometry>
#include <random>
#include <algorithm>
#include <numeric>
#include <cmath>
#include <limits>
#include <pcl/filters/voxel_grid.h>
#include <QCoreApplication>

NeuralTracker::NeuralTracker()
{
    // Fixed seed for reproducible results across runs.
    // Using deterministic seed instead of random_device so that
    // the same input always produces the same pose estimate.
    m_gen.seed(42);
}

bool NeuralTracker::isCloudUsableForInference(const NeuralCloudStats& stats) const
{
    return stats.finitePoints >= 300 &&
           std::isfinite(stats.bboxDiagonal) &&
           stats.bboxDiagonal >= 0.02f;
}

bool NeuralTracker::loadModel(const std::string& onnxPath, bool useCUDA)
{
    m_lastRejectReason.clear();
    try {
        m_env = std::make_unique<Ort::Env>(ORT_LOGGING_LEVEL_WARNING, "NeuralTracker");
        m_sessionOptions = std::make_unique<Ort::SessionOptions>();
        m_sessionOptions->SetIntraOpNumThreads(1);
        m_sessionOptions->SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED);

        if (useCUDA) {
            auto available = Ort::GetAvailableProviders();
            bool hasCUDA = std::find(available.begin(), available.end(),
                                     "CUDAExecutionProvider") != available.end();
            if (hasCUDA) {
                OrtCUDAProviderOptions cudaOpt{};
                cudaOpt.device_id = 0;
                cudaOpt.gpu_mem_limit = SIZE_MAX;
                cudaOpt.arena_extend_strategy = 0;
                m_sessionOptions->AppendExecutionProvider_CUDA(cudaOpt);
                m_usingGPU = true;
                qInfo() << "[NeuralTracker] CUDA provider enabled";
            } else {
                qWarning() << "[NeuralTracker] CUDA not available, fallback to CPU";
                m_usingGPU = false;
            }
        }

        std::wstring wpath = QString::fromStdString(onnxPath).toStdWString();
        m_session = std::make_unique<Ort::Session>(*m_env, wpath.c_str(), *m_sessionOptions);
        m_memoryInfo = std::make_unique<Ort::MemoryInfo>(
            Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault));

        qInfo() << "[NeuralTracker] Model loaded, using" << (m_usingGPU ? "CUDA" : "CPU");
        return true;
    } catch (const std::exception& e) {
        m_lastRejectReason = e.what();
        qCritical() << "[NeuralTracker] loadModel failed:" << e.what();
        return false;
    }
}

std::vector<float> NeuralTracker::preprocessCloud(
    const pcl::PointCloud<pcl::PointXYZRGB>::Ptr& cloud, int maxPoints)
{
    if (!cloud || cloud->empty()) return {};

    pcl::PointCloud<pcl::PointXYZRGB>::Ptr finiteCloud(new pcl::PointCloud<pcl::PointXYZRGB>);
    finiteCloud->reserve(cloud->size());

    // 1. Центрирование только по finite-точкам.
    Eigen::Vector3f centroid = Eigen::Vector3f::Zero();
    for (const auto& pt : cloud->points) {
        if (!std::isfinite(pt.x) || !std::isfinite(pt.y) || !std::isfinite(pt.z)) {
            continue;
        }
        centroid += pt.getVector3fMap();
        finiteCloud->push_back(pt);
    }
    if (finiteCloud->empty()) {
        return {};
    }
    centroid /= static_cast<float>(finiteCloud->size());

    // 2. Децимация (VoxelGrid + FPS) – сохраняем существующую логику,
    //    но в конце меняем формат на channels-first.
    pcl::PointCloud<pcl::PointXYZRGB>::Ptr sampled(new pcl::PointCloud<pcl::PointXYZRGB>);
    pcl::PointCloud<pcl::PointXYZRGB>::Ptr inputCloud = finiteCloud;

    // VoxelGrid если точек > 20000
    if (finiteCloud->size() > 20000) {
        pcl::VoxelGrid<pcl::PointXYZRGB> voxel;
        voxel.setLeafSize(0.015f, 0.015f, 0.015f);
        voxel.setSaveLeafLayout(false);
        voxel.setInputCloud(finiteCloud);
        pcl::PointCloud<pcl::PointXYZRGB>::Ptr down(new pcl::PointCloud<pcl::PointXYZRGB>);
        voxel.filter(*down);
        if (down->size() > 5000) {
            inputCloud = down;
        } else {
            // fallback random sampling
            pcl::PointCloud<pcl::PointXYZRGB>::Ptr randSampled(new pcl::PointCloud<pcl::PointXYZRGB>);
            int target = std::min(20000, static_cast<int>(finiteCloud->size()));
            for (int i = 0; i < target; ++i) {
                int idx = std::uniform_int_distribution<>(0, static_cast<int>(finiteCloud->size())-1)(m_gen);
                randSampled->push_back(finiteCloud->points[idx]);
            }
            inputCloud = randSampled;
        }
    }

    // Farthest point sampling заменён на random sampling + FPS для 2048 точек.
    // Полный FPS O(N*K) слишком дорог на CPU и блокирует цикл захвата.
    if (inputCloud->size() <= static_cast<size_t>(maxPoints)) {
        *sampled = *inputCloud;
        for (auto& pt : sampled->points) {
            pt.x -= centroid.x(); pt.y -= centroid.y(); pt.z -= centroid.z();
        }
        while (sampled->size() < static_cast<size_t>(maxPoints)) {
            int idx = std::uniform_int_distribution<>(0, static_cast<int>(finiteCloud->size())-1)(m_gen);
            pcl::PointXYZRGB pt = finiteCloud->points[idx];
            pt.x -= centroid.x(); pt.y -= centroid.y(); pt.z -= centroid.z();
            sampled->push_back(pt);
        }
    } else {
        // Random sampling: O(N) вместо O(N*K) FPS — в 2000x быстрее
        // для K=2048. Точности достаточно для нейросетевого трекинга.
        std::vector<int> perm(inputCloud->size());
        std::iota(perm.begin(), perm.end(), 0);
        std::shuffle(perm.begin(), perm.end(), m_gen);

        for (int i = 0; i < maxPoints; ++i) {
            pcl::PointXYZRGB pt = inputCloud->points[perm[i]];
            pt.x -= centroid.x(); pt.y -= centroid.y(); pt.z -= centroid.z();
            sampled->push_back(pt);
        }
    }

    // 3. channels-first формат (ИСПРАВЛЕНИЕ БАГА #1)
    int N = static_cast<int>(sampled->size());
    std::vector<float> result(N * 3);
    for (int i = 0; i < N; ++i) {
        result[i]       = sampled->points[i].x;
        result[i + N]   = sampled->points[i].y;
        result[i + 2*N] = sampled->points[i].z;
    }
    return result;
}

std::vector<float> NeuralTracker::runInference(const std::vector<float>& srcData,
                                               const std::vector<float>& tgtData,
                                               int maxPoints)
{
    try {
        std::vector<int64_t> shape = {1, 3, maxPoints};
        auto srcTensor = Ort::Value::CreateTensor<float>(
            *m_memoryInfo, const_cast<float*>(srcData.data()), srcData.size(),
            shape.data(), shape.size());
        auto tgtTensor = Ort::Value::CreateTensor<float>(
            *m_memoryInfo, const_cast<float*>(tgtData.data()), tgtData.size(),
            shape.data(), shape.size());

        const char* inputNames[] = {INPUT_SRC_NAME, INPUT_TGT_NAME};
        const char* outputNames[] = {OUTPUT_NAME};
        std::vector<Ort::Value> inputs;
        inputs.push_back(std::move(srcTensor));
        inputs.push_back(std::move(tgtTensor));

        auto outputs = m_session->Run(Ort::RunOptions{nullptr},
                                      inputNames, inputs.data(), 2,
                                      outputNames, 1);
        float* outData = outputs[0].GetTensorMutableData<float>();
        size_t outSize = outputs[0].GetTensorTypeAndShapeInfo().GetElementCount();
        return std::vector<float>(outData, outData + outSize);
    } catch (const std::exception& e) {
        qWarning() << "[NeuralTracker] inference error:" << e.what();
        return {};
    }
}

std::optional<Eigen::Matrix4f> NeuralTracker::tryEstimatePose(
    const pcl::PointCloud<pcl::PointXYZRGB>::Ptr& source,
    const pcl::PointCloud<pcl::PointXYZRGB>::Ptr& target,
    int maxPoints)
{
    m_lastRejectReason.clear();
    if (!m_session) {
        m_lastRejectReason = "session not loaded";
        return std::nullopt;
    }

    QMutexLocker locker(&m_mutex);

    auto srcStats = analyzeNeuralCloud(source);
    auto tgtStats = analyzeNeuralCloud(target);
    if (!srcStats || !tgtStats) {
        m_lastRejectReason = "failed to analyze cloud stats";
        return std::nullopt;
    }
    if (!isCloudUsableForInference(*srcStats) || !isCloudUsableForInference(*tgtStats)) {
        m_lastRejectReason = "cloud not usable for inference";
        return std::nullopt;
    }

    auto srcData = preprocessCloud(source, maxPoints);
    auto tgtData = preprocessCloud(target, maxPoints);
    if (srcData.empty() || tgtData.empty()) {
        m_lastRejectReason = "preprocess returned empty tensors";
        return std::nullopt;
    }

    auto output = runInference(srcData, tgtData, maxPoints);
    auto poseMetrics = poseVectorToMatrixChecked(output);
    if (!poseMetrics) {
        m_lastRejectReason = "invalid pose vector from network";
        return std::nullopt;
    }

    return poseMetrics->transform;
}

Eigen::Matrix4f NeuralTracker::estimatePose(
    const pcl::PointCloud<pcl::PointXYZRGB>::Ptr& source,
    const pcl::PointCloud<pcl::PointXYZRGB>::Ptr& target,
    int maxPoints)
{
    if (auto pose = tryEstimatePose(source, target, maxPoints)) {
        return *pose;
    }
    return Eigen::Matrix4f::Identity();
}