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.")