diff --git a/examples/frustum_pointnet/dataset_utils.py b/examples/frustum_pointnet/dataset_utils.py new file mode 100644 index 000000000..b32b39ae9 --- /dev/null +++ b/examples/frustum_pointnet/dataset_utils.py @@ -0,0 +1,21 @@ +import numpy as np + +NUM_HEADING_BIN = 12 +NUM_SIZE_CLUSTER = 8 +NUM_OBJECT_POINT = 512 + +type2class = {'Car': 0, 'Van': 1, 'Truck': 2, 'Pedestrian': 3, + 'Person_sitting': 4, 'Cyclist': 5, 'Tram': 6, 'Misc': 7} +class2type = {type2class[t]: t for t in type2class} +type2onehotclass = {'Car': 0, 'Pedestrian': 1, 'Cyclist': 2} +type_mean_size = {'Car': np.array([3.88311640418, 1.62856739989, 1.52563191]), + 'Van': np.array([5.06763659, 1.9007158, 2.20532825]), + 'Truck': np.array([10.13586957, 2.58549199, 3.2520595]), + 'Pedestrian': np.array([0.84422524, 0.66068622, 1.7625519]), + 'Person_sitting': np.array([0.80057803, 0.5983815, 1.2745]), + 'Cyclist': np.array([1.76282397, 0.59706367, 1.73698127]), + 'Tram': np.array([16.17150617, 2.53246914, 3.53079012]), + 'Misc': np.array([3.64300781, 1.54298177, 1.92320313])} +g_mean_size_arr = np.zeros((NUM_SIZE_CLUSTER, 3)) # size clustrs +for i in range(NUM_SIZE_CLUSTER): + g_mean_size_arr[i, :] = type_mean_size[class2type[i]] \ No newline at end of file diff --git a/examples/frustum_pointnet/frustum_loader.py b/examples/frustum_pointnet/frustum_loader.py new file mode 100644 index 000000000..36f4f6f35 --- /dev/null +++ b/examples/frustum_pointnet/frustum_loader.py @@ -0,0 +1,157 @@ +import tensorflow as tf +import tensorflow.python.keras.backend as K +import numpy as np + +from dataset_utils import NUM_HEADING_BIN +from paz.optimization.losses.frustumpointnet_loss import ExtractBox3DCorners + + +class frustum_data_loader(object): + def __init__(self, batch_size=32): + # g_type_object_of_interest = ['animal', 'bicycle', 'bus', 'car', + # 'emergency_vehicle', 'motorcycle', + # 'other_vehicle', 'pedestrian', 'truck'] + self.g_type_object_of_interest = ['car', 'cyclist', 'pedestrian'] + self.NUM_CLASS = len(self.g_type_object_of_interest) + self.NUM_POINT = 1024 + self.NUM_CHANNELS_OF_PC = 3 + self.batch_size = batch_size + + def parse_data(self, raw_record): + + example = self.parse_frustum_point_record(raw_record) + return example['frustum_point_cloud'], \ + tf.cast(example['one_hot_vec'], tf.float32), \ + tf.cast(example['seg_label'], tf.int32), \ + example['box3d_center'], \ + tf.cast(example['angle_class'], tf.int32), \ + example['angle_residual'], \ + tf.cast(example['size_class'], tf.int32), \ + example['size_residual'] + + def parse_frustum_point_record(self, tfexample_message: str): + + keys_to_features = { + "size_class": tf.io.FixedLenFeature((), tf.int64, + tf.zeros((), tf.int64)), + "size_residual": tf.io.FixedLenFeature((3,), tf.float32, + tf.zeros((3,), tf.float32)), + "seg_label": tf.io.FixedLenFeature((self.NUM_POINT,), tf.int64, + tf.zeros((self.NUM_POINT,), + tf.int64)), + "frustum_point_cloud": tf.io.FixedLenFeature( + (self.NUM_POINT, self.NUM_CHANNELS_OF_PC), tf.float32), + "rot_angle": tf.io.FixedLenFeature((), tf.float32, + tf.zeros((), tf.float32)), + "angle_class": tf.io.FixedLenFeature((), tf.int64, + tf.zeros((), tf.int64)), + "angle_residual": tf.io.FixedLenFeature((), tf.float32, + tf.zeros((), tf.float32)), + "one_hot_vec": tf.io.FixedLenFeature((self.NUM_CLASS,), tf.int64), + "box3d_center": tf.io.FixedLenFeature((3,), tf.float32, + tf.zeros((3,), tf.float32)), + } + parsed_example = tf.io.parse_single_example(tfexample_message, + keys_to_features) + return parsed_example + + def parse_test_data(self, raw_record): + example = self.parse_frustum_point_test_record(raw_record) + print(example) + return example['frustum_point_cloud'], tf.cast(example['one_hot_vec'], + tf.float32), \ + tf.cast(example['rot_angle'], tf.float32), tf.cast( + example['prob'], tf.float32), \ + example['type_name'], example['sample_token'], example['box_2d'] + + def parse_frustum_point_test_record(self, tfexample_message: str): + + keys_to_features = { + "frustum_point_cloud": tf.io.FixedLenFeature( + (self.NUM_POINT, self.NUM_CHANNELS_OF_PC), tf.float32), + "rot_angle": tf.io.FixedLenFeature((), tf.float32, + tf.zeros((), tf.float32)), + "one_hot_vec": tf.io.FixedLenFeature((self.NUM_CLASS,), tf.int64), + "prob": tf.io.FixedLenFeature((), tf.float32, + tf.zeros((), tf.float32)), + "type_name": tf.io.FixedLenFeature((), tf.int64), + "sample_token": tf.io.FixedLenFeature((), tf.int64), + "box_2d": tf.io.FixedLenFeature((4,), tf.float32) + } + parsed_example = tf.io.parse_single_example(tfexample_message, + keys_to_features) + return parsed_example + + def load_data(self, tfrec_path, operation): + if operation == 'train': + train_dataset = tf.data.TFRecordDataset(tfrec_path) + parsed_train_dataset = train_dataset.map(self.parse_data) + parsed_train_dataset = parsed_train_dataset.batch( + self.batch_size, drop_remainder=True) + return parsed_train_dataset + elif operation == 'validation': + val_dataset = tf.data.TFRecordDataset(tfrec_path) + parsed_val_dataset = val_dataset.map(self.parse_data) + parsed_val_dataset = parsed_val_dataset.batch( + self.batch_size, drop_remainder=True) + return parsed_val_dataset + else: + test_dataset = tf.data.TFRecordDataset(tfrec_path) + parsed_test_dataset = test_dataset.map(self.parse_test_data) + parsed_test_dataset = parsed_test_dataset.batch( + self.batch_size, drop_remainder=True) + return parsed_test_dataset + + +def make_train_iterator(dataset): + iterator = dataset.make_one_shot_iterator() + next_val = iterator.get_next() + + with K.get_session().as_default() as sess: + while True: + pointclouds_pl, one_hot_vec_pl, labels_pl, centers_pl, \ + heading_class_label_pl, heading_residual_label_pl, \ + size_class_label_pl, size_residual_label_pl = sess.run(next_val) + + hcls_onehot = tf.one_hot(tf.cast(heading_class_label_pl, tf.int64), + depth=NUM_HEADING_BIN, on_value=1, + off_value=0, + axis=-1) # BxNUM_HEADING_BIN + heading_residual_normalized_label = heading_residual_label_pl / ( + np.pi / NUM_HEADING_BIN) + + corners_3d = ExtractBox3DCorners(centers_pl, + heading_residual_label_pl, + size_residual_label_pl) + + x_train = {'frustum_point_cloud': pointclouds_pl, + 'one_hot_vector': one_hot_vec_pl} + + y_train = {'seg_logits': labels_pl, + 'center': centers_pl, + 'seg_pc_centroid': centers_pl, + 'heading_class': heading_class_label_pl, + 'heading residuals': heading_residual_label_pl, + 'size_class': size_class_label_pl, + 'size_residuals': size_residual_label_pl, + 'corners_3d_pred': corners_3d} + + yield x_train, y_train + + +def make_test_iterator(dataset): + iterator = dataset.make_one_shot_iterator() + next_val = iterator.get_next() + + with K.get_session().as_default() as sess: + while True: + pointclouds_pl, one_hot_vec_pl, rot_angle, prob_pl, class_pl, \ + tokens_pl, box2d_pl = sess.run(next_val) + data_dict = {"frustum_point_cloud": pointclouds_pl, + "one_hot_vector": one_hot_vec_pl, + "rot_angle": rot_angle, + "rgb_prob": prob_pl, "cls_type": class_pl, + "token": tokens_pl, + "box_2D": box2d_pl} + + yield data_dict diff --git a/examples/frustum_pointnet/train.py b/examples/frustum_pointnet/train.py new file mode 100644 index 000000000..e50b22640 --- /dev/null +++ b/examples/frustum_pointnet/train.py @@ -0,0 +1,83 @@ +import argparse +import math +import os + +from tensorflow.keras.callbacks import CSVLogger, ModelCheckpoint +from tensorflow.keras.callbacks import EarlyStopping, LearningRateScheduler +from tensorflow.keras.optimizers import Adam + +from frustum_loader import frustum_data_loader, make_train_iterator +from paz.models.detection.frustum_pointnet import FrustumPointNetModel + +description = 'Training script for learning 3D Object Detection' +parser = argparse.ArgumentParser(description=description) +parser.add_argument('-b', '--batch_size', default=32, type=int, + help='Batch size for training') +parser.add_argument('-lr', '--learning_rate', default=0.001, type=float, + help='Initial learning rate for Adam') +parser.add_argument('-dr', '--decay_rate', default=0.5, type=float, + help='Decay rate for scheduling the learning rate') +parser.add_argument('-e', '--max_num_epochs', default=10000, type=int, + help='Maximum number of epochs before finishing') +parser.add_argument('-sp', '--stop_patience', default=45, type=int, + help='Number of epochs before doing early stopping') +parser.add_argument('-mn', '--model_name', + default='frustumpointnet_carpedcyc', + type=str, help='Model name based on object of interest') +parser.add_argument('-s', '--save_path', + default=os.path.join( + os.path.expanduser('~'), '.keras/paz/models'), + type=str, help='Path for writing model weights and logs') +parser.add_argument('-td', '--train_tfrec_path', + default='./frustum_dataset/', + type=str, help='Path for datasets generated') +parser.add_argument('-vd', '--val_tfrec_path', + default='./frustum_dataset/', + type=str, help='Path for datasets generated') +args = parser.parse_args() + + +def step_decay(epoch): + initial_lrate = args.learning_rate + drop = args.decay_rate + epochs_drop = 10 + lrate = initial_lrate * math.pow(drop, + math.floor((1 + epoch) / epochs_drop)) + return lrate + + +# loading data +data_loader = frustum_data_loader() + +parsed_train_dataset = data_loader.load_data(args.train_tfrec_path, + operation='train') +train_iterator = make_train_iterator(parsed_train_dataset) + +parsed_validation_dataset = data_loader.load_data(args.val_tfrec_path, + operation='validation') +val_iterator = make_train_iterator(parsed_validation_dataset) + +model_name = args.model_name + +# setting callbacks +log = CSVLogger(os.path.join(args.save_path, '%s.log' % model_name)) + +stop = EarlyStopping(monitor='val_loss', min_delta=0, patience=50, verbose=0, + mode='auto', baseline=None, restore_best_weights=False) + +lrate = LearningRateScheduler(step_decay) + +model_path = os.path.join(args.save_path, '%s_weights.hdf5' % model_name) +save = ModelCheckpoint(model_path, verbose=1, save_best_only=True, + save_weights_only=True) + +# model importing +model, _ = FrustumPointNetModel() +model.compile(optimizer=Adam(args.learning_rate), + loss={'fp_loss': lambda y_true, y_pred: y_pred}) + +# model optimization +model.fit_generator(train_iterator, epochs=args.max_num_epochs, + callbacks=[stop, log, save, lrate], + validation_data=val_iterator, + verbose=1) diff --git a/paz/models/detection/frustum_pointnet.py b/paz/models/detection/frustum_pointnet.py new file mode 100644 index 000000000..65336786f --- /dev/null +++ b/paz/models/detection/frustum_pointnet.py @@ -0,0 +1,370 @@ +import numpy as np +import tensorflow as tf +from tensorflow.keras.layers import Conv2D, Dense +from tensorflow.keras.layers import Dropout +from tensorflow.keras.layers import Input +from tensorflow.keras.layers import MaxPooling2D +from tensorflow.keras.models import Model + +from examples.frustum_pointnet.dataset_utils import g_mean_size_arr +from paz.optimization.losses.frustumpointnet_loss import FrustumPointNetLoss + + +def ModelOutputToTensor(output, IntermediateOutputs, NUM_HEADING_BIN=12, + NUM_SIZE_CLUSTER=8): + """ Parse batch output to separate tensors (added to IntermediateOutputs) + Input: + output: TF tensor in shape (B,3+2*NUM_HEADING_BIN+4*NUM_SIZE_CLUSTER) + IntermediateOutputs: dict + Output: + IntermediateOutputs: dict + + """ + batch_size = output.get_shape()[0] + center = tf.slice(output, [0, 0], [-1, 3]) + IntermediateOutputs['center_boxnet'] = center + + heading_scores = tf.slice(output, [0, 3], [-1, NUM_HEADING_BIN]) + heading_residuals_normalized = tf.slice(output, [0, 3 + NUM_HEADING_BIN], + [-1, NUM_HEADING_BIN]) + IntermediateOutputs['heading_scores'] = heading_scores # BxNUM_HEADING_BIN + IntermediateOutputs['heading_residuals_normalized'] = \ + heading_residuals_normalized # BxNUM_HEADING_BIN (-1 to 1) + IntermediateOutputs['heading_residuals'] = \ + heading_residuals_normalized * (np.pi / NUM_HEADING_BIN) + + size_scores = tf.slice(output, [0, 3 + NUM_HEADING_BIN * 2], + [-1, NUM_SIZE_CLUSTER]) # BxNUM_SIZE_CLUSTER + size_residuals_normalized = tf.slice(output, + [0, + 3 + NUM_HEADING_BIN * 2 + + NUM_SIZE_CLUSTER], + [-1, NUM_SIZE_CLUSTER * 3]) + + size_residuals_normalized = tf.reshape(size_residuals_normalized, + [batch_size, NUM_SIZE_CLUSTER, + 3]) # BxNUM_SIZE_CLUSTERx3 + IntermediateOutputs['size_scores'] = size_scores + IntermediateOutputs['size_residuals_normalized'] = size_residuals_normalized + IntermediateOutputs['size_residuals'] = \ + size_residuals_normalized * tf.expand_dims( + tf.constant(g_mean_size_arr, dtype=tf.float32), 0) + + return IntermediateOutputs + + +def ObjectPointCloudMasking(point_cloud, mask, npoints=512): + """ Gather object point clouds according to predicted masks. + Input: + point_cloud: TF tensor in shape (B,N,C) + mask: TF tensor in shape (B,N) of 0 (not pick) or 1 (pick) + npoints: int scalar, maximum number of points to keep (default: 512) + Output: + object_pc: TF tensor in shape (B,npoint,C) + indices: TF int tensor in shape (B,npoint,2) + """ + + def mask_to_indices(mask): + indices = np.zeros((mask.shape[0], npoints, 2), dtype=np.int32) + for i in range(mask.shape[0]): + pos_indices = np.where(mask[i, :] > 0.5)[0] + # skip cases when pos_indices is empty + if len(pos_indices) > 0: + if len(pos_indices) > npoints: + choice = np.random.choice(len(pos_indices), + npoints, replace=False) + else: + choice = np.random.choice(len(pos_indices), + npoints - len(pos_indices), + replace=True) + choice = np.concatenate((np.arange(len(pos_indices)), + choice)) + np.random.shuffle(choice) + indices[i, :, 1] = pos_indices[choice] + indices[i, :, 0] = i + return indices + + indices = tf.py_function(mask_to_indices, [mask], tf.int32) + object_pc = tf.gather_nd(point_cloud, indices) + return object_pc, indices + +def PointCloudTranslation(point_cloud, logits, IntermediateOutputs, + NUM_OBJECT_POINT=512, xyz_only=True): + """ Select point cloud with predicted 3D mask, + translate coordinates to the masked points centroid. + + Used from Frustum PointNet + + Input: + point_cloud: TF tensor in shape (B,N,C) + logits: TF tensor in shape (B,N,2) + IntermediateOutputs: dict + xyz_only: boolean, if True only return XYZ channels + Output: + object_point_cloud: TF tensor in shape (B,M,3) + for simplicity we only keep XYZ here + M = NUM_OBJECT_POINT as a hyper-parameter + mask_xyz_mean: TF tensor in shape (B,3) + """ + batch_size = point_cloud.get_shape()[0] + num_point = point_cloud.get_shape()[1] + mask = tf.slice(logits, [0, 0, 0], [-1, -1, 1]) < \ + tf.slice(logits, [0, 0, 1], [-1, -1, 1]) + mask = tf.cast(mask, dtype=tf.float32) # BxNx1 + mask_count = tf.tile(tf.math.reduce_sum(mask, axis=1, keepdims=True), + [1, 1, 3]) # Bx1x3 + point_cloud_xyz = tf.slice(point_cloud, [0, 0, 0], [-1, -1, 3]) # BxNx3 + mask_xyz_mean = tf.math.reduce_sum( + tf.tile(mask, [1, 1, 3]) * point_cloud_xyz, axis=1, + keepdims=True) # Bx1x3 + mask = tf.squeeze(mask, axis=[2]) # BxN + IntermediateOutputs['mask'] = mask + mask_xyz_mean = mask_xyz_mean / tf.maximum(mask_count, 1) # Bx1x3 + + # Translate to masked points' centroid + point_cloud_xyz_stage1 = point_cloud_xyz - \ + tf.tile(mask_xyz_mean, [1, num_point, 1]) + + if xyz_only: + point_cloud_stage1 = point_cloud_xyz_stage1 + else: + point_cloud_features = tf.slice(point_cloud, [0, 0, 3], [-1, -1, -1]) + point_cloud_stage1 = tf.concat( + [point_cloud_xyz_stage1, point_cloud_features], axis=-1) + num_channels = point_cloud_stage1.get_shape()[2] + + object_point_cloud, indices = ObjectPointCloudMasking(point_cloud_stage1, + mask, + NUM_OBJECT_POINT) + + object_point_cloud.set_shape([batch_size, NUM_OBJECT_POINT, num_channels]) + + return object_point_cloud, tf.squeeze(mask_xyz_mean, axis=1), \ + IntermediateOutputs + + +def InstanceSegmentationNet(point_cloud, one_hot_vector): + """Instance Segmentation network with PointNet backbone. + # Arguments + point_cloud: Tensor, Frustum Point Cloud extracted using 2D object + detection and projected pointcloud. + one_hot_vector: Tensor. One hot vector of the class to which the + frustum pointcloud belongs to. + + # Reference + - [PointNet](https://arxiv.org/abs/1612.00593) + - [Frustum PointNet](https://arxiv.org/abs/1711.08488) + """ + num_points = point_cloud.get_shape().as_list()[1] + + input = tf.expand_dims(point_cloud, 2) + + x = Conv2D(64, 1, (1, 1), 'valid', activation='relu', name='conv1_1')(input) + + x = Conv2D(64, 1, (1, 1), 'valid', activation='relu', name='conv1_2')(x) + + point_feat = Conv2D(64, 1, (1, 1), 'valid', activation='relu', + name='conv1_3')(x) + + x = Conv2D(128, 1, (1, 1), 'valid', activation='relu', + name='conv1_4')(point_feat) + + x = Conv2D(1024, 1, (1, 1), 'valid', activation='relu', + name='conv1_5')(x) + + global_feat = MaxPooling2D(pool_size=[num_points, 1], padding='VALID')(x) + + global_feat = tf.concat([global_feat, tf.expand_dims(tf.expand_dims( + one_hot_vector, 1), 1)], axis=3) + + global_feat_expand = tf.tile(global_feat, [1, num_points, 1, 1]) + + concat_feat = tf.concat(axis=3, values=[point_feat, global_feat_expand]) + + x = Conv2D(512, 1, (1, 1), 'valid', activation='relu', + name='conv1_6')(concat_feat) + + x = Conv2D(256, 1, (1, 1), 'valid', activation='relu', name='conv1_7')(x) + + x = Conv2D(128, 1, (1, 1), 'valid', activation='relu', name='conv1_8')(x) + + x = Conv2D(128, 1, (1, 1), 'valid', activation='relu', name='conv1_9')(x) + + x = Dropout(rate=0.5)(x) + + logits = Conv2D(2, 1, (1, 1), 'valid', activation=None, name='conv1_10')(x) + + logits = tf.squeeze(logits, [2]) # BxNxC + + return logits + + +def BoxEstimationNetwork(object_point_cloud, one_hot_vec, NUM_HEADING_BIN=12, + NUM_SIZE_CLUSTER=8): + """Amodal Box Regression network which provides Box parameters as output + # Arguments + object_point_cloud: Tensor. Point Cloud after performing instance + segmentation. + one_hot_vector: Tensor. One hot vector of the class to which the + pointcloud belongs to. + + # Reference + - [Frustum PointNet](https://arxiv.org/abs/1711.08488) + """ + num_point = object_point_cloud.get_shape()[1] + input = tf.expand_dims(object_point_cloud, 2) + + x = Conv2D(128, 1, (1, 1), 'valid', activation='relu', + name='conv2_1')(input) + + x = Conv2D(128, 1, (1, 1), 'valid', activation='relu', name='conv2_2')(x) + + x = Conv2D(256, 1, (1, 1), 'valid', activation='relu', name='conv2_3')(x) + + x = Conv2D(512, 1, (1, 1), 'valid', activation='relu', name='conv2_4')(x) + + x = MaxPooling2D(pool_size=[num_point, 1], padding='VALID')(x) + + x = tf.squeeze(x, axis=[1, 2]) + x = tf.concat([x, one_hot_vec], axis=1) + + x = Dense(units=512, activation='relu')(x) + x = Dense(units=256, activation='relu')(x) + + output = Dense(units=3 + NUM_HEADING_BIN * 2 + NUM_SIZE_CLUSTER * 4, + activation=None)(x) + + return output + + +def SpatialTransformerNetwork(object_point_cloud, one_hot_vec): + """A Transformer Network to project object pointcloud into space + invariant frame + # Arguments + object_point_cloud: Tensor. Point Cloud after performing instance + segmentation. + one_hot_vector: Tensor. One hot vector of the class to which the + pointcloud belongs to. + + # Reference + - [Frustum PointNet](https://arxiv.org/abs/1711.08488) + """ + num_point = object_point_cloud.get_shape()[1] + input = tf.expand_dims(object_point_cloud, 2) + + x = Conv2D(128, 1, (1, 1), 'valid', activation='relu', + name='conv3_1')(input) + + x = Conv2D(128, 1, (1, 1), 'valid', activation='relu', name='conv3_2')(x) + x = Conv2D(256, 1, (1, 1), 'valid', activation='relu', name='conv3_3')(x) + x = MaxPooling2D(pool_size=[num_point, 1], padding='VALID')(x) + + x = tf.squeeze(x, axis=[1, 2]) + x = tf.concat([x, one_hot_vec], axis=1) + + x = Dense(units=256, activation='relu')(x) + x = Dense(units=128, activation='relu')(x) + predicted_center = Dense(units=3, activation=None)(x) + + return predicted_center + + +def FrustumPointNetModel(point_cloud_shape=(1024, 3), one_hot_vec_shape=(3,), + mask_label_shape=(1024,), center_label_shape=(3,), + heading_class_label_shape=(), + heading_residual_label_shape=(), + size_class_label_shape=(), + size_residual_label_shape=(3,), batch_size=32): + """ Loss functions for 3D object detection. + Input: + Frustum_Point_Cloud: TF int32 tensor in shape (B,N) + One_hot_vector: TF int32 tensor in shape (B,N) + mask_label: TF int32 tensor in shape (B,N) + center_label: TF tensor in shape (B,3) + heading_class_label: TF int32 tensor in shape (B,) + heading_residual_label: TF tensor in shape (B,) + size_class_label: TF tensor int32 in shape (B,) + size_residual_label: TF tensor tensor in shape (B,) + Output: + Training Model: Frustum PointNet model with embedded loss layer + that is used for training + Detection Model: Frustum PointNet model that can output Bounding + box parameters + References: + - [Frustum PointNets for 3D Object Detection from RGB-D Data] + (https://arxiv.org/abs/1711.08488) + """ + IntermediateOutputs = {} + + point_cloud = Input(point_cloud_shape, name="frustum_point_cloud", + batch_size=batch_size) + one_hot_vector = Input(one_hot_vec_shape, name="one_hot_vec", + batch_size=batch_size) + mask_label = Input(mask_label_shape, name="segmentation_label", + batch_size=batch_size) + center_label = Input(center_label_shape, name="box3d_center", + batch_size=batch_size) + heading_class_label = Input(heading_class_label_shape, name="angle_class", + batch_size=batch_size) + heading_residual_label = Input(heading_residual_label_shape, + name="angle_residual", batch_size=batch_size) + size_class_label = Input(size_class_label_shape, name="size_class", + batch_size=batch_size) + size_residual_label = Input(size_residual_label_shape, name="size_residual", + batch_size=batch_size) + + logits = InstanceSegmentationNet(point_cloud, one_hot_vector) # bs,n,2 + IntermediateOutputs['mask_logits'] = logits + + # Mask Point Centroid + object_point_cloud_xyz, mask_xyz_mean, IntermediateOutputs = \ + PointCloudTranslation(point_cloud, logits, IntermediateOutputs) + + # T-Net + center_delta = SpatialTransformerNetwork(object_point_cloud_xyz, + one_hot_vector) # (32,3) + + stage1_center = center_delta + mask_xyz_mean # Bx3 + IntermediateOutputs['stage1_center'] = stage1_center + # Get object point cloud in object coordinate + object_point_cloud_xyz_new = object_point_cloud_xyz - tf.expand_dims( + center_delta, 1) + + # 3D Box Estimation + box_pred = BoxEstimationNetwork(object_point_cloud_xyz_new, one_hot_vector) + + IntermediateOutputs = ModelOutputToTensor(box_pred, IntermediateOutputs) + IntermediateOutputs['center'] = IntermediateOutputs['center_boxnet'] + \ + stage1_center # Bx3 + + logits = IntermediateOutputs['mask_logits'] + heading_scores = IntermediateOutputs['heading_scores'] # BxNUM_HEADING_BIN + heading_residual = IntermediateOutputs['heading_residuals'] + size_scores = IntermediateOutputs['size_scores'] + size_residual = IntermediateOutputs['size_residuals'] + box3d_center = IntermediateOutputs['center'] + + loss = FrustumPointNetLoss().TotalLoss([mask_label, center_label, + heading_class_label, + heading_residual_label, + size_class_label, + size_residual_label, + IntermediateOutputs]) + + training_model = Model([point_cloud, one_hot_vector, mask_label, + center_label, heading_class_label, + heading_residual_label, size_class_label, + size_residual_label], loss, + name='f_pointnet_train') + + det_model = Model(inputs=[point_cloud, one_hot_vector], + outputs=[logits, box3d_center, heading_scores, + heading_residual, size_scores, size_residual], + name='f_pointnet_inference') + training_model.summary() + + return training_model, det_model + + +if __name__ == '__main__': + FrustumPointNetModel() diff --git a/paz/optimization/losses/frustumpointnet_loss.py b/paz/optimization/losses/frustumpointnet_loss.py new file mode 100644 index 000000000..57601c657 --- /dev/null +++ b/paz/optimization/losses/frustumpointnet_loss.py @@ -0,0 +1,243 @@ +import numpy as np +import tensorflow as tf +from tensorflow.keras.losses import Huber +from examples.frustum_pointnet.dataset_utils import NUM_HEADING_BIN, \ + g_mean_size_arr, NUM_SIZE_CLUSTER + +def ExtractBox3DCornersHelper(centers, headings, sizes): + """ TF layer. + Inputs: + center: (B,3) + heading_residuals: (B,NH) + size_residuals: (B,NS,3) + Outputs: + box3d_corners: (B,NH,NS,8,3) tensor + """ + batch_size = centers.get_shape()[0] + length = tf.slice(sizes, [0, 0], [-1, 1]) # (N,1) + width = tf.slice(sizes, [0, 1], [-1, 1]) # (N,1) + height = tf.slice(sizes, [0, 2], [-1, 1]) # (N,1) + x_corners = tf.concat([length / 2, length / 2, -length / 2, -length / 2, + length / 2, length / 2, -length / 2, -length / 2], + axis=1) # (N,8) + y_corners = tf.concat([height / 2, height / 2, height / 2, height / 2, + -height / 2, -height / 2, -height / 2, -height / 2], + axis=1) # (N,8) + z_corners = tf.concat([width / 2, -width / 2, -width / 2, width / 2, + width / 2, -width / 2, -width / 2, width / 2], + axis=1) # (N,8) + corners = tf.concat([tf.expand_dims(x_corners, 1), + tf.expand_dims(y_corners, 1), + tf.expand_dims(z_corners, 1)], + axis=1) # (N,3,8) + cosine_value = tf.cos(headings) + sine_value = tf.sin(headings) + ones = tf.ones([batch_size], dtype=tf.float32) + zeros = tf.zeros([batch_size], dtype=tf.float32) + row1 = tf.stack([cosine_value, zeros, sine_value], axis=1) # (N,3) + row2 = tf.stack([zeros, ones, zeros], axis=1) + row3 = tf.stack([-sine_value, zeros, cosine_value], axis=1) + R = tf.concat([tf.expand_dims(row1, 1), tf.expand_dims(row2, 1), + tf.expand_dims(row3, 1)], axis=1) # (N,3,3) + corners_3d = tf.matmul(R, corners) # (N,3,8) + corners_3d += tf.tile(tf.expand_dims(centers, 2), [1, 1, 8]) # (N,3,8) + corners_3d = tf.transpose(corners_3d, perm=[0, 2, 1]) # (N,8,3) + return corners_3d + + +def ExtractBox3DCorners(center, heading_residuals, size_residuals): + """ TF layer. + Inputs: + center: (B,3) + heading_residuals: (B,NH) + size_residuals: (B,NS,3) + Outputs: + box3d_corners: (B,NH,NS,8,3) tensor + """ + batch_size = center.get_shape()[0] + heading_bin_centers = tf.constant(np.arange(0, 2 * np.pi, 2 * np.pi / + NUM_HEADING_BIN), + dtype=tf.float32) # (NH,) + headings = heading_residuals + tf.expand_dims(heading_bin_centers, 0) + + mean_sizes = tf.expand_dims(tf.constant(g_mean_size_arr, dtype=tf.float32), + 0) + size_residuals # (B,NS,1) + sizes = mean_sizes + size_residuals # (B,NS,3) + sizes = tf.tile(tf.expand_dims(sizes, 1), [1, NUM_HEADING_BIN, 1, 1]) + headings = tf.tile(tf.expand_dims(headings, -1), [1, 1, NUM_SIZE_CLUSTER]) + centers = tf.tile(tf.expand_dims(tf.expand_dims(center, 1), 1), + [1, NUM_HEADING_BIN, NUM_SIZE_CLUSTER, 1]) + + N = batch_size * NUM_HEADING_BIN * NUM_SIZE_CLUSTER + corners_3d = ExtractBox3DCornersHelper(tf.reshape(centers, [N, 3]), + tf.reshape(headings, [N]), + tf.reshape(sizes, [N, 3])) + + return tf.reshape(corners_3d, + [batch_size, NUM_HEADING_BIN, NUM_SIZE_CLUSTER, 8, 3]) + + +class FrustumPointNetLoss: + """ Loss functions for 3D object detection. + Input: + args: dict + mask_label: TF int32 tensor in shape (B,N) + center_label: TF tensor in shape (B,3) + heading_class_label: TF int32 tensor in shape (B,) + heading_residual_label: TF tensor in shape (B,) + size_class_label: TF tensor int32 in shape (B,) + size_residual_label: TF tensor tensor in shape (B,) + end_points: dict, outputs from our model + corner_loss_weight: float scalar + box_loss_weight: float scalar + + Output: + total_loss: TF scalar tensor + the total_loss is also added to the losses collection + + References + - [Frustum PointNets for 3D Object Detection from RGB-D Data] + (https://arxiv.org/abs/1711.08488) + """ + + def __init__(self, corner_loss_weight=10.0, box_loss_weight=1.0, + mask_weight=1.0): + self.corner_loss_weight = corner_loss_weight + self.box_loss_weight = box_loss_weight + self.mask_weight = mask_weight + + def _huber_loss(self, y_true, y_pred): + """ + Huber Loss calculation for regressed output + @param y_true: Label of the loss + @param y_pred: Prediction from Frustum-PointNet + @return: Loss value + """ + h = Huber() + return h(y_true, y_pred) + + def _cross_entropy_loss(self, y_true, y_pred): + """ + Cross Entropy Loss calculation for classification output + @param y_true: Label of the loss + @param y_pred: Prediction from Frustum-PointNet + @return: Loss value + """ + loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits + (logits=y_pred, labels=tf.cast(y_true, tf.int64))) + return loss + + def to_one_hot(self, input_tensor, depth): + OneHot = tf.one_hot(tf.cast(input_tensor, tf.int64), depth=depth, + on_value=1, off_value=0, axis=-1) + return OneHot + + def TotalLoss(self, args): + mask_label, center_label, heading_class_label, heading_residual_label, \ + size_class_label, size_residual_label, end_points = args[0], args[1], \ + args[2], args[3], \ + args[4], args[5], \ + args[6] + + mask_loss = self._cross_entropy_loss(y_true=tf.cast(mask_label, + tf.int64), + y_pred=end_points['mask_logits']) + + center_loss = self._huber_loss(y_true=center_label, + y_pred=end_points['center']) + + stage1_center_loss = self._huber_loss(y_true=center_label, + y_pred=end_points[ + 'stage1_center']) + + heading_class_loss = self._cross_entropy_loss( + y_true=tf.cast(heading_class_label, tf.int64), + y_pred=end_points['heading_scores']) + + heading_cls_onehot = self.to_one_hot(heading_class_label, + NUM_HEADING_BIN) + + heading_residual_normalized_label = (heading_residual_label / + (np.pi / NUM_HEADING_BIN)) + + heading_residual_normalized_loss = self._huber_loss( + y_true=heading_residual_normalized_label, y_pred=tf.reduce_sum( + end_points['heading_residuals_normalized'] * + tf.cast(heading_cls_onehot, dtype=tf.float32), axis=1)) + + size_class_loss = self._cross_entropy_loss( + y_pred=end_points['size_scores'], y_true=tf.cast(size_class_label, + tf.int64)) + + size_cls_onehot = self.to_one_hot(tf.cast(size_class_label, tf.int64), + NUM_SIZE_CLUSTER) + + size_cls_onehot_tiled = tf.tile(tf.expand_dims(tf.cast(size_cls_onehot, + dtype=tf.float32) + ,-1), [1, 1, 3]) + predicted_size_residual_normalized = tf.reduce_sum( + end_points['size_residuals_normalized'] * size_cls_onehot_tiled, + axis=[1]) + + mean_size_arr_expand = tf.expand_dims(tf.constant(g_mean_size_arr, + dtype=tf.float32), 0) + + mean_size_label = tf.reduce_sum(size_cls_onehot_tiled * + mean_size_arr_expand, axis=[1]) + + size_residual_label_normalized = size_residual_label / mean_size_label + + size_residual_normalized_loss = self._huber_loss( + y_true=size_residual_label_normalized, + y_pred=predicted_size_residual_normalized) + + corners_3d = ExtractBox3DCorners(end_points['center'], + end_points['heading_residuals'], + end_points['size_residuals']) + + gt_mask = tf.tile(tf.expand_dims(heading_cls_onehot, 2), + [1, 1, NUM_SIZE_CLUSTER]) * tf.tile(tf.expand_dims( + size_cls_onehot, 1), [1, NUM_HEADING_BIN, 1]) + + corners_3d_pred = tf.reduce_sum(tf.cast(tf.expand_dims( + tf.expand_dims(gt_mask, -1), -1), dtype=tf.float32) * corners_3d, + axis=[1, 2]) + + heading_bin_centers = tf.constant(np.arange(0, 2 * np.pi, 2 * np.pi / + NUM_HEADING_BIN), + dtype=tf.float32) # (NH,) + + heading_label = tf.expand_dims(heading_residual_label, 1) + \ + tf.expand_dims(heading_bin_centers, 0) # (B,NH) + + heading_label = tf.reduce_sum(tf.cast(heading_cls_onehot, + dtype=tf.float32) * heading_label, + 1) + + mean_sizes = tf.expand_dims(tf.constant(g_mean_size_arr, + dtype=tf.float32),0) # (1,NS,3) + size_label = mean_sizes + tf.expand_dims(size_residual_label, 1) + size_label = tf.reduce_sum(tf.expand_dims(tf.cast(size_cls_onehot, + dtype=tf.float32), + -1) * size_label, axis=[1]) + + corners_3d_gt = ExtractBox3DCornersHelper(center_label, heading_label, + size_label) # (B,8,3) + corners_3d_gt_flip = ExtractBox3DCornersHelper(center_label, + heading_label + + np.pi, size_label) + + corners_3D_loss = self._huber_loss(y_true=corners_3d_gt, + y_pred=corners_3d_pred) + corners_3D_flip_loss = self._huber_loss(y_true=corners_3d_gt_flip, + y_pred=corners_3d_pred) + + corners_loss = tf.minimum(corners_3D_loss, corners_3D_flip_loss) + + total_loss = mask_loss * self.mask_weight + self.box_loss_weight * \ + (center_loss + heading_class_loss + size_class_loss + + heading_residual_normalized_loss * 20 + + size_residual_normalized_loss * 20 + stage1_center_loss + + self.corner_loss_weight * corners_loss) + + return total_loss