oarriaga · jaswanthbjk · Jan 2, 2021 · Jan 29, 2021 · Jan 29, 2021 · Jan 29, 2021
diff --git a/examples/frustum_pointnet/dataset_utils.py 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]]
diff --git a/examples/frustum_pointnet/frustum_loader.py 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
@@ -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)