utils/data_format_transformation_processing.py

import numpy as np
import torch
import trimesh
import os
import mesh2sdf
import time
from scipy.spatial.transform import Rotation as R

QUASI_DYN_ROOT = "/home/xueyi/diffsim/NeuS"
if not os.path.exists(QUASI_DYN_ROOT):
    QUASI_DYN_ROOT = "/root/diffsim/quasi-dyn"

ARCTIC_CANON_OBJ_SV_FOLDER = os.path.join(QUASI_DYN_ROOT, "raw_data/arctic_processed_canon_obj")



def export_canon_obj_file(kinematic_mano_gt_sv_fn, obj_name):
    
    subject_idx = kinematic_mano_gt_sv_fn.split("/")[-2] # 
    print(f"subject_idx: {subject_idx}, obj name: {obj_name}")
    sv_dict = np.load(kinematic_mano_gt_sv_fn, allow_pickle=True).item()
    
    
    object_global_orient = sv_dict["obj_rot"]  
    object_transl = sv_dict["obj_trans"] * 0.001
    obj_pcs = sv_dict["verts.object"]
    
    # obj_pcs = sv_dict['object_pc']
    obj_pcs = torch.from_numpy(obj_pcs).float().cuda()
    
    
    # self.obj_verts = obj_verts
    init_obj_verts = obj_pcs[0]
    init_obj_rot_vec = object_global_orient[0]
    init_obj_transl = object_transl[0]
    
    init_obj_transl = torch.from_numpy(init_obj_transl).float().cuda()
    init_rot_struct = R.from_rotvec(init_obj_rot_vec)
    
    init_glb_rot_mtx = init_rot_struct.as_matrix()
    init_glb_rot_mtx = torch.from_numpy(init_glb_rot_mtx).float().cuda()
    # ## reverse the global rotation matrix ##
    init_glb_rot_mtx_reversed = init_glb_rot_mtx.contiguous().transpose(1, 0).contiguous()


    
    ''' canonical object verts '''
    canon_obj_verts = torch.matmul(
        init_glb_rot_mtx_reversed.transpose(1, 0).contiguous(), (init_obj_verts - init_obj_transl.unsqueeze(0)).contiguous().transpose(1, 0).contiguous()
    ).contiguous().transpose(1, 0).contiguous()

    # ## get canon obj verts ##
    
    # /home/xueyi/diffsim/NeuS/raw_data/arctic_processed_canon_obj
    
    canon_obj_sv_folder = ARCTIC_CANON_OBJ_SV_FOLDER # "/root/diffsim/control-vae-2/assets/arctic"
    canon_obj_mesh = trimesh.Trimesh(vertices=canon_obj_verts.detach().cpu().numpy(), faces=sv_dict['f'][0])
    
    canon_obj_mesh_sv_fn = f"{subject_idx}_{obj_name}.obj"
    canon_obj_mesh_sv_fn = os.path.join(canon_obj_sv_folder, canon_obj_mesh_sv_fn)
    canon_obj_mesh.export(canon_obj_mesh_sv_fn)

    print(f"Canonical obj mesh saved to {canon_obj_mesh_sv_fn}")
    return canon_obj_mesh_sv_fn


def compute_sdf(obj_file_name):
    filename = obj_file_name

    # init_mesh_scale = 1.0
    init_mesh_scale = 0.8

    mesh_scale = 0.8
    size = 128
    level = 2 / size

    mesh = trimesh.load(filename, force='mesh')

    # normalize mesh
    vertices = mesh.vertices
    vertices = vertices * init_mesh_scale
    bbmin = vertices.min(0) # 
    bbmax = vertices.max(0) # 
    center = (bbmin + bbmax) * 0.5
    scale = 2.0 * mesh_scale / (bbmax - bbmin).max() # bounding box's max # # bbmax - bbmin # 
    vertices = (vertices - center) * scale # (vertices - center) * scale #

    scaled_bbmin = vertices.min(0)
    scaled_bbmax = vertices.max(0)
    print(f"scaled_bbmin: {scaled_bbmin}, scaled_bbmax: {scaled_bbmax}")


    t0 = time.time()
    sdf, mesh = mesh2sdf.compute( ## sdf and mesh ##
        vertices, mesh.faces, size, fix=True, level=level, return_mesh=True)
    t1 = time.time()

    print(f"sdf: {sdf.shape}, mesh: {mesh.vertices.shape}")

    mesh.vertices = mesh.vertices / scale + center
    mesh.export(filename[:-4] + '.fixed.obj') ## .fixed.obj ##
    np.save(filename[:-4] + '.npy', sdf) ## .npy ##
    print('It takes %.4f seconds to process %s' % (t1-t0, filename))



def convert_tot_states_to_data_ref_format(tot_states_fn, sv_gt_ref_data_fn):
    tot_states_data = np.load(tot_states_fn, allow_pickle=True).item()
    tot_states = tot_states_data['tot_states']
    tot_mano_rot = []
    tot_mano_glb_trans = []
    tot_mano_states = []
    
    tot_obj_rot = []
    tot_obj_trans = []
    
    for i_fr in range(len(tot_states)):
        cur_state = tot_states[i_fr]
        cur_mano_state = cur_state[:-7]
        cur_obj_state  = cur_state[-7:]
        
        tot_mano_glb_trans.append(cur_mano_state[:3])
        cur_mano_rot_vec = cur_mano_state[3:6]
        cur_mano_rot_euler_zyx = [cur_mano_rot_vec[2], cur_mano_rot_vec[1], cur_mano_rot_vec[0]]
        cur_mano_rot_euler_zyx = np.array(cur_mano_rot_euler_zyx, dtype=np.float32)
        cur_mano_rot_struct = R.from_euler('zyx', cur_mano_rot_euler_zyx, degrees=False)
        cur_mano_rot_quat_xyzw = cur_mano_rot_struct.as_quat()
        cur_mano_rot_quat_wxyz = cur_mano_rot_quat_xyzw[[3, 0, 1, 2]]
        tot_mano_rot.append(cur_mano_rot_quat_wxyz.astype(np.float32))
        
        tot_mano_states.append(cur_mano_state[4:])
        
        tot_obj_rot.append(cur_obj_state[-4:][[3, 0, 1, 2]])
        tot_obj_trans.append(cur_obj_state[:3])
    tot_obj_rot = np.stack(tot_obj_rot, axis=0)
    tot_obj_trans = np.stack(tot_obj_trans, axis=0)
    tot_mano_rot =  np.stack(tot_mano_rot, axis=0)
    tot_mano_glb_trans = np.stack(tot_mano_glb_trans, axis=0)
    tot_mano_states = np.stack(tot_mano_states, axis=0)
    
    gt_ref_data = {
        'obj_rot': tot_obj_rot, 
        'obj_trans': tot_obj_trans,
        'mano_states': tot_mano_states,
        'mano_glb_trans': tot_mano_glb_trans,
        'mano_glb_rot': tot_mano_rot
    }
    
    np.save(sv_gt_ref_data_fn, gt_ref_data)
    print(f"gt ref data svaed to {sv_gt_ref_data_fn}")


# spoon --> how to sue the spoon #
# 

if __name__=='__main__':
    
    tot_states_fn = "/home/xueyi/diffsim/raisim/dgrasp/raisimGymTorch/bullet_env/obj_mouse_idx_102_tracking_2/2024-02-27-05-44-09/sv_info_800.npy"
    sv_gt_ref_data_fn = "/home/xueyi/diffsim/Control-VAE/ReferenceData/shadow_grab_mouse_102_dgrasptracking.npy"
    
    tot_states_fn = "/home/xueyi/diffsim/raisim/dgrasp/raisimGymTorch/bullet_env/obj_bunny_idx_85_tracking_2/2024-02-26-08-33-58/sv_info_600.npy"
    sv_gt_ref_data_fn = "/home/xueyi/diffsim/Control-VAE/ReferenceData/shadow_grab_bunny_85_dgrasptracking.npy"
    
    tot_states_fn = "/home/xueyi/diffsim/raisim/dgrasp/raisimGymTorch/bullet_env/obj_spoon2_idx_20231105_067_tracking_2/2024-02-29-09-49-32/sv_info_300.npy"
    sv_gt_ref_data_fn = "/home/xueyi/diffsim/Control-VAE/ReferenceData/shadow_taco_spoon2_idx_20231105_067_dgrasptracking.npy"
    
    tot_states_fn= "/home/xueyi/diffsim/raisim/dgrasp/raisimGymTorch/bullet_env/obj_20231027_086_idx_20231027_086_tracking_2/2024-03-09-13-01-24/sv_info_best.npy"
    sv_gt_ref_data_fn = "/home/xueyi/diffsim/Control-VAE/ReferenceData/shadow_taco_20231027_086_idx_20231027_086_dgrasptracking.npy"
    
    # /home/xueyi/diffsim/raisim/dgrasp/raisimGymTorch/bullet_env/obj_20231024_044_idx_20231024_044_tracking_2/2024-03-09-13-20-07/sv_info_best.npy
    tot_states_fn= "/home/xueyi/diffsim/raisim/dgrasp/raisimGymTorch/bullet_env/obj_20231024_044_idx_20231024_044_tracking_2/2024-03-09-13-20-07/sv_info_best.npy"
    sv_gt_ref_data_fn = "/home/xueyi/diffsim/Control-VAE/ReferenceData/shadow_taco_20231024_044_idx_20231024_044_dgrasptracking.npy"
    
    tot_states_fn = "/home/xueyi/diffsim/raisim/dgrasp/raisimGymTorch/bullet_env/obj_20231027_130_idx_20231027_130_tracking_2/2024-03-09-13-40-52/sv_info_best.npy"
    sv_gt_ref_data_fn = "/home/xueyi/diffsim/Control-VAE/ReferenceData/shadow_taco_20231027_130_idx_20231027_130_dgrasptracking.npy" ## gt 130 ##
    
    tot_states_fn = "/home/xueyi/diffsim/raisim/dgrasp/raisimGymTorch/bullet_env/obj_20231020_199_idx_20231020_199_tracking_2/2024-03-09-14-01-13/sv_info_best.npy"
    sv_gt_ref_data_fn = "/home/xueyi/diffsim/Control-VAE/ReferenceData/shadow_taco_20231020_199_idx_20231020_199_dgrasptracking.npy" ## gt 130 ##
    
    tot_states_fn = "/home/xueyi/diffsim/raisim/dgrasp/raisimGymTorch/bullet_env/obj_20231026_016_idx_20231026_016_tracking_2/2024-03-11-10-06-43/sv_info_best.npy"
    
    sv_gt_ref_data_fn = "/home/xueyi/diffsim/Control-VAE/ReferenceData/shadow_taco_20231026_016_idx_20231026_016_dgrasptracking.npy" ## gt 130 ##
    
    tot_states_fn = "/home/xueyi/diffsim/raisim/dgrasp/raisimGymTorch/bullet_env/obj_20231027_114_idx_20231027_114_tracking_2/2024-03-11-10-07-36/sv_info_best.npy"
    sv_gt_ref_data_fn = "/home/xueyi/diffsim/Control-VAE/ReferenceData/shadow_taco_20231027_114_idx_20231027_114_dgrasptracking.npy" ## gt 130 ##
    
    convert_tot_states_to_data_ref_format(tot_states_fn, sv_gt_ref_data_fn)
    exit(0)
    
    
    kinematic_mano_gt_sv_fn = '/data/xueyi/sim/arctic_processed_data/processed_seqs/s01/box_grab_01.npy'
    
    kinematic_mano_gt_sv_fn = '/data/xueyi/sim/arctic_processed_data/processed_seqs/s01/ketchup_grab_01.npy'
    
    # kinematic_mano_gt_sv_fn = '/data/xueyi/sim/arctic_processed_data/processed_seqs/s01/mixer_grab_01.npy'
    
    # kinematic_mano_gt_sv_fn = '/data/xueyi/sim/arctic_processed_data/processed_seqs/s01/laptop_grab_01.npy'
    
    # kinematic_mano_gt_sv_fn = '/data/xueyi/sim/arctic_processed_data/processed_seqs/s01/box_grab_01.npy'
    
    kinematic_mano_gt_sv_fn = '/data/xueyi/sim/arctic_processed_data/processed_seqs/s02/waffleiron_grab_01.npy'
    kinematic_mano_gt_sv_fn = '/data/xueyi/sim/arctic_processed_data/processed_seqs/s02/ketchup_grab_01.npy'
    kinematic_mano_gt_sv_fn = '/data/xueyi/sim/arctic_processed_data/processed_seqs/s02/phone_grab_01.npy'
    kinematic_mano_gt_sv_fn = '/data/xueyi/sim/arctic_processed_data/processed_seqs/s02/box_grab_01.npy'
    
    kinematic_mano_gt_sv_folder = "/data/xueyi/sim/arctic_processed_data/processed_seqs/s02"
    kinematic_mano_gt_sv_folder = "/data/xueyi/sim/arctic_processed_data/processed_seqs/s04"
    kinematic_mano_gt_sv_folder = "/data/xueyi/sim/arctic_processed_data/processed_seqs/s05"
    kinematic_mano_gt_sv_folder = "/data/xueyi/sim/arctic_processed_data/processed_seqs/s06"
    kinematic_mano_gt_sv_folder = "/data/xueyi/sim/arctic_processed_data/processed_seqs/s07"
    kinematic_mano_gt_sv_folder = "/data/xueyi/sim/arctic_processed_data/processed_seqs/s10"
    
    kinematic_mano_gt_sv_fn_all = os.listdir(kinematic_mano_gt_sv_folder)
    kinematic_mano_gt_sv_fn_all = [fn for fn in kinematic_mano_gt_sv_fn_all if fn.endswith(".npy") and "grab" in fn]
    kinematic_mano_gt_sv_fn_all = [os.path.join(kinematic_mano_gt_sv_folder, fn) for fn in kinematic_mano_gt_sv_fn_all]
    for kinematic_mano_gt_sv_fn in kinematic_mano_gt_sv_fn_all:
        obj_name = kinematic_mano_gt_sv_fn.split("/")[-1].split(".")[0].split("_")[0]
        print(f"obj_name: {obj_name}")
        
        ##### process and export canonical object file #####
        canon_obj_mesh_sv_fn = export_canon_obj_file(kinematic_mano_gt_sv_fn, obj_name)
        
        
        ##### compute sdf of the canonical object mesh #####
        compute_sdf(canon_obj_mesh_sv_fn)
        
    
    # obj_name = kinematic_mano_gt_sv_fn.split("/")[-1].split(".")[0].split("_")[0]
    # print(f"obj_name: {obj_name}")
    
    # ##### process and export canonical object file #####
    # canon_obj_mesh_sv_fn = export_canon_obj_file(kinematic_mano_gt_sv_fn, obj_name)
    
    
    # ##### compute sdf of the canonical object mesh #####
    # compute_sdf(canon_obj_mesh_sv_fn)