3D Visualizations

vispy_3d.py

"""
Get cylinder points in 3D grid given start and end point of cylinder axis
"""

import sys
import pdb
import copy
import numpy as np
from scipy.spatial import distance

import vispy
from vispy import visuals, scene
# print(vispy.sys_info())

def get_vispy_canvas():

    # Step 1 - Create canvas and a view
    canvas = vispy.scene.SceneCanvas(keys='interactive', bgcolor='w')
    view = canvas.central_widget.add_view()
    # view.camera = vispy.scene.TurntableCamera(up='z', fov=60)
    view.camera = ['perspective', 'panzoom', 'fly', 'arcball', 'base', 'turntable', None][2]
    axis    = vispy.scene.visuals.XYZAxis(parent=view.scene)
    print (' - Vispy scene setup ...')

    # Step 2 - Different kind of axis
    if 1:
        modes = ((1,2), (10,5), (100,10))
        mode = modes[0]
        scale = mode[0]
        font_size = np.power(2,mode[1]) # [(scale, pow) = [1=2, 10=5, 100=10]]
        xax = vispy.scene.Axis(pos=[[0, 0], [1*scale, 0]], tick_direction=(0, -1*scale), axis_color='r', tick_color='r', text_color='r', font_size=font_size, parent=view.scene)
        yax = vispy.scene.Axis(pos=[[0, 0], [0, 1*scale]], tick_direction=(-1*scale, 0), axis_color='g', tick_color='g', text_color='g', font_size=font_size, parent=view.scene)
        zax = vispy.scene.Axis(pos=[[0, 0], [-1*scale, 0]], tick_direction=(0, -1*scale), axis_color='b', tick_color='b', text_color='b', font_size=font_size, parent=view.scene)
        zax.transform = vispy.scene.transforms.MatrixTransform()  # its acutally an inverted xaxis
        zax.transform.rotate(90, (0, 1, 0))  # rotate cw around yaxis
        zax.transform.rotate(-45, (0, 0, 1))  # tick direction towards (-1,-1)
        print (' - Vispy axis setup ... ')

        # Step3 - Markers (that indicate scale)
        Scatter3D = vispy.scene.visuals.create_visual_node(vispy.visuals.MarkersVisual)
        p1 = Scatter3D(parent=view.scene)
        pos = np.random.random((100,3))
        pos = np.array([
            [0,0,0]
            , [0.5,0.5,0]
            , [1,1,0]
            , [1,0,0]
            , [0,1,0]
        ])
        pos = pos * scale
        p1.set_data(pos, symbol='x', size=20)
        print (' - Vispy Marker setup ...')
    
    return canvas, view
   
   
def plot3D(data, markers):

    # Step 1 - Get canvas
    canvas, view = get_vispy_canvas()

    # Step 2 - Plot data
    print (' - Plotting data ... ')
    Scatter3D = vispy.scene.visuals.create_visual_node(vispy.visuals.MarkersVisual)
    p2 = Scatter3D(parent=view.scene)
    p2.set_data(data, symbol='o', size=2)
    # p2.set_gl_state('translucent', blend=True, depth_test=True)
    # p2.set_data(pos, face_color=colors, symbol='o', size=10,
    #             edge_width=0.5, edge_color='blue')
    
    # Step 3 - Plot markers
    Scatter3D = vispy.scene.visuals.create_visual_node(vispy.visuals.MarkersVisual)
    p3_1 = Scatter3D(parent=view.scene)
    p3_1.set_data(np.array([markers[0], markers[0]]), symbol='o', size=15, edge_color='blue', face_color='blue',edge_width=1)

    Scatter3D = vispy.scene.visuals.create_visual_node(vispy.visuals.MarkersVisual)
    p3_2 = Scatter3D(parent=view.scene)
    p3_2.set_data(np.array([markers[1], markers[1]]), symbol='o', size=10, edge_color='red', face_color='red',edge_width=1)
    
    # Step 4 - Arrows
    Arrow = vispy.scene.visuals.create_visual_node(vispy.visuals.ArrowVisual)
    p4 = Arrow(parent=view.scene)
    p4.set_data(markers, color='blue', width=4)
    
    canvas.show()
    vispy.app.run()
    
def get_paf_grid(joint2, joint1, grid, limb_width=2, title=None, show=False, verbose=False):
    """
     - 2D PAF: https://arxiv.org/abs/1812.08008
     - (for 3D extension) Ref: https://stackoverflow.com/questions/27161533/find-the-shortest-distance-between-a-point-and-line-segments-not-line
    """
    
    # Step 1 - Basic Lambda functions
    unit_vector = lambda x: x / np.linalg.norm(x)
    dist_line = lambda seg, jt1, x: seg.dot(x - jt1)
    dist_segment = lambda jt2, jt1, x: np.linalg.norm(np.cross(jt2 - jt1, jt1 - x)) / np.linalg.norm(jt2 - jt1)
    is_body_part = lambda dist_line_param, thresh_line, dist_radial_param, thresh_radial : \
        1 if (0 <= dist_line_param <= thresh_line) and (0 <= dist_radial_param <= thresh_radial) else 0
    # dist_radial = lambda seg_perp, jt1, x: abs(seg_perp.dot(x - jt1))
    
    # Step 2 - Unit Vectors
    limb_uv = unit_vector(joint2 - joint1)
    tmp = np.random.randn(3)  # take a random vector
    limb_uv_perp = unit_vector(tmp - tmp.dot(limb_uv) * limb_uv)  # make it orthogonal to uv_limb
    
    # Step 3 - Thresholds
    limb_length = distance.euclidean(joint2, joint1)
    print (' - Limb Length: ', limb_length, ' || limb_uv: ', limb_uv)

    # Step 4 - Loop over the grid
    grid_new = copy.deepcopy(grid)
    grid_pts = []
    for id_x in range(grid.shape[0]):
        for id_y in range(grid.shape[1]):
            for id_z in range(grid.shape[2]):
                x = np.array([id_x, id_y, id_z])
                dist_line_val = dist_line(limb_uv, joint1, x)
                dist_segment_val = dist_segment(joint2, joint1, x)
                # dist_radial_val = dist_radial(limb_uv_perp, joint1, x)
                # decision_bool = is_body_part(dist_line_val, limb_length, dist_radial_val, limb_width)
                decision_bool = is_body_part(dist_line_val, limb_length, dist_segment_val, limb_width)
                if decision_bool:
                    if verbose:
                        print (x, ' || dist_line_val: ', round(dist_line_val,2), ' || dist_segment_val: ', round(dist_segment_val,5), ' || decision_bool: ', decision_bool)
                    grid_pts.append(list(x))
                grid_new[id_x, id_y, id_z] = decision_bool

    grid_pts = np.array(grid_pts)

    return grid_pts
 
 if __name__ == '__main__':
    grid = np.zeros((48,48,24))
    joint1 = np.array([30,12,11]) # parent
    joint2 = np.array([35,11,9])  # child
    # joint1 = np.array([25,25,10])
    # joint2 = np.array([25,25,15])
    data = get_paf_grid(joint2, joint1, grid, show=True, verbose=True)
    markers = np.array([joint1, joint2])
    
    if len(data):
        scale = 100.0
        plot3D(data/scale, markers/scale)