import os os.environ["CDF_LIB"] = "/is/cluster/scratch/stripathi/data/cdf37_1-dist/src/lib" import cv2 import pandas as pd # import json # import glob # import h5py import torch # import trimesh import numpy as np import pickle as pkl # from xml.dom import minidom # import xml.etree.ElementTree as ET from tqdm import tqdm from spacepy import pycdf # from .read_openpose import read_openpose import sys sys.path.append('../../') # from models import hmr, SMPL # import config # import constants import argparse # import shutil import smplx import pytorch3d.transforms as p3dt from utils.geometry import batch_rodrigues, batch_rot2aa, ea2rm # from vis_utils.world_vis import overlay_mesh, vis_smpl_with_ground model_type = 'smplx' model_folder = '/ps/project/common/smplifyx/models/' body_model_params = dict(model_path=model_folder, model_type=model_type, create_global_orient=True, create_body_pose=True, create_betas=True, num_betas=10, create_left_hand_pose=True, create_right_hand_pose=True, create_expression=True, create_jaw_pose=True, create_leye_pose=True, create_reye_pose=True, create_transl=True, use_pca=False) body_model = smplx.create(gender='neutral', **body_model_params).to('cuda') def rich_extract(img_dataset_path, out_path, scene_indx, split=None, vis_path=None, visualize=False, downsample_factor=1): # structs we use imgnames_ = [] # scales_, centers_, parts_, Ss_, Ss_world_, openposes_ = [], [], [], [], [], [] poses_, shapes_, transls_ = [], [], [] # poses_world_, transls_world_, cams_r_, cams_t_ = [], [], [], [] cams_k_ = [] # ground_offset_ = [] # in_bos_label_ = [] contact_label_ = [] scene_seg_, part_seg_ = [], [] # # seqs in validation set # if split == 'val': # seq_list = ['2021-06-15_Multi_IOI_ID_00176_Yoga2', # '2021-06-15_Multi_IOI_ID_00228_Yoga1', # '2021-06-15_Multi_IOI_ID_03588_Yoga1', # '2021-06-15_Multi_IOI_ID_00176_Yoga1'] # # seqs in testing set # if split == 'test': # seq_list = ['2021-06-15_Multi_IOI_ID_00186_Yoga1', # '2021-06-15_Multi_IOI_ID_03588_Yoga2', # 'MultiIOI_201019_ID03581_parkingLot_Calibration06_Settings06_PushUp__2', # 'Multi-IOI_ID00227_Scene_ParkingLot_Calibration_03_CameraSettings_4_pushup_1'] scenes = sorted(os.listdir('/ps/project/datasets/RICH/rich_toolkit/data/images/test')) # scene = scenes[scene_indx] scene = 'ParkingLot2_009_impro5' for scene_name in [scene]: # for scene_name in ['LectureHall_003_wipingchairs1']: out_file = os.path.join(out_path, f'rich_{scene_name}_smplx.npz') if os.path.exists(out_file): return print(scene_name) for i, fl in tqdm(enumerate(sorted(os.listdir(os.path.join(img_dataset_path, 'images', split)))), dynamic_ncols=True): if not scene_name in fl: continue ind = fl.index('cam') location = fl[:ind-1] cam_num = fl[ind:ind+6] img = fl[ind+7:-3] + 'jpeg' imgname = os.path.join(location, cam_num, img) mask_name = fl sp = mask_name.split('_') indx = mask_name.index('cam') st = mask_name[indx-1:indx+7] mask_name = mask_name.replace(st, '/') mask_name = mask_name[:-7] new_p = mask_name.split('/') mask_name = new_p[0] + '/' + new_p[1] + '/' + sp[1] + '.pkl' mask_path = os.path.join(img_dataset_path, 'labels', split, mask_name) df = pd.read_pickle(mask_path) mask = df['contact'] scene_path = os.path.join('/ps/scratch/ps_shared/stripathi/deco/4agniv/rich/seg_masks_new', split, fl[:-3] + 'png') part_path = os.path.join('/ps/scratch/ps_shared/stripathi/deco/4agniv/rich/part_masks_new', split, fl[:-3] + 'png') dataset_path = '/ps/project/datasets/RICH' ind = fl.index('cam') frame_id = fl[:ind-1] location = frame_id.split('_')[0] if location == 'LectureHall': if 'chair' in frame_id: cam2world_location = location + '_' + 'chair' else: cam2world_location = location + '_' + 'yoga' else: cam2world_location = location img_num = fl.split('_')[-2] cam_num = int(fl.split('_')[-1][:2]) # get ioi2scan transformation per sequence ioi2scan_fn = os.path.join(dataset_path, 'website_release/multicam2world', cam2world_location + '_multicam2world.json') try: camera_fn = os.path.join(dataset_path, 'rich_toolkit/data/scan_calibration', location, f'calibration/{cam_num:03d}.xml') focal_length_x, focal_length_y, camC, camR, camT, _, _, _ = extract_cam_param_xml(camera_fn) except: print(f'camera calibration file not found: {camera_fn}') continue # print('X: ', focal_length_x) # print('Y: ', focal_length_y) # path to smpl params smplx_param = os.path.join(dataset_path, 'rich_toolkit/data/bodies', split, frame_id, str(img_num), frame_id.split('_')[1] + '.pkl') # # path to GT bounding boxes # bbox_path = os.path.join(dataset_path, 'preprocessed', split, frame_id, img_num, frame_id.split('_')[1], 'bbox_refine', f'{img_num}_{cam_num:02d}.json') # # path with 2D openpose keypoints # openpose_path = os.path.join(dataset_path, 'preprocessed', split, frame_id, img_num, frame_id.split('_')[1], 'keypoints_refine', f'{img_num}_{str(cam_num).zfill(2)}_keypoints.json') # # path to image crops # img_path = os.path.join(dataset_path, 'preprocessed', split, frame_id, img_num, frame_id.split('_')[1], 'images_refine', f'{img_num}_{cam_num:02d}.png') # # bbox file # try: # with open(bbox_path, 'r') as f: # bbox_dict = json.load(f) # except: # print(f'bbox file not found: {bbox_path}') # continue # # read GT bounding box # x1_ul = bbox_dict['x1'] // downsample_factor # y1_ul = bbox_dict['y1'] // downsample_factor # x2_br = bbox_dict['x2'] // downsample_factor # y2_br = bbox_dict['y2'] // downsample_factor # bbox = np.array([x1_ul, y1_ul, x2_br, y2_br]) # center = [(bbox[2]+bbox[0])/2, (bbox[3]+bbox[1])/2] # scale = 0.9 * max(bbox[2] - bbox[0], bbox[3] - bbox[1]) / 200. # get smpl parameters ## body resides in multi-ioi coordidate, where camera 0 is world zero. with open(smplx_param, 'rb') as f: body_params = pkl.load(f) # in ioi coordinates: cam 0 beta = body_params['betas'] pose_aa = body_params['body_pose'] pose_rotmat = p3dt.axis_angle_to_matrix(torch.FloatTensor(pose_aa.reshape(-1,3))).numpy() transl = body_params['transl'] global_orient = body_params['global_orient'] global_orient = p3dt.axis_angle_to_matrix(torch.FloatTensor(global_orient.reshape(-1,3))).numpy() smpl_body_cam0 = body_model(betas=torch.FloatTensor(beta).to('cuda')) # canonical body with shape vertices_cam0 = smpl_body_cam0.vertices.detach().cpu().numpy().squeeze() joints_cam0 = smpl_body_cam0.joints.detach().cpu().numpy() pelvis_cam0 = joints_cam0[:, 0, :] # ## rigid transformation between multi-ioi and Leica scan (world) # with open(ioi2scan_fn, 'r') as f: # ioi2scan_dict = json.load(f) # R_ioi2world = np.array(ioi2scan_dict['R']) # Note: R is transposed # t_ioi2world= np.array(ioi2scan_dict['t']).reshape(1, 3) # # # get SMPL params in world coordinates # # # import ipdb; ipdb.set_trace() # global_orient_world = np.matmul(R_ioi2world.T, global_orient) # transl_world = np.matmul((pelvis_cam0+transl), R_ioi2world) + t_ioi2world - pelvis_cam0 # right multiplication to avoid transpose # full_pose_rotmat_world = np.concatenate((global_orient_world, pose_rotmat), axis=0).squeeze() # theta_world = batch_rot2aa(torch.FloatTensor(full_pose_rotmat_world)).reshape(-1, 66).cpu().numpy() # smpl_body_world = body_model(betas=torch.FloatTensor(beta).to('cuda'), # body_pose=torch.FloatTensor(theta_world[:, 3:]).to('cuda'), # transl=torch.FloatTensor(transl_world).to('cuda'), # global_orient=torch.FloatTensor(theta_world[:, :3]).to('cuda')) # vertices_world = smpl_body_world.vertices.detach().cpu().numpy().squeeze() # joints3d_world = smpl_body_world.joints[:, 25:, :].detach().cpu().numpy().squeeze() # # get SMPL params in camera coordinates global_orient_cam = np.matmul(np.array(camR), global_orient) transl_cam = np.matmul(camR, (pelvis_cam0 + transl).T).T + camT - pelvis_cam0 full_pose_rotmat_cam = np.concatenate((global_orient_cam, pose_rotmat), axis=0).squeeze() theta_cam = batch_rot2aa(torch.FloatTensor(full_pose_rotmat_cam)).reshape(-1, 66).cpu().numpy() # read GT 2D keypoints K = np.eye(3, dtype=np.float64) K[0, 0] = focal_length_x / downsample_factor K[1, 1] = focal_length_y / downsample_factor K[:2, 2:] = camC.T / downsample_factor # # get openpose 2D keypoints # try: # with open(openpose_path, 'r') as f: # openpose = json.load(f) # openpose = np.array(openpose['people'][0]['pose_keypoints_2d']).reshape([-1, 3]) # except: # print(f'No openpose !! Missing {openpose_path}') # continue # get camera parameters wrt to scan # R_worldtocam = np.matmul(camR, R_ioi2world) # Note: R_ioi2world is transposed # T_worldtocam = -t_ioi2world + camT # ground offset # ground_offset = ground_eq[2] # store data jpg_img_path = os.path.join('/ps/project/datasets/RICH_JPG', split, imgname) bmp_img_path = jpg_img_path.replace('/ps/project/datasets/RICH_JPG', '/ps/project/datasets/RICH/rich_toolkit/data/images') bmp_img_path = bmp_img_path.replace('.jpeg', '.bmp') if not os.path.exists(bmp_img_path): bmp_img_path = bmp_img_path.replace('.bmp', '.png') imgnames_.append(bmp_img_path) contact_label_.append(mask) scene_seg_.append(scene_path) part_seg_.append(part_path) # centers_.append(center) # scales_.append(scale) # openposes_.append(openpose) poses_.append(theta_cam.squeeze()) transls_.append(transl_cam.squeeze()) # poses_world_.append(theta_world.squeeze()) # transls_world_.append(transl_world.squeeze()) shapes_.append(beta.squeeze()) # cams_r_.append(R_worldtocam.tolist()) # # Todo: note that T_worldtocam here is (1,3) whereas in h36m T_worldtocam is (1,3) # cams_t_.append(T_worldtocam.tolist()) cams_k_.append(K.tolist()) # ground_offset_.append(ground_offset) # for seq_i in tqdm(seq_list): # print(f'Processing sequence: {seq_i}') # # path with GT bounding boxes # params_path = os.path.join(dataset_path, seq_i, 'params') # # path to metadata for files # md_path = os.path.join(dataset_path, seq_i, 'data') # # glob all folders in params path # frame_param_paths = sorted(glob.glob(os.path.join(params_path, '*'))) # frame_param_paths = [p for p in frame_param_paths if '.yaml' not in p] # # get ioi2scan transformation per sequence # ioi2scan_fn = os.path.join(dataset_path, seq_i, 'cam2scan.json') # ## ground resides in Leica scan coordinate, which is (roughly) axis aligned. # # ground_mesh = trimesh.load(os.path.join(dataset_path, seq_i, 'ground_mesh.ply'), process=False) # # ground_eq = np.mean(ground_mesh.vertices, axis=0) # # list all files in the folder # cam_files = os.listdir(os.path.join(dataset_path, seq_i, f'calibration')) # cam_list = sorted([int(os.path.splitext(f)[0]) for f in cam_files if '.xml' in f]) # # if split == 'val': # # cam_list = cam_list[1:] # remove first camera in val # for cam_num in cam_list[:1]: # camera_fn = os.path.join(dataset_path, seq_i, f'calibration/{cam_num:03d}.xml') # focal_length_x, focal_length_y, camC, camR, camT, _, _, _ = extract_cam_param_xml(camera_fn) # for frame_param_path in tqdm(frame_param_paths): # frame_id = os.path.basename(frame_param_path) # frame_num = int(frame_id) # # path to smpl params # try: # smplx_param = os.path.join(frame_param_path, '00', 'results/000.pkl') # except: # import ipdb; ipdb.set_trace() # # path to GT bounding boxes # bbox_path = os.path.join(md_path, frame_id, '00', 'bbox_refine', f'{frame_id}_{cam_num:02d}.json') # # path with 2D openpose keypoints # openpose_path = os.path.join(md_path, frame_id, '00', 'keypoints_refine', f'{frame_id}_{str(cam_num).zfill(2)}_keypoints.json') # # path to image crops # if downsample_factor == 1: # img_path = os.path.join(md_path, frame_id, '00', 'images_orig', f'{frame_id}_{cam_num:02d}.png') # else: # img_path = os.path.join(md_path, frame_id, '00', 'images_orig_720p', f'{frame_id}_{cam_num:02d}.png') # if not os.path.isfile(img_path): # print(f'image not found: {img_path}') # continue # # raise FileNotFoundError # # bbox file # try: # with open(bbox_path, 'r') as f: # bbox_dict = json.load(f) # except: # print(f'bbox file not found: {bbox_path}') # continue # # read GT bounding box # x1_ul = bbox_dict['x1'] // downsample_factor # y1_ul = bbox_dict['y1'] // downsample_factor # x2_br = bbox_dict['x2'] // downsample_factor # y2_br = bbox_dict['y2'] // downsample_factor # bbox = np.array([x1_ul, y1_ul, x2_br, y2_br]) # center = [(bbox[2]+bbox[0])/2, (bbox[3]+bbox[1])/2] # scale = 0.9 * max(bbox[2] - bbox[0], bbox[3] - bbox[1]) / 200. # # get smpl parameters # ## body resides in multi-ioi coordidate, where camera 0 is world zero. # with open(smplx_param, 'rb') as f: # body_params = pkl.load(f) # # in ioi coordinates: cam 0 # beta = body_params['betas'] # pose_aa = body_params['body_pose'] # pose_rotmat = p3dt.axis_angle_to_matrix(torch.FloatTensor(pose_aa.reshape(-1,3))).numpy() # transl = body_params['transl'] # global_orient = body_params['global_orient'] # global_orient = p3dt.axis_angle_to_matrix(torch.FloatTensor(global_orient.reshape(-1,3))).numpy() # smpl_body_cam0 = body_model(betas=torch.FloatTensor(beta).to('cuda')) # canonical body with shape # vertices_cam0 = smpl_body_cam0.vertices.detach().cpu().numpy().squeeze() # joints_cam0 = smpl_body_cam0.joints.detach().cpu().numpy() # pelvis_cam0 = joints_cam0[:, 0, :] # ## rigid transformation between multi-ioi and Leica scan (world) # with open(ioi2scan_fn, 'r') as f: # ioi2scan_dict = json.load(f) # R_ioi2world = np.array(ioi2scan_dict['R']) # Note: R is transposed # t_ioi2world= np.array(ioi2scan_dict['t']).reshape(1, 3) # # get SMPL params in world coordinates # # import ipdb; ipdb.set_trace() # global_orient_world = np.matmul(R_ioi2world.T, global_orient) # transl_world = np.matmul((pelvis_cam0+transl), R_ioi2world) + t_ioi2world - pelvis_cam0 # right multiplication to avoid transpose # full_pose_rotmat_world = np.concatenate((global_orient_world, pose_rotmat), axis=0).squeeze() # theta_world = batch_rot2aa(torch.FloatTensor(full_pose_rotmat_world)).reshape(-1, 66).cpu().numpy() # smpl_body_world = body_model(betas=torch.FloatTensor(beta).to('cuda'), # body_pose=torch.FloatTensor(theta_world[:, 3:]).to('cuda'), # transl=torch.FloatTensor(transl_world).to('cuda'), # global_orient=torch.FloatTensor(theta_world[:, :3]).to('cuda')) # vertices_world = smpl_body_world.vertices.detach().cpu().numpy().squeeze() # joints3d_world = smpl_body_world.joints[:, 25:, :].detach().cpu().numpy().squeeze() # mesh = trimesh.Trimesh(vertices_world, body_model.faces, # process=False, # maintain_order=True) # mesh.export('gt_mesh_world_smplx.obj') # # smpl_body_world = body_model(betas=torch.FloatTensor(beta).to('cuda'), # # body_pose=torch.FloatTensor(pose_rotmat[None, ...]).to('cuda'), # # transl=torch.FloatTensor(transl_world[None, ...]).to('cuda'), # # global_orient=torch.FloatTensor(global_orient_world[None, ...]).to('cuda'), # # left_hand_pose=torch.eye(3).reshape(1, 1, 3, 3).expand(1, 15, -1, -1).to('cuda'), # # right_hand_pose=torch.eye(3).reshape(1, 1, 3, 3).expand(1, 15, -1, -1).to('cuda'), # # leye_pose= torch.eye(3).reshape(1, 1, 3, 3).expand(batch_size, -1, -1, -1), # # # # pose2rot=False) # # vertices_world = smpl_body_world.vertices.detach().cpu().numpy().squeeze() # # joints3d_world = smpl_body_world.joints[:, 25:, :].detach().cpu().numpy().squeeze() # # mesh = trimesh.Trimesh(vertices_world, body_model.faces, # # process=False, # # maintain_order=True) # # mesh.export('gt_mesh_world_smplx.obj') # # get SMPL params in camera coordinates # global_orient_cam = np.matmul(camR, global_orient) # transl_cam = np.matmul(camR, (pelvis_cam0 + transl).T).T + camT - pelvis_cam0 # full_pose_rotmat_cam = np.concatenate((global_orient_cam, pose_rotmat), axis=0).squeeze() # theta_cam = batch_rot2aa(torch.FloatTensor(full_pose_rotmat_cam)).reshape(-1, 66).cpu().numpy() # # smpl_body_cam = body_model(betas=torch.FloatTensor(beta).to('cuda'), # # body_pose=torch.FloatTensor(pose_rotmat).to('cuda'), # # transl=torch.FloatTensor(transl_cam).to('cuda'), # # global_orient=torch.FloatTensor(global_orient_cam).to('cuda'), # # pose2rot=False) # # vertices_cam = smpl_body_cam.vertices.detach().cpu().numpy().squeeze() # # joints3d_cam = smpl_body_cam.joints[:, 25:, :].detach().cpu().numpy().squeeze() # # # # mesh = trimesh.Trimesh(vertices_cam, body_model.faces, # # process=False, # # maintain_order=True) # # mesh.export('mesh_in_cam0.obj') # # read GT 2D keypoints # K = np.eye(3, dtype=np.float) # K[0, 0] = focal_length_x / downsample_factor # K[1, 1] = focal_length_y / downsample_factor # K[:2, 2:] = camC.T / downsample_factor # # projected_points = (K @ joints3d_cam.T).T # # joints2d = projected_points[:, :2] / np.hstack((projected_points[:, 2:], projected_points[:, 2:])) # # part = np.hstack((joints2d, np.ones((joints2d.shape[0], 1)))) # # get openpose 2D keypoints # try: # with open(openpose_path, 'r') as f: # openpose = json.load(f) # openpose = np.array(openpose['people'][0]['pose_keypoints_2d']).reshape([-1, 3]) # except: # print(f'No openpose !! Missing {openpose_path}') # continue # # get camera parameters wrt to scan # R_worldtocam = np.matmul(camR, R_ioi2world) # Note: R_ioi2world is transposed # T_worldtocam = -t_ioi2world + camT # # ground offset # # ground_offset = ground_eq[2]/ # # # get stability labels: 1: stable, 0: unstable but in contact, -1: unstable and not `in contact # # in_bos_label, contact_label, contact_mask = vis_smpl_with_ground(theta_world, transl_world, beta, seq_i, # # vis_path, # # start_idx=frame_num, # # sub_sample=1, # # ground_offset=ground_offset, # # smpl_batch_size=1, # # visualize=False) # # in_bos_label = in_bos_label.detach().cpu().numpy() # # contact_label = contact_label.detach().cpu().numpy() # # contact_mask = contact_mask.detach().cpu().numpy() # # visualize world smpl on ground plane # # if visualize: # # if cam_num == 0: # # vis_smpl_with_ground(theta_world, transl_world, beta, split+'_'+seq_i, vis_path, # # start_idx=frame_num, # # sub_sample=1, # # ground_offset=ground_offset, # # smpl_batch_size=1, # # visualize=True) # # ## visualize projected points # # img = cv2.imread(img_path) # # joints2d = joints2d.astype(np.int) # # img[joints2d[:, 1], joints2d[:, 0], :] = [0, 255, 0] # # read GT 3D pose in cam coordinates # # S24 = joints3d_cam # # pelvis_cam = (S24[[2], :] + S24[[3], :]) / 2 # # S24 -= pelvis_cam # # S24 = np.hstack([S24, np.ones((S24.shape[0], 1))]) # # read GT 3D pose in world coordinates # # S24_world = joints3d_world # # S24_world = np.hstack([S24_world, np.ones((S24_world.shape[0], 1))]) # # store data # imgnames_.append(img_path) # centers_.append(center) # scales_.append(scale) # # parts_.append(part) # # Ss_.append(S24) # # Ss_world_.append(S24_world) # openposes_.append(openpose) # poses_.append(theta_cam.squeeze()) # transls_.append(transl.squeeze()) # poses_world_.append(theta_world.squeeze()) # transls_world_.append(transl_world.squeeze()) # shapes_.append(beta.squeeze()) # cams_r_.append(R_worldtocam.tolist()) # # Todo: note that T_worldtocam here is (1,3) whereas in h36m T_worldtocam is (1,3) # cams_t_.append(T_worldtocam.tolist()) # cams_k_.append(K.tolist()) # # in_bos_label_.append(in_bos_label) # # contact_label_.append(contact_label) # ground_offset_.append(ground_offset) # store the data struct if not os.path.isdir(out_path): os.makedirs(out_path) out_file = os.path.join(out_path, f'rich_{scene_name}_smplx.npz') np.savez(out_file, imgname=imgnames_, # center=centers_, # scale=scales_, # part=parts_, # S=Ss_, # S_world=Ss_world_, pose=poses_, transl=transls_, shape=shapes_, # openpose=openposes_, # pose_world=poses_world_, # transl_world=transls_world_, # cam_r=cams_r_, # cam_t=cams_t_, cam_k=cams_k_, # in_bos_label=in_bos_label_, contact_label=contact_label_, # ground_offset=ground_offset_ scene_seg=scene_seg_, part_seg=part_seg_ ) print('Saved to ', out_file) def rectify_pose(camera_r, body_aa): body_r = batch_rodrigues(body_aa).reshape(-1,3,3) final_r = camera_r @ body_r body_aa = batch_rot2aa(final_r) return body_aa def extract_cam_param_xml(xml_path: str = '', dtype=float): import xml.etree.ElementTree as ET tree = ET.parse(xml_path) extrinsics_mat = [float(s) for s in tree.find('./CameraMatrix/data').text.split()] intrinsics_mat = [float(s) for s in tree.find('./Intrinsics/data').text.split()] distortion_vec = [float(s) for s in tree.find('./Distortion/data').text.split()] focal_length_x = intrinsics_mat[0] focal_length_y = intrinsics_mat[4] center = np.array([[intrinsics_mat[2], intrinsics_mat[5]]], dtype=dtype) rotation = np.array([[extrinsics_mat[0], extrinsics_mat[1], extrinsics_mat[2]], [extrinsics_mat[4], extrinsics_mat[5], extrinsics_mat[6]], [extrinsics_mat[8], extrinsics_mat[9], extrinsics_mat[10]]], dtype=dtype) translation = np.array([[extrinsics_mat[3], extrinsics_mat[7], extrinsics_mat[11]]], dtype=dtype) # t = -Rc --> c = -R^Tt cam_center = [-extrinsics_mat[0] * extrinsics_mat[3] - extrinsics_mat[4] * extrinsics_mat[7] - extrinsics_mat[8] * extrinsics_mat[11], -extrinsics_mat[1] * extrinsics_mat[3] - extrinsics_mat[5] * extrinsics_mat[7] - extrinsics_mat[9] * extrinsics_mat[11], -extrinsics_mat[2] * extrinsics_mat[3] - extrinsics_mat[6] * extrinsics_mat[7] - extrinsics_mat[10] * extrinsics_mat[11]] cam_center = np.array([cam_center], dtype=dtype) k1 = np.array([distortion_vec[0]], dtype=dtype) k2 = np.array([distortion_vec[1]], dtype=dtype) return focal_length_x, focal_length_y, center, rotation, translation, cam_center, k1, k2 # rich_extract(img_dataset_path='/is/cluster/work/achatterjee/rich', out_path='/is/cluster/work/achatterjee/rich/npzs', split='train') # rich_extract(img_dataset_path='/is/cluster/work/achatterjee/rich', out_path='/is/cluster/fast/achatterjee/rich/scene_npzs', split='val') # rich_extract(img_dataset_path='/is/cluster/work/achatterjee/rich', out_path='/is/cluster/fast/achatterjee/rich/scene_npzs/test', split='test') # rich_extract(dataset_path='/ps/scratch/ps_shared/stripathi/4yogi/RICH/val/', out_path='/home/achatterjee/rich_ext', split='val') if __name__=='__main__': parser = argparse.ArgumentParser() parser.add_argument('--index', type=int) args = parser.parse_args() rich_extract(img_dataset_path='/is/cluster/work/achatterjee/rich', out_path='/is/cluster/fast/achatterjee/rich/scene_npzs/test', scene_indx=args.index, split='test')