segmentator's Introduction

A simple cpp lib for 3d unsupervised segmentation

It provides point cloud segmentation (using graph) and mesh segmentation by offering a simple pytorch binding.

Our code is written based on this great work: Segmentator

Build Example:

cd csrc && mkdir build && cd build

cmake .. \
-DCMAKE_PREFIX_PATH=`python -c 'import torch;print(torch.utils.cmake_prefix_path)'` \
-DPYTHON_INCLUDE_DIR=$(python -c "from distutils.sysconfig import get_python_inc; print(get_python_inc())")  \
-DPYTHON_LIBRARY=$(python -c "import distutils.sysconfig as sysconfig; print(sysconfig.get_config_var('LIBDIR'))") \
-DCMAKE_INSTALL_PREFIX=`python -c 'from distutils.sysconfig import get_python_lib; print(get_python_lib())'` 

make && make install # after install, please do not delete this folder (as we only create a symbolic link)

Some Helpful Notes:

Our code is expected to be run with g++ and has not been tested on MSVC
If you cannot compile the code, please ensure your compilation setup is correct

Examples

import trimesh
import numpy as np
from segmentator import segment_mesh, segment_point, compute_vn

MESH_PATH = ??? # here is the path to your triangle mesh
mesh = trimesh.load_mesh(MESH_PATH)

# segment on mesh
vertices = torch.from_numpy(mesh.vertices.astype(np.float32))
faces = torch.from_numpy(mesh.faces.astype(np.int64))
ind = segment_mesh(vertices, faces) 
color_table = torch.randint(0, 256, size=(1 + ind.max(), 3))
np.savetxt("result_pc_mesh.txt", torch.cat([vertices, color_table[ind]], dim=1).numpy())

# segment on pointcloud
normals = torch.from_numpy(compute_vn(mesh).astype(np.float32))
edges = torch.from_numpy(mesh.edges.astype(np.int64)) # NOTE the edges can actually be obtained from knn graph or radius graph
ind = segment_point(vertices, normals, edges) 
np.savetxt("result_pc_point.txt", torch.cat([vertices, color_table[ind]], dim=1).numpy())

Recommend Projects