trimesh to plotly

visualize.py

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()

visualize_circle.py

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")