in lib/datasets/patch_dataset.py [0:0]
def __getitem__(self, index):
''' Get index-th element from the picked file dataset. '''
# ------------------------------ INPUTS ----------------------------
rot_angle = self.get_center_view_rot_angle(index)
patch = self.patch_xyz_list[index]
if self.add_rgb:
rgb = self.patch_rgb_list[index]
patch = np.concatenate((patch, rgb), axis=-1)
self.rotate_to_center = True
if self.rotate_to_center:
patch[:, :, 0:3] = self.rotato_patch_to_center(patch[:, :, 0:3], rot_angle)
# transport patch to fixed size
# and change shape to C * H * W from H * W * C
#print('===============>',self.id_list)
patch = torch.from_numpy(patch) # H * W * C
patch = patch.unsqueeze(0) # 1 * H * W * C
patch = patch.transpose(2, 3).transpose(1, 2) # 1 * H * W * C -> 1 * H * C * W -> 1 * C * H * W
patch = F.interpolate(patch, self.patch_size, mode='bilinear', align_corners=True).squeeze(0).numpy()
#print('=======1========>',self.id_list[index])
cls_type = self.type_list[index]
assert(cls_type in ['Car', 'Pedestrian', 'Cyclist'])
one_hot_vec = np.zeros((3), dtype=np.float32)
one_hot_vec[self.dataset_helper.type2onehot[cls_type]] = 1
patch_feat = self.rgb_feat_list[index]
if self.from_rgb_detection:
return patch, rot_angle, self.prob_list[index], self.id_list[index], \
self.type_list[index], self.box2d_list[index], one_hot_vec, patch_feat
# ------------------------------ LABELS ----------------------------
# size labels
size_class, size_residual = self.dataset_helper.size2class(self.box3d_size_list[index], self.type_list[index])
# center labels
center = self.box3d_center_list[index]
if self.rotate_to_center:
center = rotate_pc_along_y(np.expand_dims(center,0), self.get_center_view_rot_angle(index)).squeeze()
# heading labels
heading_angle = self.heading_list[index]
if self.rotate_to_center:
heading_angle = heading_angle - rot_angle
#patch_feat = self.rgb_feat_list[index]
#print('=========>', patch_feat.shape)
if self.random_flip:
if np.random.random() > 0.5: # 50% chance flipping
patch[0, :, :] *= -1
center[0] *= -1
heading_angle = np.pi - heading_angle
# print('================1>', patch_feat.transpose(1,2,0).shape)
patch_feat_tmp = np.fliplr(patch_feat.transpose(1,2,0))
# print('================2>', patch_feat_tmp.shape)
patch_feat_new = np.ascontiguousarray(patch_feat_tmp.transpose(2,0,1))
# print('================3>', patch_feat_new.shape)
else:
patch_feat_new = patch_feat
if self.random_shift: # random shift object center
dist = np.sqrt(np.sum(center[0]**2 + center[2]**2))
shift = np.clip(np.random.randn() * dist * 0.2, -dist * 0.2, dist * 0.2)
patch[2, :, :] += shift
center[2] += shift
angle_class, angle_residual = self.dataset_helper.angle2class(heading_angle)
return patch, center, angle_class, angle_residual, size_class, size_residual, rot_angle, one_hot_vec, patch_feat_new#, self.level_list[index]