def __init__(self, data_root, human, ann_file, split):
super(Dataset, self).__init__()
self.data_root = data_root
self.human = human
self.split = split
annots = np.load(ann_file, allow_pickle=True).item()
self.cams = annots['cams']
K, RT = render_utils.load_cam(ann_file)
render_w2c = RT
self.ts = np.arange(0, np.pi * 2, np.pi / 72)
self.nt = len(self.ts)
i = 0
i = i + cfg.begin_i
self.ims = np.array([
np.array(ims_data['ims'])[cfg.training_view]
for ims_data in annots['ims'][i:i + cfg.ni * cfg.i_intv]
])
self.K = K[0]
self.render_w2c = render_w2c
img_root = 'data/render/{}'.format(cfg.exp_name)
# base_utils.write_K_pose_inf(self.K, self.render_w2c, img_root)
self.Ks = np.array(K)[cfg.training_view].astype(np.float32)
self.RT = np.array(RT)[cfg.training_view].astype(np.float32)
self.center_rayd = [
render_utils.get_center_rayd(K_, RT_)
for K_, RT_ in zip(self.Ks, self.RT)
]
self.Ds = np.array(self.cams['D'])[cfg.training_view].astype(
np.float32)
self.nrays = cfg.N_rand
def get_nearest_cam(self, K, RT):
center_rayd = render_utils.get_center_rayd(K, RT)
sim = [np.dot(center_rayd, cam_rayd) for cam_rayd in self.center_rayd]
return self.RT[np.argmax(sim)]
