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 tf_bboxes_select_layer(classifications_layer, localizations_layer,
select_threshold=None,
ignore_class=0,
scope=None):
"""Extract classes, scores and bounding boxes from features in one layer.
Batch-compatible: inputs are supposed to have batch-type shapes.
Args:
predictions_layer: A SSD prediction layer;
localizations_layer: A SSD localization layer;
select_threshold: Classification threshold for selecting a box. All boxes
under the threshold are set to 'zero'. If None, no threshold applied.
Return:
d_scores, d_bboxes: Dictionary of scores and bboxes Tensors of
size Batches X N x 1 | 4. Each key corresponding to a class.
"""
select_threshold = 0.0 if select_threshold is None else select_threshold
with tf.name_scope(scope, 'bboxes_select_layer',
[classifications_layer, localizations_layer]):
# Reshape features: Batches x N x N_labels | 4
c_shape = get_shape(classifications_layer)
classifications_layer = tf.reshape(classifications_layer,
tf.stack([c_shape[0], -1, c_shape[-1]]))
l_shape = get_shape(localizations_layer)
localizations_layer = tf.reshape(localizations_layer,
tf.stack([l_shape[0], -1, l_shape[-1]]))
d_scores = {}
d_bboxes = {}
for c in range(c_shape[-1]):
if c != ignore_class:
# Remove boxes under the threshold.
scores = classifications_layer[:, :, c]
fmask = tf.cast(tf.greater_equal(scores, select_threshold), scores.dtype)
scores = scores * fmask
bboxes = localizations_layer * tf.expand_dims(fmask, axis=-1)
# Append to dictionary.
d_scores[c] = scores
d_bboxes[c] = bboxes
return d_scores, d_bboxes
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 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
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 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_name = tensor_dict['frame_name']
points_xyz = tensor_dict['lidars']['points_xyz']
points_feature = tensor_dict['lidars']['points_feature']
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']
num_valid_points = tf_util.get_shape(points_xyz)[0]
points_xyz = tf_util.pad_or_trim_to(points_xyz, [max_num_points, 3])
points_feature = tf_util.pad_or_trim_to(points_feature, [max_num_points, 2])
points_mask = tf.sequence_mask(num_valid_points,
maxlen=max_num_points)
points_mask = tf.cast(points_mask, dtype=tf.dtypes.float32)
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 = augment(
points_xyz=points_xyz,
points_mask=points_mask,
bboxes=bboxes)
(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])
return {
'frame_name': frame_name,
'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,
}