def call(self, points_xyz, points_feature, points_mask, points_voxel):
batch_size, _ = tf_util.get_shape(points_feature, 2)
points_feature = self.pointnet(points_feature, points_mask)
voxels = tf_util.batched_unsorted_segment_max(
batched_data=points_feature,
batched_segment_ids=points_voxel['indices'],
num_segments=points_voxel['num_voxels'],
batched_padding=points_voxel['paddings'])
_, _, nc = tf_util.get_shape(voxels)
voxels = tf.reshape(voxels, [-1])
voxels = tf.where(voxels <= voxels.dtype.min, tf.zeros_like(voxels), voxels)
voxels_in = tf.reshape(voxels, [batch_size] + self.grid_size + [nc])
voxels_out1 = self.res1(voxels_in)
voxels_out2 = self.res2(voxels_in)
voxels_out3 = self.res3(voxels_out2)
voxels_out2 = self.deconv2(voxels_out2)
voxels_out3 = self.deconv3(voxels_out3)
voxels_out = tf.concat([voxels_out1, voxels_out2, voxels_out3], axis=-1)
voxels_out = self.conv(voxels_out)
nc = tf_util.get_shape(voxels_out)[-1]
# w/ bilinear interpolation
points_out = tfa.image.interpolate_bilinear(
voxels_out, points_voxel['voxel_xyz'][:, :, :2])
return points_out
def call(self, inputs):
features = self.pillar_net(inputs)
cls_features = self.cls_mlp(features)
loc_features = self.loc_mlp(features)
cls_preds = self.cls_fc(cls_features)
reg_preds = self.loc_fc(loc_features)
batch_size, nx, ny, _ = tf_util.get_shape(reg_preds)
cls_logits = tf.reshape(cls_preds, [batch_size, nx*ny, 1])
reg_logits = tf.reshape(reg_preds, [batch_size, nx*ny, -1])
preds = {
'cls_logits': cls_logits,
'reg_logits': reg_logits,
}
return preds
def decode_fn(value,
data_aug=False,
max_num_points=245760,
max_num_bboxes=100,
class_id=1,
difficulty=1,
pillar_map_size=(256, 256),
pillar_map_range=(75.2, 75.2)):
"""Decode function."""
tensor_dict = waymo_decoder.decode_tf_example(
serialized_example=value, features=waymo_decoder.FEATURE_SPEC)
frame_valid = tensor_dict['frame_valid']
points_xyz = tensor_dict['lidars']['points_xyz']
points_feature = tensor_dict['lidars']['points_feature']
points_mask = tensor_dict['lidars']['points_mask']
all_points_xyz = tensor_dict['lidars']['all_points_xyz']
all_points_xyz_transformed = (
tensor_dict['lidars']['all_points_xyz_transformed'])
all_points_feature = tensor_dict['lidars']['all_points_feature']
all_points_mask = tensor_dict['lidars']['all_points_mask']
bboxes = tensor_dict['objects']['box']
bboxes_label = tensor_dict['objects']['label']
bboxes_speed = tensor_dict['objects']['speed']
bboxes_difficulty = tensor_dict['objects']['combined_difficulty_level']
bboxes_detection_difficulty = (
tensor_dict['objects']['combined_difficulty_level'])
bboxes_difficulty = bboxes_difficulty <= difficulty
bboxes_mask = tf.equal(bboxes_label, class_id)
bboxes_mask = tf.math.logical_and(bboxes_mask, bboxes_difficulty)
bboxes_mask = tf.cast(bboxes_mask, dtype=tf.dtypes.float32)
num_valid_bboxes = tf_util.get_shape(bboxes)[0]
bboxes_index = tf.math.top_k(bboxes_mask,
k=tf.math.minimum(max_num_bboxes,
num_valid_bboxes))[1]
bboxes_mask = tf.gather(bboxes_mask, bboxes_index)
bboxes_label = tf.gather(bboxes_label, bboxes_index)
bboxes = tf.gather(bboxes, bboxes_index)
bboxes_speed = tf.gather(bboxes_speed, bboxes_index)
bboxes = tf_util.pad_or_trim_to(bboxes, [max_num_bboxes, 7])
bboxes_label = tf_util.pad_or_trim_to(bboxes_label, [max_num_bboxes])
bboxes_speed = tf_util.pad_or_trim_to(bboxes_speed, [max_num_bboxes, 2])
bboxes_difficulty = tf_util.pad_or_trim_to(bboxes_difficulty,
[max_num_bboxes])
bboxes_mask = tf_util.pad_or_trim_to(bboxes_mask, [max_num_bboxes])
if data_aug:
(points_xyz, points_mask, bboxes, all_points_xyz,
all_points_xyz_transformed, all_points_mask) = augment(
points_xyz=points_xyz,
points_mask=points_mask,
bboxes=bboxes,
all_points_xyz=all_points_xyz,
all_points_xyz_transformed=all_points_xyz_transformed,
all_points_mask=all_points_mask)
(pillar_map_xyz, pillar_map_bboxes, pillar_map_bboxes_label,
pillar_map_if_in_bboxes, pillar_map_centerness,
pillar_map_bboxes_index) = (assign_bboxes(
pillar_map_size=pillar_map_size,
pillar_map_range=pillar_map_range,
bboxes=bboxes,
bboxes_label=bboxes_label,
bboxes_mask=bboxes_mask))
pillar_map_xyz = tf.reshape(pillar_map_xyz, [-1, 3])
pillar_map_bboxes = tf.reshape(pillar_map_bboxes, [-1, 7])
pillar_map_bboxes_label = tf.reshape(pillar_map_bboxes_label, [-1])
pillar_map_if_in_bboxes = tf.reshape(pillar_map_if_in_bboxes, [-1])
pillar_map_centerness = tf.reshape(pillar_map_centerness, [-1])
pillar_map_bboxes_index = tf.reshape(pillar_map_bboxes_index, [-1])
all_points_mask = tf.expand_dims(all_points_mask, axis=-1)
all_points = tf.concat([
all_points_xyz, all_points_xyz_transformed, all_points_feature,
all_points_mask
],
axis=-1)
num_frames, num_points, num_features = tf_util.get_shape(all_points)
all_points = tf.reshape(all_points,
[num_frames * num_points, num_features])
return {
'points_xyz': points_xyz,
'points_feature': points_feature,
'points_mask': points_mask,
'bboxes': bboxes,
'bboxes_label': bboxes_label,
'bboxes_mask': bboxes_mask,
'bboxes_speed': bboxes_speed,
'pillar_map_xyz': pillar_map_xyz,
'pillar_map_bboxes': pillar_map_bboxes,
'pillar_map_if_in_bboxes': pillar_map_if_in_bboxes,
}
def call(self, inputs):
points_xyz = inputs['points_xyz']
points_feature = inputs['points_feature']
points_mask = inputs['points_mask']
batch_size, num_points = tf_util.get_shape(points_xyz, 2)
xy_view_voxels = tf_util.points_to_voxels(points_xyz,
points_mask,
self.xy_view_grid_size,
self.xy_view_grid_range_x,
self.xy_view_grid_range_y,
self.xy_view_grid_range_z)
xy_view_voxels_stats = tf_util.points_to_voxels_stats(points_xyz,
xy_view_voxels)
xy_view_points_xyz = points_xyz - xy_view_voxels['centers']
points_cylinder = tf_util.points_xyz_to_cylinder(points_xyz)
cylinder_view_voxels = tf_util.points_to_voxels(
points_cylinder, points_mask, self.cylinder_view_grid_size,
self.cylinder_view_grid_range_x,
self.cylinder_view_grid_range_y,
self.cylinder_view_grid_range_z)
cylinder_view_voxels_stats = tf_util.points_to_voxels_stats(
points_cylinder, cylinder_view_voxels)
cylinder_view_points = points_cylinder - cylinder_view_voxels['centers']
points_feature = tf.concat([
points_xyz,
xy_view_points_xyz,
tf.cast(tf.reshape(xy_view_voxels['voxel_point_count'],
[batch_size, num_points, 1]),
tf.dtypes.float32),
xy_view_voxels_stats['centered_xyz'],
xy_view_voxels_stats['points_covariance'],
xy_view_voxels_stats['centroids'],
points_cylinder,
cylinder_view_points,
tf.cast(tf.reshape(cylinder_view_voxels['voxel_point_count'],
[batch_size, num_points, 1]),
tf.dtypes.float32),
cylinder_view_voxels_stats['centered_xyz'],
cylinder_view_voxels_stats['points_covariance'],
cylinder_view_voxels_stats['centroids'],
points_feature], axis=-1)
x = self.pointnet1(points_feature, points_mask)
x_xy_view = self.xy_view(points_xyz,
x,
points_mask,
xy_view_voxels)
x_cylinder_view = self.cylinder_view(points_cylinder,
x,
points_mask,
cylinder_view_voxels)
x_pointwise = self.pointnet2(x, points_mask)
x = tf.concat([
x_xy_view,
x_cylinder_view,
x_pointwise], axis=-1)
x = self.pointnet3(x, points_mask)
pillars = tf_util.batched_unsorted_segment_max(
batched_data=x,
batched_segment_ids=xy_view_voxels['indices'],
num_segments=xy_view_voxels['num_voxels'],
batched_padding=xy_view_voxels['paddings'])
_, _, nc = tf_util.get_shape(pillars)
pillars = tf.reshape(pillars, [-1])
pillars = tf.where(pillars <= pillars.dtype.min,
tf.zeros_like(pillars),
pillars)
nx, ny, nz = self.xy_view_grid_size
pillars = tf.reshape(pillars, [batch_size, nx, ny, nz * nc])
out1 = self.block1(pillars)
out2 = self.block2(out1)
out3 = self.block3(out2)
out1 = self.up1(out1)
out2 = self.up2(out2)
out3 = self.up3(out3)
out = tf.concat([out1, out2, out3], axis=-1)
out = self.conv(out)
return out
def call(self, points_feature, points_mask): batch_size, num_points = tf_util.get_shape(points_feature, 2) points_mask = tf.reshape(points_mask, [batch_size, num_points, 1]) points_feature = self.pointnet(points_feature) * points_mask return points_feature
def main(_):
batch_size = FLAGS.test_batch_size
# Fake optimizer
optimizer = tf.keras.optimizers.Adam(FLAGS.lr, clipnorm=10.0)
# Make a model
model = builder.PillarModel(
class_id=FLAGS.class_id,
norm_type=FLAGS.norm_type,
act_type=FLAGS.act_type,
nms_iou_threshold=FLAGS.nms_iou_threshold,
nms_score_threshold=FLAGS.nms_score_threshold,
max_nms_boxes=FLAGS.max_nms_boxes,
use_oriented_per_class_nms=FLAGS.use_oriented_per_class_nms)
# Create summary writers
model_dir = FLAGS.model_dir
summary_dir = os.path.join(model_dir, 'summaries')
eval_summary_writer = tf.summary.create_file_writer(
os.path.join(summary_dir, 'eval'))
# Make a dataset
dataset_val = waymo_loader.waymo_open_dataset(
data_path=FLAGS.data_path,
batch_size=batch_size,
split='valid',
cycle_length=FLAGS.cycle_length,
shuffle_buffer_size=FLAGS.shuffle_buffer_size,
num_parallel_calls=FLAGS.num_parallel_calls,
percentile=FLAGS.percentile,
max_num_points=FLAGS.max_num_points,
max_num_bboxes=FLAGS.max_num_bboxes,
class_id=FLAGS.class_id,
difficulty=FLAGS.difficulty,
pillar_map_size=(FLAGS.pillar_map_size, FLAGS.pillar_map_size),
pillar_map_range=(FLAGS.pillar_map_range, FLAGS.pillar_map_range))
checkpoint_file = None
while True:
# Validation loop starts here.
checkpoint = tf.train.Checkpoint(ema_model=model, optimizer=optimizer)
if FLAGS.ckpt_path and FLAGS.eval_once:
latest_checkpoint_file = FLAGS.ckpt_path
else:
latest_checkpoint_file = tf.train.latest_checkpoint(model_dir)
if latest_checkpoint_file == checkpoint_file:
time.sleep(60)
continue
else:
logging.info(
'Checkpoint file %s found and restoring from '
'checkpoint', latest_checkpoint_file)
checkpoint.restore(latest_checkpoint_file)
logging.info('Loading from checkpoint file completed')
checkpoint_file = latest_checkpoint_file
current_step = optimizer.iterations.numpy()
total_loss = 0
total_cls_loss = 0
total_loc_loss = 0
total_example = 0
bboxes_pred = []
bboxes_pred_score = []
bboxes_pred_mask = []
bboxes = []
bboxes_mask = []
bboxes_speed = []
for inputs in dataset_val:
preds = model(inputs, training=False)
outputs = model.infer(inputs, preds)
bboxes_pred.append(outputs['loc_preds'])
bboxes_pred_score.append(outputs['cls_preds'])
bboxes_pred_mask.append(outputs['loc_mask'])
bboxes.append(inputs['bboxes'])
bboxes_mask.append(inputs['bboxes_mask'])
bboxes_speed.append(inputs['bboxes_speed'])
batch_size = tf_util.get_shape(inputs['points_xyz'])[0]
cls_loss, loc_loss = model.compute_loss(inputs, preds)
cls_loss = tf.reduce_sum(cls_loss)
loc_loss = tf.reduce_sum(loc_loss)
total_loss += cls_loss.numpy() + loc_loss.numpy()
total_cls_loss += cls_loss.numpy()
total_loc_loss += loc_loss.numpy()
total_example += batch_size
if total_example % 100 == 0:
logging.info('finished decoding %d examples', total_example)
decoded_outputs = {
'bboxes_pred': tf.concat(bboxes_pred, axis=0),
'bboxes_pred_score': tf.concat(bboxes_pred_score, axis=0),
'bboxes_pred_mask': tf.concat(bboxes_pred_mask, axis=0),
'bboxes': tf.concat(bboxes, axis=0),
'bboxes_mask': tf.concat(bboxes_mask, axis=0),
'bboxes_speed': tf.concat(bboxes_speed, axis=0),
}
metrics = tf_util.compute_ap(decoded_outputs, FLAGS.class_id)
val_status = ('Val Step: %d / loc_loss = %s, cls_loss = %s.') % (
current_step, total_loc_loss / total_example,
total_cls_loss / total_example)
if eval_summary_writer:
with eval_summary_writer.as_default():
tf.summary.scalar('loc_loss',
total_loc_loss / total_example,
step=current_step)
tf.summary.scalar('cls_loss',
total_cls_loss / total_example,
step=current_step)
tf.summary.scalar('total_example',
total_example,
step=current_step)
for metric in metrics:
for key in metric:
tf.summary.scalar(key, metric[key], step=current_step)
metric_status = ('step: %s, %s: %s') % (
current_step, key, metric[key])
logging.info(metric_status)
eval_summary_writer.flush()
logging.info(val_status)
if FLAGS.eval_once:
break