SkeletonDiffusion_demo/plot_several_meshes.py

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