RigNet/utils/vis_utils.py

#-------------------------------------------------------------------------------
# Name:        vis_utils.py
# Purpose:     utilize functions for visualization, highly relied on Open3D (0.9.0)
# RigNet Copyright 2020 University of Massachusetts
# RigNet is made available under General Public License Version 3 (GPLv3), or under a Commercial License.
# Please see the LICENSE README.txt file in the main directory for more information and instruction on using and licensing RigNet.
#-------------------------------------------------------------------------------

import os
import cv2
import glob
import numpy as np
import open3d as o3d
import matplotlib.pyplot as plt


def drawSphere(center, radius, color=[0.0,0.0,0.0]):
    mesh_sphere = o3d.geometry.TriangleMesh.create_sphere(radius=radius)
    transform_mat = np.eye(4)
    transform_mat[0:3, -1] = center
    mesh_sphere.transform(transform_mat)
    mesh_sphere.paint_uniform_color(color)
    return mesh_sphere


def drawCone(bottom_center, top_position, color=[0.6, 0.6, 0.9]):
    cone = o3d.geometry.TriangleMesh.create_cone(radius=0.007, height=np.linalg.norm(top_position - bottom_center)+1e-6)
    line1 = np.array([0.0, 0.0, 1.0])
    line2 = (top_position - bottom_center) / (np.linalg.norm(top_position - bottom_center)+1e-6)
    v = np.cross(line1, line2)
    c = np.dot(line1, line2) + 1e-8
    k = np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])
    R = np.eye(3) + k + np.matmul(k, k) * (1 / (1 + c))
    if np.abs(c + 1.0) < 1e-4: # the above formula doesn't apply when cos(∠(𝑎,𝑏))=−1
        R = np.array([[-1, 0, 0], [0, 1, 0], [0, 0, -1]])
    T = bottom_center + 5e-3 * line2
    #print(R)
    cone.transform(np.concatenate((np.concatenate((R, T[:, np.newaxis]), axis=1), np.array([[0.0, 0.0, 0.0, 1.0]])), axis=0))
    cone.paint_uniform_color(color)
    return cone


def show_obj_skel(mesh_name, root):
    vis = o3d.visualization.Visualizer()
    vis.create_window()
    ctr = vis.get_view_control()

    # draw mesh
    mesh = o3d.io.read_triangle_mesh(mesh_name)
    mesh_ls = o3d.geometry.LineSet.create_from_triangle_mesh(mesh)
    mesh_ls.colors = o3d.utility.Vector3dVector([[0.8, 0.8, 0.8] for i in range(len(mesh_ls.lines))])
    vis.add_geometry(mesh_ls)

    vis.add_geometry(drawSphere(root.pos, 0.01, color=[0.1, 0.1, 0.1]))
    this_level = root.children
    while this_level:
        next_level = []
        for p_node in this_level:
            vis.add_geometry(drawSphere(p_node.pos, 0.008, color=[1.0, 0.0, 0.0])) # [0.3, 0.1, 0.1]
            vis.add_geometry(drawCone(np.array(p_node.parent.pos), np.array(p_node.pos)))
            next_level+=p_node.children
        this_level = next_level

    #param = o3d.io.read_pinhole_camera_parameters('sideview.json')
    #ctr.convert_from_pinhole_camera_parameters(param)

    vis.run()
    #vis.update_geometry()
    #vis.poll_events()
    #vis.update_renderer()

    #param = ctr.convert_to_pinhole_camera_parameters()
    #o3d.io.write_pinhole_camera_parameters('sideview.json', param)

    image = vis.capture_screen_float_buffer()
    vis.destroy_window()
    image = np.asarray(image) * 255
    image = image.astype(np.uint8)
    return image


def draw_shifted_pts(mesh_name, pts, weights=None):
    mesh = o3d.io.read_triangle_mesh(mesh_name)
    mesh_ls = o3d.geometry.LineSet.create_from_triangle_mesh(mesh)
    mesh_ls.colors = o3d.utility.Vector3dVector([[0.8, 0.8, 0.8] for i in range(len(mesh_ls.lines))])
    pred_joints = o3d.geometry.PointCloud()
    pred_joints.points = o3d.utility.Vector3dVector(pts)
    if weights is None:
        color_joints = [[1.0, 0.0, 0.0] for i in range(len(pts))]
    else:
        import matplotlib.pyplot as plt
        cmap = plt.cm.get_cmap('YlOrRd')
        #weights = (weights - np.min(weights)) / (np.max(weights) - np.min(weights))
        #weights = 1 / (1 + np.exp(-weights))
        color_joints = cmap(weights.squeeze())
        color_joints = color_joints[:, :-1]
    pred_joints.colors = o3d.utility.Vector3dVector(color_joints)

    vis = o3d.visualization.Visualizer()
    vis.create_window()
    ctr = vis.get_view_control()
    vis.add_geometry(mesh_ls)
    vis.add_geometry(pred_joints)

    param = o3d.io.read_pinhole_camera_parameters('sideview.json')
    ctr.convert_from_pinhole_camera_parameters(param)

    #vis.run()
    vis.update_geometry()
    vis.poll_events()
    vis.update_renderer()

    #param = ctr.convert_to_pinhole_camera_parameters()
    #o3d.io.write_pinhole_camera_parameters('sideview.json', param)

    image = vis.capture_screen_float_buffer()
    vis.destroy_window()
    image = np.asarray(image) * 255
    image = image.astype(np.uint8)
    return image


def draw_joints(mesh_name, pts):
    mesh = o3d.io.read_triangle_mesh(mesh_name)
    mesh_ls = o3d.geometry.LineSet.create_from_triangle_mesh(mesh)
    mesh_ls.colors = o3d.utility.Vector3dVector([[0.8, 0.8, 0.8] for i in range(len(mesh_ls.lines))])
    vis = o3d.visualization.Visualizer()
    vis.create_window()
    ctr = vis.get_view_control()
    vis.add_geometry(mesh_ls)
    for joint_pos in pts:
        vis.add_geometry(drawSphere(joint_pos, 0.006, color=[1.0, 0.0, 0.0]))

    param = o3d.io.read_pinhole_camera_parameters('sideview.json')
    ctr.convert_from_pinhole_camera_parameters(param)

    #vis.run()
    vis.update_geometry()
    vis.poll_events()
    vis.update_renderer()

    image = vis.capture_screen_float_buffer()
    vis.destroy_window()
    image = np.asarray(image) * 255
    image = image.astype(np.uint8)
    return image


def volume_to_cubes(volume, threshold=0, dim=[1., 1., 1.]):
    #o = np.array([-dim[0]/2., -dim[1]/2., -dim[2]/2.])
    o = np.array([0, 0, 0])
    step = np.array([dim[0]/volume.shape[0], dim[1]/volume.shape[1], dim[2]/volume.shape[2]])
    points = []
    lines = []
    for x in range(1, volume.shape[0]-1):
        for y in range(1, volume.shape[1]-1):
            for z in range(1, volume.shape[2]-1):
                pos = o + np.array([x, y, z]) * step
                if volume[x, y, z] > threshold:
                    vidx = len(points)
                    POS = pos + step*0.95
                    xx = pos[0]
                    yy = pos[1]
                    zz = pos[2]
                    XX = POS[0]
                    YY = POS[1]
                    ZZ = POS[2]

                    points.append(np.array([xx, yy, zz])[np.newaxis, :])
                    points.append(np.array([xx, YY, zz])[np.newaxis, :])
                    points.append(np.array([XX, YY, zz])[np.newaxis, :])
                    points.append(np.array([XX, yy, zz])[np.newaxis, :])
                    points.append(np.array([xx, yy, ZZ])[np.newaxis, :])
                    points.append(np.array([xx, YY, ZZ])[np.newaxis, :])
                    points.append(np.array([XX, YY, ZZ])[np.newaxis, :])
                    points.append(np.array([XX, yy, ZZ])[np.newaxis, :])

                    lines.append(np.array([vidx + 1, vidx + 2]))
                    lines.append(np.array([vidx + 2, vidx + 6]))
                    lines.append(np.array([vidx + 6, vidx + 5]))
                    lines.append(np.array([vidx + 1, vidx + 5]))

                    lines.append(np.array([vidx + 1, vidx + 1]))
                    lines.append(np.array([vidx + 3, vidx + 3]))
                    lines.append(np.array([vidx + 7, vidx + 7]))
                    lines.append(np.array([vidx + 5, vidx + 5]))

                    lines.append(np.array([vidx + 0, vidx + 3]))
                    lines.append(np.array([vidx + 0, vidx + 4]))
                    lines.append(np.array([vidx + 4, vidx + 7]))
                    lines.append(np.array([vidx + 7, vidx + 3]))

    return points, lines


def show_mesh_vox(mesh_filename, vox):
    vis = o3d.visualization.Visualizer()
    vis.create_window()

    vox_pts, vox_lines = volume_to_cubes(vox.data)
    vox_pts = np.concatenate(vox_pts, axis=0)
    line_set_vox = o3d.geometry.LineSet()
    line_set_vox.points = o3d.utility.Vector3dVector(vox_pts+np.array(vox.translate)[np.newaxis, :])
    line_set_vox.lines = o3d.utility.Vector2iVector(vox_lines)
    colors = [[0.0, 0.0, 1.0] for i in range(len(vox_lines))]
    line_set_vox.colors = o3d.utility.Vector3dVector(colors)
    vis.add_geometry(line_set_vox)

    mesh = o3d.io.read_triangle_mesh(mesh_filename)
    mesh_ls = o3d.geometry.LineSet.create_from_triangle_mesh(mesh)
    mesh_ls.colors = o3d.utility.Vector3dVector([[1.0, 0.0, 0.0] for i in range(len(mesh_ls.lines))])
    vis.add_geometry(mesh_ls)

    vis.run()
    vis.destroy_window()

    return