Skip to content

Instantly share code, notes, and snippets.

@thomasweng15
Last active December 21, 2022 22:48
Show Gist options
  • Save thomasweng15/2f636a5ffa6e0d191e42e138738cf58a to your computer and use it in GitHub Desktop.
Save thomasweng15/2f636a5ffa6e0d191e42e138738cf58a to your computer and use it in GitHub Desktop.
trimesh to plotly
import os
import trimesh
import plotly.graph_objects as go
import plotly.figure_factory as ff
import numpy as np
from voxelized_pointcloud_sampling import create_grid_points_from_bounds
import pickle
"""Load all meshes and inputs into the same plot to make sure everything is aligned. """
def get_voxel_grid(path):
vdata = np.load(
path
) # ['point_cloud', 'compressed_occupancies', 'bb_min', 'bb_max', 'res']
vdata[
"compressed_occupancies"
].shape # Need to do some post processing to get the voxels, check if there's existing code
occupancies = np.unpackbits(vdata["compressed_occupancies"])
input_res = 256
input = np.reshape(occupancies, (256,) * 3) # 256, 256, 256
bb_min = vdata["bb_min"]
bb_max = vdata["bb_max"]
grid_points = create_grid_points_from_bounds(bb_min, bb_max, input_res)
occupied_idxs = np.where(occupancies == 1)[0]
occupied_pts = grid_points[occupied_idxs]
return input, occupied_pts
data_dir = "/media/ExtraDrive1/IDM/data"
t = 4
exp_name = "3-5layer_trajs5"
# Load cloth and sphere meshes
path = f"{data_dir}/{exp_name}/0-{t}/obs_curr_cloth_norm.off"
with open(path, "r") as f:
off_data = trimesh.exchange.off.load_off(f)
mesh = trimesh.Trimesh(vertices=off_data["vertices"], faces=off_data["faces"])
x, z, y = mesh.vertices.T
size = (mesh.vertices.max(axis=0) - mesh.vertices.min(axis=0)).T
path = f"{data_dir}/{exp_name}/0-{t}/obs_curr_sphere_l_norm.off"
with open(path, "r") as f:
off_data = trimesh.exchange.off.load_off(f)
mesh_l = trimesh.Trimesh(vertices=off_data["vertices"], faces=off_data["faces"])
x_l, z_l, y_l = mesh_l.vertices.T
path = f"{data_dir}/{exp_name}/0-{t}/obs_curr_sphere_r_norm.off"
with open(path, "r") as f:
off_data = trimesh.exchange.off.load_off(f)
mesh_r = trimesh.Trimesh(vertices=off_data["vertices"], faces=off_data["faces"])
x_r, z_r, y_r = mesh_r.vertices.T
# Load sampled query pairs
ppath = f"{data_dir}/{exp_name}/0-{t}/labels_gifs_curr.npz"
labels = np.load(ppath) # ['pairs', 'df', 'grid_coords', 'labels']
pairs = labels["pairs"] # 100k x 6
px, pz, py = pairs[:, 0], pairs[:, 1], pairs[:, 2]
pu, pw, pv = (
pairs[:, 3] - pairs[:, 0],
pairs[:, 4] - pairs[:, 1],
pairs[:, 5] - pairs[:, 2],
)
# Get scaled action
with open(
f"{data_dir}/{exp_name}/scale.pkl",
"rb",
) as f:
scale = pickle.load(f)
scale_centroid = scale["centroid"]
scale_bounds = scale["bounds"]
with open(
f"{data_dir}/{exp_name}/0-{t}/obs_curr.pkl",
"rb",
) as f:
obs_curr_data = pickle.load(f)
center_l_tip = obs_curr_data["center_l_tip"]
center_l_scaled = (center_l_tip - scale_centroid) / scale_bounds
with open(
f"{data_dir}/{exp_name}/0-{t}/obs_next.pkl",
"rb",
) as f:
obs_next_data = pickle.load(f)
center_l_tip_next = obs_next_data["center_l_tip"]
center_l_scaled_next = (center_l_tip_next - scale_centroid) / scale_bounds
action = (
center_l_scaled_next - center_l_scaled
) # action vector should be the same whether we use sphere_l or sphere_r to compute this.
# Get voxel occupancy grid
vpath = f"{data_dir}/{exp_name}/0-{t}/voxelized_point_cloud_256res_10000points_obs_curr_cloth_norm.npz"
cloth_voxel, cloth_occupied_pts = get_voxel_grid(vpath)
vpath = f"{data_dir}/{exp_name}/0-{t}/voxelized_point_cloud_256res_10000points_obs_curr_sphere_l_norm.npz"
sphere_l_voxel, sphere_l_occupied_pts = get_voxel_grid(vpath)
vpath = f"{data_dir}/{exp_name}/0-{t}/voxelized_point_cloud_256res_10000points_obs_curr_sphere_r_norm.npz"
sphere_r_voxel, sphere_r_occupied_pts = get_voxel_grid(vpath)
# create input_voxel shaped 256 x 256 x 256 x F, where F are the input features
# F: [cloth occupancy, flow xyz (only at action pts), ]
flow_voxel = np.tile(
sphere_l_voxel.astype(np.float32)[:, :, :, np.newaxis], (1, 1, 1, 3)
)
flow_voxel[:, :, :, 0] *= action[0]
flow_voxel[:, :, :, 1] *= action[1]
flow_voxel[:, :, :, 2] *= action[2]
input_voxel = np.concatenate([cloth_voxel[:, :, :, np.newaxis], flow_voxel], axis=-1)
# visualize the input voxel grid
input_voxel_flat = np.reshape(input_voxel, (256**3, 4))
cloth_idxs = np.where(input_voxel_flat[:, 0] == 1)[0]
grid_points = create_grid_points_from_bounds(-0.5, 0.5, 256)
cloth_pts = grid_points[cloth_idxs]
sphere_idxs = np.where(input_voxel_flat[:, 1] != 0)[0]
sphere_pts = grid_points[sphere_idxs]
sx, sz, sy = sphere_pts[:, 0], sphere_pts[:, 1], sphere_pts[:, 2]
flows = input_voxel_flat[sphere_idxs, -3:]
su, sw, sv = flows[:, 0], flows[:, 1], flows[:, 2]
# Plot cloth mesh
fig = ff.create_trisurf(
x=x,
y=y,
z=z,
simplices=mesh.faces,
aspectratio={"x": 1.0, "y": 1.0, "z": 1.0},
show_colorbar=False,
)
# Plot cloth voxel points
fig.add_trace(
go.Scatter3d(
x=cloth_pts[:, 0],
y=cloth_pts[:, 2],
z=cloth_pts[:, 1],
mode="markers",
marker=dict(size=1, color="orange"),
)
)
# Plot action arrows
skip = 50
fig.add_trace(
go.Cone(
x=sx[::skip],
y=sy[::skip],
z=sz[::skip],
u=su[::skip],
v=sv[::skip],
w=sw[::skip],
colorscale="Reds",
sizemode="absolute",
sizeref=0.05
# mode="markers",
# marker=dict(size=2, color='red')
)
)
fig.show()
# fig.write_html("test.html")
# also visualize sampled points
fig.add_trace(
go.Scatter3d(
x=pairs[:, 0],
y=pairs[:, 2],
z=pairs[:, 1],
mode="markers",
marker=dict(size=1, color="blue", opacity=0.3),
)
)
fig.add_trace(
go.Scatter3d(
x=pairs[:, 3],
y=pairs[:, 5],
z=pairs[:, 4],
mode="markers",
marker=dict(size=1, color="blue", opacity=0.3),
)
)
fig.show()
import os
import trimesh
import plotly.graph_objects as go
import plotly.figure_factory as ff
import numpy as np
# Make a sphere
mesh = trimesh.primitives.Sphere(radius=10, center=(0, 0, 0))
x, z, y = mesh.vertices.T
# Plot cloth mesh
fig = ff.create_trisurf(
x=x,
y=y,
z=z,
simplices=mesh.faces,
aspectratio={"x": 1.0, "y": 1.0, "z": 1.0},
show_colorbar=False,
)
# fig.show()
fig.write_html("test.html")
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment