Spaces:
Sleeping
Sleeping
| import numpy as np | |
| import trimesh | |
| import os | |
| # try: | |
| # import mesh2sdf | |
| # except: | |
| # pass | |
| import mesh2sdf | |
| import time | |
| from scipy.spatial.transform import Rotation as R | |
| 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}") | |
| 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 get_grab_data_dict(data_fn): | |
| grab_data_dict = np.load(data_fn, allow_pickle=True).item() | |
| print(f"grab_data_keys: { grab_data_dict.keys()}") | |
| if __name__=='__main__': | |
| ## data_fn ## ### that we can use in the GRAB | |
| data_fn = "/data1/xueyi/GRAB_extracted_test/train/363_sv_dict.npy" | |
| get_grab_data_dict(data_fn) | |
| exit(0) | |
| # compute_sdf('/data/xueyi/taco/processed_data/20230917/right_20230917_032.obj') | |
| # exit(0) | |
| 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" | |
| pkl_root_folder = "/data3/datasets/xueyi/taco/processed_data/20231105" | |
| pkl_root_folder = "/data/xueyi/taco/processed_data/20230917" ## pkl folder | |
| pkl_root_folder = "/data3/datasets/xueyi/taco/processed_data/20231102" | |
| pkl_root_folder = "/data3/datasets/xueyi/taco/processed_data/20230923" | |
| pkl_root_folder = "/data3/datasets/xueyi/taco/processed_data/20230926" | |
| pkl_root_folder = "/data3/datasets/xueyi/taco/processed_data/20230927" | |
| pkl_root_folder = "/data/xueyi/taco/processed_data/20230930" | |
| pkl_root_folder = "/data2/datasets/xueyi/taco/processed_data/20231031" | |
| # pkl_root_folder = "/data2/datasets/xueyi/taco/processed_data/20230929" | |
| # pkl_root_folder = "/data2/datasets/xueyi/taco/processed_data/20230930" | |
| tot_pkl_fns = os.listdir(pkl_root_folder) | |
| tot_pkl_fns = [fn for fn in tot_pkl_fns if fn.endswith(".pkl")] | |
| # tot_pkl_fns = ['right_20230930_001.pkl'] | |
| # tot_pkl_fns = ['/data/xueyi/taco/processed_data/20230917/right_20230917_032.obj'] | |
| 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) | |
| # compute_sdf(cur_full_fn) | |
| # exit(0) | |
| # pkl_root_folder = "/data3/datasets/xueyi/taco/processed_data/20231104" | |
| # pkl_root_folder = "/data3/datasets/xueyi/taco/processed_data/20231105" | |
| # pkl_root_folder = "/data/xueyi/taco/processed_data/20230917" | |
| # pkl_root_folder = "/data3/datasets/xueyi/taco/processed_data/20231102" | |
| tot_fns = os.listdir(pkl_root_folder) | |
| tot_fns = [fn for fn in tot_fns if fn.endswith(".obj")] | |
| for cur_fn in tot_fns: | |
| cur_full_fn = os.path.join(pkl_root_folder, cur_fn) | |
| compute_sdf(cur_full_fn) | |
| # obj_mesh_fn = "/data3/datasets/xueyi/taco/processed_data/20231104/right_20231104_017.obj" | |
| # compute_sdf(obj_mesh_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) | |