SkeletonDiffusion_demo/visual.py

from vedo import Mesh, Points, Plotter, screenshot, Arrow, Light, Plane
import imageio
import os
import numpy as np
from scipy.interpolate import Rbf, NearestNDInterpolator as Near
import re

screenshot_scale = 3

cam = dict(
        position=(-3.85922, -4.78140, 0.689328),
        focal_point=(0.192544, 4.77379e-3, 0.0127248),
        viewup=(0.0724348, 0.109097, 0.991388),
        distance=5.25119,
        clipping_range=(4.38228, 6.36775),
    )

_red = [0.9, 0.1, 0.1]
_green= [0.4, 1.0, 0.8]
_blue =[0.7, 0.9, 1.0]
_dark_blue =[0.03, 0.4, 0.7]

def sort_list(l):
    try:
        return list(sorted(l, key=lambda x: int(re.search(r'\d+(?=\.)', x).group())))
    except AttributeError:
        return sorted(l)


def get_lights(light_color=[255, 255, 245], offset=[0, 0, 0], intensity_mult=1):
    # Add light sources at the given positions
    # (grab the position and color of the arrow object)
    orig = np.array([0, 0, 0]) + np.array(offset)
    phl = Arrow(np.array([0.1, 0.1, 10]) + offset, orig, c=light_color).scale(0.2)
    pfl = Arrow(np.array([1.5, 0.1, 0.3]) + offset, orig, c=light_color).scale(0.2)
    pbl = Arrow(np.array([-1.5, 0.1, 0.3]) + offset, orig, c=light_color).scale(0.2)
    prl = Arrow(np.array([0.1, -1.5, 0.3]) + offset, orig, c=light_color).scale(0.2)
    pll = Arrow(np.array([0.1, 1.5, 0.3]) + offset, orig, c=light_color).scale(0.2)
    hl = Light(phl, intensity=0.7 * intensity_mult, angle=180, )
    rl = Light(pfl, intensity=0.6 * intensity_mult, angle=180, )
    ll = Light(pbl, intensity=0.6 * intensity_mult, angle=180, )
    bl = Light(pll, intensity=0.6 * intensity_mult, angle=180, )
    fl = Light(prl, intensity=1 * intensity_mult, angle=180, )
    lights = [hl, fl, bl, ll, rl]
    return lights

def get_mesh(mesh_name, mesh_opacity=1.0, mesh_lighting='default', rotation=[0, 0, 0], offset=[0, 0, 0], scale=0):
    offset = np.array(offset)
    #styles = ['default', 'metallic', 'plastic', 'shiny', 'glossy', 'ambient', 'off']
    mesh = Mesh(mesh_name, c=_blue, alpha=mesh_opacity).lighting(mesh_lighting)
    if scale == 0:
        bnd = np.array(mesh.bounds())
        scale = 1 / max(bnd[1]-bnd[0], bnd[3]-bnd[2], bnd[5]-bnd[4])
    mesh.scale(scale)
    mesh.pos(offset)
    mesh.rotate_x(rotation[0])
    mesh.rotate_y(rotation[1])
    mesh.rotate_z(rotation[2])
    # mesh.phong()
    # mesh2.phong()
    return mesh, scale

def get_ptc(ptc_name, rotation=[0, 0, 0], offset=[0, 0, 0], scale=0, color=_dark_blue):
    offset = np.array(offset)
    #styles = ['default', 'metallic', 'plastic', 'shiny', 'glossy', 'ambient', 'off']
    mesh = Points(Mesh(ptc_name), r=1).c([0.5,0.5,0.5]).alpha(0.8)
    if scale == 0:
        bnd = np.array(mesh.bounds())
        scale = 1 / max(bnd[1]-bnd[0], bnd[3]-bnd[2], bnd[5]-bnd[4])
    mesh.scale(scale)
    mesh.pos(offset)
    mesh.rotate_x(rotation[0])
    mesh.rotate_y(rotation[1])
    mesh.rotate_z(rotation[2])
    # mesh.phong()
    # mesh2.phong()
    return mesh, scale

def get_mesh_shadow(mesh, offset=[0, 0, 0], plane_normal=(0, 0, 1), direction=[0.1, -1.8, 3]):
    shadow = []
    shad_col = np.array([0.8, 0.8, 0.8])
    min_z = mesh.points().min(axis=0)[2]
    # mesh.add_shadow(plane='z', point=min_z, alpha=1, c=shad_col, culling=0.9,)
    plane = Plane(pos=np.array([0, 0, min_z]) + np.array(offset), normal=plane_normal, s=[7, 7]).alpha(0.2)
    shad = mesh.clone().project_on_plane(plane, direction=-np.array(direction) + np.array(offset))
    shad.c(shad_col).alpha(1).lighting("off").use_bounds(False)
    shadow = shad
    return shadow




def create_visual(mesh_files, directory, x,y,z, type='mesh'):
    meshes, scales =[], []
    # plotter = Plotter(offscreen=True)
    plotter = Plotter(bg=[255, 255, 255], offscreen=True)

    frame_files = []
    # lights = get_lights()
    for file in mesh_files:
        if type == 'mesh':
            mesh, scale = get_mesh(file, rotation=[x,y,z], mesh_lighting='metallic')
        else:
            mesh, scale = get_ptc(file)
        plotter.clear()
        
        shadow = get_mesh_shadow(mesh)
        # plotter.add([mesh, shadow])
        plotter.show(mesh, shadow, __doc__, interactive=False, camera=cam, resetcam=True, zoom=2)
        frame_path =file[:-4] + ".png"
        frame_files.append(frame_path)
        screenshot(frame_path, scale=1)
    
    # Compile saved frames into a GIF
    filename = os.path.join(directory, type+'_animation.gif')
    with imageio.get_writer(filename, mode='I', duration=0.2) as writer:
        for filename in frame_files:
            image = imageio.imread(filename)
            writer.append_data(image)
            
        

def create_mesh_animation(mesh_files, directory, x,y,z):
    
    # Load all meshes
    meshes = [Mesh(f) for f in mesh_files]

    # Create a plotter
    plotter = Plotter(bg=[255, 255, 255], offscreen=True)
    # plotter = vedo.Plotter()# Use offscreen rendering to save frames

    frame_files = []  # To store the path of each frame image
    
    # plotter.camera.SetViewUp(0,0,1)
    
    
    # shadow = mesh.clone().projectOnPlane(direction=(0,0,-1))

    # Color the shadow mesh dark and make it slightly transparent
    
    for frame_id, mesh in enumerate(meshes):
        mesh.rotate_x(x).rotate_y(y).rotate_z(z).lighting('metallic')
        shadow = get_mesh_shadow(mesh)
        scalars = 1- mesh.vertices[:, 2]
        mesh.cmap('Blues', scalars)
        # mesh.add_shadow('x', 0.95)
        plotter.clear()
        # plotter.add(mesh)
        plotter.show(mesh, shadow, __doc__, interactive=False, camera=cam, resetcam=True, zoom=2)
        frame_path = mesh_files[frame_id][:-4] + ".png"
        frame_files.append(frame_path)
        screenshot(frame_path, scale=1)  # Save frame as image
        
        
        #crop the images
        # os.system("""convert -trim {0} "{1}" """.format(frame_path, frame_path))


    # Compile saved frames into a GIF
    filename = os.path.join(directory, 'mesh_animation.gif')
    with imageio.get_writer(filename, mode='I', duration=0.2) as writer:
        for f in frame_files:
            image = imageio.imread(f)
            writer.append_data(image)
            os.system("""convert -trim {0} "{1}" """.format(f, f))
    
    



def create_point_cloud_animation(files, directory, x,y,z):
    # Load all meshes
    meshes = [Points(Mesh(f), r=1) for f in files]

    # Create a plotter
    plotter = Plotter(bg=[255, 255, 255], offscreen=True)
    # plotter = vedo.Plotter()# Use offscreen rendering to save frames

    frame_files = []  # To store the path of each frame image
    
    # plotter.camera.SetViewUp(0,0,1)
    
    
    # shadow = mesh.clone().projectOnPlane(direction=(0,0,-1))

    # Color the shadow mesh dark and make it slightly transparent
    
    for frame_id, mesh in enumerate(meshes):
        mesh.rotate_x(x).rotate_y(y).rotate_z(z).c([0.5,0.5,0.5])
        shadow = get_mesh_shadow(mesh)
        # scalars = 1- mesh.vertices[:, 2]
        # mesh.cmap('binary', scalars)
        plotter.clear()
        plotter.show(mesh, shadow, __doc__, interactive=False, camera=cam, resetcam=True, zoom=2)

        frame_path = files[frame_id][:-4] + ".png"
        frame_files.append(frame_path)
        screenshot(frame_path, scale=1)  # Save frame as image
        
        
        #crop the images
        # os.system("""convert -trim {0} "{1}" """.format(frame_path, frame_path))


    # Compile saved frames into a GIF
    filename = os.path.join(directory, 'pct_animation.gif')
    with imageio.get_writer(filename, mode='I', duration=0.2) as writer:
        for f in frame_files:
            image = imageio.imread(f)
            writer.append_data(image)
            os.system("""convert -trim {0} "{1}" """.format(f, f))



def point_cloud_colormap(gt_ptc, ptc, corr_num):
    
    s1, _ = get_ptc(gt_ptc, color=_red)
    s2, _ = get_ptc(ptc, color=_blue)
    
    landmarks1 = s1.vertices[:corr_num,:]
    landmarks2 = s2.vertices[:corr_num,:]
    
    x, y, z = np.split(landmarks1, 3, axis=1)
    desc = x**2 + y**2 + z**2
    
    # itr = Rbf(x, y, z, desc)          # Radial Basis Function interpolator
    itr = Near(landmarks1, desc)       # Nearest-neighbour interpolator

    # interpolate descriptor on the full set of mesh vertices
    xi, yi, zi = np.split(s2.vertices, 3, axis=1)
    interpolated_desc = itr(xi, yi, zi)
    
    s2.cmap('coolwarm', interpolated_desc.squeeze()).add_scalarbar(title='error')
    # s2.point_colors(interpolated_desc, cmap='coolwarm', vmin=min(interpolated_desc), vmax=max(interpolated_desc)).addScalarBar(title='$\error$')
    
    plotter = Plotter()
    
    plotter.show(s2, __doc__, axes=True)


def mesh_colormap(gt_mesh, mesh, corr_num):
    
    s1, _ = get_ptc(gt_ptc, color=_red)
    s2, _ = get_ptc(ptc, color=_blue)
    
    landmarks1 = s1.vertices[:corr_num,:]
    landmarks2 = s2.vertices[:corr_num,:]
    
    x, y, z = np.split(landmarks1, 3, axis=1)
    desc = x**2 + y**2 + z**2
    
    # itr = Rbf(x, y, z, desc)          # Radial Basis Function interpolator
    itr = Near(landmarks1, desc)       # Nearest-neighbour interpolator

    # interpolate descriptor on the full set of mesh vertices
    xi, yi, zi = np.split(s2.vertices, 3, axis=1)
    interpolated_desc = itr(xi, yi, zi)
    
    s2.cmap('coolwarm', interpolated_desc.squeeze()).add_scalarbar(title='error')
    # s2.point_colors(interpolated_desc, cmap='coolwarm', vmin=min(interpolated_desc), vmax=max(interpolated_desc)).addScalarBar(title='$\error$')
    
    plotter = Plotter()
    
    plotter.show(s2, __doc__, axes=True)
    
    
def plot_two_mesh(meshfile1, meshfile2, x,y,z):
    plotter = Plotter(bg=[255, 255, 255], offscreen=True)
    mesh1, scale = get_ptc(meshfile1, rotation=[x,y,z])
    mesh1.alpha(1.0).c([0.2,0.2,0.2])
    # scalars = 1- mesh1.vertices[:, 2]
    # mesh1.cmap('Oranges', scalars)
    # mesh1.c(_dark_blue)
    mesh2, scale = get_mesh(meshfile2, rotation=[x,y,z], mesh_lighting='shiny')
    scalars = 1- mesh2.vertices[:, 2]
    mesh2.cmap('Blues', scalars)
    plotter.show(mesh1, mesh2, __doc__, interactive=False, camera=cam, resetcam=True, zoom=2)
    frame_path =meshfile2[:-4] + ".png"
    
    screenshot(frame_path, scale=1)
    os.system("""convert -trim {0} "{1}" """.format(frame_path, frame_path))
    

if __name__ == "__main__":
    
    folder = '/home/wiss/sang/git/implicit_neuro_morph/jax_implicit_neuro_morph/exp/smal/2024_04_18_15_11_59/reconstructions'
    prefix = 'shape_0_to_4000_'
    
    # filenames = [os.path.join(folder, prefix + str(i) + '_mesh.ply') for i in range(0,11)]
    # create_mesh_animation(filenames, folder, -45, -45, 90)
    
    
    # filenames = [os.path.join(folder, prefix + str(i) + '_ptc.ply') for i in range(0,11)]
    
    # gt_ptc = os.path.join(folder, prefix + 'gt_pointcloud_y.ply')
    # ptc = os.path.join(folder, prefix + 'epoch_10000_time_10_ptc.ply')
    
    # create_point_cloud_animation(filenames, folder, -45, -45, 90)
    
    folder = '/home/wiss/sang/git/implicit_neuro_morph/jax_implicit_neuro_morph/exp/smal/2024_04_18_15_11_59/eval'
    # mesh_files = [os.path.join(folder, f) for f in os.listdir(folder) if f.endswith('mesh.ply')]
    # mesh_files = sort_list(mesh_files)
    
    # create_visual(mesh_files, folder, 0,0,-90, type='mesh')
    
    # create_mesh_animation(mesh_files, folder, 0,0, -90)
    
    #  folder = '/home/wiss/sang/git/implicit_neuro_morph/jax_implicit_neuro_morph/exp/fraust_r/2024_04_22_14_01_05/eval'
    # mesh_files = [os.path.join(folder, f) for f in os.listdir(folder) if f.endswith('ptc.ply')]
    # mesh_files = sort_list(mesh_files)
    # create_point_cloud_animation(mesh_files, folder, 0,0,-90)
    
    prefix = 'shape_0_to_4000_'
    mesh1 = 'gt_pointcloud_y.ply'
    mesh2 = 'step_10_mesh.ply'
    
    meshfile1 = os.path.join(folder, prefix + mesh1)
    meshfile2 = os.path.join(folder, prefix + mesh2)
    
    plot_two_mesh(meshfile1, meshfile2, 0, 0, -30)