Spaces:
Running
on
Zero
Running
on
Zero
| import numpy as np | |
| import cv2 | |
| from vedo import Mesh, Plane, Plotter, screenshot | |
| import imageio.v2 as imageio | |
| # import imageio as imageio | |
| import re | |
| import os | |
| from .visual import get_ptc # probaly vpython | |
| from PIL import Image | |
| from tqdm import tqdm | |
| screenshot_scale = 3 | |
| PLAY_OBS_BEFORE_PRED = True | |
| obs_dir = "" | |
| cam = dict( | |
| position=(-3.85922, -4.78140, 2.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), | |
| ) | |
| _green= [0.4, 0.7, 0.5] | |
| _blue =[0.7, 0.9, 1.0] | |
| _dark_blue =[0.03, 0.4, 0.7] | |
| _orange = [0.9, 0.7, 0.2] | |
| def get_mesh(mesh_name, color='darksalmon', 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=color, 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 save_png_rm_bg(filename, im, bg_color_min=[0, 255, 100], bg_color_max=None, rgb=False): | |
| # https://stackoverflow.com/questions/72062001/remove-everything-of-a-specific-color-with-a-color-variation-tolerance-from-an | |
| # for transperancy: | |
| # https://stackoverflow.com/questions/55582117/efficiently-converting-color-to-transparency-in-python | |
| if bg_color_max is None: | |
| bg_color_max = bg_color_min | |
| # bg_color_min = [bg_color_min[2], bg_color_min[1], bg_color_min[0]] | |
| # bg_color_max = [bg_color_max[2], bg_color_max[1], bg_color_max[0]] | |
| # Load image | |
| # im = cv2.imread('test.png') | |
| # Define lower and upper limits of our blue | |
| min = np.array(bg_color_min, np.uint8) | |
| max = np.array(bg_color_max, np.uint8) | |
| # Go to HSV colourspace and get mask of blue pixels | |
| # HSV = cv2.cvtColor(im, cv2.COLOR_BGR2HSV) | |
| mask = cv2.inRange(im, min, max) | |
| # Try dilating (enlarging) mask with 3x3 structuring element | |
| # SE = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3)) | |
| # mask = cv2.dilate(mask, kernel, iterations=1) | |
| # create alpha channel | |
| alpha = (im[:, :, 0] * 0 + 255) | |
| # Make all pixels in mask white | |
| alpha[mask > 0] = 0 | |
| # add alpha channel | |
| if rgb: | |
| im = cv2.merge((im[:, :, 0], im[:, :, 1], im[:, :, 2], alpha)) | |
| else: | |
| im = cv2.merge((im[:, :, 2], im[:, :, 1], im[:, :, 0], alpha)) | |
| # crop to visible part | |
| mask = im[:, :, 3] != 0. | |
| coords = np.argwhere(mask) | |
| y0, x0 = coords.min(axis=0) | |
| y1, x1 = coords.max(axis=0) + 1 | |
| # im = im[y0:y1, x0:x1, :] | |
| # write to file | |
| cv2.imwrite(filename, im) | |
| def save_plot_as_transparent_png(plt, out_filename): | |
| # if not out_filename is None: | |
| bg_col = [255, 255, 255] | |
| plt.background(c1=bg_col) | |
| save_png_rm_bg(out_filename, plt.screenshot(asarray=True), bg_color_min=bg_col) | |
| 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.vertices.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 plot_meshes(mesh_files, cam, mesh_color=[1, 1, 1], rotation=[0, 0, 0]): | |
| mesh_opacity = 0.8 | |
| plot_shadows = True | |
| pass | |
| def create_visual(mesh_files, color='coral', dest_name="", x=0,y=0,z=0, type='mesh', mesh_lighting='plastic', obs_dir=""): | |
| meshes, scales =[], [] | |
| bg_col = [255, 255, 255]# plotter = Plotter(offscreen=True) | |
| plotter = Plotter(bg=[255, 255, 255], offscreen=True) | |
| frame_files = [] | |
| # lights = get_lights() | |
| for file in tqdm(mesh_files): | |
| mesh, scale = get_mesh(file, color=color, rotation=[x,y,z], mesh_lighting=mesh_lighting, scale=0.6) | |
| # scalars = 1- mesh.vertices[:, 2] | |
| # mesh.cmap('Oranges', scalars) | |
| plotter.clear() | |
| shadow = get_mesh_shadow(mesh) | |
| # plotter.add([mesh, shadow]) | |
| # plotter.show(mesh, __doc__, interactive=False, camera=cam, resetcam=True, zoom=2) | |
| 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) | |
| save_png_rm_bg(file[:-4]+"new.png", plotter.screenshot(asarray=True), bg_color_min=bg_col) | |
| # Compile saved frames into a GIF | |
| filename = os.path.join(dest_name + '.gif') | |
| create_transparent_gif(frame_files, dest_name, obs_dir) | |
| # with imageio.get_writer(filename, mode='I', duration=2) as writer: #duration=0.2) as writer: | |
| # for filename in frame_files: | |
| # image = imageio.imread(filename) | |
| # writer.append_data(image) | |
| def create_transparent_gif(frame_files, dest_name, obs_dir): | |
| images = [] | |
| if "obs" not in dest_name and PLAY_OBS_BEFORE_PRED: | |
| # visit all file under obs_dir that end with new.png | |
| obs_files = [os.path.join(obs_dir, f) for f in os.listdir(obs_dir) if f.endswith('new.png')] | |
| # sort the files | |
| obs_files = sort_list(obs_files) | |
| for i, frame in enumerate(obs_files): | |
| im = Image.open(frame) | |
| images.append(im) | |
| # Optimize: Load all frames at once and process them in batch | |
| frame_files = [f.replace('.png', 'new.png') for f in frame_files] | |
| frame_files = sort_list(frame_files) | |
| # Use a more efficient way to load and process images | |
| for frame in frame_files: | |
| im = Image.open(frame) | |
| images.append(im) | |
| # Save with optimized settings | |
| images[0].save( | |
| os.path.join(dest_name + '_tranp.gif'), | |
| format="GIF", | |
| save_all=True, | |
| loop=1, | |
| append_images=images[1:], # Skip first image as it's already saved | |
| duration=0.033, # Reduced from 0.016 to 0.033 (30fps) | |
| disposal=2, | |
| optimize=True, # Enable optimization | |
| quality=85, # Reduced quality for better performance | |
| colors=64 # Reduce color palette | |
| ) | |
| # Clean up | |
| for im in images: | |
| im.close() | |
| 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 visual_gt(mesh_files, color, x,y,z, type='mesh'): | |
| bg_col = [255, 255, 255]# 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, color=color, rotation=[x,y,z], mesh_lighting='plastic') | |
| # scalars = 1- mesh.vertices[:, 2] | |
| # mesh.cmap('Blues', scalars) | |
| else: | |
| mesh, scale = get_ptc(file, rotation=[x,y,z]) | |
| mesh.c([0.5,0.5,0.5]) | |
| 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) | |
| save_png_rm_bg(file[:-4]+"new.png", plotter.screenshot(asarray=True), bg_color_min=bg_col) | |
| def plot_two_mesh(meshfile1, meshfile2, color, 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.5,0.5,0.5]) | |
| # scalars = 1- mesh1.vertices[:, 2] | |
| # mesh1.cmap('Oranges', scalars) | |
| # mesh1.c(_dark_blue) | |
| mesh2, scale = get_mesh(meshfile2, color=color, rotation=[x,y,z], mesh_lighting='plastic') | |
| shadow = get_mesh_shadow(mesh2) | |
| plotter.show(mesh1, mesh2, shadow, __doc__, interactive=False, camera=cam, resetcam=True, zoom=2) | |
| frame_path =meshfile2[:-4] + ".png" | |
| screenshot(frame_path, scale=1) | |
| save_png_rm_bg(meshfile2[:-4]+"both.png", plotter.screenshot(asarray=True), bg_color_min=[255, 255, 255]) | |
| import argparse | |
| def main(parent_folder, displayed_preds=3): | |
| file_type = '.obj' #'mesh.ply' | |
| x,y,z = 0, 0, -90 # Changed from 0,0,0 to rotate around x-axis | |
| dest_dir = os.path.join(os.path.dirname(parent_folder), 'shadow_gif') | |
| obs_dir = os.path.join(os.path.dirname(parent_folder), 'obs_obj') | |
| os.makedirs(dest_dir, exist_ok=True) | |
| # Check if we already have enough GIFs | |
| existing_gifs = [f for f in os.listdir(dest_dir) if f.endswith('_tranp.gif') and not f.startswith('obs')] | |
| if len(existing_gifs) >= displayed_preds: | |
| print(f"Found {len(existing_gifs)} existing GIFs, which is enough for {displayed_preds} predictions") | |
| return | |
| print(f"Found {len(existing_gifs)} existing GIFs, need {displayed_preds}, proceeding with generation") | |
| # If no existing GIFs, proceed with generation | |
| subfolders = sorted([f for f in os.listdir(parent_folder) if os.path.isdir(os.path.join(parent_folder, f)) and 'obj' in f]) | |
| for f in subfolders: | |
| folder = os.path.join(parent_folder, f) | |
| dest_name = os.path.join(dest_dir, f) | |
| mesh_files = [os.path.join(folder, f) for f in os.listdir(folder) if f.endswith(file_type)] | |
| mesh_files = sort_list(mesh_files) | |
| os.makedirs(dest_dir, exist_ok=True) | |
| #color "lightblue" "navajowhite" | |
| if 'obs' in f: | |
| create_visual(mesh_files, _green, dest_name, x,y,z, type='mesh', mesh_lighting='plastic', obs_dir=obs_dir) | |
| else: | |
| create_visual(mesh_files, _dark_blue, dest_name, x,y,z, type='mesh', mesh_lighting='plastic', obs_dir=obs_dir) | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser() | |
| #conda activate hmp_visualize | |
| parser.add_argument('-f', '--mesh_parent_folder', type=str, default="final_output/hmp/visuals/amass/SkeletonDiffusion/test/952_WalkTogether/obj") | |
| args = parser.parse_args() | |
| main(args.mesh_parent_folder) | |
| #################################################################################################################################### | |
| # Old code snippet, kept just for reference. Jialin here there may be somehting you will need, or maybe not :) If you do not need you can delete | |
| #################################################################################################################################### | |
| # other functions not needed, from older code | |
| # visual_gt(gt_mesh_files, 0,0,-90, type='mesh') # this line was commented out by Lu | |
| # visual_gt(gt_obj_files, 'salmon', x,y,z)#--> .png and new.png (transparent) per frame | |
| # gt_ptc_files = [os.path.join(folder, prefix +'gt_pointcloud_x.ply'),os.path.join(folder, prefix +'gt_pointcloud_y.ply')] | |
| # gt_obj_files = sort_list([os.path.join(target_folder, f) for f in os.listdir(target_folder) if f.endswith(file_type)]) | |
| # plot_two_mesh(gt_ptc_files[1], mesh_files[-1], "lightblue", x,y,z) | |
| # older code | |
| # final_folder | |
| # gif_folder | |
| # filename = os.path.join(dest_name + '.gif') | |
| # frame_files = [os.path.join(gif_folder, f"pred{i}.gif") for i in range(10)] | |
| # obs_image = imageio.imread(os.path.join(gif_folder, f"obs_salmon.png")) | |
| # for i,filename in enumerate(frame_files): | |
| # image = imageio.imread(filename) | |
| # frames = np.vstack([obs_image, image]) | |
| # imageio.imwrite(os.path.join(final_folder, f"pred{i}.gif"), frames, fps=50) | |
| # image = imageio.imread(os.path.join(gif_folder, f"gt_salmon.gif")) | |
| # frames = np.vstack([obs_image, image]) | |
| # imageio.imwrite(os.path.join(final_folder, f"gt.gif"), frames, fps=50) | |
| # gifs= [imageio.imread(os.path.join(final_folder, f"pred{i}.gif")) for i in range(10)] | |
| # #Create writer object | |
| # new_gif = imageio.get_writer(os.path.join(final_folder, f"output.gif")) | |
| # for frame_number in range(len(gifs[0])): | |
| # img1 = gif1.get_next_data() | |
| # img2 = gif2.get_next_data() | |
| # #here is the magic | |
| # upper_row = np.hstack((gifs[:5, frame_number])) | |
| # lower_row = np.hstack((gifs[5:, frame_number])) | |
| # new_image = np.vstack((upper_row, lower_row)) | |
| # new_gif.append_data(new_image) | |
| print("done.") | |