00_gen_dataset.cpp

This is an example of how to use the netdem library.

#include "distribution_uniform.hpp"
#include "igl_wrapper.hpp"
#include "particle.hpp"
#include "shape_ellipsoid.hpp"
#include "shape_level_set.hpp"
#include "shape_poly_super_ellipsoid.hpp"
#include "shape_spherical_harmonics.hpp"
#include "shape_trimesh.hpp"
#include "spherical_voronoi.hpp"
#include "utils_math.hpp"
#include <filesystem>
#include <fstream>
#include <iostream>
#include <random>
#include <sstream>
#include <string>
 
using namespace netdem;
using namespace std;
 
void SaveDatasetSDF(int num_samples, double (*ds_inputs)[3],
                    double (*ds_SDF)[1], string filename) {
  stringbuf buf;
  ostream os(&buf);
  int os_width = 24;
  os.setf(ios::scientific);
  os.precision(15);
  for (int i = 0; i < num_samples; i++) {
    // write the data as one sample one row
    for (int i_inputs = 0; i_inputs < 3; i_inputs++) {
      os.width(os_width);
      os << ds_inputs[i][i_inputs] << ", ";
    }
    os.width(os_width);
    os << ds_SDF[i][0] << endl;
  }
 
  ofstream outfile;
  outfile.open(filename);
  if (!outfile) {
    cout << "cannot open file: " << filename << endl;
  }
  outfile << buf.str();
  outfile.close();
  cout << "data saved to: " << filename << endl;
}
 
void SaveDatasetSPF(int num_samples, double (*ds_inputs)[3],
                    double (*ds_SPF)[3], string filename) {
  stringbuf buf;
  ostream os(&buf);
  int os_width = 24;
  os.setf(ios::scientific);
  os.precision(15);
  for (int i = 0; i < num_samples; i++) {
    // skip the non-contact data
    if (ds_SPF[i][0] == 10)
      continue;
 
    // write the data as one sample one row
    for (int i_inputs = 0; i_inputs < 3; i_inputs++) {
      os.width(os_width);
      os << ds_inputs[i][i_inputs] << ", ";
    }
    for (int i_outputs = 0; i_outputs < 2; i_outputs++) {
      os.width(os_width);
      os << ds_SPF[i][i_outputs] << ", ";
    }
    os.width(os_width);
    os << ds_SPF[i][2] << endl;
  }
 
  ofstream outfile;
  outfile.open(filename);
  if (!outfile) {
    cout << "cannot open file: " << filename << endl;
  }
  outfile << buf.str();
  outfile.close();
  cout << "data saved to: " << filename << endl;
}
 
void GenDataset(int num_samples = 100) {
  // load particle
  TriMesh trimesh;
  trimesh.InitFromSTL("shape_template.stl");
  trimesh.Decimate(500);
  trimesh.AlignAxes();
  trimesh.SetSize(1.0);
  Particle obj_p = Particle(&trimesh);
  obj_p.need_update_stl_model = true;
  cout << "particle created ... " << endl;
 
  // allocate memory
  double(*ds_inputs)[3] = new double[num_samples][3];
  double(*ds_SDF)[1] = new double[num_samples][1];
  double(*ds_SPF)[3] = new double[num_samples][3];
 
  // use bound sphere to narrow down the random space
  double dist_max = obj_p.shape->GetBoundSphereRadius() * 1.1;
  double dist_min{dist_max};
  for (auto &vert : trimesh.vertices) {
    dist_min = min(Math::NormL2(vert), dist_min);
  }
  dist_min *= 0.9;
  double dist_range = dist_max - dist_min;
 
  // random generator
  UniformDistribution uniform_dist(0.0, 1.0);
 
  // use spherical centroidal voronoi to sample uniform unit VecXT
  VecXT<Vec3d> vertices = SphericalVoronoi::Solve(1000, 10000, 1.0e-4);
  VecXT<Vec3i> facets;
  IGLWrapper::ConvexHull(vertices, &vertices, &facets);
 
  // random cases
  int count = 0;
  for (int trial = 0, i = 0; trial < num_samples * 100; trial++) {
    // random direction
    int id_facet = floor(uniform_dist.Get() * facets.size());
    auto vert_0 = vertices[facets[id_facet][0]];
    auto vert_1 = vertices[facets[id_facet][1]];
    auto vert_2 = vertices[facets[id_facet][2]];
 
    double u_vert = uniform_dist.Get();
    double v_vert = uniform_dist.Get() * (1 - u_vert);
    double w_vert = 1 - u_vert - v_vert;
 
    Vec3d tmp_dir;
    tmp_dir[0] = u_vert * vert_0[0] + v_vert * vert_1[0] + w_vert * vert_2[0];
    tmp_dir[1] = u_vert * vert_0[1] + v_vert * vert_1[1] + w_vert * vert_2[1];
    tmp_dir[2] = u_vert * vert_0[2] + v_vert * vert_1[2] + w_vert * vert_2[2];
    Math::Normalize(&tmp_dir);
 
    // random radial distance
    double tmp_dist = dist_min + uniform_dist.Get() * dist_range;
 
    // query point
    Vec3d query_p;
    query_p[0] = tmp_dist * tmp_dir[0];
    query_p[1] = tmp_dist * tmp_dir[1];
    query_p[2] = tmp_dist * tmp_dir[2];
 
    double SDF_P = trimesh.SignedDistance(query_p);
 
    if (SDF_P < 0 && SDF_P > -0.1) {
      // use query point as inputs
      ds_inputs[i][0] = query_p[0];
      ds_inputs[i][1] = query_p[1];
      ds_inputs[i][2] = query_p[2];
 
      // use signed distance field as output
      ds_SDF[i][0] = SDF_P;
 
      ds_SPF[i][0] = 10;
      ds_SPF[i][1] = 10;
      ds_SPF[i][2] = 10;
 
      i++;
    } else if (SDF_P > 0 && SDF_P < 0.1) {
      auto peoject_q = trimesh.SurfacePoint(query_p);
 
      // use query point as inputs
      ds_inputs[i][0] = query_p[0];
      ds_inputs[i][1] = query_p[1];
      ds_inputs[i][2] = query_p[2];
 
      ds_SDF[i][0] = SDF_P;
 
      // use surface project point as output
      ds_SPF[i][0] = peoject_q[0];
      ds_SPF[i][1] = peoject_q[1];
      ds_SPF[i][2] = peoject_q[2];
 
      i++;
    }
 
    count = i;
 
    if (i >= num_samples) {
      break;
    }
  }
 
  trimesh.SaveAsSTL("shape_standardized.stl");
  cout << "generated " << count << " samples." << endl;
  SaveDatasetSDF(count, ds_inputs, ds_SDF, "dataset_SDF.txt");
  SaveDatasetSPF(count, ds_inputs, ds_SPF, "dataset_SPF.txt");
 
  delete[] ds_inputs;
  delete[] ds_SDF;
  delete[] ds_SPF;
}