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, }