hoho/vis.py

import trimesh
import numpy as np
from copy import deepcopy

def line(p1, p2, c=(255,0,0), resolution=10, radius=0.05):
    '''draws a 3d cylinder along the line (p1, p2)'''
    # check colors
    if len(c) == 1:
        c = [c[0]]*4
    elif len(c) == 3:
        c = [*c, 255]
    elif len(c) != 4:
        raise ValueError(f'{c} is not a valid color (must have 1,3, or 4 elements).')
        
    # compute length and direction of segment
    p1, p2 = np.asarray(p1), np.asarray(p2)
    l = np.linalg.norm(p2-p1)
    
    direction = (p2 - p1) / l
    
    # point z along direction of segment
    T = np.eye(4)
    T[:3, 2] = direction
    T[:3, 3] = (p1+p2)/2
    
    #reorthogonalize basis
    b0, b1 = T[:3, 0], T[:3, 1]
    if np.abs(np.dot(b0, direction)) < np.abs(np.dot(b1, direction)):
        T[:3, 1] = -np.cross(b0, direction)
    else:
        T[:3, 0] = np.cross(b1, direction)
    
    # generate and transform mesh
    mesh = trimesh.primitives.Cylinder(radius=radius, height=l, transform=T)
    
    # apply uniform color
    mesh.visual.vertex_colors = np.ones_like(mesh.visual.vertex_colors)*c
         
    return mesh


def show_grid(edges, meshes=None, row_length=5):
    '''
        edges: list of list of meshes
        meshes: optional corresponding list of meshes
        row_length: number of meshes per row
  
        returns trimesh.Scene()
    '''
    
    T = np.eye(4)
    out = []
    edges = [sum(e[1:], e[0]) for e in edges]
    row_height = 1.1 * max((e.extents for e in edges), key=lambda e: e[1])[2]
    col_width = 1.1 * max((e.extents for e in edges), key=lambda e: e[0])[0]
    # print(row_height, col_width)
    
    if meshes is None:
        meshes = [None]*len(edges)

    for i, (gt, mesh) in enumerate(zip(edges, meshes), start=0):
        mesh = deepcopy(mesh)
        gt = deepcopy(gt)

        if i%row_length != 0:
            T[0, 3] += col_width

        else:
            T[0, 3] = 0
            T[1, 3] += row_height

        # print(T[0,3]/col_width, T[2,3]/row_height)
        
        if mesh is not None:
            mesh.apply_transform(T)
            out.append(mesh)
                            
        gt.apply_transform(T)
        out.append(gt)
        
                            
        out.extend([mesh, gt])

            
    return trimesh.Scene(out)