Nuscenes Warping Test Code

import matplotlib.pyplot as plt
from nuscenes.utils import data_classes
seq_sample_recs = data['seq_records']

h, w = self.cfg['bev']['h'], self.cfg['bev']['w']
h_meters, w_meters = self.cfg['bev']['h_meters'], self.cfg['bev']['w_meters']
h_scale, w_scale = float(h) / h_meters, float(w) / w_meters
h_shift, w_shift = int(h/2), int(w/2)

lidar_record = self.nusc.get('sample_data', seq_sample_recs[0]['data']['LIDAR_TOP'])
egolidar = self.nusc.get('ego_pose', lidar_record['ego_pose_token'])
w2e_ref = get_pose(egolidar['rotation'], egolidar['translation'], inv=True, flat=True)

ins_map = {}
topview_ref = np.zeros(shape=(h, w)).astype('uint8')
for t in range(0, len(seq_sample_recs)):

    lidar_record = self.nusc.get('sample_data', seq_sample_recs[t]['data']['LIDAR_TOP'])
    egolidar = self.nusc.get('ego_pose', lidar_record['ego_pose_token'])
    w2e = get_pose(egolidar['rotation'], egolidar['translation'], inv=True, flat=True)

    topview = np.zeros(shape=(h, w)).astype('uint8')
    for ann_token in seq_sample_recs[t]['anns']:
        ann = self.nusc.get('sample_annotation', ann_token)
        if ('vehicle' in  ann['category_name']):

            if ann['instance_token'] not in ins_map:
                ins_map[ann['instance_token']] = len(ins_map) + 1
            ins_id = ins_map[ann['instance_token']]

            box = data_classes.Box(ann['translation'], ann['size'], Quaternion(ann['rotation']))
            corners_w = box.bottom_corners() # 3 x 4
            corners_w = np.concatenate([corners_w, np.ones(4).reshape(1, 4)], axis=0) # 4 x 4
            conrners_e_ref = (w2e_ref @ corners_w)[:2] # 2 x 4
            conrners_e = (w2e @ corners_w)[:2]

            col_pels = -(conrners_e_ref[1, :] * w_scale).astype(np.int32) + int(w_shift)
            row_pels = -(conrners_e_ref[0, :] * h_scale).astype(np.int32) + int(h_shift)
            corners_tv = np.concatenate([col_pels.reshape(4, 1), row_pels.reshape(4, 1)], axis=1)
            topview_ref = cv2.fillConvexPoly(topview_ref, corners_tv, ins_id)

            col_pels = -(conrners_e[1, :] * w_scale).astype(np.int32) + int(w_shift)
            row_pels = -(conrners_e[0, :] * h_scale).astype(np.int32) + int(h_shift)
            corners_tv = np.concatenate([col_pels.reshape(4, 1), row_pels.reshape(4, 1)], axis=1)
            topview = cv2.fillConvexPoly(topview, corners_tv, ins_id)


    # plt.imshow(topview.astype('float'))
    # plt.show()
plt.imshow(topview_ref.astype('float'))
plt.show()

seq_sample_recs = data['seq_records']
# t, c, h, w -> b, t, c, h, w
instance = torch.flip(data['instance'], dims=(2, 3)).unsqueeze(0).permute(0, 1, 2, 4, 3)
# instance = data['instance'].unsqueeze(0).permute(0, 1, 2, 4, 3)

e2w, w2e = data['e2w'], data['w2e']

instance_acc = toNP(instance[:, 0].squeeze())
for t in range(0, len(seq_sample_recs)):
    RT_2to0 = w2e[0].numpy() @ e2w[t].numpy()
    R = RT_2to0[:2, :2]
    T = RT_2to0[:2, -1].reshape(2, 1) / 50.0
    T[0, 0] *= -1

    transformation = np.concatenate([R, T], axis=1)[None]
    transformation = torch.from_numpy(transformation)

    x = instance[:, t]
    grid = torch.nn.functional.affine_grid(transformation, size=x.shape, align_corners=False)
    warped_x = torch.nn.functional.grid_sample(x, grid.float(), mode='bilinear', padding_mode='zeros',
                                               align_corners=False)

    instance_t = toNP(warped_x.squeeze())
    chk = (instance_t > 0)
    instance_acc[chk] = instance_t[chk]

instance_acc = toNP(torch.flip(torch.from_numpy(instance_acc).permute(1, 0), dims=(0, 1)))
plt.imshow(instance_acc)
plt.show()