def generate_labels(self, center, dimension, angle, ref_xyz):
box_corner1 = compute_box_3d(center, dimension * 0.5, angle)
box_corner2 = compute_box_3d(center, dimension, angle)
labels = np.zeros(len(ref_xyz))
inside1 = extract_pc_in_box3d(ref_xyz, box_corner1)
inside2 = extract_pc_in_box3d(ref_xyz, box_corner2)
labels[inside2] = -1
labels[inside1] = 1
if inside1.sum() == 0:
dis = np.sqrt(((ref_xyz - center)**2).sum(1))
argmin = np.argmin(dis)
labels[argmin] = 1
return labels
def check_box_frustum(box, P, center, dimension, angle):
x1, y1, x2, y2 = box
box_corner = compute_box_3d(center, dimension, angle, P) # 8, 3
z1 = np.arange(0, 70, 0.1)
xyz1 = np.zeros((len(z1), 3))
xyz1[:, 0] = x1
xyz1[:, 1] = y1
xyz1[:, 2] = z1
xyz1_rect = project_image_to_rect(xyz1, P)
xyz1[:, 0] = x2
xyz1[:, 1] = y2
xyz1[:, 2] = z1
xyz2_rect = project_image_to_rect(xyz1, P)
xyz1[:, 0] = x1
xyz1[:, 1] = y2
xyz1[:, 2] = z1
xyz3_rect = project_image_to_rect(xyz1, P)
xyz1[:, 0] = x2
xyz1[:, 1] = y1
xyz1[:, 2] = z1
xyz4_rect = project_image_to_rect(xyz1, P)
fig = plt.figure()
ax = fig.gca(projection='3d')
draw_box3d(box_corner, ax)
draw_points(xyz1_rect, ax)
draw_points(xyz2_rect, ax)
draw_points(xyz3_rect, ax)
draw_points(xyz4_rect, ax)
plt.show()
def __getitem__(self, index):
rotate_to_center = cfg.DATA.RTC
with_extra_feat = cfg.DATA.WITH_EXTRA_FEAT
point_set = self.input_list[index].copy()
pred_box3d = self.pred_box3d_list[index].copy()
pred_box3d_center = (pred_box3d[0, :] + pred_box3d[6, :]) / 2
pred_box3d_angle = self.pred_box3d_angle_list[index]
pred_box3d_size = self.pred_box3d_size_list[index].copy()
cls_type = self.type_list[index]
# assert cls_type in KITTICategory.CLASSES, cls_type
# size_class = KITTICategory.CLASSES.index(cls_type)
assert cls_type in self.category_info.CLASSES, '%s not in category_info' % cls_type
size_class = self.category_info.CLASSES.index(cls_type)
# Compute one hot vector
if self.one_hot:
one_hot_vec = np.zeros((len(self.category_info.CLASSES)))
one_hot_vec[size_class] = 1
if rotate_to_center:
point_set[:, :3] = self.get_center_view_point(
point_set[:, :3], pred_box3d_center, pred_box3d_angle)
if not with_extra_feat:
point_set = point_set[:, :3]
# Resample
if self.npoints > 0:
# choice = np.random.choice(point_set.shape[0], self.npoints, replace=True)
choice = np.random.choice(point_set.shape[0], self.npoints,
point_set.shape[0] < self.npoints)
else:
choice = np.random.permutation(len(point_set.shape[0]))
point_set = point_set[choice, :]
# P = self.calib_list[index]['P2'].reshape(3, 4)
if rotate_to_center:
pred_box3d_center_rot, pred_box3d_angle_rot = self.get_center_view_box3d(
pred_box3d_center, pred_box3d_angle, pred_box3d_center,
pred_box3d_angle)
else:
pred_box3d_center_rot = pred_box3d_center
pred_box3d_angle_rot = pred_box3d_angle
pred_box3d = compute_box_3d(pred_box3d_center_rot, pred_box3d_size,
pred_box3d_angle_rot)
ref1, ref2, ref3, ref4 = self.generate_ref(pred_box3d)
if self.from_rgb_detection:
data_inputs = {
'point_cloud': np.transpose(point_set,
(1, 0)).astype(np.float32),
'center_ref1': np.transpose(ref1, (1, 0)).astype(np.float32),
'center_ref2': np.transpose(ref2, (1, 0)).astype(np.float32),
'center_ref3': np.transpose(ref3, (1, 0)).astype(np.float32),
'center_ref4': np.transpose(ref4, (1, 0)).astype(np.float32),
'rgb_prob':
np.array([self.prob_list[index]]).astype(np.float32),
'rot_angle': np.array([pred_box3d_angle]).astype(np.float32),
'ref_center': pred_box3d_center.astype(np.float32),
}
if not rotate_to_center:
data_inputs.update({'rot_angle': torch.zeros(1)})
data_inputs.update({'ref_center': torch.zeros(3)})
if self.one_hot:
data_inputs.update({'one_hot': torch.FloatTensor(one_hot_vec)})
return data_inputs
# mask_label = self.label_list[index].copy()
# mask_label = mask_label[choice]
box3d = self.box3d_list[index].copy()
heading_angle = self.heading_list[index]
box3d_size = self.size_list[index].copy()
box3d_center = (box3d[0, :] + box3d[6, :]) / 2
if rotate_to_center:
box3d_center, heading_angle = self.get_center_view_box3d(
box3d_center, heading_angle, pred_box3d_center,
pred_box3d_angle)
# box3d_rot = compute_box_3d(box3d_center, box3d_size, heading_angle)
# self.check(point_set, ref1, box3d_rot, pred_box3d)
if self.random_flip:
# note: rot_angle won't be correct if we have random_flip
# so do not use it in case of random flipping.
if np.random.random() > 0.5: # 50% chance flipping
point_set[:, 0] *= -1
box3d_center[0] *= -1
heading_angle = np.pi - heading_angle
ref1[:, 0] *= -1
ref2[:, 0] *= -1
ref3[:, 0] *= -1
ref4[:, 0] *= -1
if self.random_shift:
s1 = cfg.DATA.STRIDE[0]
l, w, h = self.size_list[index]
dist = np.sqrt(np.sum(l**2 + w**2))
shift = np.clip(np.random.randn() * dist * 0.1, -s1 * 2, 2 * s1)
point_set[:, 2] += shift
box3d_center[2] += shift
# angle_class, angle_residual = angle2class(heading_angle, NUM_HEADING_BIN)
# self.check(box, P, box3d_center, self.size_list[index], heading_angle)
labels = self.generate_labels(box3d_center, box3d_size, heading_angle,
ref2)
data_inputs = {
'point_cloud': np.transpose(point_set, (1, 0)).astype(np.float32),
'cls_label': labels.astype(np.int64),
'box3d_center': box3d_center.astype(np.float32),
'box3d_heading': np.array([heading_angle], dtype=np.float32),
'box3d_size': box3d_size.astype(np.float32),
'size_class': np.array([size_class], dtype=np.int64),
'center_ref1': np.transpose(ref1, (1, 0)).astype(np.float32),
'center_ref2': np.transpose(ref2, (1, 0)).astype(np.float32),
'center_ref3': np.transpose(ref3, (1, 0)).astype(np.float32),
'center_ref4': np.transpose(ref4, (1, 0)).astype(np.float32),
'rot_angle': np.array([pred_box3d_angle]).astype(np.float32),
'ref_center': pred_box3d_center.astype(np.float32),
}
if not rotate_to_center:
data_inputs.update({'rot_angle': torch.zeros(1)})
if self.one_hot:
data_inputs.update({'one_hot': torch.FloatTensor(one_hot_vec)})
return data_inputs