Spaces:
Sleeping
Sleeping
File size: 2,174 Bytes
99a05f0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 |
'''
Combine cropped npz given a folder of npz
'''
import argparse
import os
import numpy as np
def combine_npz(npz_dir, out_npz):
# get all npz files in the folder
npz_files = [os.path.join(npz_dir, out_dir, f) for f in os.listdir(npz_dir, out_dir) if f.endswith('.npz')]
print('Found {} npz files'.format(len(npz_files)))
# combine all the values in all keys in all npz files
c_imgname = []
c_pose = []
c_transl = []
c_shape = []
c_cam_k = []
c_contact_label = []
c_scene_seg = []
c_part_seg = []
for npz_file in npz_files:
npz= np.load(npz_file)
c_imgname.extend(npz['imgname'])
c_pose.extend(npz['pose'])
c_transl.extend(npz['transl'])
c_shape.extend(npz['shape'])
c_cam_k.extend(npz['cam_k'])
c_contact_label.extend(npz['contact_label'])
c_scene_seg.extend(npz['scene_seg'])
c_part_seg.extend(npz['part_seg'])
# convert to numpy arrays
c_imgname = np.concatenate(c_imgname, axis=0)
c_pose = np.concatenate(c_pose, axis=0)
c_transl = np.concatenate(c_transl, axis=0)
c_shape = np.concatenate(c_shape, axis=0)
c_cam_k = np.concatenate(c_cam_k, axis=0)
c_contact_label = np.concatenate(c_contact_label, axis=0)
c_scene_seg = np.concatenate(c_scene_seg, axis=0)
c_part_seg = np.concatenate(c_part_seg, axis=0)
# save the new npz
np.savez(out_npz,
imgname=c_imgname,
pose=c_pose,
transl=c_transl,
shape=c_shape,
cam_k=c_cam_k,
contact_label=c_contact_label,
scene_seg=c_scene_seg,
part_seg=c_part_seg)
print('Saved combined npz to {}'.format(out_npz))
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--orig_npz_dir', type=str, default='/is/cluster/fast/achatterjee/rich/scene_npzs/train')
parser.add_argument('--out_npz', type=str, required=True,default='/is/cluster/fast/achatterjee/rich/scene_npzs/train_combined.npz')
args = parser.parse_args()
combine_npz(npz_dir = args.orig_npz_dir, out_npz = args.out_npz) |