Octree/tests/search_assign_test/main.cpp

/**
 * ------------------------------------
 * @author: Weipeng Kong
 * @date: 2021/11/17
 * @email: yjxkwp@foxmail.com
 * @site: https://donot.fit
 * @description: 
 * ------------------------------------
**/

#include <iostream>
#include <igl/marching_cubes.h>
#include <igl/voxel_grid.h>
#include <igl/writeOBJ.h>
#include <igl/in_element.h>
#include <pMesh/mesh/TriangleMesh.h>
#include <pMesh/mesh/HexahedronMesh.h>
#include <pMesh/io/reader/OBJReader.h>
#include <pMesh/io/writer/VTKWriter.h>
#include <pMesh/io/adapter/DefaultReadAdapter.h>
#include <pMesh/io/adapter/DefaultWriteAdapter.h>
#include <boost/timer.hpp>
#include <boost/log/trivial.hpp>
#include <pMesh/io/reader/BaseReader.h>

#include "Octree/udf/UDFOctree.h"
#include "Octree/OctreeBuilder.h"
#include "Octree/OctreeTraverser.h"
#include "Octree/udf/UDFTraversalSampler.h"
#include "test-path.h"
#include "search_assign.h"


int main() {
    pMesh::io::fs_path data_base_path = TEST_DATA_BASE_PATH;
    BOOST_LOG_TRIVIAL(debug) << "base data path is " << boost::filesystem::absolute(data_base_path);
    bool use_udf = true;
    auto out_mc_path = data_base_path / (use_udf? "sphere-udf-mc.obj": "sphere-sdf-mc.obj");

    Eigen::AlignedBox<double, 3> aabb(Eigen::Vector3d(0,0,0), Eigen::Vector3d(1, 1, 1));

    int aabb_res_x = 400;
    int aabb_res_y = 400;
    int aabb_res_z = 400;
    int aabb_len = aabb_res_x * aabb_res_y * aabb_res_z;

    double x0 = 0;
    double y0 = 0;
    double z0 = 0;
    double R = 1.2;

    auto sphere_shape = [=](double x, double y, double z) {
        return sqrt(pow(x - x0, 2) + pow(y - y0, 2) + pow(z - z0, 2)) - R;
    };
    Octree::UDFOctree octree(0, Octree::AABB(aabb.min(), aabb.max()));
    octree.get_root()->udf = std::make_shared<Octree::UDFData>(std::vector<double>(aabb_len), aabb_res_x, aabb_res_y,
                                                               aabb_res_z);
    auto &udf_data = *octree.get_root()->udf;
    BOOST_LOG_TRIVIAL(debug) << "Calculating UDF";

    for (int x = 0; x < aabb_res_x; ++x) {
        double pos_x = ((aabb_res_x - x) * aabb.min().x() + x * aabb.max().x()) / (1.0 * aabb_res_x);
        for (int y = 0; y < aabb_res_y; ++y) {
            double pos_y = ((aabb_res_y - y) * aabb.min().y() + y * aabb.max().y()) / (1.0 * aabb_res_y);
            for (int z = 0; z < aabb_res_z; ++z) {
                double pos_z = ((aabb_res_z - z) * aabb.min().z() + z * aabb.max().z()) / (1.0 * aabb_res_z);

                int ind = udf_data.get_array_index(x, y, z);
                if(use_udf)
                    udf_data.value[ind] = abs(sphere_shape(pos_x, pos_y, pos_z));
                else
                    udf_data.value[ind] = sphere_shape(pos_x, pos_y, pos_z);
            }
        }
    }
    BOOST_LOG_TRIVIAL(debug) << "End computing UDF";

    if(use_udf) {
        int start[3] = {aabb_res_x/2, aabb_res_y/2, aabb_res_z/2};
        double block_size[3] = {
                aabb.sizes().x() / (aabb_res_x - 1),
                aabb.sizes().y() / (aabb_res_y - 1),
                aabb.sizes().z() / (aabb_res_z - 1)
        };
        // TODO:
        double threshold = std::max({block_size[0], block_size[1], block_size[2]});

        boost::timer t;
        search_assign(udf_data, start, threshold);
        BOOST_LOG_TRIVIAL(debug) << "[search_assign]: " << t.elapsed();
    }

    if (1) {
        Eigen::MatrixXd GV;
        Eigen::RowVector3i res;
        const int s = 10;
        igl::voxel_grid(aabb, s, 1, GV, res);

        // compute values
        std::cout << "Computing distances..." << std::endl;
        Eigen::VectorXd S = Eigen::VectorXd(GV.rows()), B;

        {
            for (int i = 0; i < GV.rows(); ++i) {
                auto node = octree.map_node(GV.row(i));
                if (node == nullptr) {
                    S[i] = 1000;
                } else {
                    S[i] = node->get_udf(GV.row(i));
                }
            }
        }
        std::cout << "Marching cubes..." << std::endl;
        Eigen::MatrixXd SV, BV;
        Eigen::MatrixXi SF, BF;

        igl::marching_cubes(S, GV, res(0), res(1), res(2), 0, SV, SF);
        igl::writeOBJ(out_mc_path.string(), SV, SF);
    }
    return 0;
}