utils/extract_obj_meshes.py

import numpy as np
import trimesh
import os 


def extract_obj_meshes(sv_dict_fn):
    # sv_fn = "/data2/datasets/sim/arctic_processed_data/processed_sv_dicts/s01/box_grab_01_extracted_dict.npy"
    # if not os.path.exists(sv_fn):
    #     sv_fn = "/data/xueyi/arctic/processed_sv_dicts/box_grab_01_extracted_dict.npy"
    active_passive_sv_dict = np.load(sv_dict_fn, allow_pickle=True).item() # 

    obj_verts = active_passive_sv_dict['obj_verts'] # object orientation #
    obj_faces = active_passive_sv_dict['obj_faces']
    
    obj_mesh = trimesh.Trimesh(obj_verts[0], obj_faces)
    obj_mesh_sv_fn = os.path.join(
        "/home/xueyi/diffsim/NeuS/utils", "init_box.obj"
    )
    obj_mesh.export(obj_mesh_sv_fn)
    
    obj_verts_reversed = obj_verts[:, :, [1, 0, 2]]
    obj_verts_reversed[:, :, 1] = -obj_verts_reversed[:, :, 1]
    obj_mesh_reversed = trimesh.Trimesh(obj_verts_reversed[0], obj_faces)
    obj_mesh_sv_fn_reversed = os.path.join(
        "/home/xueyi/diffsim/NeuS/utils", "init_box_reversed.obj"
    )
    obj_mesh_reversed.export(obj_mesh_sv_fn_reversed)


def extract_obj_meshes_boundingbox(sv_dict_fn):
    # sv_fn = "/data2/datasets/sim/arctic_processed_data/processed_sv_dicts/s01/box_grab_01_extracted_dict.npy"
    # if not os.path.exists(sv_fn):
    #     sv_fn = "/data/xueyi/arctic/processed_sv_dicts/box_grab_01_extracted_dict.npy"
    active_passive_sv_dict = np.load(sv_dict_fn, allow_pickle=True).item() # 

    obj_verts = active_passive_sv_dict['obj_verts'] # object orientation #
    obj_faces = active_passive_sv_dict['obj_faces']
    
    init_obj_verts = obj_verts[0]
    minn_box = np.min(init_obj_verts, axis=0, keepdims=True)
    maxx_box = np.max(init_obj_verts, axis=0, keepdims=True)
    # get the minn box and maxx box #
    
    
    box_triangle_mesh_faces = np.array([
        [1, 2, 3],  # Left face (triangle 1)
        [2, 3, 4],  # Left face (triangle 2)
        [5, 6, 7],  # Right face (triangle 1)
        [6, 7, 8],  # Right face (triangle 2)
        [1, 3, 5],  # Bottom face (triangle 1)
        [3, 5, 7],  # Bottom face (triangle 2)
        [2, 4, 6],  # Top face (triangle 1)
        [4, 6, 8],  # Top face (triangle 2)
        [1, 2, 5],  # Front face (triangle 1)
        [2, 5, 6],  # Front face (triangle 2)
        [3, 4, 7],  # Back face (triangle 1)
        [4, 7, 8]   # Back face (triangle 2)
    ], dtype=np.int32) - 1
    
    box_vertices = np.array([
        [-1, -1, -1],
        [-1, -1, 1],
        [-1, 1, -1],
        [-1, 1, 1],
        [1, -1, -1],
        [1, -1, 1],
        [1, 1, -1],
        [1, 1, 1]
    ], dtype=np.float32)
    
    box_vertices = (box_vertices -  (-1)) / 2
    
    
    box_vertices = box_vertices * (maxx_box - minn_box) + minn_box
    
    box_mesh=  trimesh.Trimesh(box_vertices, box_triangle_mesh_faces)
    # 
    
    # obj_mesh = trimesh.Trimesh(obj_verts[0], obj_faces)
    obj_mesh_sv_fn = os.path.join(
        "/home/xueyi/diffsim/NeuS/utils", "init_bounding_box.obj"
    )
    box_mesh.export(obj_mesh_sv_fn)
    
    # obj_verts_reversed = obj_verts[:, :, [1, 0, 2]]
    # obj_verts_reversed[:, :, 1] = -obj_verts_reversed[:, :, 1]
    # obj_mesh_reversed = trimesh.Trimesh(obj_verts_reversed[0], obj_faces)
    # obj_mesh_sv_fn_reversed = os.path.join(
    #     "/home/xueyi/diffsim/NeuS/utils", "init_box_reversed.obj"
    # )
    # obj_mesh_reversed.export(obj_mesh_sv_fn_reversed)
    
def extract_obj_meshes_taco(pkl_fn):
    import pickle as pkl
    
    sv_dict = pkl.load(open(pkl_fn, "rb"))
    
    print(f"sv_dict: {sv_dict.keys()}")
    
    # maxx_ws = min(maxx_ws, len(sv_dict['obj_verts']) - start_idx)
    
    obj_pcs = sv_dict['obj_verts'] # [start_idx: start_idx + maxx_ws]
    # obj_pcs = torch.from_numpy(obj_pcs).float().cuda()
    
    # self.obj_pcs = obj_pcs
    # # obj_vertex_normals = sv_dict['obj_vertex_normals']
    # # obj_vertex_normals = torch.from_numpy(obj_vertex_normals).float().cuda()
    # self.obj_normals = torch.zeros_like(obj_pcs[0]) ### get the obj naormal vectors ##
    
    object_pose = sv_dict['obj_pose'] # [start_idx: start_idx + maxx_ws]
    # object_pose = torch.from_numpy(object_pose).float().cuda() ### nn_frames x 4 x 4 ###
    object_global_orient_mtx = object_pose[:, :3, :3 ] ## nn_frames x 3 x 3 ##
    object_transl = object_pose[:, :3, 3] ## nn_frmaes x 3 ##
    
    obj_faces = sv_dict['obj_faces']
    # obj_faces = torch.from_numpy(obj_faces).long().cuda()
    # self.obj_faces = obj_faces # [0] ### obj faces ##
    
    # obj_verts = sv_dict['obj_verts']
    # minn_verts = np.min(obj_verts, axis=0)
    # maxx_verts = np.max(obj_verts, axis=0)
    # extent = maxx_verts - minn_verts
    # center_ori = (maxx_verts + minn_verts) / 2
    # scale_ori = np.sqrt(np.sum(extent ** 2))
    # obj_verts = torch.from_numpy(obj_verts).float().cuda()
    
    init_obj_verts = obj_pcs[0]
    init_obj_ornt_mtx = object_global_orient_mtx[0]
    init_obj_transl = object_transl[0]
    
    canon_obj_verts = np.matmul(
        init_obj_ornt_mtx.T, (init_obj_verts - init_obj_transl[None]).T
    ).T
    # self.obj_verts = canon_obj_verts.clone()
    # obj_verts = canon_obj_verts.clone()
    canon_obj_mesh = trimesh.Trimesh(vertices=canon_obj_verts, faces=obj_faces)
    canon_obj_mesh_export_dir = "/".join(pkl_fn.split("/")[:-1])
    pkl_name = pkl_fn.split("/")[-1].split(".")[0]
    canon_obj_mesh_sv_fn = f"{pkl_name}.obj"
    canon_obj_mesh.export(os.path.join(canon_obj_mesh_export_dir, canon_obj_mesh_sv_fn))
    print(f"canon_obj_mesh_sv_fn: {canon_obj_mesh_sv_fn}")
        

if __name__=='__main__':
    
    pkl_fn = "/data3/datasets/xueyi/taco/processed_data/20230917/right_20230917_004.pkl"
    pkl_fn = "/data3/datasets/xueyi/taco/processed_data/20230917/right_20230917_030.pkl"
    pkl_fn = "/data3/datasets/xueyi/taco/processed_data/20230917/right_20230917_037.pkl"
    pkl_fn = "/data/xueyi/taco/processed_data/20230917/right_20230917_037.pkl"
    pkl_root_folder = "/data3/datasets/xueyi/taco/processed_data/20230917"
    pkl_root_folder = "/data3/datasets/xueyi/taco/processed_data/20231010"
    pkl_root_folder = "/data3/datasets/xueyi/taco/processed_data/20230919"
    pkl_root_folder = "/data3/datasets/xueyi/taco/processed_data/20231104"
    tot_pkl_fns = os.listdir(pkl_root_folder)
    tot_pkl_fns  = [fn for fn in tot_pkl_fns if fn.endswith(".pkl")]
    for i_fn, cur_pkl_fn in enumerate(tot_pkl_fns):
        cur_full_pkl_fn = os.path.join(pkl_root_folder, cur_pkl_fn)
        extract_obj_meshes_taco(cur_full_pkl_fn)
    exit(0)
    
    extract_obj_meshes_taco(pkl_fn)
    exit(0)
    
    sv_dict_fn = "/data2/datasets/sim/arctic_processed_data/processed_sv_dicts/s01/box_grab_01_extracted_dict.npy"
    # extract_obj_meshes(sv_dict_fn=sv_dict_fn)
    
    extract_obj_meshes_boundingbox(sv_dict_fn)