Add Faiss ANN option + small bug fix in fullDataMaterial toggle

Author: Rsnelllenberg

DVRViewPlugin/res/shaders/FullDataMaterialBlending.frag

@@ -20,6 +20,8 @@ uniform vec2 invMatTexSize;     // Pre-divided matTexSize (1.0 / matTexSize)
 uniform int numRays;            // The number of rays in the batch (this is the same for all rays in the batch)
 uniform float stepSize;         // The step size used for the ray marching, this is used to compensate for the alpha blending in the shader
 
+uniform bool useClutterRemover; // If true, the clutter remover is used to smoothen the visualization
+
 // holds the start index for the meanPositions array for each rayID (rayIDs are decided per batch and have no relation to the pixelPos) 
 // The values it holds are the samplePosition, since each meanPos is a vec2, the start index needs to be multiplied by 2 in the shader
 // It also holds that total number of samples at the end such that the final ray length can also be calculated
@@ -41,7 +43,7 @@ float getMaterialID(inout float[5] materials, vec3 samplePos) {
     float nextMaterial = materials[3];
     float lastMaterial = materials[4];
 
-    if(firstMaterial == previousMaterial && nextMaterial == lastMaterial && currentMaterial != previousMaterial && currentMaterial != nextMaterial){
+    if(useClutterRemover && firstMaterial == previousMaterial && nextMaterial == lastMaterial && currentMaterial != previousMaterial && currentMaterial != nextMaterial){
         currentMaterial = texture(materialVolumeData, samplePos).r + 0.5f;
 
         // Update the materials array with the new material

DVRViewPlugin/res/shaders/MaterialTransition2D.frag

@@ -23,7 +23,7 @@ uniform vec3 u_maxClippingPlane;
 uniform bool useShading;
 uniform bool useClutterRemover;
 
-float getMaterialID(inout float[5] materials, inout vec3[5] samplePositions) {
+float getMaterialID(inout float[5] materials, vec3[5] samplePositions) {
     float firstMaterial = materials[0];
     float previousMaterial = materials[1];
 
@@ -39,7 +39,6 @@ float getMaterialID(inout float[5] materials, inout vec3[5] samplePositions) {
 
         // Update the materials array with the new material
         materials[2] = currentMaterial;
-        samplePositions[2] = samplePos;
     }
 
     return currentMaterial;

DVRViewPlugin/res/shaders/NNMaterialTransition.frag

@@ -57,11 +57,9 @@ void updateArrays(
 // Get the current material using the sliding window (like MaterialTransition2D.frag)
 // Refinement step: copy previous value instead of re-sampling
 float getMaterialID(inout float[5] materials, inout vec3[5] samplePositions) {
-    float firstMaterial = materials[0];
     float previousMaterial = materials[1];
     float currentMaterial = materials[2];
     float nextMaterial = materials[3];
-    float lastMaterial = materials[4];
 
     // If a transition is detected, copy the previous value instead of re-sampling
     if (currentMaterial != previousMaterial && currentMaterial != nextMaterial

DVRViewPlugin/src/VolumeRenderer.cpp

@@ -736,86 +736,53 @@ void VolumeRenderer::renderDirections()
 }
 
 
-void VolumeRenderer::prepareHNSW()
+void VolumeRenderer::prepareANN()
 {
     if (!_volumeDataset.isValid()) {
-        qCritical() << "Volume dataset is not valid. Cannot prepare HNSW.";
+        qCritical() << "Volume dataset is not valid. Cannot prepare ANN.";
         return;
     }
 
-    // Get the number of voxels and dimensions from the dataset
     uint32_t numVoxels = _volumeDataset->getNumberOfVoxels();
     uint32_t dimensions = _volumeDataset->getComponentsPerVoxel();
 
-    qDebug() << "Preparing HNSW with" << numVoxels << "voxels and" << dimensions << "dimensions.";
-
-    // Initialize HNSW space and index
-    _hnswSpace = std::make_unique<hnswlib::L2Space>(dimensions); //If we use a local parameter here instead of a member variable we get a crash later on in the program when calling the hwnsIndex again
-    _hnswIndex = std::make_unique<hnswlib::HierarchicalNSW<float>>(
-        _hnswSpace.get(),
-        numVoxels,
-        _hnswM,
-        _hnswEfConstruction
-    );
-
-    // Populate HNSW index with volume data.
+    // Populate ANN index with volume data.
     std::vector<float> voxelData(dimensions * numVoxels);
     QPair<float, float> scalarDataRange;
     _volumeDataset->getVolumeData(_compositeIndices, voxelData, scalarDataRange);
 
-    // Generate random data in the range [0, 1]
-    //std::random_device rd;
-    //std::mt19937 gen(rd());
-    //std::uniform_real_distribution<float> dis(0.0f, 1.0f);
-
-    //// Parallelize data generation using OpenMP.
-    //#pragma omp parallel for schedule(static)
-    //for (int32_t i = 0; i < numVoxels; i++) {
-    //    for (int32_t j = 0; j < dimensions; j++) {
-    //        // Each element is a random float in [0, 1]
-    //        voxelData[i * dimensions + j] = dis(gen);
-    //    }
-    //}
-
-    // Optionally, print the first few voxel vectors for inspection.
-    //for (uint32_t i = 10000; i < std::min(numVoxels, (uint32_t)100000); i++) {
-    //    QString s = "Voxel " + QString::number(i) + ":";
-    //    for (uint32_t j = 0; j < dimensions; j++) {
-    //        s += " " + QString::number(voxelData[i * dimensions + j]);
-    //    }
-    //    qDebug() << s;
-    //}
-
-    //for (uint32_t i = 0; i < numVoxels; ++i) {
-    //    // Normalize the voxel data to the range [0, 1]
-    //    for (uint32_t j = 0; j < dimensions; ++j) {
-    //        voxelData[i * dimensions + j] = (voxelData[i * dimensions + j] - scalarDataRange.first) / (scalarDataRange.second - scalarDataRange.first);
-    //    }
-    //}
-
-    // Add points to the HNSW index
-    for (uint32_t i = 0; i < numVoxels; ++i) {
-        _hnswIndex->addPoint(voxelData.data() + i * dimensions, i);
+    if (_useFaissANN) {
+        _nlist = std::clamp(static_cast<int>(numVoxels / 1000), 32, 4096); // nlist is the number of clusters in Faiss
+        //_nprobe = std::clamp(static_cast<int>(numVoxels / 1000000), 1, 64); // nprobe is the number of clusters to search in Faiss
+
+        // IVF index for large datasets
+        _faissIndexFlat = std::make_unique<faiss::IndexFlatL2>(dimensions);
+        _faissIndexIVF = std::make_unique<faiss::IndexIVFFlat>(_faissIndexFlat.get(), dimensions, _nlist, faiss::METRIC_L2);
+        _faissIndexIVF->train(numVoxels, voxelData.data());
+        _faissIndexIVF->add(numVoxels, voxelData.data());
+
+        //_faissIndexIVF->nprobe = _nprobe; // Set the number of clusters to search
+
     }
+    else {
 
-    // Set a high ef for query-time (improves recall, at the expense of query latency)
-    //_hnswIndex->setEf(_hwnsEfSearch);
-
-    // Test recall: query each point for its nearest neighbor.
-    float correct = 0;
-#pragma omp parallel for schedule(guided)
-    for (int32_t i = 0; i < numVoxels; i++) {
-        const float* queryPoint = voxelData.data() + i * dimensions;
-        std::priority_queue<std::pair<float, hnswlib::labeltype>> result = _hnswIndex->searchKnn(queryPoint, 1);
-        float distance = result.top().first;
-#pragma omp critical
-        {
-            if (distance == 0)
-                correct++;
+        // Initialize HNSW space and index
+        _hnswSpace = std::make_unique<hnswlib::L2Space>(dimensions); //If we use a local parameter here instead of a member variable we get a crash later on in the program when calling the hwnsIndex again
+        _hnswIndex = std::make_unique<hnswlib::HierarchicalNSW<float>>(
+            _hnswSpace.get(),
+            numVoxels,
+            _hnswM,
+            _hnswEfConstruction
+        );
+
+        // Add points to the HNSW index
+        for (uint32_t i = 0; i < numVoxels; ++i) {
+            _hnswIndex->addPoint(voxelData.data() + i * dimensions, i);
         }
+
+        // Set a high ef for query-time (improves recall, at the expense of query latency)
+        _hnswIndex->setEf(_hwnsEfSearch);
     }
-    float recall = correct / numVoxels;
-    std::cout << "Recall: " << recall << "\n";
 }
 
 
@@ -835,34 +802,56 @@ void VolumeRenderer::batchSearch(
         qCritical() << "Query data size is not a multiple of dimensions.";
     }
 
-    // Check that the index is valid.
-    if (!_hnswIndex) {
-        qCritical() << "HNSW index is not initialized.";
-    }
-
     int64_t numQueries = static_cast<int64_t>(queryData.size() / dimensions);
 
-    // Prepare a container to hold result vectors (one per query) with k elements each containing a pair of distance and label.
-    std::vector<std::vector<std::pair<float, hnswlib::labeltype>>> batchResults(numQueries);
-    qDebug() << "Batch search size: " << batchResults.size();
-
-    #pragma omp parallel for schedule(guided)
-    for (int64_t i = 0; i < numQueries; i++) { // it is important to use int64_t here to avoid overflow crashes
-        // Find pointer to the start of the i-th query.
-        const float* query = queryData.data() + static_cast<int64_t>(i * dimensions);
-        std::priority_queue<std::pair<float, hnswlib::labeltype>> resultQueue = _hnswIndex->searchKnn(query, k);
-
-        // Convert the priority queue to a vector.
-        std::vector<std::pair<float, hnswlib::labeltype>> answers;
-        while (!resultQueue.empty()) {
-            answers.push_back(resultQueue.top());
-            resultQueue.pop();
+    if (_useFaissANN) {
+        if (!_faissIndexIVF->is_trained) {
+            qCritical() << "Faiss IVF index is not trained!";
+            return;
         }
 
-        QVector2D meanPos = ComputeMeanOfNN(answers, k, positionData);
-        meanPositionData[i * 2] = meanPos.x();
-        meanPositionData[i * 2 + 1] = meanPos.y();
+        std::vector<faiss::idx_t> labels(numQueries * k);
+        std::vector<float> distances(numQueries * k);
 
+        qDebug() << "Searching for" << k << "nearest neighbors for" << numQueries << "queries using Faiss IVF index.";
+        _faissIndexIVF->search(numQueries, queryData.data(), k, distances.data(), labels.data());
+        qDebug() << "Faiss IVF search completed.";
+
+        for (int64_t i = 0; i < numQueries; i++) {
+            std::vector<std::pair<float, int64_t>> answers;
+            for (int j = 0; j < k; j++) {
+                answers.emplace_back(distances[i * k + j], labels[i * k + j]);
+            }
+            // Compute the mean position of the nearest neighbors.
+            QVector2D meanPos = ComputeMeanOfNN(answers, k, positionData);
+            // Store the mean position in the output vector.
+            meanPositionData[i * 2] = meanPos.x();
+            meanPositionData[i * 2 + 1] = meanPos.y();
+        }
+    }
+    else {
+        // Check that the index is valid.
+        if (!_hnswIndex) {
+            qCritical() << "HNSW index is not initialized.";
+        }
+
+        #pragma omp parallel for schedule(guided)
+        for (int64_t i = 0; i < numQueries; i++) { // it is important to use int64_t here to avoid overflow crashes
+            // Find pointer to the start of the i-th query.
+            const float* query = queryData.data() + static_cast<int64_t>(i * dimensions);
+            std::priority_queue<std::pair<float, hnswlib::labeltype>> resultQueue = _hnswIndex->searchKnn(query, k);
+
+            // Convert the priority queue to a vector.
+            std::vector<std::pair<float, int64_t>> answers;
+            while (!resultQueue.empty()) {
+                answers.push_back({ resultQueue.top().first, static_cast<int64_t>(resultQueue.top().second) });
+                resultQueue.pop();
+            }
+            QVector2D meanPos = ComputeMeanOfNN(answers, k, positionData);
+            meanPositionData[i * 2] = meanPos.x();
+            meanPositionData[i * 2 + 1] = meanPos.y();
+
+        }
     }
 }
 
@@ -1248,6 +1237,7 @@ void VolumeRenderer::renderBatchToScreen(int batchIndex, uint32_t sampleDim, std
         _fullDataCompositeShader.uniform2f("invFaceTexSize", 1.0f / float(width), 1.0f / float(height));
         _fullDataCompositeShader.uniform2f("invTfTexSize", 1.0f / _tfDataset->getImageSize().width(), 1.0f / _tfDataset->getImageSize().height());
         _fullDataCompositeShader.uniform1f("stepSize", _stepSize);
+        _fullDataCompositeShader.uniform1i("useClutterRemover", _useClutterRemover);
 
         // Render a full-screen quad to composite the current batch's results over prevFullCompositeTexture.
         drawDVRQuad(_fullDataCompositeShader);
@@ -1311,18 +1301,17 @@ void VolumeRenderer::renderBatchToScreen(int batchIndex, uint32_t sampleDim, std
 // @param neighbours: A vector of pairs containing the distance and label of each neighbour.
 // @param k: The number of neighbours to consider.
 // @param positionData: A vector containing the 2D positions of the neighbours.
-QVector2D VolumeRenderer::ComputeMeanOfNN(const std::vector<std::pair<float, hnswlib::labeltype>>& neighbors, int k, const std::vector<float>& positionData) {
+QVector2D VolumeRenderer::ComputeMeanOfNN(const std::vector<std::pair<float, int64_t>>& neighbors, int k, const std::vector<float>& positionData) {
     float epsilon = 1.0f;  // To avoid division by zero and limit the impact of very close neighbours.
     float clusterSlack = 0.1f; // Slack for cluster thresholding, can be adjusted based on the dataset.
 
     std::vector<float> weights(k, 0.0f);
     std::vector<QVector2D> candidatePositions(k, QVector2D(0.0f, 0.0f));
     int j = 0;
     for (const auto& entry : neighbors) {
-        int32_t index = static_cast<int32_t>(entry.second);
         weights[j] = (1.0f / (entry.first + epsilon));  // Inverse distance is used as weight
-        float posX = positionData[index * 2];
-        float posY = positionData[index * 2 + 1];
+        float posX = positionData[entry.second * 2];
+        float posY = positionData[entry.second * 2 + 1];
         candidatePositions[j] = QVector2D(posX, posY);
         j++;
     }
@@ -1479,8 +1468,9 @@ void VolumeRenderer::renderFullData()
 
     // Make sure the ANN (e.g. hnswlib) is prepared for the dataset.
     if (!_ANNAlgorithmTrained) {
-        prepareHNSW();
+        prepareANN();
         _ANNAlgorithmTrained = true;
+        qDebug() << "ANN algorithm trained for full data mode.";
     }
 
     // Initialize the GPU full data mode parameters if not already done.

DVRViewPlugin/src/VolumeRenderer.h

@@ -22,6 +22,10 @@
 #include "MCArrays.h"
 
 #include <hnswlib.h>
+#include <faiss/IndexIVFFlat.h>
+#include <faiss/IndexFlat.h>
+#include <faiss/Index.h>
+
 #include <QOpenGLFunctions_4_3_Core>
 
 namespace mv {
@@ -117,13 +121,13 @@ class VolumeRenderer : protected QOpenGLFunctions_4_3_Core
     void drawDVRQuad(mv::ShaderProgram& shader);
 
     // Full data render mode methods
-    void prepareHNSW();
+    void prepareANN();
     void batchSearch(const std::vector<float>& queryData, std::vector<float>& positionData, uint32_t dimensions, int k, bool useWeightedMean, std::vector<float>& meanPositionData);
     void getFacesTextureData(std::vector<float>& frontfacesData, std::vector<float>& backfacesData);
     void getGPUFullDataModeBatches(std::vector<float>& frontfacesData, std::vector<float>& backfacesData);
     void retrieveBatchFullData(std::vector<float>& cpuOutput, int batchIndex, bool deleteBuffers);
     void renderBatchToScreen(int batchIndex, uint32_t sampleDim, std::vector<float>& meanPositions);
-    QVector2D ComputeMeanOfNN(const std::vector<std::pair<float, hnswlib::labeltype>>& neighbors, int k, const std::vector<float>& positionData);
+    QVector2D ComputeMeanOfNN(const std::vector<std::pair<float, int64_t>>& neighbors, int k, const std::vector<float>& positionData);
     void updateRenderModeParameters();
 
     void renderFullData();
@@ -230,14 +234,22 @@ class VolumeRenderer : protected QOpenGLFunctions_4_3_Core
     mv::Vector3f _cameraPos;
 
     size_t _fullDataMemorySize = 0; // The size of the full data in bytes
-    size_t _fullGPUMemorySize = static_cast<size_t>(2 * 1024 * 1024) * 1024; // The size of the full data in bytes on the GPU if we use normal int it causes a overflow; The SSBOs are limited to 2GB, so even if the GPU has more VRAM we limit the size to 2GB for the full data mode.
+    size_t _fullGPUMemorySize = static_cast<size_t>(14 * 1024 * 1024) * 1024; // The size of the full data in bytes on the GPU if we use normal int it causes a overflow; The SSBOs are limited to 2GB, so even if the GPU has more VRAM we limit the size to 2GB for the full data mode.
 
-    // HNSWLib-related members  
+    // ANN-related members  
     std::unique_ptr<hnswlib::L2Space> _hnswSpace;
     std::unique_ptr<hnswlib::HierarchicalNSW<float>> _hnswIndex;
-    int _hnswM = 4;
-    int _hnswEfConstruction = 50;
+    std::unique_ptr<faiss::IndexIVFFlat> _faissIndexIVF;
+    std::unique_ptr<faiss::IndexFlatL2> _faissIndexFlat;
+    int _hnswM = 8;
+    int _hnswEfConstruction = 100;
     int _hwnsEfSearch = 50;
+
+    int _nlist = 1000;
+    int _nprobe = 100; // Number of probes for Faiss IVF index
+
+    // Boolean to select ANN library
+    bool _useFaissANN = true;
 
     // Full Data Rendermode Parameters
     std::vector<std::vector<int>> _GPUBatches; // Batches of pixel indices for the full data mode as it is not always possible to fit all pixels in one batch