bb87928bce
Upstream: https://github.com/Rhoban/onshape-to-robot Upstream-Commit: 80e710700aac9573a2230f74f7ce9e094833a0bc Upstream-Branch: master
121 lines
4.1 KiB
Python
121 lines
4.1 KiB
Python
import hashlib
|
|
import os
|
|
from pathlib import Path
|
|
from .message import bright, info, error, warning
|
|
from .processor import Processor
|
|
from .config import Config
|
|
from .robot import Robot, Part
|
|
from .geometry import Mesh
|
|
import numpy as np
|
|
import pickle
|
|
|
|
|
|
class ProcessorConvexDecomposition(Processor):
|
|
"""
|
|
Convex decomposition processor. Runs CoACD algorithm on collision meshes to use a convex approximation.
|
|
"""
|
|
|
|
is_safe: bool = False
|
|
|
|
def __init__(self, config: Config):
|
|
super().__init__(config)
|
|
|
|
# Enable convex decomposition
|
|
self.convex_decomposition: bool = config.get("convex_decomposition", False)
|
|
self.rainbow_colors: bool = config.get("rainbow_colors", False)
|
|
|
|
self.check_coacd()
|
|
|
|
def get_cache_path(self) -> Path:
|
|
"""
|
|
Return the path to the user cache.
|
|
"""
|
|
path = Path.home() / ".cache" / "onshape-to-robot-convex-decomposition"
|
|
path.mkdir(parents=True, exist_ok=True)
|
|
return path
|
|
|
|
def check_coacd(self):
|
|
if self.convex_decomposition:
|
|
print(bright("* Checking CoACD presence..."))
|
|
try:
|
|
import coacd
|
|
import trimesh
|
|
except ImportError:
|
|
print(bright("Can't import CoACD, disabling convex decomposition."))
|
|
print(info("TIP: consider installing CoACD:"))
|
|
print(info("pip install coacd trimesh"))
|
|
self.convex_decomposition = False
|
|
|
|
def process(self, robot: Robot):
|
|
if self.convex_decomposition:
|
|
os.makedirs(self.config.asset_path("convex_decomposition"), exist_ok=True)
|
|
getcwd = os.getcwd()
|
|
os.chdir(self.config.output_directory)
|
|
|
|
for link in robot.links:
|
|
for part in link.parts:
|
|
self.convex_decompose(part)
|
|
|
|
os.chdir(getcwd)
|
|
|
|
def convex_decompose(self, part: Part):
|
|
import coacd
|
|
import trimesh
|
|
|
|
collision_meshes = [mesh for mesh in part.meshes if mesh.collision]
|
|
if len(collision_meshes) > 0:
|
|
if len(collision_meshes) > 1:
|
|
print(
|
|
warning(
|
|
f"* Skipping convex decomposition for part {part.name} as it already has multiple collision meshes."
|
|
)
|
|
)
|
|
|
|
collision_mesh = collision_meshes[0]
|
|
|
|
# Retrieving file SHA1
|
|
sha1 = hashlib.sha1(open(collision_mesh.filename, "rb").read()).hexdigest()
|
|
cache_filename = f"{self.get_cache_path()}/{sha1}.pkl"
|
|
|
|
if os.path.exists(cache_filename):
|
|
print(
|
|
info(
|
|
f"* Loading cached CoACD decomposition cache for part {part.name}"
|
|
)
|
|
)
|
|
with open(cache_filename, "rb") as f:
|
|
meshes = pickle.load(f)
|
|
else:
|
|
mesh = trimesh.load(collision_mesh.filename, force="mesh")
|
|
mesh = coacd.Mesh(mesh.vertices, mesh.faces)
|
|
meshes = coacd.run_coacd(mesh, max_convex_hull=16)
|
|
with open(cache_filename, "wb") as f:
|
|
pickle.dump(meshes, f)
|
|
|
|
part.collision_meshes = []
|
|
filename = self.config.asset_path(
|
|
f"convex_decomposition/{part.name}_%05d.stl"
|
|
)
|
|
for k, mesh in enumerate(meshes):
|
|
mesh = trimesh.Trimesh(vertices=mesh[0], faces=mesh[1])
|
|
mesh.export(filename % k)
|
|
color = (
|
|
np.concatenate([np.random.rand(3), [1.0]]) if self.rainbow_colors else collision_mesh.color
|
|
)
|
|
part.meshes.append(
|
|
Mesh(
|
|
filename % k,
|
|
color,
|
|
visual=False,
|
|
collision=True,
|
|
)
|
|
)
|
|
part.collision_meshes.append(filename % k)
|
|
|
|
collision_mesh.collision = False
|
|
part.prune_unused_geometry()
|
|
|
|
print(
|
|
info(f"* Decomposed part {part.name} into {len(meshes)} convex shapes.")
|
|
)
|