cpp_ext.cpp

#include <torch/extension.h>
#include "voxels.hpp"

namespace py = pybind11;

PYBIND11_MODULE(_metapedia_cpp_ext, m) {
  using namespace metapedia;

  m.doc() = "Metapedia CPP extensions";
  m.attr("__version__") = "0.0.1";

  // Bind the VoxelConfig class.
  py::class_<VoxelConfig>(m, "VoxelConfig")
      .def(py::init<>())
      .def(
          "__setitem__",
          [](VoxelConfig& config, VoxelConfig::Key key, VoxelDef def) {
            config.defs.emplace(key, def);
          })
      .def_static(
          "from_dict",
          [](const std::unordered_map<VoxelConfig::Key, VoxelDef>& dict) {
            VoxelConfig config;
            for (const auto& [key, def] : dict) {
              config.defs.emplace(key, def);
            }
            return config;
          });
  py::class_<EmptyVoxel>(m, "EmptyVoxel").def(py::init<>());
  py::class_<ColorVoxel>(m, "ColorVoxel").def(py::init<int, int, int>());

  // Bind the VoxelMesh class.
  py::class_<VoxelMesh>(m, "VoxelMesh")
      .def(py::init<>())
      .def_readwrite("positions", &VoxelMesh::positions)
      .def_readwrite("normals", &VoxelMesh::normals)
      .def_readwrite("colors", &VoxelMesh::colors)
      .def_readwrite("triangles", &VoxelMesh::triangles);

  // Bind the mesh generation routine.
  m.def("generate_mesh", generate_mesh);
};

utils.hpp

#pragma once 
#include <array>
#include <functional>
#include <tuple>

namespace metapedia {
// Enables lambda-based static visitor pattern for std::variant
template <class... Ts>
struct Overloaded : Ts... {
  using Ts::operator()...;
};

// explicit deduction guide for Overloaded visitor pattern
template <class... Ts>
explicit Overloaded(Ts...)->Overloaded<Ts...>;

// Basic geometric types.
using Vec3i = std::array<int, 3>;
using Vec3f = std::array<float, 3>;

template <typename Vec3>
inline auto add(Vec3 u, Vec3 v) {
  return Vec3{u[0] + v[0], u[1] + v[1], u[2] + v[2]};
}

template <typename S, typename Vec3>
inline auto scale(S a, Vec3 u) {
  return Vec3{a * u[0], a * u[1], a * u[2]};
}

// Masks for each cube face direction
enum DirMask {
  X_NEG = 0b000001,
  X_POS = 0b000010,
  Y_NEG = 0b000100,
  Y_POS = 0b001000,
  Z_NEG = 0b010000,
  Z_POS = 0b100000,
};

// Integer log base 2
inline constexpr uint32_t lg2(uint32_t x) {
  return x < 2 ? 0 : 1 + lg2(x >> 1);
}

// Generic routine for applying a function to each component of a tuple.
template <typename... Args, typename Fn, std::size_t... Idx>
inline constexpr void for_each(
    const std::tuple<Args...>& t, Fn&& f, std::index_sequence<Idx...>) {
  (f(std::get<Idx>(t)), ...);
}

// Provides efficient traversal over tensor cells via a boolean mask.
template <typename Mask, typename... Arrays, typename Fn>
inline void array_walk(Fn&& fn, Mask&& mask, Arrays&&... a) {
  CHECK_ARGUMENT(((mask.len() == a.len()) && ...));

  auto store = mask.store();
  auto iters = std::make_tuple(std::tuple(0, a.store())...);

  int pos = 0;
  auto advance = [&](auto& it) {
    auto& index = std::get<0>(it);
    auto& store = std::get<1>(it);
    while (store->ends[index] <= pos) {
      ++index;
    }
    return store->vals[index];
  };

  for (int i = 0; i < store->size; i += 1) {
    auto end = store->ends[i];
    auto val = store->vals[i];
    if (!val) {
      pos = end;
      continue;
    }
    for (; pos < end; ++pos) {
      // Advance tensor iterators.
      auto vals = std::apply(
          [&](auto&... it) { return std::tuple(advance(it)...); }, iters);

      // Invoke the traversal fn.
      std::apply(fn, std::tuple_cat(std::tuple(pos), vals));
    }
  }
}

auto pack_bytes(uint8_t a, uint8_t b, uint8_t c, uint8_t d) {
  uint32_t ret = static_cast<uint32_t>(d);
  ret |= static_cast<uint32_t>(c) << 8;
  ret |= static_cast<uint32_t>(b) << 16;
  ret |= static_cast<uint32_t>(a) << 24;
  return ret;
}

auto unpack_bytes(uint32_t bytes) {
  uint8_t a = (bytes & 0xFF000000) >> 24;
  uint8_t b = (bytes & 0x00FF0000) >> 16;
  uint8_t c = (bytes & 0x0000FF00) >> 8;
  uint8_t d = (bytes & 0x000000FF);
  return std::tuple(a, b, c, d);
}

}  // namespace metapedia

voxel_renderer.py

import torch
import random
import numpy as np
from _metapedia_cpp_ext import (
    generate_mesh,
    VoxelConfig,
    VoxelMesh,
    EmptyVoxel,
    ColorVoxel,
)
from pythreejs import (
    BufferGeometry,
    BufferAttribute,
    Mesh,
    MeshLambertMaterial,
    PerspectiveCamera,
    PointLight,
    AmbientLight,
    Renderer,
    OrbitControls,
    Scene,
)
from IPython.display import display


def random_color_config(n: int) -> VoxelConfig:
    config = VoxelConfig.from_dict({0: EmptyVoxel(),})
    for i in range(1, n):
        config[i] = ColorVoxel(
            random.randint(50, 255), random.randint(50, 255), random.randint(50, 255),
        )
    return config


def to_mesh(tensor: torch.Tensor, config: VoxelConfig) -> VoxelMesh:
    return generate_mesh(config, tensor)


def compute_mesh(tensor: torch.Tensor) -> Mesh:
    config = VoxelConfig()
    config[0] = EmptyVoxel()
    config[1] = ColorVoxel(255, 255, 255)
    config[2] = ColorVoxel(0, 130, 0)
    config[3] = ColorVoxel(255, 105, 180)
    mesh = generate_mesh(config, tensor)
    geometry = BufferGeometry(
        attributes={
            "position": BufferAttribute(
                np.array(mesh.positions, dtype="float32").reshape(-1, 3),
                normalized=False,
            ),
            "index": BufferAttribute(
                np.array(mesh.triangles, dtype="uint32"), normalized=False,
            ),
            "color": BufferAttribute(
                np.array(mesh.colors, dtype="float32").reshape(-1, 4) / 255.0,
            ),
        },
    )
    geometry.exec_three_obj_method("computeVertexNormals")
    js_mesh = Mesh(
        geometry=geometry,
        material=MeshLambertMaterial(vertexColors="VertexColors"),
        position=[0, 0, 0],
    )
    return js_mesh


def ipython_render(
    tensor: torch.Tensor,
    width=600,
    height=500,
    camera_pos=None,
    look_at=None,
    controls=None,
    light_position=None,
):
    js_mesh = compute_mesh(tensor)
    if camera_pos is None:
        camera_pos = tuple([e * 1.5 for e in tensor.shape])
    if look_at is None:
        look_at = (0, 0, 0)
    camera = PerspectiveCamera(position=tuple(camera_pos), fov=75)
    up = [0, 1, 0]
    camera.up = tuple(up)
    camera.lookAt(tuple(look_at))
    light_position = light_position or tuple(tensor.shape)
    point_light = PointLight(color="#ffffff", position=light_position)
    global_light = AmbientLight(color="#333333")
    scene = Scene(
        children=[js_mesh, camera, point_light, global_light], background="black"
    )

    if controls is None:
        controls = [OrbitControls(controlling=camera)]

    renderer = Renderer(
        camera=camera,
        background="#b0b0b0",
        background_opacity=1,
        scene=scene,
        controls=controls,
        width=width,
        height=height,
    )

    display(renderer)
    return renderer

voxels.hpp

#pragma once

#include <torch/extension.h>
#include "utils.hpp"

namespace metapedia {

using namespace torch::indexing;

struct EmptyVoxel {};

struct ColorVoxel {
  uint32_t rgba;

  ColorVoxel(uint8_t r, uint8_t g, uint8_t b, uint8_t a)
      : rgba(pack_bytes(r, g, b, a)) {}
  ColorVoxel(int8_t r, int8_t g, int8_t b) : ColorVoxel(r, g, b, 255) {}
};

struct ThreeSliceSpec {
  Slice zero;
  Slice one;
  Slice two;

  auto index_into(torch::Tensor &tensor) {
    return tensor.index({zero, one, two});
  }
};

using VoxelDef = std::variant<EmptyVoxel, ColorVoxel>;

struct VoxelConfig {
  using Key = int32_t;
  std::unordered_map<Key, VoxelDef> defs;
};

using VoxelTensor = torch::Tensor;

struct VoxelMesh {
  std::vector<float> positions;
  std::vector<float> normals;
  std::vector<uint8_t> colors;
  std::vector<uint32_t> triangles;

  auto vertex_count() const {
    return positions.size() / 3;
  }

  auto triangle_count() const {
    return triangles.size() / 3;
  }
};

auto compute_adjacency_mask(const VoxelTensor& tensor) {
  assert(tensor.dim() == 3);

  auto w = tensor.size(0);
  auto h = tensor.size(1);
  auto d = tensor.size(2);
  auto sl = [&](int x_0, int x_1, int y_0, int y_1, int z_0, int z_1) {
    x_0 = x_0 < 0 ? w + x_0 : x_0;
    y_0 = y_0 < 0 ? h + y_0 : y_0;
    z_0 = z_0 < 0 ? d + z_0 : z_0;
    x_1 = x_1 <= 0 ? w + x_1 : x_1;
    y_1 = y_1 <= 0 ? h + y_1 : y_1;
    z_1 = z_1 <= 0 ? d + z_1 : z_1;
    return ThreeSliceSpec(
      {Slice(x_0, x_1, 1), Slice(y_0, y_1, 1), Slice(z_0, z_1, 1)}
    );
  };

  auto update = [](torch::Tensor& tensor, ThreeSliceSpec slice, torch::Tensor input) {
    auto merged = torch::bitwise_or(tensor.index({slice.zero, slice.one, slice.two}), input);
    tensor.index_put_({slice.zero, slice.one, slice.two}, merged);
  };

  auto empty = (tensor == 0) * ((int32_t)0xFFFFFFFF);
  auto adjacencies = torch::zeros_like(tensor);

  constexpr auto x_neg_mask = static_cast<int>(DirMask::X_NEG);
  constexpr auto x_pos_mask = static_cast<int>(DirMask::X_POS);
  constexpr auto y_neg_mask = static_cast<int>(DirMask::Y_NEG);
  constexpr auto y_pos_mask = static_cast<int>(DirMask::Y_POS);
  constexpr auto z_neg_mask = static_cast<int>(DirMask::Z_NEG);
  constexpr auto z_pos_mask = static_cast<int>(DirMask::Z_POS);

  update(
    adjacencies,
    sl(0, 1, 0, h, 0, d),
    torch::bitwise_and(torch::bitwise_not(
      sl(0, 1, 0, h, 0, d).index_into(empty)), 
      x_neg_mask));

  update(
    adjacencies,
    sl(-1, w, 0, h, 0, d),
    torch::bitwise_and(torch::bitwise_not(
      sl(-1, w, 0, h, 0, d).index_into(empty)),
      x_pos_mask));


  update(
    adjacencies,
    sl(0, w, 0, 1, 0, d),
    torch::bitwise_and(torch::bitwise_not(
      sl(0, w, 0, 1, 0, d).index_into(empty)),
      y_neg_mask
    ));

  update(
      adjacencies,
      sl(0, w, -1, h, 0, d),
      torch::bitwise_and(torch::bitwise_not(
        sl(0, w, -1, h, 0, d).index_into(empty)),
        y_pos_mask
      ));

  update(
      adjacencies,
      sl(0, w, 0, h, 0, 1),
      torch::bitwise_and(torch::bitwise_not(
        sl(0, w, 0, h, 0, 1).index_into(empty)),
        z_neg_mask
      ));

  update(
      adjacencies,
      sl(0, w, 0, h, -1, d),
      torch::bitwise_and(torch::bitwise_not(
        sl(0, w, 0, h, -1, d).index_into(empty)),
        z_pos_mask
      ));

  // Initialize the adjacency tensor interior values.
  auto shift = [&](int dx, int dy, int dz) {
    int x_0 = std::max(0, dx), x_1 = std::min(w, w + dx);
    int y_0 = std::max(0, dy), y_1 = std::min(h, h + dy);
    int z_0 = std::max(0, dz), z_1 = std::min(d, d + dz);
    return sl(x_0, x_1, y_0, y_1, z_0, z_1);
  };

  auto edge = [&](int dx, int dy, int dz) {
    auto s_1 = shift(dx, dy, dz);
    auto s_2 = shift(-dx, -dy, -dz);
    return torch::bitwise_and(
      torch::bitwise_not(s_1.index_into(empty)),
      s_2.index_into(empty)
    );
  };

  update(adjacencies, shift(1, 0, 0), torch::bitwise_and(edge(1, 0, 0), x_neg_mask));
  update(adjacencies, shift(-1, 0, 0), torch::bitwise_and(edge(-1, 0, 0), x_pos_mask));
  update(adjacencies, shift(0, 1, 0), torch::bitwise_and(edge(0, 1, 0), y_neg_mask));
  update(adjacencies, shift(0, -1, 0), torch::bitwise_and(edge(0, -1, 0), y_pos_mask));
  update(adjacencies, shift(0, 0, 1), torch::bitwise_and(edge(0, 0, 1), z_neg_mask));
  update(adjacencies, shift(0, 0, -1), torch::bitwise_and(edge(0, 0, -1), z_pos_mask));

  return adjacencies;
}

auto generate_mesh(const VoxelConfig& config, const VoxelTensor& tensor) {
  assert(tensor.dim() == 3);

  VoxelMesh mesh;

  // Adds 4 vertices and 2 triangles (for one cube face) to the output mesh.
  auto emit_face_geometry = [&](Vec3i origin, DirMask dir) {
    static auto positions = [] {
      using T = std::array<Vec3i, 4>;
      std::vector<T> ret;
      ret.push_back(T{{{0, 0, 0}, {0, 1, 0}, {0, 1, 1}, {0, 0, 1}}});  // X_NEG
      ret.push_back(T{{{1, 1, 0}, {1, 0, 0}, {1, 0, 1}, {1, 1, 1}}});  // X_POS
      ret.push_back(T{{{1, 0, 0}, {0, 0, 0}, {0, 0, 1}, {1, 0, 1}}});  // Y_NEG
      ret.push_back(T{{{0, 1, 0}, {1, 1, 0}, {1, 1, 1}, {0, 1, 1}}});  // Y_POS
      ret.push_back(T{{{0, 0, 0}, {1, 0, 0}, {1, 1, 0}, {0, 1, 0}}});  // Z_NEG
      ret.push_back(T{{{0, 1, 1}, {1, 1, 1}, {1, 0, 1}, {0, 0, 1}}});  // Z_POS
      return ret;
    }();

    static auto normals = [] {
      std::vector<Vec3i> ret;
      ret.push_back({-1, 0, 0});  // X_NEG
      ret.push_back({1, 0, 0});   // X_POS
      ret.push_back({0, -1, 0});  // Y_NEG
      ret.push_back({0, 1, 0});   // Y_POS
      ret.push_back({0, 0, -1});  // Y_NEG
      ret.push_back({0, 0, 1});   // Y_POS
      return ret;
    }();

    // Emit vertex indices for the two new triangles.
    auto base = mesh.vertex_count();
    mesh.triangles.push_back(base);
    mesh.triangles.push_back(base + 2);
    mesh.triangles.push_back(base + 1);
    mesh.triangles.push_back(base + 3);
    mesh.triangles.push_back(base + 2);
    mesh.triangles.push_back(base);

    // Emit vertex positions and normals for the 4 new vertices.
    auto dir_index = lg2(static_cast<int>(dir));
    for (const auto& pos_vec : positions.at(dir_index)) {
      Vec3i out = add(origin, pos_vec);
      mesh.positions.push_back(out[0]);
      mesh.positions.push_back(out[1]);
      mesh.positions.push_back(out[2]);
    }
    for (int i = 0; i < 4; i += 1) {
      Vec3i out = normals.at(dir_index);
      mesh.normals.push_back(out[0]);
      mesh.normals.push_back(out[1]);
      mesh.normals.push_back(out[2]);
    }
  };

  // Adds 4 vertex color attributes (for one cube face) to the output mesh.
  auto emit_face_colors = [&](int32_t rgba) {
    for (int i = 0; i < 4; i += 1) {
      auto [r, g, b, a] = unpack_bytes(rgba);
      mesh.colors.push_back(r);
      mesh.colors.push_back(g);
      mesh.colors.push_back(b);
      mesh.colors.push_back(a);
    }
  };

  // Create a tensor of the adjacency bitmask at each non-empty tensor cell.
  auto adjacencies = compute_adjacency_mask(tensor);

  // TODO: Pre-allocate mesh vectors to fit the exact output.

  // Iterate over each boundary voxel and emit its mesh attributes.
  auto w = tensor.size(0);
  auto h = tensor.size(1);
  auto d = tensor.size(2);
  for (int x = 0; x < w; x++) {
    for (int y = 0; y < h; y++) {
      for (int z = 0; z < d; z++) {
        Vec3i origin = {x, y, z};
        int adj = adjacencies.index({x, y, z}).item<int>();
        int key = tensor.index({x, y, z}).item<int>();

        if (adj == 0) {
          continue;
        }

        auto v = std::get<ColorVoxel>(config.defs.at(key));
        if (adj & static_cast<int>(DirMask::X_NEG)) {
          emit_face_geometry(origin, DirMask::X_NEG);
          emit_face_colors(v.rgba);
        }
        if (adj & static_cast<int>(DirMask::X_POS)) {
          emit_face_geometry(origin, DirMask::X_POS);
          emit_face_colors(v.rgba);
        }
        if (adj & static_cast<int>(DirMask::Y_NEG)) {
          emit_face_geometry(origin, DirMask::Y_NEG);
          emit_face_colors(v.rgba);
        }
        if (adj & static_cast<int>(DirMask::Y_POS)) {
          emit_face_geometry(origin, DirMask::Y_POS);
          emit_face_colors(v.rgba);
        }
        if (adj & static_cast<int>(DirMask::Z_NEG)) {
          emit_face_geometry(origin, DirMask::Z_NEG);
          emit_face_colors(v.rgba);
        }
        if (adj & static_cast<int>(DirMask::Z_POS)) {
          emit_face_geometry(origin, DirMask::Z_POS);
          emit_face_colors(v.rgba);
        }
      }
    }
  }
  return mesh;
}

} // namespace metapedia