Spaces:
Running
on
Zero
Running
on
Zero
| 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) | |