def Mem2Ref(xyz_mem, Z, Y, X):
# xyz is B x N x 3, in mem coordinates
# transforms mem coordinates into ref coordinates
B, N, C = list(xyz_mem.shape)
ref_T_mem = get_ref_T_mem(B, Z, Y, X)
xyz_ref = utils_disco.apply_4x4(ref_T_mem, xyz_mem)
return xyz_ref
def Ref2Mem(xyz, Z, Y, X):
# xyz is B x N x 3, in ref coordinates
# transforms velo coordinates into mem coordinates
B, N, C = list(xyz.shape)
mem_T_ref = get_mem_T_ref(B, Z, Y, X)
xyz = utils_disco.apply_4x4(mem_T_ref, xyz)
return xyz
def __getitem__(self, idx):
hyp_B = 1
hyp_S = 1
hyp_N = self.N
__p = lambda x: utils_disco.pack_seqdim(x, hyp_B)
__u = lambda x: utils_disco.unpack_seqdim(x, hyp_B)
__pb = lambda x: utils_disco.pack_boxdim(x, hyp_N)
__ub = lambda x: utils_disco.unpack_boxdim(x, hyp_N)
Z, Y, X = 144, 144, 144
Z2, Y2, X2 = int(Z / 2), int(Y / 2), int(X / 2)
Z4, Y4, X4 = int(Z / 4), int(Y / 4), int(X / 4)
######## Get query and key index #########################
scene_num = idx // self.views
query_idx = idx % self.views
scene_path = self.all_files[scene_num]
data = pickle.load(open(scene_path, "rb"))
###### Take necessary matrices from data dict ############
tids = torch.from_numpy(
np.reshape(np.arange(hyp_B * hyp_N), [hyp_B, hyp_N]))
pix_T_cams = torch.from_numpy(
data["pix_T_cams_raw"][query_idx]).reshape(hyp_B, hyp_S, 4,
4).cuda()
camRs_T_origin = data['camR_T_origin_raw'][query_idx]
camRs_T_origin = torch.from_numpy(camRs_T_origin).reshape(
hyp_B, hyp_S, 4, 4).cuda()
origin_T_camRs = __u(utils_disco.safe_inverse(__p(camRs_T_origin)))
origin_T_camXs = torch.from_numpy(
data['origin_T_camXs_raw'][query_idx]).reshape(hyp_B, hyp_S, 4,
4).cuda()
camX0_T_camXs = utils_disco.get_camM_T_camXs(origin_T_camXs, ind=0)
camRs_T_camXs = __u(
torch.matmul(utils_disco.safe_inverse(__p(origin_T_camRs)),
__p(origin_T_camXs)))
print(origin_T_camXs)
camXs_T_camRs = __u(utils_disco.safe_inverse(__p(camRs_T_camXs)))
camX0_T_camRs = camXs_T_camRs[:, 0]
camR_T_camX0 = utils_disco.safe_inverse(camX0_T_camRs)
rgb_camXs = data["rgb_camXs_raw"][:, :, :, :3]
rgb_camX0 = torch.from_numpy(rgb_camXs[query_idx]).permute(
2, 0, 1
).reshape(
hyp_B, 3, 256, 256
) #torch.from_numpy(np.fliplr(rgb_camXs[0,0])).reshape(1, 256, 256, 3).permute(0,3,1,2)
############ load tree file ##############################
tree_path = data['tree_seq_filename'].replace("shamitl", "mprabhud")
tree_path = tree_path.replace("datasets", "dataset")
tree_file = pickle.load(open(os.path.join(tree_path), "rb"))
gt_boxesR, scores, classes = self.trees_rearrange([tree_file])
########### set up gt boxes #############################
gt_boxesR = torch.from_numpy(
gt_boxesR).cuda().float() # torch.Size([2, 3, 6])
gt_boxesR_end = torch.reshape(gt_boxesR, [hyp_B, hyp_N, 2, 3])
gt_boxesR_theta = utils_disco.get_alignedboxes2thetaformat(
gt_boxesR_end)
gt_boxesR_corners = utils_disco.transform_boxes_to_corners(
gt_boxesR_theta)
gt_boxesRMem_corners = __ub(
utils_vox.Ref2Mem(__pb(gt_boxesR_corners), Z2, Y2, X2))
gt_boxesRMem_end = utils_disco.get_ends_of_corner(gt_boxesRMem_corners)
gt_boxesRMem_theta = utils_disco.transform_corners_to_boxes(
gt_boxesRMem_corners)
gt_boxesRUnp_corners = __ub(
utils_vox.Ref2Mem(__pb(gt_boxesR_corners), Z, Y, X))
gt_boxesRUnp_end = utils_disco.get_ends_of_corner(gt_boxesRUnp_corners)
gt_boxesX0_corners = __ub(
utils_disco.apply_4x4(camX0_T_camRs, __pb(gt_boxesR_corners)))
gt_boxesXs_corners = __u(
__ub(
utils_disco.apply_4x4(
__p(camXs_T_camRs),
__p(
__pb(gt_boxesR_corners).unsqueeze(1).repeat(
1, hyp_S, 1, 1)))))
gt_boxesXs_end = __u(
utils_disco.get_ends_of_corner(__p(gt_boxesXs_corners)))
gt_boxesX0Mem_corners = __ub(
utils_vox.Ref2Mem(__pb(gt_boxesX0_corners), Z2, Y2, X2))
gt_boxesX0Mem_theta = utils_disco.transform_corners_to_boxes(
gt_boxesX0Mem_corners)
gt_boxesX0Mem_end = utils_disco.get_ends_of_corner(
gt_boxesX0Mem_corners)
gt_boxesX0_end = utils_disco.get_ends_of_corner(gt_boxesX0_corners)
gt_cornersX0_pix = __ub(
utils_disco.apply_pix_T_cam(pix_T_cams[:, 0],
__pb(gt_boxesX0_corners)))
########### gt egomotion #############################################
################# get bbobex on image ##############################
boxes_vis_q, corners_pix_q = summ_box_by_corners(
rgb_camX0, gt_boxesX0_corners, torch.from_numpy(scores), tids,
pix_T_cams[:, 0])
boxes_q = torch.zeros([hyp_N, 4])
for n in range(hyp_N):
boxes_q[n][0] = torch.min(corners_pix_q[0, n, :, 0])
boxes_q[n][1] = torch.min(corners_pix_q[0, n, :, 1])
boxes_q[n][2] = torch.max(corners_pix_q[0, n, :, 0])
boxes_q[n][3] = torch.max(corners_pix_q[0, n, :, 1])
return (rgb_camX0 /
255.).squeeze(), hyp_N, boxes_q, boxes_vis_q.squeeze(
), scene_num, query_idx, pix_T_cams.squeeze(
), origin_T_camXs.squeeze(), idx, scene_path
def __getitem__(self, idx):
hyp_B = 1
hyp_S = 2
hyp_N = self.N
__p = lambda x: utils_disco.pack_seqdim(x, hyp_B)
__u = lambda x: utils_disco.unpack_seqdim(x, hyp_B)
__pb = lambda x: utils_disco.pack_boxdim(x, hyp_N)
__ub = lambda x: utils_disco.unpack_boxdim(x, hyp_N)
Z, Y, X = 144, 144, 144
Z2, Y2, X2 = int(Z / 2), int(Y / 2), int(X / 2)
Z4, Y4, X4 = int(Z / 4), int(Y / 4), int(X / 4)
######## Get query and key index #########################
scene_num = idx
if self.q_IDX is not None:
query_idx, key_idx = self.q_IDX, np.random.randint(
0, self.views, 1)[0]
else:
query_idx, key_idx = np.random.randint(0, self.views, 2)
# query_idx=0
# query_idx, key_idx = 0, random.sample(range(0, self.views), 1)[0]#; print(key_idx)
# print(self.views, random.sample(range(0, self.views), 1), key_idx)
index = scene_num * self.views + query_idx
scene_path = self.all_files[scene_num]
data = pickle.load(open(scene_path, "rb"))
###### Take necessary matrices from data dict ############
tids = torch.from_numpy(
np.reshape(np.arange(hyp_B * hyp_N), [hyp_B, hyp_N]))
pix_T_cams = torch.from_numpy(
np.stack((data["pix_T_cams_raw"][query_idx],
data["pix_T_cams_raw"][key_idx]))).reshape(
hyp_B, hyp_S, 4, 4).cuda()
camRs_T_origin = np.stack((data['camR_T_origin_raw'][query_idx],
data['camR_T_origin_raw'][key_idx]))
camRs_T_origin = torch.from_numpy(camRs_T_origin).reshape(
hyp_B, hyp_S, 4, 4).cuda()
origin_T_camRs = __u(utils_disco.safe_inverse(__p(camRs_T_origin)))
origin_T_camXs = torch.from_numpy(
np.stack((data['origin_T_camXs_raw'][query_idx],
data['origin_T_camXs_raw'][key_idx]))).reshape(
hyp_B, hyp_S, 4, 4).cuda()
camX0_T_camXs = utils_disco.get_camM_T_camXs(origin_T_camXs, ind=0)
camRs_T_camXs = __u(
torch.matmul(utils_disco.safe_inverse(__p(origin_T_camRs)),
__p(origin_T_camXs)))
camXs_T_camRs = __u(utils_disco.safe_inverse(__p(camRs_T_camXs)))
camX0_T_camRs = camXs_T_camRs[:, 0]
camX1_T_camRs = camXs_T_camRs[:, 1]
camR_T_camX0 = utils_disco.safe_inverse(camX0_T_camRs)
rgb_camXs = data["rgb_camXs_raw"][:, :, :, :3]
rgb_camX0 = torch.from_numpy(rgb_camXs[query_idx]).permute(
2, 0, 1
).reshape(
hyp_B, 3, 256, 256
) #torch.from_numpy(np.fliplr(rgb_camXs[0,0])).reshape(1, 256, 256, 3).permute(0,3,1,2)
rgb_camX1 = torch.from_numpy(rgb_camXs[key_idx]).permute(
2, 0, 1).reshape(hyp_B, 3, 256, 256)
############ load tree file ##############################
tree_path = data['tree_seq_filename'].replace("shamitl", "mprabhud")
tree_path = tree_path.replace("datasets", "dataset")
tree_file = pickle.load(open(os.path.join(tree_path), "rb"))
gt_boxesR, scores, classes = self.trees_rearrange([tree_file])
########### set up gt boxes #############################
gt_boxesR = torch.from_numpy(
gt_boxesR).cuda().float() # torch.Size([2, 3, 6])
gt_boxesR_end = torch.reshape(gt_boxesR, [hyp_B, hyp_N, 2, 3])
gt_boxesR_theta = utils_disco.get_alignedboxes2thetaformat(
gt_boxesR_end)
gt_boxesR_corners = utils_disco.transform_boxes_to_corners(
gt_boxesR_theta)
gt_boxesRMem_corners = __ub(
utils_vox.Ref2Mem(__pb(gt_boxesR_corners), Z2, Y2, X2))
gt_boxesRMem_end = utils_disco.get_ends_of_corner(gt_boxesRMem_corners)
gt_boxesRMem_theta = utils_disco.transform_corners_to_boxes(
gt_boxesRMem_corners)
gt_boxesRUnp_corners = __ub(
utils_vox.Ref2Mem(__pb(gt_boxesR_corners), Z, Y, X))
gt_boxesRUnp_end = utils_disco.get_ends_of_corner(gt_boxesRUnp_corners)
gt_boxesX0_corners = __ub(
utils_disco.apply_4x4(camX0_T_camRs, __pb(gt_boxesR_corners)))
gt_boxesX1_corners = __ub(
utils_disco.apply_4x4(camX1_T_camRs, __pb(gt_boxesR_corners)))
gt_boxesXs_corners = __u(
__ub(
utils_disco.apply_4x4(
__p(camXs_T_camRs),
__p(
__pb(gt_boxesR_corners).unsqueeze(1).repeat(
1, hyp_S, 1, 1)))))
gt_boxesXs_end = __u(
utils_disco.get_ends_of_corner(__p(gt_boxesXs_corners)))
gt_boxesX0Mem_corners = __ub(
utils_vox.Ref2Mem(__pb(gt_boxesX0_corners), Z2, Y2, X2))
gt_boxesX0Mem_theta = utils_disco.transform_corners_to_boxes(
gt_boxesX0Mem_corners)
gt_boxesX0Mem_end = utils_disco.get_ends_of_corner(
gt_boxesX0Mem_corners)
gt_boxesX0_end = utils_disco.get_ends_of_corner(gt_boxesX0_corners)
gt_cornersX0_pix = __ub(
utils_disco.apply_pix_T_cam(pix_T_cams[:, 0],
__pb(gt_boxesX0_corners)))
gt_cornersX1_pix = __ub(
utils_disco.apply_pix_T_cam(pix_T_cams[:, 1],
__pb(gt_boxesX1_corners)))
########### gt egomotion #############################################
'''
matmul(query_camX_T_origin, target_camX_T_origin.inverse())
'''
query_camX_T_origin = torch.from_numpy(
data['origin_T_camXs_raw'][query_idx]).reshape(1, 4, 4)
origin_T_camX_target = utils_disco.safe_inverse(
torch.from_numpy(data['origin_T_camXs_raw'][key_idx]).reshape(
1, 4, 4))
viewpoints = torch.matmul(query_camX_T_origin, origin_T_camX_target)
rx, ry, rz = utils_disco.rotm2eul(viewpoints)
rx, ry, rz = rx.unsqueeze(1), ry.unsqueeze(1), rz.unsqueeze(1)
xyz = viewpoints[:, :3, -1]
view_vector = [
xyz,
torch.cos(rx),
torch.sin(rx),
torch.cos(rz),
torch.sin(rz)
]
rel_viewpoint = torch.cat(view_vector, dim=-1)
################# get bbobex on image ##############################
boxes_vis_q, corners_pix_q = summ_box_by_corners(
rgb_camX0, gt_boxesX0_corners, torch.from_numpy(scores), tids,
pix_T_cams[:, 0])
boxes_vis_k, corners_pix_k = summ_box_by_corners(
rgb_camX1, gt_boxesX1_corners, torch.from_numpy(scores), tids,
pix_T_cams[:, 1])
boxes_q = torch.zeros([hyp_N, 4])
for n in range(hyp_N):
boxes_q[n][0] = torch.min(corners_pix_q[0, n, :, 0])
boxes_q[n][1] = torch.min(corners_pix_q[0, n, :, 1])
boxes_q[n][2] = torch.max(corners_pix_q[0, n, :, 0])
boxes_q[n][3] = torch.max(corners_pix_q[0, n, :, 1])
boxes_k = torch.zeros([hyp_N, 4])
for n in range(hyp_N):
boxes_k[n][0] = torch.min(corners_pix_k[0, n, :, 0])
boxes_k[n][1] = torch.min(corners_pix_k[0, n, :, 1])
boxes_k[n][2] = torch.max(corners_pix_k[0, n, :, 0])
boxes_k[n][3] = torch.max(corners_pix_k[0, n, :, 1])
return (rgb_camX0 /
255.).squeeze(), hyp_N, boxes_q, boxes_vis_q.squeeze(), (
rgb_camX1 /
255.).squeeze(), hyp_N, boxes_k, boxes_vis_k.squeeze(
), scene_num, query_idx, key_idx, pix_T_cams.squeeze(
), camRs_T_origin.squeeze(), origin_T_camXs.squeeze(
), rel_viewpoint, index, scene_path