You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
243 lines
7.5 KiB
243 lines
7.5 KiB
5 months ago
|
#pragma once
|
||
|
|
|
||
|
|
#include <torch/torch.h>
|
||
|
|
#include <torch/script.h>
|
||
|
|
#include <omp.h>
|
||
|
|
#include <cmath>
|
||
|
|
#include <algorithm>
|
||
|
|
#include "ThirdParty/IsoSurfGen/ImplicitFunc.h"
|
||
|
|
|
||
|
|
template<typename Real>
|
||
|
|
class MyUtil
|
||
|
|
{
|
||
|
|
public:
|
||
|
|
MyUtil()
|
||
|
|
{
|
||
|
|
max_iter = 50; threshold = 1.0e-7f;
|
||
|
|
buffer = 131072;// 262144;
|
||
|
|
use_GPU = true;
|
||
|
|
dim_type = 2;
|
||
|
|
verbose = true;
|
||
|
|
}
|
||
|
|
public:
|
||
|
|
int max_iter;
|
||
|
|
int buffer;
|
||
|
|
float threshold;
|
||
|
|
bool use_GPU;
|
||
|
|
int dim_type;
|
||
|
|
at::Tensor direct_input, direct_output, direct_grad_input, batch_point_input, batch_output;
|
||
|
|
std::vector<TinyVector<Real, 3>> world_start, world_end;
|
||
|
|
std::vector<Real> fstart, fend, value;
|
||
|
|
std::vector<Real> values;
|
||
|
|
std::vector<int> point_id_total_cache, point_id_cache;
|
||
|
|
bool verbose;
|
||
|
|
};
|
||
|
|
|
||
|
|
template <typename Real>
|
||
|
|
class MyImplicitFunc : public ImplicitFunc<Real>
|
||
|
|
{
|
||
|
|
public:
|
||
|
|
torch::jit::script::Module* net;
|
||
|
|
MyUtil<Real>* util;
|
||
|
|
public:
|
||
|
|
MyImplicitFunc(torch::jit::script::Module* _net, MyUtil<Real>* _util)
|
||
|
|
:ImplicitFunc<Real>(), net(_net), util(_util)
|
||
|
|
{
|
||
|
|
}
|
||
|
|
/////////////////////////////////////////
|
||
|
|
~MyImplicitFunc() {}
|
||
|
|
/// //////////////////////////////////////
|
||
|
|
bool is_inside(const TinyVector<Real, 3>& p, const Real _isovalue = 0) const
|
||
|
|
{
|
||
|
|
return scalar_value(p) < _isovalue;
|
||
|
|
}
|
||
|
|
/////////////////////////////////////////
|
||
|
|
Real scalar_value(const TinyVector<Real, 3>& p) const
|
||
|
|
{
|
||
|
|
std::vector<torch::jit::IValue> inputs;
|
||
|
|
at::Tensor m;
|
||
|
|
if (util->dim_type == 3)
|
||
|
|
m = torch::zeros({ 1, 1, 3 });
|
||
|
|
else
|
||
|
|
m = torch::zeros({ 1, 3 });
|
||
|
|
|
||
|
|
auto mpstorage = static_cast<float*>(m.storage().data());
|
||
|
|
mpstorage[0] = p[0], mpstorage[1] = p[1], mpstorage[2] = p[2];
|
||
|
|
if (util->use_GPU)
|
||
|
|
inputs.push_back(m.to(at::kCUDA));
|
||
|
|
else
|
||
|
|
inputs.push_back(m);
|
||
|
|
|
||
|
|
at::Tensor output = net->forward(inputs).toTensor().cpu();
|
||
|
|
auto outputstorage = static_cast<float*>(output.storage().data());
|
||
|
|
return outputstorage[0];
|
||
|
|
}
|
||
|
|
/////////////////////////////////////////
|
||
|
|
void gradient(const TinyVector<Real, 3>& p, TinyVector<Real, 3>& gradient) const
|
||
|
|
{
|
||
|
|
std::vector<torch::jit::IValue> inputs;
|
||
|
|
at::Tensor m;
|
||
|
|
if (util->dim_type == 3)
|
||
|
|
m = torch::zeros({ 1, 1, 3 }, torch::requires_grad(true));
|
||
|
|
else
|
||
|
|
m = torch::zeros({ 1, 3 }, torch::requires_grad(true));
|
||
|
|
|
||
|
|
auto mpstorage = static_cast<float*>(m.storage().data());
|
||
|
|
mpstorage[0] = p[0], mpstorage[1] = p[1], mpstorage[2] = p[2];
|
||
|
|
if (util->use_GPU)
|
||
|
|
inputs.push_back(m.to(at::kCUDA));
|
||
|
|
else
|
||
|
|
inputs.push_back(m);
|
||
|
|
|
||
|
|
at::Tensor output = net->forward(inputs).toTensor();
|
||
|
|
output.sum().backward();
|
||
|
|
auto gradstorage = static_cast<float*>(m.grad().cpu().storage().data());
|
||
|
|
gradient[0] = gradstorage[0], gradient[1] = gradstorage[1], gradient[2] = gradstorage[2];
|
||
|
|
}
|
||
|
|
/// ////////////////////////////////////////
|
||
|
|
void scalar_value(const std::vector<TinyVector<Real, 3>>& pvec, std::vector<Real>& vec_values) const
|
||
|
|
{
|
||
|
|
if (util->dim_type == 3)
|
||
|
|
util->batch_point_input = torch::zeros({ 1, (int)pvec.size(), 3 });
|
||
|
|
else
|
||
|
|
util->batch_point_input = torch::zeros({ (int)pvec.size(), 3 });
|
||
|
|
|
||
|
|
auto mpstorage = static_cast<float*>(util->batch_point_input.storage().data());
|
||
|
|
#pragma omp parallel for
|
||
|
|
for (int i = 0; i < (int)pvec.size(); i++)
|
||
|
|
{
|
||
|
|
mpstorage[3 * i] = pvec[i][0], mpstorage[3 * i + 1] = pvec[i][1], mpstorage[3 * i + 2] = pvec[i][2];
|
||
|
|
}
|
||
|
|
std::vector<torch::jit::IValue> inputs;
|
||
|
|
if (util->use_GPU)
|
||
|
|
inputs.push_back(util->batch_point_input.to(at::kCUDA));
|
||
|
|
else
|
||
|
|
inputs.push_back(util->batch_point_input);
|
||
|
|
|
||
|
|
util->batch_output = net->forward(inputs).toTensor().cpu();
|
||
|
|
auto outputstorage = static_cast<float*>(util->batch_output.storage().data());
|
||
|
|
vec_values.resize(pvec.size());
|
||
|
|
#pragma omp parallel for
|
||
|
|
for (int i = 0; i < (int)pvec.size(); i++)
|
||
|
|
{
|
||
|
|
vec_values[i] = outputstorage[i];
|
||
|
|
}
|
||
|
|
}
|
||
|
|
//////////////////////////////////////////////////
|
||
|
|
void directed_distance(
|
||
|
|
const std::vector<TinyVector<Real, 3>>& _p0vec,
|
||
|
|
const std::vector<TinyVector<Real, 3>>& _p1vec,
|
||
|
|
const std::vector<Real>& _p0vec_values,
|
||
|
|
const std::vector<Real>& _p1vec_values,
|
||
|
|
std::vector<TinyVector<Real, 3>>& _point,
|
||
|
|
std::vector<TinyVector<Real, 3>>& _normal,
|
||
|
|
const Real _isovalue = 0) const
|
||
|
|
{
|
||
|
|
int size_edges = (int)_p0vec.size();
|
||
|
|
_point.resize(size_edges);
|
||
|
|
_normal.resize(size_edges);
|
||
|
|
|
||
|
|
if (util->dim_type == 3)
|
||
|
|
{
|
||
|
|
util->direct_input = torch::zeros({ 1, size_edges, 3 });
|
||
|
|
util->direct_output = torch::zeros({ 1, size_edges });
|
||
|
|
util->direct_grad_input = torch::zeros({ 1, size_edges, 3 }, torch::requires_grad(true));
|
||
|
|
}
|
||
|
|
else
|
||
|
|
{
|
||
|
|
util->direct_input = torch::zeros({ size_edges, 3 });
|
||
|
|
util->direct_output = torch::zeros({ size_edges, 1 });
|
||
|
|
util->direct_grad_input = torch::zeros({ size_edges, 3 }, torch::requires_grad(true));
|
||
|
|
}
|
||
|
|
|
||
|
|
util->world_start.assign(_p0vec.begin(), _p0vec.end());
|
||
|
|
util->world_end.assign(_p1vec.begin(), _p1vec.end());
|
||
|
|
util->fstart.assign(_p0vec_values.begin(), _p0vec_values.end());
|
||
|
|
util->fend.assign(_p1vec_values.begin(), _p1vec_values.end());
|
||
|
|
util->value.resize(size_edges);
|
||
|
|
|
||
|
|
if (_isovalue != (Real)0.0)
|
||
|
|
{
|
||
|
|
#pragma omp parallel for
|
||
|
|
for (int i = 0; i < (int)util->fstart.size(); i++)
|
||
|
|
{
|
||
|
|
util->fstart[i] -= _isovalue;
|
||
|
|
util->fend[i] -= _isovalue;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
auto mpstorage = static_cast<float*>(util->direct_input.storage().data());
|
||
|
|
Real tiny_value = 1.0e-12;
|
||
|
|
for (int iter = 0; iter < util->max_iter; iter++)
|
||
|
|
{
|
||
|
|
#pragma omp parallel for
|
||
|
|
for (int i = 0; i < size_edges; i++)
|
||
|
|
{
|
||
|
|
Real s0 = fabs(util->fstart[i]);
|
||
|
|
Real s1 = fabs(util->fend[i]);
|
||
|
|
Real t = s0 / std::max(tiny_value, s0 + s1);
|
||
|
|
_point[i] = ((Real)1.0 - t) * util->world_start[i] + t * util->world_end[i];
|
||
|
|
mpstorage[3 * i] = _point[i][0], mpstorage[3 * i + 1] = _point[i][1], mpstorage[3 * i + 2] = _point[i][2];
|
||
|
|
}
|
||
|
|
if (iter < util->max_iter - 1)
|
||
|
|
{
|
||
|
|
std::vector<torch::jit::IValue> inputs;
|
||
|
|
if (util->use_GPU)
|
||
|
|
inputs.push_back(util->direct_input.to(at::kCUDA));
|
||
|
|
else
|
||
|
|
inputs.push_back(util->direct_input);
|
||
|
|
|
||
|
|
util->direct_output = net->forward(inputs).toTensor().cpu();
|
||
|
|
auto outputstorage = static_cast<float*>(util->direct_output.storage().data());
|
||
|
|
#pragma omp parallel for
|
||
|
|
for (int i = 0; i < size_edges; i++)
|
||
|
|
{
|
||
|
|
util->value[i] = outputstorage[i] - _isovalue;
|
||
|
|
if (util->value[i] * util->fstart[i] < 0)
|
||
|
|
{
|
||
|
|
util->world_end[i] = _point[i];
|
||
|
|
util->fend[i] = util->value[i];
|
||
|
|
}
|
||
|
|
else if (util->value[i] * util->fend[i] < 0)
|
||
|
|
{
|
||
|
|
util->world_start[i] = _point[i];
|
||
|
|
util->fstart[i] = util->value[i];
|
||
|
|
}
|
||
|
|
}
|
||
|
|
float vmax = 0;
|
||
|
|
for (int i = 0; i < size_edges; i++)
|
||
|
|
if ((_p0vec_values[i] - _isovalue) * (_p1vec_values[i] - _isovalue) < 0 && fabs(util->value[i]) > vmax)
|
||
|
|
vmax = fabs(util->value[i]);
|
||
|
|
if (vmax < util->threshold)
|
||
|
|
break;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
auto gradstorage = static_cast<float*>(util->direct_grad_input.storage().data());
|
||
|
|
#pragma omp parallel for
|
||
|
|
for (int i = 0; i < size_edges; i++)
|
||
|
|
{
|
||
|
|
gradstorage[3 * i] = _point[i][0]; gradstorage[3 * i + 1] = _point[i][1]; gradstorage[3 * i + 2] = _point[i][2];
|
||
|
|
}
|
||
|
|
std::vector<torch::jit::IValue> inputs;
|
||
|
|
if (util->use_GPU)
|
||
|
|
inputs.push_back(util->direct_grad_input.to(at::kCUDA));
|
||
|
|
else
|
||
|
|
inputs.push_back(util->direct_grad_input);
|
||
|
|
|
||
|
|
at::Tensor output = net->forward(inputs).toTensor();
|
||
|
|
output.sum().backward();
|
||
|
|
auto gradoutputstorage = static_cast<float*>(util->direct_grad_input.grad().cpu().storage().data());
|
||
|
|
|
||
|
|
#pragma omp parallel for
|
||
|
|
for (int i = 0; i < size_edges; i++)
|
||
|
|
{
|
||
|
|
_normal[i][0] = gradoutputstorage[3 * i];
|
||
|
|
_normal[i][1] = gradoutputstorage[3 * i + 1];
|
||
|
|
_normal[i][2] = gradoutputstorage[3 * i + 2];
|
||
|
|
_normal[i].Normalize();
|
||
|
|
}
|
||
|
|
}
|
||
|
|
//////////////////////////////////////////////////
|
||
|
|
};
|