├── LICENSE
├── README.md
├── data
├── README.md
└── data_prepare.py
├── data_gen.py
├── docs
└── cover.png
├── eval.py
├── eval_cmp.py
├── models
├── __init__.py
├── baseline.py
├── bmwnet.py
├── dronet.py
├── flowdrivenet.py
├── pilotnet.py
└── pointnet.py
├── tf_ops
├── 3d_interpolation
│ ├── __pycache__
│ │ └── tf_interpolate.cpython-35.pyc
│ ├── interpolate.cpp
│ ├── tf_interpolate.cpp
│ ├── tf_interpolate.py
│ ├── tf_interpolate_compile.sh
│ ├── tf_interpolate_op_test.py
│ ├── tf_interpolate_so.so
│ └── visu_interpolation.py
├── grouping
│ ├── __pycache__
│ │ └── tf_grouping.cpython-35.pyc
│ ├── test
│ │ ├── compile.sh
│ │ ├── query_ball_point.cpp
│ │ ├── query_ball_point.cu
│ │ ├── query_ball_point_block.cu
│ │ ├── query_ball_point_grid.cu
│ │ ├── selection_sort.cpp
│ │ ├── selection_sort.cu
│ │ └── selection_sort_const.cu
│ ├── tf_grouping.cpp
│ ├── tf_grouping.py
│ ├── tf_grouping_compile.sh
│ ├── tf_grouping_g.cu
│ ├── tf_grouping_g.cu.o
│ ├── tf_grouping_op_test.py
│ └── tf_grouping_so.so
└── sampling
│ ├── __pycache__
│ └── tf_sampling.cpython-35.pyc
│ ├── tf_sampling.cpp
│ ├── tf_sampling.py
│ ├── tf_sampling_compile.sh
│ ├── tf_sampling_g.cu
│ ├── tf_sampling_g.cu.o
│ └── tf_sampling_so.so
├── train.py
├── train_cmp.py
├── train_multi_gpus.py
└── utils
├── README.md
├── __init__.py
├── compile_render_balls_so.sh
├── pc_util.py
├── pointnet_util.py
├── provider.py
├── render_balls_so.cpp
├── render_balls_so.so
├── show3d_balls.py
└── tf_util.py
/LICENSE:
--------------------------------------------------------------------------------
1 | Apache License
2 | Version 2.0, January 2004
3 | http://www.apache.org/licenses/
4 |
5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
6 |
7 | 1. Definitions.
8 |
9 | "License" shall mean the terms and conditions for use, reproduction,
10 | and distribution as defined by Sections 1 through 9 of this document.
11 |
12 | "Licensor" shall mean the copyright owner or entity authorized by
13 | the copyright owner that is granting the License.
14 |
15 | "Legal Entity" shall mean the union of the acting entity and all
16 | other entities that control, are controlled by, or are under common
17 | control with that entity. For the purposes of this definition,
18 | "control" means (i) the power, direct or indirect, to cause the
19 | direction or management of such entity, whether by contract or
20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the
21 | outstanding shares, or (iii) beneficial ownership of such entity.
22 |
23 | "You" (or "Your") shall mean an individual or Legal Entity
24 | exercising permissions granted by this License.
25 |
26 | "Source" form shall mean the preferred form for making modifications,
27 | including but not limited to software source code, documentation
28 | source, and configuration files.
29 |
30 | "Object" form shall mean any form resulting from mechanical
31 | transformation or translation of a Source form, including but
32 | not limited to compiled object code, generated documentation,
33 | and conversions to other media types.
34 |
35 | "Work" shall mean the work of authorship, whether in Source or
36 | Object form, made available under the License, as indicated by a
37 | copyright notice that is included in or attached to the work
38 | (an example is provided in the Appendix below).
39 |
40 | "Derivative Works" shall mean any work, whether in Source or Object
41 | form, that is based on (or derived from) the Work and for which the
42 | editorial revisions, annotations, elaborations, or other modifications
43 | represent, as a whole, an original work of authorship. For the purposes
44 | of this License, Derivative Works shall not include works that remain
45 | separable from, or merely link (or bind by name) to the interfaces of,
46 | the Work and Derivative Works thereof.
47 |
48 | "Contribution" shall mean any work of authorship, including
49 | the original version of the Work and any modifications or additions
50 | to that Work or Derivative Works thereof, that is intentionally
51 | submitted to Licensor for inclusion in the Work by the copyright owner
52 | or by an individual or Legal Entity authorized to submit on behalf of
53 | the copyright owner. For the purposes of this definition, "submitted"
54 | means any form of electronic, verbal, or written communication sent
55 | to the Licensor or its representatives, including but not limited to
56 | communication on electronic mailing lists, source code control systems,
57 | and issue tracking systems that are managed by, or on behalf of, the
58 | Licensor for the purpose of discussing and improving the Work, but
59 | excluding communication that is conspicuously marked or otherwise
60 | designated in writing by the copyright owner as "Not a Contribution."
61 |
62 | "Contributor" shall mean Licensor and any individual or Legal Entity
63 | on behalf of whom a Contribution has been received by Licensor and
64 | subsequently incorporated within the Work.
65 |
66 | 2. Grant of Copyright License. Subject to the terms and conditions of
67 | this License, each Contributor hereby grants to You a perpetual,
68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable
69 | copyright license to reproduce, prepare Derivative Works of,
70 | publicly display, publicly perform, sublicense, and distribute the
71 | Work and such Derivative Works in Source or Object form.
72 |
73 | 3. Grant of Patent License. Subject to the terms and conditions of
74 | this License, each Contributor hereby grants to You a perpetual,
75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable
76 | (except as stated in this section) patent license to make, have made,
77 | use, offer to sell, sell, import, and otherwise transfer the Work,
78 | where such license applies only to those patent claims licensable
79 | by such Contributor that are necessarily infringed by their
80 | Contribution(s) alone or by combination of their Contribution(s)
81 | with the Work to which such Contribution(s) was submitted. If You
82 | institute patent litigation against any entity (including a
83 | cross-claim or counterclaim in a lawsuit) alleging that the Work
84 | or a Contribution incorporated within the Work constitutes direct
85 | or contributory patent infringement, then any patent licenses
86 | granted to You under this License for that Work shall terminate
87 | as of the date such litigation is filed.
88 |
89 | 4. Redistribution. You may reproduce and distribute copies of the
90 | Work or Derivative Works thereof in any medium, with or without
91 | modifications, and in Source or Object form, provided that You
92 | meet the following conditions:
93 |
94 | (a) You must give any other recipients of the Work or
95 | Derivative Works a copy of this License; and
96 |
97 | (b) You must cause any modified files to carry prominent notices
98 | stating that You changed the files; and
99 |
100 | (c) You must retain, in the Source form of any Derivative Works
101 | that You distribute, all copyright, patent, trademark, and
102 | attribution notices from the Source form of the Work,
103 | excluding those notices that do not pertain to any part of
104 | the Derivative Works; and
105 |
106 | (d) If the Work includes a "NOTICE" text file as part of its
107 | distribution, then any Derivative Works that You distribute must
108 | include a readable copy of the attribution notices contained
109 | within such NOTICE file, excluding those notices that do not
110 | pertain to any part of the Derivative Works, in at least one
111 | of the following places: within a NOTICE text file distributed
112 | as part of the Derivative Works; within the Source form or
113 | documentation, if provided along with the Derivative Works; or,
114 | within a display generated by the Derivative Works, if and
115 | wherever such third-party notices normally appear. The contents
116 | of the NOTICE file are for informational purposes only and
117 | do not modify the License. You may add Your own attribution
118 | notices within Derivative Works that You distribute, alongside
119 | or as an addendum to the NOTICE text from the Work, provided
120 | that such additional attribution notices cannot be construed
121 | as modifying the License.
122 |
123 | You may add Your own copyright statement to Your modifications and
124 | may provide additional or different license terms and conditions
125 | for use, reproduction, or distribution of Your modifications, or
126 | for any such Derivative Works as a whole, provided Your use,
127 | reproduction, and distribution of the Work otherwise complies with
128 | the conditions stated in this License.
129 |
130 | 5. Submission of Contributions. Unless You explicitly state otherwise,
131 | any Contribution intentionally submitted for inclusion in the Work
132 | by You to the Licensor shall be under the terms and conditions of
133 | this License, without any additional terms or conditions.
134 | Notwithstanding the above, nothing herein shall supersede or modify
135 | the terms of any separate license agreement you may have executed
136 | with Licensor regarding such Contributions.
137 |
138 | 6. Trademarks. This License does not grant permission to use the trade
139 | names, trademarks, service marks, or product names of the Licensor,
140 | except as required for reasonable and customary use in describing the
141 | origin of the Work and reproducing the content of the NOTICE file.
142 |
143 | 7. Disclaimer of Warranty. Unless required by applicable law or
144 | agreed to in writing, Licensor provides the Work (and each
145 | Contributor provides its Contributions) on an "AS IS" BASIS,
146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
147 | implied, including, without limitation, any warranties or conditions
148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
149 | PARTICULAR PURPOSE. You are solely responsible for determining the
150 | appropriateness of using or redistributing the Work and assume any
151 | risks associated with Your exercise of permissions under this License.
152 |
153 | 8. Limitation of Liability. In no event and under no legal theory,
154 | whether in tort (including negligence), contract, or otherwise,
155 | unless required by applicable law (such as deliberate and grossly
156 | negligent acts) or agreed to in writing, shall any Contributor be
157 | liable to You for damages, including any direct, indirect, special,
158 | incidental, or consequential damages of any character arising as a
159 | result of this License or out of the use or inability to use the
160 | Work (including but not limited to damages for loss of goodwill,
161 | work stoppage, computer failure or malfunction, or any and all
162 | other commercial damages or losses), even if such Contributor
163 | has been advised of the possibility of such damages.
164 |
165 | 9. Accepting Warranty or Additional Liability. While redistributing
166 | the Work or Derivative Works thereof, You may choose to offer,
167 | and charge a fee for, acceptance of support, warranty, indemnity,
168 | or other liability obligations and/or rights consistent with this
169 | License. However, in accepting such obligations, You may act only
170 | on Your own behalf and on Your sole responsibility, not on behalf
171 | of any other Contributor, and only if You agree to indemnify,
172 | defend, and hold each Contributor harmless for any liability
173 | incurred by, or claims asserted against, such Contributor by reason
174 | of your accepting any such warranty or additional liability.
175 |
176 | END OF TERMS AND CONDITIONS
177 |
178 | APPENDIX: How to apply the Apache License to your work.
179 |
180 | To apply the Apache License to your work, attach the following
181 | boilerplate notice, with the fields enclosed by brackets "[]"
182 | replaced with your own identifying information. (Don't include
183 | the brackets!) The text should be enclosed in the appropriate
184 | comment syntax for the file format. We also recommend that a
185 | file or class name and description of purpose be included on the
186 | same "printed page" as the copyright notice for easier
187 | identification within third-party archives.
188 |
189 | Copyright [yyyy] [name of copyright owner]
190 |
191 | Licensed under the Apache License, Version 2.0 (the "License");
192 | you may not use this file except in compliance with the License.
193 | You may obtain a copy of the License at
194 |
195 | http://www.apache.org/licenses/LICENSE-2.0
196 |
197 | Unless required by applicable law or agreed to in writing, software
198 | distributed under the License is distributed on an "AS IS" BASIS,
199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
200 | See the License for the specific language governing permissions and
201 | limitations under the License.
202 |
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 |
8 | # FlowDriveNet
9 |
10 | ## 1. Introduction
11 |

12 |
13 | Learning driving policies using an end-to-end network has been proved a promising solution for autonomous driving. Due to the lack of a benchmark driver behavior dataset that contains both the visual and the LiDAR data, existing works solely focus on learning driving from visual sensors. Besides, most works are limited to predict steering angle yet neglect the more challenging vehicle speed control problem.
14 |
15 | In this project, we propose release our __Code__ and __Dataset__ for training our FlowDriveNet, which is a novel end-to-end network and takes advantages of sequential visual data and LiDAR data jointly to predict steering angle and vehicle speed.
16 |
17 | ## 2. Requirements
18 | * Python 3.x
19 | * Tensorflow 1.x.x
20 | * Python Libraries: numpy, scipy and __laspy__
21 |
22 | ## 3. Dataset
23 | ### 3.1 Udacity CH2-LiDAR
24 | The dataset used in this project is created from the Udacity `CH2` dataset, which is a popular open source driving dataset and can be used for the vision-based driving learning methods. However, the original dataset does not contain LiDAR data, then we extract the raw point cloud data and remove distortions from the original ROS bag file.
25 |
26 | **Pipeline for creating this data:**
27 | - Extracting images (center, right, left) and corresponding label file from bag file. (6311,6311,6288)
28 | - Extracting Point Cloud data from bag file and save to seperate pcd file. (3096)
29 | - fix calibration file
30 | - covert raw velodyne point packet topic to PointCloud2 topic
31 | - save to pcd file using ros node
32 | - Registering the point cloud with images and labels using timestamps.(3096)
33 |
34 | **Related tools:**
35 | - Data extraction tool: [udacity launch](https://github.com/wsustcid/self-driving-car/tree/master/datasets/udacity_launch)
36 | - Optical Flow and Point Flow extraction tool: './data/data_prepare.py' for more implementation details.
37 |
38 |
39 | ## 4. Train & Evaluation
40 | ```python
41 | # single GPU
42 | python train.py --data_root xx --input_config xx ...
43 |
44 | # Multiple GPU
45 | python train_multi-gpus.py
46 |
47 | # Evaluation
48 | python eval.py
49 | ```
50 |
51 | ## 5. Citation
52 | ```
53 | @article{wang,
54 | title={FlowDriveNet: An End-to-End Network for Learning Driving Policies from Image Optical Flow and LiDAR Point Flow},
55 | author={Shuai Wang, Jiahu Qin, Menglin Li and Yaonan Wang},
56 | journal={IEEE International Conference on Robotics and Automation (ICRA)},
57 | year={2021}
58 | }
59 | ```
60 |
61 | ## 6. License
62 | This Project is released under the [Apache licenes](LICENSE).
63 |
64 |
--------------------------------------------------------------------------------
/data/README.md:
--------------------------------------------------------------------------------
1 |
9 | # DATASET
10 |
11 | ## Udacity CH2-LiDAR
12 |
13 | ```python
14 | CH2 folder tree:
15 | └── CH2_001 # test set
16 | │ ├── center
17 | │ ├── center.csv
18 | │ ├── left
19 | │ ├── left.csv
20 | │ ├── points
21 | │ ├── right
22 | │ ├── right.csv
23 | └── CH2_002 # train set
24 | ├── HMB_1
25 | │ ├── center
26 | │ ├── center.csv
27 | │ ├── left
28 | │ ├── left.csv
29 | │ ├── points
30 | │ ├── right
31 | │ ├── right.csv
32 | ├── HMB_2; _4, _5, _6
33 | ```
34 |
35 | ## LiVi-Set
--------------------------------------------------------------------------------
/data/data_prepare.py:
--------------------------------------------------------------------------------
1 | '''
2 | @Author: Shuai Wang
3 | @Github: https://github.com/wsustcid
4 | @Version: 1.0.0
5 | @Date: 2020-09-07 21:08:55
6 | @LastEditTime: 2020-09-09 15:19:37
7 | @Description:
8 | '''
9 | import os, sys
10 | import time
11 | import numpy as np
12 | base_dir = os.path.dirname(os.path.abspath(__file__))
13 | sys.path.append(os.path.join(base_dir, '../tools/'))
14 |
15 | from label_loading import load_image_path, load_image_path_all, load_clouds_path
16 | from data_processor import ImageProcessor, CloudProcessor
17 |
18 | def create_optical_flow(data_root, set_type, dtype):
19 | '''
20 | Processing Procedure:
21 | 1. Load two consecutive image with grayscale format (uint8,480,640)
22 | 2. compute the optical flow (float32,480,640)
23 | 3. save the optical flow to bin file
24 | '''
25 | # load image path
26 | if set_type=='train':
27 | data_path = os.path.join(data_root, 'CH2_002')
28 | images_path = load_image_path_all(data_path, dtype)
29 | elif set_type == 'test':
30 | HMB_path = os.path.join(data_root, "CH2_001")
31 | images_path = load_image_path(HMB_path, dtype)
32 |
33 | imp = ImageProcessor()
34 | start = time.time()
35 | for i in range(len(images_path)):
36 | # load image
37 | if i == 0:
38 | pre_img_path = images_path[0]
39 | pre_img = imp.load_image(pre_img_path)
40 |
41 | end_img_path = images_path[i]
42 | end_img = imp.load_image(end_img_path)
43 |
44 | # compute flow
45 | flow = imp.compute_flow(pre_img, end_img)
46 |
47 | # save to bin
48 | flow_name = end_img_path.split('/')[-1]
49 | bin_path = os.path.join(os.path.dirname(end_img_path), '../'+dtype+'_bin', flow_name[:-4]+'.bin')
50 | imp.save_flow(flow, bin_path)
51 |
52 | pre_img = end_img
53 |
54 | end = time.time()
55 | print("Procedure Duration: {} s".format(end-start))
56 |
57 |
58 | def create_point_flow(HMB_path):
59 | """
60 | Input: the HMB path
61 | Pipeline:
62 | 1. load point cloud from pcd file
63 | 2. cloud process
64 | 3. compute flow and save to .bin file
65 | """
66 | start = time.time()
67 |
68 | # load cloud path
69 | clouds_path = load_clouds_path(HMB_path)
70 |
71 | clp = CloudProcessor()
72 |
73 | for i in range(len(clouds_path)):
74 | if i == 0:
75 | cloud_pre_path = clouds_path[0]
76 | cloud_pre = o3d.io.read_point_cloud(cloud_pre_path)
77 | cloud_pre = clp.cloud_process(cloud_pre)
78 |
79 | cloud_path = clouds_path[i]
80 | cloud = o3d.io.read_point_cloud(cloud_path)
81 | cloud = clp.cloud_process(cloud)
82 |
83 | # save to bin
84 | cloud_name = cloud_path.split('/')[-1]
85 | bin_path = os.path.join(os.path.dirname(cloud_path), '../points_bin',cloud_name[:-4]+'.bin')
86 | clp.save_cloud_flow(cloud, cloud_pre, bin_path)
87 |
88 | cloud_pre = o3d.geometry.PointCloud()
89 | cloud_pre.points = cloud.points
90 |
91 | end = time.time()
92 | print("Procedure Duration: {} s".format(end-start))
93 |
94 | if __name__ == '__main__':
95 | data_root = '/media/ubuntu16/Documents/Datasets/Udacity/CH2'
96 |
97 | create_optical_flow(data_root, set_type='train', dtype='right')
98 |
99 |
100 |
--------------------------------------------------------------------------------
/docs/cover.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/wsustcid/FlowDriveNet/3604495269ae45e5b43964046104f685ec66e383/docs/cover.png
--------------------------------------------------------------------------------
/eval.py:
--------------------------------------------------------------------------------
1 | '''
2 | @Author: Shuai Wang
3 | @Github: https://github.com/wsustcid
4 | @Version: 1.0.0
5 | @Date: 2020-09-11 23:42:23
6 | @LastEditTime: 2020-10-13 22:32:20
7 | '''
8 |
9 | import os
10 | import sys
11 | import argparse
12 | from datetime import datetime
13 | import time
14 | from tqdm import tqdm
15 | import time
16 |
17 | import numpy as np
18 | import tensorflow as tf
19 |
20 | base_dir = os.path.dirname(os.path.abspath(__file__))
21 | sys.path.append(base_dir)
22 | from utils.tf_util import log_string
23 |
24 | from data_gen import DataLoader
25 | from models.flowdrivenet import FlowDriveNet
26 |
27 | parser = argparse.ArgumentParser()
28 | parser.add_argument('--data_root', default='/media/ubuntu16/Documents/Datasets/Udacity/CH2',
29 | help='data_root path [default: local path]')
30 | parser.add_argument('--input_cfg', default='GRAY',
31 | help='Input type: GRAY, GRAYF, GRAYF-T, XYZ, XYZF, XYZF-T, GRAYF-XYZF-T')
32 | parser.add_argument('--model_cfg', default='VFE',
33 | help='Model type: VFE, VFE-TFP, PFE, PFE-TFP, VFE-PFE-TFP')
34 | parser.add_argument('--height', type=int, default=200, help='img height')
35 | parser.add_argument('--width', type=int, default=200, help='img width')
36 | parser.add_argument('--seq_len', type=int, default=5, help='sel length')
37 | parser.add_argument('--aug_cfg', default='None', help='None, IA, RP, SC, BA, BS')
38 | #parser.add_argument('--use_side_cam', default=False, action='store_true')
39 | parser.add_argument('--num_point', type=int, default=10000, help='Point N')
40 | parser.add_argument('--log_dir', default='test',
41 | help='Log dir [default: test]')
42 | parser.add_argument('--batch_size', type=int, default=1,
43 | help='Batch Size during training [default: 16]')
44 | parser.add_argument('--decay_steps', type=int, default=300000,
45 | help='Decay step for lr decay [default: 200000]') # decay_steps = n_train * epochs
46 | parser.add_argument('--decay_rate', type=float, default=0.7,
47 | help='Decay rate for lr decay [default: 0.7]')
48 | parser.add_argument('--model_file', default='/media/ubuntu16/F/FlowDriveNet/logs/VFE/gray_base/model_best.ckpt',
49 | help='the model path to be evaluated')
50 |
51 |
52 | FLAGS = parser.parse_args()
53 |
54 | BATCH_SIZE = FLAGS.batch_size
55 |
56 | log_dir = os.path.join(base_dir, 'logs', FLAGS.log_dir)
57 | os.makedirs(log_dir, exist_ok=True)
58 | test_log_dir = os.path.join(log_dir, 'log_test.txt')
59 | log_string(test_log_dir, str(FLAGS)+'\n')
60 |
61 | #
62 | dataloader = DataLoader(FLAGS.data_root, FLAGS.input_cfg,
63 | FLAGS.height, FLAGS.width,
64 | FLAGS.seq_len,
65 | FLAGS.num_point,
66 | FLAGS.aug_cfg)
67 | model = FlowDriveNet(FLAGS.input_cfg, FLAGS.model_cfg,
68 | FLAGS.height, FLAGS.width, FLAGS.seq_len, FLAGS.num_point)
69 |
70 | def get_bn_decay(batch):
71 | bn_momentum = tf.train.exponential_decay(
72 | 0.5,
73 | batch*BATCH_SIZE,
74 | float(FLAGS.decay_steps),
75 | 0.5,
76 | staircase=True)
77 | bn_decay = tf.minimum(0.99, 1 - bn_momentum)
78 | return bn_decay
79 |
80 | def eval():
81 | with tf.Graph().as_default():
82 | image_pl, points_pl, _ = model.get_inputs_pl(BATCH_SIZE)
83 | is_training_pl = tf.placeholder(tf.bool, shape=())
84 | # define global_step; optimizer will increase it in every training loop
85 | batch = tf.get_variable('batch', [],
86 | initializer=tf.constant_initializer(0),
87 | trainable=False)
88 | bn_decay = get_bn_decay(batch)
89 |
90 | pred = model.get_model(image_pl, points_pl, is_training_pl, bn_decay)
91 |
92 | # Create a session
93 | config = tf.ConfigProto()
94 | config.gpu_options.allow_growth = True
95 | config.allow_soft_placement = True
96 | config.log_device_placement = False
97 | sess = tf.Session(config=config)
98 |
99 | # Init variables
100 | init = tf.global_variables_initializer()
101 | sess.run(init)
102 |
103 | # restore model
104 | saver = tf.train.Saver()
105 | saver.restore(sess, FLAGS.model_file)
106 |
107 | # save all tensor
108 | ops = {'image_pl': image_pl,
109 | 'points_pl': points_pl,
110 | 'is_training_pl': is_training_pl,
111 | 'pred': pred}
112 |
113 | ## evaluation
114 | is_training = False
115 | num_batches = dataloader.num_test // BATCH_SIZE
116 | rmse_angle_sum = 0.0
117 | rmse_speed_sum = 0.0
118 | result_all = np.zeros((0,4)) # pred_a, pred_s, label_a, label_s
119 |
120 | time_sum = 0.0
121 | for i in tqdm(range(num_batches)):
122 | X_image_batch, X_cloud_batch, y_batch = dataloader.load_test_batch(BATCH_SIZE)
123 |
124 | feed_dict = {ops['image_pl']: X_image_batch,
125 | ops['points_pl']: X_cloud_batch,
126 | ops['is_training_pl']: is_training}
127 | t1 = time.time()
128 | pred_batch = sess.run(ops['pred'],feed_dict=feed_dict)
129 | t2 = time.time()
130 | time_sum += (t2-t1)
131 | result_batch = np.hstack((pred_batch, y_batch))
132 | result_all = np.concatenate((result_all, result_batch), axis=0)
133 |
134 |
135 | np.savetxt(os.path.join(log_dir, 'results.csv'), result_all, delimiter=",")
136 | # b = np.loadtxt("temp.csv", delimiter=",")
137 |
138 | rmse_angle = np.sqrt(np.mean(np.square(result_all[:,0] - result_all[:,2])))
139 | rmse_speed = np.sqrt(np.mean(np.square(result_all[:,1] - result_all[:,3])))
140 | log_string(test_log_dir, 'Test rmse_angle: %f' % (rmse_angle))
141 | log_string(test_log_dir, 'Test rmse_speed: %f' % (rmse_speed))
142 | log_string(test_log_dir, 'Test rmse_average: %f' % ((rmse_angle+rmse_speed)/2))
143 | log_string(test_log_dir, 'Test FPS: %f' % (1/(time_sum/num_batches)))
144 |
145 |
146 | if __name__ == "__main__":
147 | eval()
148 |
--------------------------------------------------------------------------------
/eval_cmp.py:
--------------------------------------------------------------------------------
1 | '''
2 | @Author: Shuai Wang
3 | @Github: https://github.com/wsustcid
4 | @Version: 1.0.0
5 | @Date: 2020-09-11 23:42:23
6 | @LastEditTime: 2020-09-28 22:43:54
7 | '''
8 |
9 | import os
10 | import sys
11 | import argparse
12 | from datetime import datetime
13 | import time
14 | from tqdm import tqdm
15 |
16 | import numpy as np
17 | import tensorflow as tf
18 |
19 | base_dir = os.path.dirname(os.path.abspath(__file__))
20 | sys.path.append(base_dir)
21 | from utils.tf_util import log_string
22 |
23 | from data_gen import DataLoader
24 | from models.flowdrivenet import FlowDriveNet
25 |
26 | parser = argparse.ArgumentParser()
27 | parser.add_argument('--data_root', default='/media/ubuntu16/Documents/Datasets/Udacity/CH2',
28 | help='data_root path [default: local path]')
29 | parser.add_argument('--input_cfg', default='BGR',
30 | help='Input type: BGR, GRAYF-T, XYZ, GRAY')
31 | parser.add_argument('--model_cfg', default='PilotNet',
32 | help='Model type: PilotNet, BMWNet, PointNet, DroNet')
33 | parser.add_argument('--use_side_cam', default=False, action='store_true')
34 | parser.add_argument('--log_dir', default='test',
35 | help='Log dir [default: test]')
36 | parser.add_argument('--batch_size', type=int, default=1,
37 | help='Batch Size during training [default: 16]')
38 | parser.add_argument('--decay_steps', type=int, default=300000,
39 | help='Decay step for lr decay [default: 200000]') # decay_steps = n_train * epochs
40 | parser.add_argument('--decay_rate', type=float, default=0.7,
41 | help='Decay rate for lr decay [default: 0.7]')
42 | parser.add_argument('--model_file', default='/media/ubuntu16/F/FlowDriveNet/logs/VFE/gray_base/model_best.ckpt',
43 | help='the model path to be evaluated')
44 |
45 |
46 | FLAGS = parser.parse_args()
47 |
48 | BATCH_SIZE = FLAGS.batch_size
49 |
50 | log_dir = os.path.join(base_dir, 'logs', FLAGS.log_dir)
51 | os.makedirs(log_dir, exist_ok=True)
52 | test_log_dir = os.path.join(log_dir, 'log_test.txt')
53 | log_string(test_log_dir, str(FLAGS)+'\n')
54 |
55 | #
56 | if FLAGS.model_cfg == 'PilotNet':
57 | from models.pilotnet import PilotNet
58 | dataloader = DataLoader(FLAGS.data_root, "BGR",
59 | height=66, width=200,
60 | seq_len=None,
61 | num_point=None,
62 | use_side_cam=FLAGS.use_side_cam)
63 | model = PilotNet()
64 | elif FLAGS.model_cfg == 'BMWNet':
65 | from models.bmwnet import BMWNet
66 | # TODO add seq_len
67 | dataloader = DataLoader(FLAGS.data_root, 'GRAYF-T',
68 | height=66, width=200,
69 | seq_len=10,
70 | num_point=None,
71 | use_side_cam=FLAGS.use_side_cam)
72 | model = BMWNet()
73 | elif FLAGS.model_cfg == 'PointNet':
74 | from models.pointnet import PointNet
75 | dataloader = DataLoader(FLAGS.data_root, 'XYZ',
76 | height=None, width=None,
77 | seq_len=None,
78 | num_point=10000,
79 | use_side_cam=FLAGS.use_side_cam)
80 | model = PointNet(num_point=10000)
81 | elif FLAGS.model_cfg == 'DroNet':
82 | from models.dronet import DroNet
83 | dataloader = DataLoader(FLAGS.data_root, 'GRAY',
84 | height=200, width=200,
85 | seq_len=None,
86 | num_point=None,
87 | use_side_cam=FLAGS.use_side_cam)
88 | model = DroNet()
89 | else:
90 | raise TypeError
91 |
92 | def get_bn_decay(batch):
93 | bn_momentum = tf.train.exponential_decay(
94 | 0.5,
95 | batch*BATCH_SIZE,
96 | float(FLAGS.decay_steps),
97 | 0.5,
98 | staircase=True)
99 | bn_decay = tf.minimum(0.99, 1 - bn_momentum)
100 | return bn_decay
101 |
102 | def eval():
103 | with tf.Graph().as_default():
104 | feature_pl, _ = model.get_inputs_pl(BATCH_SIZE)
105 | is_training_pl = tf.placeholder(tf.bool, shape=())
106 | # define global_step; optimizer will increase it in every training loop
107 | batch = tf.get_variable('batch', [],
108 | initializer=tf.constant_initializer(0),
109 | trainable=False)
110 | bn_decay = get_bn_decay(batch)
111 |
112 | pred = model.get_model(feature_pl, is_training_pl, bn_decay)
113 |
114 | # Create a session
115 | config = tf.ConfigProto()
116 | config.gpu_options.allow_growth = True
117 | config.allow_soft_placement = True
118 | config.log_device_placement = False
119 | sess = tf.Session(config=config)
120 |
121 | # Init variables
122 | init = tf.global_variables_initializer()
123 | sess.run(init)
124 |
125 | # restore model
126 | saver = tf.train.Saver()
127 | saver.restore(sess, FLAGS.model_file)
128 |
129 | # save all tensor
130 | ops = {'feature_pl': feature_pl,
131 | 'is_training_pl': is_training_pl,
132 | 'pred': pred}
133 |
134 | ## evaluation
135 | is_training = False
136 | num_batches = dataloader.num_test // BATCH_SIZE
137 | rmse_angle_sum = 0.0
138 | rmse_speed_sum = 0.0
139 | if FLAGS.model_cfg in ['PilotNet', 'BMWNet', 'PointNet']:
140 | result_all = np.zeros((0,2))
141 | elif FLAGS.model_cfg == 'DroNet':
142 | result_all = np.zeros((0,4)) # pred_a, pred_s, label_a, label_s
143 | else:
144 | raise TypeError
145 |
146 | time_sum = 0.0
147 | for i in tqdm(range(num_batches)):
148 | if FLAGS.model_cfg in ['PilotNet', 'BMWNet']:
149 | X_batch, y = dataloader.load_image_test_batch(BATCH_SIZE)
150 | y_batch = y[:,0:1]
151 | elif FLAGS.model_cfg == 'PointNet':
152 | X_batch, y = dataloader.load_cloud_test_batch(BATCH_SIZE)
153 | y_batch = y[:,1:2]
154 | elif FLAGS.model_cfg == 'DroNet':
155 | X_batch, y_batch = dataloader.load_image_test_batch(BATCH_SIZE)
156 | else:
157 | raise TypeError
158 |
159 | feed_dict = {ops['feature_pl']: X_batch,
160 | ops['is_training_pl']: is_training}
161 | t1 = time.time()
162 | pred_batch = sess.run(ops['pred'],feed_dict=feed_dict)
163 | t2 = time.time()
164 | time_sum += (t2-t1)
165 | result_batch = np.hstack((pred_batch, y_batch))
166 | result_all = np.concatenate((result_all, result_batch), axis=0)
167 |
168 |
169 |
170 | np.savetxt(os.path.join(log_dir, 'results.csv'), result_all, delimiter=",")
171 | # b = np.loadtxt("temp.csv", delimiter=",")
172 | if FLAGS.model_cfg in ['PilotNet', 'BMWNet']:
173 | rmse_angle = np.sqrt(np.mean(np.square(result_all[:,0] - result_all[:,1])))
174 | log_string(test_log_dir, 'Test rmse_angle: %f' % (rmse_angle))
175 | log_string(test_log_dir, 'Test FPS: %f' % (1/(time_sum/num_batches)))
176 | elif FLAGS.model_cfg == 'PointNet':
177 | rmse_speed = np.sqrt(np.mean(np.square(result_all[:,0] - result_all[:,1])))
178 | log_string(test_log_dir, 'Test rmse_speed: %f' % (rmse_speed))
179 | log_string(test_log_dir, 'Test FPS: %f' % (1/(time_sum/num_batches)))
180 | elif FLAGS.model_cfg == 'DroNet':
181 | rmse_angle = np.sqrt(np.mean(np.square(result_all[:,0] - result_all[:,2])))
182 | rmse_speed = np.sqrt(np.mean(np.square(result_all[:,1] - result_all[:,3])))
183 | log_string(test_log_dir, 'Test rmse_angle: %f' % (rmse_angle))
184 | log_string(test_log_dir, 'Test rmse_speed: %f' % (rmse_speed))
185 | log_string(test_log_dir, 'Test rmse_average: %f' % ((rmse_angle+rmse_speed)/2))
186 | log_string(test_log_dir, 'Test FPS: %f' % (1/(time_sum/num_batches)))
187 | else:
188 | raise TypeError
189 |
190 |
191 | if __name__ == "__main__":
192 | eval()
193 |
--------------------------------------------------------------------------------
/models/__init__.py:
--------------------------------------------------------------------------------
1 | '''
2 | @Author: Shuai Wang
3 | @Github: https://github.com/wsustcid
4 | @Version: 1.0.0
5 | @Date: 1970-01-01 08:00:00
6 | @LastEditTime: 2020-06-09 17:43:19
7 | @Description:
8 | '''
9 |
--------------------------------------------------------------------------------
/models/baseline.py:
--------------------------------------------------------------------------------
1 | '''
2 | @Author: Shuai Wang
3 | @Github: https://github.com/wsustcid
4 | @Version: 1.0.0
5 | @Date: 2020-08-04 19:43:43
6 | @LastEditTime: 2020-09-13 22:39:59
7 | @Description:
8 | '''
9 | import os
10 | import sys
11 | import numpy as np
12 |
13 | base_dir = os.path.dirname(os.path.abspath(__file__))
14 | sys.path.append(os.path.join(base_dir, '../'))
15 |
16 | from data_gen import DataLoader
17 |
18 |
19 | if __name__ == '__main__':
20 | data_root='/media/ubuntu16/Documents/Datasets/Udacity/CH2'
21 | dataloader = DataLoader(data_root, input_cfg='GRAY', height=200, width=200, use_side_cam=False)
22 |
23 | label_train = dataloader.y_train
24 | label_val = dataloader.y_val
25 | label_test = dataloader.y_test
26 |
27 | pred_train = np.zeros(label_train.shape, dtype=np.float32)
28 | pred_val = np.zeros(label_val.shape, dtype=np.float32)
29 | pred_test = np.zeros(label_test.shape, dtype=np.float32)
30 |
31 | rmse_train = np.sqrt(np.mean(np.square(label_train-pred_train), axis=0))
32 | rmse_val = np.sqrt(np.mean(np.square(label_val-pred_val), axis=0))
33 | rmse_test = np.sqrt(np.mean(np.square(label_test-pred_test), axis=0))
34 |
35 | print("Train: rmse_angle: {}; rmse_speed: {}; rmse_all: {}".format(rmse_train[0], rmse_train[1], np.mean(rmse_train)))
36 | print("Val: rmse_angle: {}; rmse_speed: {}; rmse_all: {}".format(rmse_val[0], rmse_val[1], np.mean(rmse_val)))
37 | print("Test: rmse_angle: {}; rmse_speed: {}; rmse_all: {}".format(rmse_test[0], rmse_test[1], np.mean(rmse_test)))
--------------------------------------------------------------------------------
/models/bmwnet.py:
--------------------------------------------------------------------------------
1 | '''
2 | @Author: Shuai Wang
3 | @Github: https://github.com/wsustcid
4 | @Version: 1.0.0
5 | @Date: 2020-05-28 22:19:51
6 | @LastEditTime: 2020-09-28 22:02:40
7 | '''
8 |
9 | import os
10 | import sys
11 | import numpy as np
12 | import tensorflow as tf
13 | BASE_DIR = os.path.dirname(os.path.abspath(__file__))
14 | sys.path.append(os.path.join(BASE_DIR, '..'))
15 |
16 | from utils import tf_util
17 |
18 |
19 | class BMWNet(object):
20 | def __init__(self, height=66, width=200, seq_len=10):
21 | self.height = height
22 | self.width = width
23 | self.seq_len = seq_len
24 |
25 | self.output_dim = 1
26 |
27 | self.drop_rate = 0.5
28 |
29 |
30 | def get_inputs_pl(self, batch_size):
31 | image_pl = tf.placeholder(tf.float32, shape=(batch_size, self.seq_len, self.height, self.width, 3))
32 | label_pl = tf.placeholder(tf.float32, shape=(batch_size, self.output_dim))
33 |
34 | return image_pl, label_pl
35 |
36 |
37 | def get_model(self, image_pl, is_training, bn_decay=None):
38 | """ BMWNet model
39 | Args:
40 | - input: input Tensor: (B,T,66,200,3)
41 | - is_training: A bool flag used for dropout
42 | Return:
43 | - output: output Tensor (B,2)
44 | """
45 |
46 | img_shape = image_pl.get_shape().as_list()[2:]
47 | input_img = tf.reshape(image_pl, (-1,*img_shape)) # (B*T,H,W,C)
48 |
49 | X = tf.layers.conv2d(input_img, 24, [5,5], strides=[2,2],
50 | activation=tf.nn.relu,
51 | padding='VALID', name='conv1') # (B*T,31,98,24)
52 | X = tf.layers.conv2d(X, 36, [5,5], strides=[2,2],
53 | activation=tf.nn.relu,
54 | padding='VALID', name='conv2') # (B*T,14,47,36)
55 | X = tf.layers.conv2d(X, 48, [5,5], strides=[2,2],
56 | activation=tf.nn.relu,
57 | padding='VALID', name='conv3') # (B*T,5,22,48)
58 | X = tf.layers.conv2d(X, 64, [3,3], strides=[1,1],
59 | activation=tf.nn.relu,
60 | padding='VALID', name='conv4') # (B*T,3,20,64)
61 | X = tf.layers.conv2d(X, 64, [3,3], strides=[1,1],
62 | activation=tf.nn.relu,
63 | padding='VALID', name='conv5') # (B*T,1,18,64)
64 |
65 | X = tf.contrib.layers.flatten(X) # (B*T, 1152)
66 |
67 | X = tf.layers.dense(X, 1152, activation=tf.nn.relu, name='fc1')
68 | X = tf.layers.dropout(X, rate=self.drop_rate, training=is_training, name='dp1')
69 |
70 | X = tf.layers.dense(X, 512, activation=tf.nn.relu, name='fc2')
71 | X = tf.layers.dropout(X, rate=self.drop_rate, training=is_training, name='dp2')
72 |
73 | # add lstm
74 | X = tf.reshape(X, (-1, self.seq_len, X.get_shape().as_list()[-1])) # (B,T,D)
75 | X = tf_util.lstm(X, hidden_size=128, scope='lstm')
76 |
77 | output = tf.layers.dense(X, self.output_dim, activation=None, name='output')
78 | # y = tf.mul(tf.atan(tf.matmul(h_fc4_drop, W_fc5) + b_fc5), 2, name='y') # [-pi,pi]
79 |
80 |
81 | return output
82 |
83 | def get_loss(self, pred, label, name='Loss'):
84 | """ Return loss for a batch of data
85 | Args:
86 | Return: MSE loss
87 | """
88 |
89 | loss = tf.reduce_mean(tf.square(tf.subtract(pred, label)))
90 |
91 | tf.summary.scalar(name, loss)
92 |
93 | return loss
94 |
95 | def get_rmse(self, pred, label, name='rmse'):
96 | """Return rmse as evalation metrics for a batch of data
97 | """
98 | rmse = tf.sqrt(tf.reduce_mean(tf.square(tf.subtract(pred, label))))
99 |
100 | tf.summary.scalar(name, rmse)
101 |
102 | return rmse
103 |
104 |
105 | if __name__ == '__main__':
106 | with tf.Graph().as_default():
107 | is_training_pl = tf.placeholder(tf.bool, shape=())
108 |
109 | model = BMWNet()
110 |
111 | img_pl, label_pl = model.get_inputs_pl(batch_size=32)
112 | outputs = model.get_model(img_pl, is_training=is_training_pl)
113 |
114 | tf_util.model_summary()
115 | #Total size of variables: 2378261
116 | # Total bytes of variables: 9513044
--------------------------------------------------------------------------------
/models/dronet.py:
--------------------------------------------------------------------------------
1 | '''
2 | @Author: Shuai Wang
3 | @Github: https://github.com/wsustcid
4 | @Version: 1.0.0
5 | @Date: 2020-09-15 19:34:29
6 | @LastEditTime: 2020-09-28 22:17:54
7 | '''
8 |
9 | import os
10 | import sys
11 | import numpy as np
12 | import tensorflow as tf
13 |
14 | base_dir = os.path.dirname(os.path.abspath(__file__))
15 | sys.path.append(os.path.join(base_dir, '..'))
16 | from utils import tf_util
17 |
18 | class DroNet(object):
19 | def __init__(self, height=200, width=200, channels=1):
20 | self.height = height
21 | self.width = width
22 | self.channels = channels
23 | self.output_dim = 2
24 |
25 | self.drop_rate = 0.5
26 |
27 | def get_inputs_pl(self, batch_size):
28 |
29 | image_pl = tf.placeholder(tf.float32, shape=(batch_size, self.height, self.width, self.channels))
30 | label_pl = tf.placeholder(tf.float32, shape=(batch_size, self.output_dim))
31 |
32 | return image_pl, label_pl
33 |
34 |
35 | def get_model(self, image_pl, is_training, bn_decay=None):
36 | """ DroNet model (ResNet8)
37 | Args:
38 | - input: input Tensor: (B,200,200,1)
39 | - is_training: A bool flag used for dropout
40 | Return:
41 | - output: output Tensor (B,1)
42 | TODO: in conv2d set
43 | kernel_initializer="he_normal",
44 | kernel_regularizer=regularizers.l2(1e-4)
45 | """
46 | # Input
47 | X1 = tf.layers.conv2d(image_pl, 32, [5,5], strides=[2,2],
48 | padding='same', name='conv1')
49 | X1 = tf.layers.max_pooling2d(X1, pool_size=(3,3), strides=[2,2])
50 |
51 | # First residual block
52 | X2 = tf.layers.batch_normalization(X1, training=is_training, name='res1_bn1')
53 | X2 = tf.nn.relu(X2, name='res1_relu1')
54 | X2 = tf.layers.conv2d(X2, 32, [3,3], strides=[2,2],
55 | padding='same', name='res1_conv1')
56 |
57 | X2 = tf.layers.batch_normalization(X2, training=is_training, name='res1_bn2')
58 | X2 = tf.nn.relu(X2, name='res1_relu2')
59 | X2 = tf.layers.conv2d(X2, 32, [3,3], strides=[1,1],
60 | padding='same', name='res1_conv2')
61 |
62 | X1 = tf.layers.conv2d(X1, 32, [1,1], strides=[2,2],
63 | padding='same', name='res1_skip')
64 | X3 = tf.add(X1, X2, name='res1_add')
65 |
66 | # Second residual block
67 | X4 = tf.layers.batch_normalization(X3, training=is_training, name='res2_bn1')
68 | X4 = tf.nn.relu(X4, name='res2_relu1')
69 | X4 = tf.layers.conv2d(X4, 64, [3,3], strides=[2,2],
70 | padding='same', name='res2_conv1')
71 |
72 | X4 = tf.layers.batch_normalization(X4, training=is_training, name='res2_bn2')
73 | X4 = tf.nn.relu(X4, name='res2_relu2')
74 | X4 = tf.layers.conv2d(X4, 64, [3,3], strides=[1,1],
75 | padding='same', name='res2_conv2')
76 |
77 | X3 = tf.layers.conv2d(X3, 64, [1,1], strides=[2,2],
78 | padding='same', name='res2_skip')
79 | X5 = tf.add(X3, X4, name='res2_add')
80 |
81 | # Third residual block
82 | X6 = tf.layers.batch_normalization(X5, training=is_training, name='res3_bn1')
83 | X6 = tf.nn.relu(X6, name='res3_relu1')
84 | X6 = tf.layers.conv2d(X6, 128, [3,3], strides=[2,2],
85 | padding='same', name='res3_conv1')
86 |
87 | X6 = tf.layers.batch_normalization(X6, training=is_training, name='res3_bn2')
88 | X6 = tf.nn.relu(X6, name='res3_relu2')
89 | X6 = tf.layers.conv2d(X6, 128, [3,3], strides=[1,1],
90 | padding='same', name='res3_conv2')
91 |
92 | X5 = tf.layers.conv2d(X5, 128, [1,1], strides=[2,2],
93 | padding='same', name='res3_skip')
94 | X7 = tf.add(X5, X6, name='res3_add')
95 |
96 | X = tf.contrib.layers.flatten(X7) # (B, 1152)
97 | X = tf.nn.relu(X)
98 | X = tf.layers.dropout(X, rate=self.drop_rate)
99 |
100 | output = tf.layers.dense(X, self.output_dim, activation=None, name='output')
101 |
102 | return output
103 |
104 | def get_loss(self, pred, label, name='Loss'):
105 | """ Return loss for a batch of data
106 | Args:
107 | Return: MSE loss
108 | """
109 |
110 | loss = tf.reduce_mean(tf.square(tf.subtract(pred, label)))
111 |
112 | tf.summary.scalar(name, loss)
113 |
114 | return loss
115 |
116 | def get_rmse(self, pred, label, name='rmse'):
117 | """Return rmse as evalation metrics for a batch of data
118 | """
119 | rmse = tf.sqrt(tf.reduce_mean(tf.square(tf.subtract(pred, label)), axis=0))
120 |
121 | tf.summary.scalar(name+'_angle', rmse[0])
122 | tf.summary.scalar(name+'_speed', rmse[1])
123 |
124 | return rmse[0], rmse[1]
125 |
126 | def get_acc(self, pred, label, delta_a=5.0, delta_s=1.0, name='acc'):
127 | """
128 | TODO: Tolerance is to large! angle < 0.01;
129 | - delta_a: tolerance of angle: 5 degrees
130 | - delta_s: tolerance of speed: 1 ft/s = 0.3048 m/s = 1.09728 km/h
131 | """
132 | acc_a = tf.abs(tf.subtract(pred[:, 0], label[:, 0])) < (delta_a/180*np.pi)
133 | acc_a = tf.reduce_mean(tf.cast(acc_a,tf.float32))
134 |
135 | acc_s = tf.abs(tf.subtract(pred[:, 1], label[:, 1])) < delta_s
136 | acc_s = tf.reduce_mean(tf.cast(acc_s, tf.float32))
137 |
138 | tf.summary.scalar(name+'_angle', acc_a)
139 | tf.summary.scalar(name+'_speed', acc_s)
140 |
141 | return acc_a, acc_s
142 |
143 | if __name__ == '__main__':
144 | with tf.Graph().as_default():
145 | is_training = tf.cast(True, tf.bool)
146 | batch_size = 16
147 |
148 | model = DroNet()
149 |
150 | image_pl, label_pl = model.get_inputs_pl(batch_size)
151 | pred = model.get_model(image_pl, is_training)
152 | loss = model.get_loss(pred, label_pl)
153 |
154 | tf_util.model_summary()
155 | # Total size of variables: 320930
156 | # Total bytes of variables: 1283720
157 |
158 |
--------------------------------------------------------------------------------
/models/pilotnet.py:
--------------------------------------------------------------------------------
1 | '''
2 | @Author: Shuai Wang
3 | @Github: https://github.com/wsustcid
4 | @Version: 1.0.0
5 | @Date: 2020-05-28 22:19:51
6 | @LastEditTime: 2020-09-15 23:32:47
7 | '''
8 |
9 | import os
10 | import sys
11 | import numpy as np
12 | import tensorflow as tf
13 |
14 | base_dir = os.path.dirname(os.path.abspath(__file__))
15 | sys.path.append(os.path.join(base_dir, '..'))
16 | from utils import tf_util
17 |
18 | class PilotNet(object):
19 | def __init__(self, height=66, width=200):
20 | self.height = height
21 | self.width = width
22 | self.output_dim = 1
23 |
24 | self.drop_rate = 0.5
25 |
26 | def get_inputs_pl(self, batch_size):
27 | """ Create placeholders for the PilotNet
28 | """
29 | image_pl = tf.placeholder(tf.float32, shape=(batch_size, self.height, self.width, 3))
30 | label_pl = tf.placeholder(tf.float32, shape=(batch_size, self.output_dim))
31 |
32 | return image_pl, label_pl
33 |
34 |
35 | def get_model(self, image_pl, is_training, bn_decay=None):
36 | """ Nvidia PilotNet model
37 | Args:
38 | - input: input Tensor: (B,66,200,3)
39 | - is_training: A bool flag used for dropout
40 | Return:
41 | - output: output Tensor (B,1)
42 | """
43 | X = tf.layers.conv2d(image_pl, 24, [5,5], strides=[2,2],
44 | activation=tf.nn.relu,
45 | padding='VALID', name='conv1') # (B,31,98,24)
46 | X = tf.layers.conv2d(X, 36, [5,5], strides=[2,2],
47 | activation=tf.nn.relu,
48 | padding='VALID', name='conv2') # (B,14,47,36)
49 | X = tf.layers.conv2d(X, 48, [5,5], strides=[2,2],
50 | activation=tf.nn.relu,
51 | padding='VALID', name='conv3') # (B,5,22,48)
52 | X = tf.layers.conv2d(X, 64, [3,3], strides=[1,1],
53 | activation=tf.nn.relu,
54 | padding='VALID', name='conv4') # (B,3,20,64)
55 | X = tf.layers.conv2d(X, 64, [3,3], strides=[1,1],
56 | activation=tf.nn.relu,
57 | padding='VALID', name='conv5') # (B,1,18,64)
58 |
59 | X = tf.contrib.layers.flatten(X) # (B, 1152)
60 |
61 | X = tf.layers.dense(X, 1164, activation=tf.nn.relu, name='fc1')
62 | X = tf.layers.dropout(X, rate=self.drop_rate, training=is_training, name='dp1')
63 |
64 | X = tf.layers.dense(X, 100, activation=tf.nn.relu, name='fc2')
65 | X = tf.layers.dropout(X, rate=self.drop_rate, training=is_training, name='dp2')
66 |
67 | X = tf.layers.dense(X, 50, activation=tf.nn.relu, name='fc3')
68 | X = tf.layers.dropout(X, rate=self.drop_rate, training=is_training, name='dp4')
69 |
70 | X = tf.layers.dense(X, 10, activation=tf.nn.relu, name='fc4')
71 | X = tf.layers.dropout(X, rate=self.drop_rate, training=is_training, name='dp4')
72 |
73 | output = tf.layers.dense(X, self.output_dim, activation=None, name='output')
74 |
75 | # y = tf.mul(tf.atan(tf.matmul(h_fc4_drop, W_fc5) + b_fc5), 2, name='y') # [-pi,pi]
76 |
77 | return output
78 |
79 | def get_loss(self, pred, label, name='Loss'):
80 | """ Return loss for a batch of data
81 | Args:
82 | Return: MSE loss
83 | """
84 |
85 | loss = tf.reduce_mean(tf.square(tf.subtract(pred, label)))
86 |
87 | tf.summary.scalar(name, loss)
88 |
89 | return loss
90 |
91 | def get_rmse(self, pred, label, name='rmse'):
92 | """Return rmse as evalation metrics for a batch of data
93 | """
94 | rmse = tf.sqrt(tf.reduce_mean(tf.square(tf.subtract(pred, label))))
95 |
96 | tf.summary.scalar(name, rmse)
97 |
98 | return rmse
99 |
100 | def get_acc(self, pred, label, delta_a=5.0, delta_s=1.0, name='acc'):
101 | """
102 | TODO: Tolerance is to large! angle < 0.01;
103 | - delta_a: tolerance of angle: 5 degrees
104 | - delta_s: tolerance of speed: 1 ft/s = 0.3048 m/s = 1.09728 km/h
105 | """
106 | acc_a = tf.abs(tf.subtract(pred[:, 0], label[:, 0])) < (delta_a/180*np.pi)
107 | acc_a = tf.reduce_mean(tf.cast(acc_a,tf.float32))
108 |
109 | acc_s = tf.abs(tf.subtract(pred[:, 1], label[:, 1])) < delta_s
110 | acc_s = tf.reduce_mean(tf.cast(acc_s, tf.float32))
111 |
112 | tf.summary.scalar(name+'_angle', acc_a)
113 | tf.summary.scalar(name+'_speed', acc_s)
114 |
115 | return acc_a, acc_s
116 |
117 | if __name__ == '__main__':
118 | with tf.Graph().as_default():
119 | is_training = tf.cast(True, tf.bool)
120 | batch_size = 16
121 |
122 | model = PilotNet()
123 |
124 | image_pl, label_pl = model.get_inputs_pl(batch_size)
125 | pred = model.get_model(image_pl, is_training)
126 | loss = model.get_loss(pred, label_pl)
127 | rmse = model.get_rmse(pred, label_pl)
128 |
129 | tf_util.model_summary()
130 |
131 | print('loss:', loss)
132 | print('rmse:', rmse)
133 | # Total size of variables: 1595511
134 | # Total bytes of variables: 6382044
135 |
--------------------------------------------------------------------------------
/models/pointnet.py:
--------------------------------------------------------------------------------
1 | '''
2 | @Author: Shuai Wang
3 | @Github: https://github.com/wsustcid
4 | @Version: 1.0.0
5 | @Date: 2020-06-28 11:01:02
6 | @LastEditTime: 2020-09-28 22:06:02
7 | @Description:
8 | '''
9 |
10 | '''
11 | PointNet version 1 Model
12 | Reference: https://github.com/charlesq34/pointnet
13 | '''
14 |
15 | import os
16 | import sys
17 | import numpy as np
18 | import tensorflow as tf
19 | BASE_DIR = os.path.dirname(os.path.abspath(__file__))
20 | sys.path.append(os.path.join(BASE_DIR, '../utils'))
21 | import tf_util
22 |
23 |
24 | class PointNet(object):
25 | def __init__(self, num_point):
26 |
27 | self.num_point = num_point
28 | self.output_dim = 1
29 |
30 |
31 | def get_inputs_pl(self, batch_size):
32 | points_pl = tf.placeholder(tf.float32, shape=(batch_size, self.num_point, 3))
33 | label_pl = tf.placeholder(tf.float32, shape=(batch_size, self.output_dim))
34 |
35 | return points_pl, label_pl
36 |
37 |
38 | def get_model(self, points_pl, is_training, bn_decay=None):
39 | """ Regression PointNet, input is BxNx3, output Bx1 """
40 | input_image = tf.expand_dims(points_pl, -1) # (B,N,3,1)
41 |
42 | # Point functions (MLP implemented as conv2d)
43 | net = tf_util.conv2d(input_image, 64, [1,3],
44 | padding='VALID', stride=[1,1],
45 | bn=True, is_training=is_training,
46 | scope='conv1', bn_decay=bn_decay)
47 | net = tf_util.conv2d(net, 64, [1,1],
48 | padding='VALID', stride=[1,1],
49 | bn=True, is_training=is_training,
50 | scope='conv2', bn_decay=bn_decay)
51 | net = tf_util.conv2d(net, 64, [1,1],
52 | padding='VALID', stride=[1,1],
53 | bn=True, is_training=is_training,
54 | scope='conv3', bn_decay=bn_decay)
55 | net = tf_util.conv2d(net, 128, [1,1],
56 | padding='VALID', stride=[1,1],
57 | bn=True, is_training=is_training,
58 | scope='conv4', bn_decay=bn_decay)
59 | net = tf_util.conv2d(net, 1024, [1,1],
60 | padding='VALID', stride=[1,1],
61 | bn=True, is_training=is_training,
62 | scope='conv5', bn_decay=bn_decay)
63 |
64 | # Symmetric function: max pooling
65 | net = tf_util.max_pool2d(net, [self.num_point, 1],
66 | padding='VALID', scope='maxpool') # (B,1,1,1024)
67 |
68 | # MLP on global point cloud vector
69 | #net = tf.reshape(net, [batch_size, -1])
70 | net = tf.squeeze(net, axis=[1,2]) # (B,1024)
71 |
72 | net = tf_util.fully_connected(net, 512, bn=True, is_training=is_training,
73 | scope='fc1', bn_decay=bn_decay)
74 | net = tf_util.fully_connected(net, 256, bn=True, is_training=is_training,
75 | scope='fc2', bn_decay=bn_decay)
76 | net = tf_util.dropout(net, keep_prob=0.7, is_training=is_training,
77 | scope='dp1')
78 | net = tf_util.fully_connected(net, self.output_dim, activation_fn=None, scope='fc3')
79 |
80 | return net
81 |
82 |
83 | def get_loss(self, pred, label, name='Loss'):
84 | """ Return mse loss for a batch of data
85 | """
86 | loss = tf.reduce_mean(tf.square(tf.subtract(pred, label)))
87 |
88 | tf.summary.scalar(name, loss)
89 |
90 | return loss
91 |
92 |
93 | def get_rmse(self,pred, label, name='rmse'):
94 | """Return rmse as evalation metrics for a batch of data
95 | """
96 | rmse = tf.sqrt(tf.reduce_mean(tf.square(tf.subtract(pred, label))))
97 |
98 | tf.summary.scalar(name, rmse)
99 |
100 | return rmse
101 |
102 | if __name__=='__main__':
103 | with tf.Graph().as_default():
104 | batch_size = 16
105 | is_training = tf.constant(True)
106 |
107 | model = PointNet(num_point=20000)
108 |
109 | points_pl, label_pl = model.get_inputs_pl(batch_size)
110 | pred = model.get_model(points_pl, is_training)
111 | loss = model.get_loss(pred, label_pl)
112 | print(pred)
113 | print(loss)
114 | tf_util.model_summary()
115 | # Total size of variables: 809601
116 | # Total bytes of variables: 3238404
117 |
--------------------------------------------------------------------------------
/tf_ops/3d_interpolation/__pycache__/tf_interpolate.cpython-35.pyc:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/wsustcid/FlowDriveNet/3604495269ae45e5b43964046104f685ec66e383/tf_ops/3d_interpolation/__pycache__/tf_interpolate.cpython-35.pyc
--------------------------------------------------------------------------------
/tf_ops/3d_interpolation/interpolate.cpp:
--------------------------------------------------------------------------------
1 | #include
2 | #include
3 | #include // memset
4 | #include // rand, RAND_MAX
5 | #include // sqrtf
6 | #include
7 | #include
8 | using namespace std;
9 | float randomf(){
10 | return (rand()+0.5)/(RAND_MAX+1.0);
11 | }
12 | static double get_time(){
13 | timespec tp;
14 | clock_gettime(CLOCK_MONOTONIC,&tp);
15 | return tp.tv_sec+tp.tv_nsec*1e-9;
16 | }
17 |
18 | // Find three nearest neigbors with square distance
19 | // input: xyz1 (b,n,3), xyz2(b,m,3)
20 | // output: dist (b,n,3), idx (b,n,3)
21 | void threenn_cpu(int b, int n, int m, const float *xyz1, const float *xyz2, float *dist, int *idx) {
22 | for (int i=0;i
2 | #include
3 | #include // memset
4 | #include // rand, RAND_MAX
5 | #include // sqrtf
6 | #include "tensorflow/core/framework/op.h"
7 | #include "tensorflow/core/framework/op_kernel.h"
8 | #include "tensorflow/core/framework/shape_inference.h"
9 | #include "tensorflow/core/framework/common_shape_fns.h"
10 | using namespace tensorflow;
11 |
12 | REGISTER_OP("ThreeNN")
13 | .Input("xyz1: float32")
14 | .Input("xyz2: float32")
15 | .Output("dist: float32")
16 | .Output("idx: int32")
17 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
18 | c->set_output(0, c->input(0));
19 | c->set_output(1, c->input(0));
20 | return Status::OK();
21 | });
22 | REGISTER_OP("ThreeInterpolate")
23 | .Input("points: float32")
24 | .Input("idx: int32")
25 | .Input("weight: float32")
26 | .Output("out: float32")
27 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
28 | ::tensorflow::shape_inference::ShapeHandle dims1; // (b,m,c)
29 | c->WithRank(c->input(0), 3, &dims1);
30 | ::tensorflow::shape_inference::ShapeHandle dims2; // (b,n,3)
31 | c->WithRank(c->input(1), 3, &dims2);
32 | // (b,n,c)
33 | ::tensorflow::shape_inference::ShapeHandle output = c->MakeShape({c->Dim(dims1, 0), c->Dim(dims2, 1), c->Dim(dims1, 2)});
34 | c->set_output(0, output);
35 | return Status::OK();
36 | });
37 | REGISTER_OP("ThreeInterpolateGrad")
38 | .Input("points: float32")
39 | .Input("idx: int32")
40 | .Input("weight: float32")
41 | .Input("grad_out: float32")
42 | .Output("grad_points: float32")
43 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
44 | c->set_output(0, c->input(0));
45 | return Status::OK();
46 | });
47 |
48 | float randomf(){
49 | return (rand()+0.5)/(RAND_MAX+1.0);
50 | }
51 | static double get_time(){
52 | timespec tp;
53 | clock_gettime(CLOCK_MONOTONIC,&tp);
54 | return tp.tv_sec+tp.tv_nsec*1e-9;
55 | }
56 |
57 | // Find three nearest neigbors with square distance
58 | // input: xyz1 (b,n,3), xyz2(b,m,3)
59 | // output: dist (b,n,3), idx (b,n,3)
60 | void threenn_cpu(int b, int n, int m, const float *xyz1, const float *xyz2, float *dist, int *idx) {
61 | for (int i=0;iinput(0);
163 | OP_REQUIRES(context, xyz1_tensor.dims()==3 && xyz1_tensor.shape().dim_size(2)==3, errors::InvalidArgument("ThreeNN expects (b,n,3) xyz1 shape."));
164 | int b = xyz1_tensor.shape().dim_size(0);
165 | int n = xyz1_tensor.shape().dim_size(1);
166 |
167 | const Tensor& xyz2_tensor = context->input(1);
168 | OP_REQUIRES(context, xyz2_tensor.dims()==3 && xyz2_tensor.shape().dim_size(2)==3, errors::InvalidArgument("ThreeNN expects (b,m,3) xyz2 shape."));
169 | int m = xyz2_tensor.shape().dim_size(1);
170 |
171 | Tensor *dist_tensor = nullptr;
172 | OP_REQUIRES_OK(context, context->allocate_output(0, TensorShape{b,n,3}, &dist_tensor));
173 | Tensor *idx_tensor = nullptr;
174 | OP_REQUIRES_OK(context, context->allocate_output(1, TensorShape{b,n,3}, &idx_tensor));
175 |
176 | auto xyz1_flat = xyz1_tensor.flat();
177 | const float *xyz1 = &(xyz1_flat(0));
178 | auto xyz2_flat = xyz2_tensor.flat();
179 | const float *xyz2 = &(xyz2_flat(0));
180 | auto dist_flat = dist_tensor->flat();
181 | float *dist = &(dist_flat(0));
182 | auto idx_flat = idx_tensor->flat();
183 | int *idx = &(idx_flat(0));
184 | threenn_cpu(b,n,m,xyz1,xyz2,dist,idx);
185 | }
186 | };
187 | REGISTER_KERNEL_BUILDER(Name("ThreeNN").Device(DEVICE_CPU), ThreeNNOp);
188 |
189 |
190 |
191 | class ThreeInterpolateOp: public OpKernel{
192 | public:
193 | explicit ThreeInterpolateOp(OpKernelConstruction * context):OpKernel(context){}
194 |
195 | void Compute(OpKernelContext * context) override {
196 | const Tensor& points_tensor=context->input(0);
197 | OP_REQUIRES(context, points_tensor.dims()==3, errors::InvalidArgument("ThreeInterpolate expects (b,m,c) points shape"));
198 | int b = points_tensor.shape().dim_size(0);
199 | int m = points_tensor.shape().dim_size(1);
200 | int c = points_tensor.shape().dim_size(2);
201 |
202 | const Tensor& idx_tensor=context->input(1);
203 | OP_REQUIRES(context,idx_tensor.dims()==3 && idx_tensor.shape().dim_size(0)==b && idx_tensor.shape().dim_size(2)==3, errors::InvalidArgument("ThreeInterpolate expects (b,n,3) idx shape"));
204 | int n = idx_tensor.shape().dim_size(1);
205 | const Tensor& weight_tensor=context->input(2);
206 | OP_REQUIRES(context,weight_tensor.dims()==3 && weight_tensor.shape().dim_size(0)==b && weight_tensor.shape().dim_size(1)==n && weight_tensor.shape().dim_size(2)==3, errors::InvalidArgument("ThreeInterpolate expects (b,n,3) weight shape"));
207 |
208 | Tensor * out_tensor = nullptr;
209 | OP_REQUIRES_OK(context, context->allocate_output(0,TensorShape{b,n,c}, &out_tensor));
210 |
211 | auto points_flat = points_tensor.flat();
212 | const float *points = &(points_flat(0));
213 | auto idx_flat = idx_tensor.flat();
214 | const int *idx = &(idx_flat(0));
215 | auto weight_flat = weight_tensor.flat();
216 | const float *weight = &(weight_flat(0));
217 | auto out_flat = out_tensor->flat();
218 | float *out = &(out_flat(0));
219 | threeinterpolate_cpu(b,m,c,n,points,idx,weight,out);
220 | }
221 | };
222 | REGISTER_KERNEL_BUILDER(Name("ThreeInterpolate").Device(DEVICE_CPU),ThreeInterpolateOp);
223 |
224 |
225 | class ThreeInterpolateGradOp: public OpKernel{
226 | public:
227 | explicit ThreeInterpolateGradOp(OpKernelConstruction * context):OpKernel(context){}
228 |
229 | void Compute(OpKernelContext * context) override {
230 | const Tensor& points_tensor=context->input(0);
231 | OP_REQUIRES(context, points_tensor.dims()==3, errors::InvalidArgument("ThreeInterpolateGrad expects (b,m,c) points shape"));
232 | int b = points_tensor.shape().dim_size(0);
233 | int m = points_tensor.shape().dim_size(1);
234 | int c = points_tensor.shape().dim_size(2);
235 |
236 | const Tensor& idx_tensor=context->input(1);
237 | OP_REQUIRES(context,idx_tensor.dims()==3 && idx_tensor.shape().dim_size(0)==b, errors::InvalidArgument("ThreeInterpolateGrad expects (b,n,3) idx shape"));
238 | int n = idx_tensor.shape().dim_size(1);
239 | const Tensor& weight_tensor=context->input(2);
240 | OP_REQUIRES(context,weight_tensor.dims()==3 && weight_tensor.shape().dim_size(0)==b && weight_tensor.shape().dim_size(1)==n && weight_tensor.shape().dim_size(2)==3, errors::InvalidArgument("ThreeInterpolateGrad expects (b,n,3) weight shape"));
241 |
242 | const Tensor& grad_out_tensor=context->input(3);
243 | OP_REQUIRES(context,grad_out_tensor.dims()==3 && grad_out_tensor.shape().dim_size(0)==b && grad_out_tensor.shape().dim_size(1)==n && grad_out_tensor.shape().dim_size(2)==c, errors::InvalidArgument("ThreeInterpolateGrad expects (b,n,c) grad_out shape"));
244 |
245 | Tensor * grad_points_tensor = nullptr;
246 | OP_REQUIRES_OK(context, context->allocate_output(0,TensorShape{b,m,c}, &grad_points_tensor));
247 |
248 | auto points_flat = points_tensor.flat();
249 | const float *points = &(points_flat(0));
250 | auto idx_flat = idx_tensor.flat();
251 | const int *idx = &(idx_flat(0));
252 | auto weight_flat = weight_tensor.flat();
253 | const float *weight = &(weight_flat(0));
254 | auto grad_out_flat = grad_out_tensor.flat();
255 | const float *grad_out = &(grad_out_flat(0));
256 | auto grad_points_flat = grad_points_tensor->flat();
257 | float *grad_points = &(grad_points_flat(0));
258 | memset(grad_points, 0, sizeof(float)*b*m*c);
259 | threeinterpolate_grad_cpu(b,n,c,m,grad_out,idx,weight,grad_points);
260 | }
261 | };
262 | REGISTER_KERNEL_BUILDER(Name("ThreeInterpolateGrad").Device(DEVICE_CPU),ThreeInterpolateGradOp);
263 |
264 |
265 |
--------------------------------------------------------------------------------
/tf_ops/3d_interpolation/tf_interpolate.py:
--------------------------------------------------------------------------------
1 | '''
2 | @Author: Shuai Wang
3 | @Github: https://github.com/wsustcid
4 | @Version: 1.0.0
5 | @Date: 1970-01-01 08:00:00
6 | @LastEditTime: 2020-06-16 22:22:43
7 | @Description:
8 | '''
9 | import tensorflow as tf
10 | from tensorflow.python.framework import ops
11 | import sys
12 | import os
13 | BASE_DIR = os.path.dirname(__file__)
14 | sys.path.append(BASE_DIR)
15 | interpolate_module=tf.load_op_library(os.path.join(BASE_DIR, 'tf_interpolate_so.so'))
16 | def three_nn(xyz1, xyz2):
17 | '''
18 | Input:
19 | xyz1: (b,n,3) float32 array, unknown points
20 | xyz2: (b,m,3) float32 array, known points
21 | Output:
22 | dist: (b,n,3) float32 array, distances to known points
23 | idx: (b,n,3) int32 array, indices to known points
24 | '''
25 | return interpolate_module.three_nn(xyz1, xyz2)
26 | ops.NoGradient('ThreeNN')
27 | def three_interpolate(points, idx, weight):
28 | '''
29 | Input:
30 | points: (b,m,c) float32 array, known points
31 | idx: (b,n,3) int32 array, indices to known points
32 | weight: (b,n,3) float32 array, weights on known points
33 | Output:
34 | out: (b,n,c) float32 array, interpolated point values
35 | '''
36 | return interpolate_module.three_interpolate(points, idx, weight)
37 | @tf.RegisterGradient('ThreeInterpolate')
38 | def _three_interpolate_grad(op, grad_out):
39 | points = op.inputs[0]
40 | idx = op.inputs[1]
41 | weight = op.inputs[2]
42 | return [interpolate_module.three_interpolate_grad(points, idx, weight, grad_out), None, None]
43 |
44 | if __name__=='__main__':
45 | import numpy as np
46 | import time
47 | np.random.seed(100)
48 | pts = np.random.random((32,128,64)).astype('float32')
49 | tmp1 = np.random.random((32,512,3)).astype('float32')
50 | tmp2 = np.random.random((32,128,3)).astype('float32')
51 | with tf.device('/cpu:0'):
52 | points = tf.constant(pts)
53 | xyz1 = tf.constant(tmp1)
54 | xyz2 = tf.constant(tmp2)
55 | dist, idx = three_nn(xyz1, xyz2)
56 | weight = tf.ones_like(dist)/3.0
57 | interpolated_points = three_interpolate(points, idx, weight)
58 | with tf.Session('') as sess:
59 | now = time.time()
60 | for _ in range(100):
61 | ret = sess.run(interpolated_points)
62 | print(time.time() - now)
63 | print(ret.shape, ret.dtype)
64 | #print ret
65 |
66 |
67 |
68 |
--------------------------------------------------------------------------------
/tf_ops/3d_interpolation/tf_interpolate_compile.sh:
--------------------------------------------------------------------------------
1 | # TF1.2
2 | #g++ -std=c++11 tf_interpolate.cpp -o tf_interpolate_so.so -shared -fPIC -I /usr/local/lib/python2.7/dist-packages/tensorflow/include -I /usr/local/cuda-8.0/include -lcudart -L /usr/local/cuda-8.0/lib64/ -O2 -D_GLIBCXX_USE_CXX11_ABI=0
3 |
4 | # TF1.4
5 | #g++ -std=c++11 tf_interpolate.cpp -o tf_interpolate_so.so -shared -fPIC -I /usr/local/lib/python2.7/dist-packages/tensorflow/include -I /usr/local/cuda-8.0/include -I /usr/local/lib/python2.7/dist-packages/tensorflow/include/external/nsync/public -lcudart -L /usr/local/cuda-8.0/lib64/ -L/usr/local/lib/python2.7/dist-packages/tensorflow -ltensorflow_framework -O2 -D_GLIBCXX_USE_CXX11_ABI=0
6 |
7 | # TF1.12 (local venv)
8 | g++ -std=c++11 tf_interpolate.cpp -o tf_interpolate_so.so -shared -fPIC -I /home/ubuntu16/venv/lib/python3.5/site-packages/tensorflow/include -I /usr/local/cuda-9.0/include -I /home/ubuntu16/venv/lib/python3.5/site-packages/tensorflow/include/external/nsync/public -lcudart -L /usr/local/cuda-9.0/lib64/ -L/home/ubuntu16/venv/lib/python3.5/site-packages/tensorflow -ltensorflow_framework -O2 -D_GLIBCXX_USE_CXX11_ABI=0
--------------------------------------------------------------------------------
/tf_ops/3d_interpolation/tf_interpolate_op_test.py:
--------------------------------------------------------------------------------
1 | import tensorflow as tf
2 | import numpy as np
3 | from tf_interpolate import three_nn, three_interpolate
4 |
5 | class GroupPointTest(tf.test.TestCase):
6 | def test(self):
7 | pass
8 |
9 | def test_grad(self):
10 | with self.test_session():
11 | points = tf.constant(np.random.random((1,8,16)).astype('float32'))
12 | print points
13 | xyz1 = tf.constant(np.random.random((1,128,3)).astype('float32'))
14 | xyz2 = tf.constant(np.random.random((1,8,3)).astype('float32'))
15 | dist, idx = three_nn(xyz1, xyz2)
16 | weight = tf.ones_like(dist)/3.0
17 | interpolated_points = three_interpolate(points, idx, weight)
18 | print interpolated_points
19 | err = tf.test.compute_gradient_error(points, (1,8,16), interpolated_points, (1,128,16))
20 | print err
21 | self.assertLess(err, 1e-4)
22 |
23 | if __name__=='__main__':
24 | tf.test.main()
25 |
--------------------------------------------------------------------------------
/tf_ops/3d_interpolation/tf_interpolate_so.so:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/wsustcid/FlowDriveNet/3604495269ae45e5b43964046104f685ec66e383/tf_ops/3d_interpolation/tf_interpolate_so.so
--------------------------------------------------------------------------------
/tf_ops/3d_interpolation/visu_interpolation.py:
--------------------------------------------------------------------------------
1 | ''' Visualize part segmentation '''
2 | import os
3 | import sys
4 | ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
5 | sys.path.append('/home/rqi/Projects/toolkits/visualization')
6 | from show3d_balls import showpoints
7 | import numpy as np
8 | from tf_interpolate import three_nn, three_interpolate
9 | import tensorflow as tf
10 |
11 |
12 | pts2 = np.array([[0,0,1],[1,0,0],[0,1,0],[1,1,0]]).astype('float32')
13 | xyz1 = np.random.random((100,3)).astype('float32')
14 | xyz2 = np.array([[0,0,0],[1,0,0],[0,1,0],[1,1,1]]).astype('float32')
15 |
16 | def fun(xyz1,xyz2,pts2):
17 | with tf.device('/cpu:0'):
18 | points = tf.constant(np.expand_dims(pts2,0))
19 | xyz1 = tf.constant(np.expand_dims(xyz1,0))
20 | xyz2 = tf.constant(np.expand_dims(xyz2,0))
21 | dist, idx = three_nn(xyz1, xyz2)
22 | #weight = tf.ones_like(dist)/3.0
23 | dist = tf.maximum(dist, 1e-10)
24 | norm = tf.reduce_sum((1.0/dist),axis=2,keep_dims=True)
25 | norm = tf.tile(norm, [1,1,3])
26 | print norm
27 | weight = (1.0/dist) / norm
28 | interpolated_points = three_interpolate(points, idx, weight)
29 | with tf.Session('') as sess:
30 | tmp,pts1,d,w = sess.run([xyz1, interpolated_points, dist, weight])
31 | #print w
32 | pts1 = pts1.squeeze()
33 | return pts1
34 |
35 | pts1 = fun(xyz1,xyz2,pts2)
36 | all_pts = np.zeros((104,3))
37 | all_pts[0:100,:] = pts1
38 | all_pts[100:,:] = pts2
39 | all_xyz = np.zeros((104,3))
40 | all_xyz[0:100,:]=xyz1
41 | all_xyz[100:,:]=xyz2
42 | showpoints(xyz2, pts2, ballradius=8)
43 | showpoints(xyz1, pts1, ballradius=8)
44 | showpoints(all_xyz, all_pts, ballradius=8)
45 |
--------------------------------------------------------------------------------
/tf_ops/grouping/__pycache__/tf_grouping.cpython-35.pyc:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/wsustcid/FlowDriveNet/3604495269ae45e5b43964046104f685ec66e383/tf_ops/grouping/__pycache__/tf_grouping.cpython-35.pyc
--------------------------------------------------------------------------------
/tf_ops/grouping/test/compile.sh:
--------------------------------------------------------------------------------
1 | g++ query_ball_point.cpp -o query_ball_point
2 | nvcc query_ball_point.cu -o query_ball_point_cuda
3 | nvcc query_ball_point_block.cu -o query_ball_point_block
4 | nvcc query_ball_point_grid.cu -o query_ball_point_grid
5 | g++ -Wall selection_sort.cpp -o selection_sort
6 | nvcc selection_sort.cu -o selection_sort_cuda
7 |
--------------------------------------------------------------------------------
/tf_ops/grouping/test/query_ball_point.cpp:
--------------------------------------------------------------------------------
1 | #include
2 | #include
3 | #include // memset
4 | #include // rand, RAND_MAX
5 | #include // sqrtf
6 | #include
7 | #include
8 | using namespace std;
9 | float randomf(){
10 | return (rand()+0.5)/(RAND_MAX+1.0);
11 | }
12 | static double get_time(){
13 | timespec tp;
14 | clock_gettime(CLOCK_MONOTONIC,&tp);
15 | return tp.tv_sec+tp.tv_nsec*1e-9;
16 | }
17 | // input: radius (1), nsample (1), xyz1 (b,n,3), xyz2 (b,m,3)
18 | // output: idx (b,m,nsample)
19 | void query_ball_point_cpu(int b, int n, int m, float radius, int nsample, const float *xyz1, const float *xyz2, int *idx) {
20 | for (int i=0;i
2 | #include
3 | #include // memset
4 | #include // rand, RAND_MAX
5 | #include // sqrtf
6 | #include
7 | #include
8 | using namespace std;
9 | float randomf(){
10 | return (rand()+0.5)/(RAND_MAX+1.0);
11 | }
12 | static double get_time(){
13 | timespec tp;
14 | clock_gettime(CLOCK_MONOTONIC,&tp);
15 | return tp.tv_sec+tp.tv_nsec*1e-9;
16 | }
17 | // input: radius (1), nsample (1), xyz1 (b,n,3), xyz2 (b,m,3)
18 | // output: idx (b,m,nsample)
19 | __global__ void query_ball_point_gpu(int b, int n, int m, float radius, int nsample, const float *xyz1, const float *xyz2, int *idx) {
20 | for (int i=0;i>>(b,n,m,radius,nsample,xyz1,xyz2,idx);
113 | cudaDeviceSynchronize();
114 | printf("query_ball_point gpu time %f\n",get_time()-t0);
115 |
116 | t0=get_time();
117 | group_point_gpu<<<1,1>>>(b,n,c,m,nsample,points,idx,out);
118 | cudaDeviceSynchronize();
119 | printf("grou_point gpu time %f\n",get_time()-t0);
120 |
121 | t0=get_time();
122 | group_point_grad_gpu<<<1,1>>>(b,n,c,m,nsample,grad_out,idx,grad_points);
123 | cudaDeviceSynchronize();
124 | printf("grou_point_grad gpu time %f\n",get_time()-t0);
125 |
126 | cudaFree(xyz1);
127 | cudaFree(xyz2);
128 | cudaFree(points);
129 | cudaFree(idx);
130 | cudaFree(out);
131 | cudaFree(grad_out);
132 | cudaFree(grad_points);
133 | return 0;
134 | }
135 |
--------------------------------------------------------------------------------
/tf_ops/grouping/test/query_ball_point_block.cu:
--------------------------------------------------------------------------------
1 | #include
2 | #include
3 | #include // memset
4 | #include // rand, RAND_MAX
5 | #include // sqrtf
6 | #include
7 | #include
8 | using namespace std;
9 | float randomf(){
10 | return (rand()+0.5)/(RAND_MAX+1.0);
11 | }
12 | static double get_time(){
13 | timespec tp;
14 | clock_gettime(CLOCK_MONOTONIC,&tp);
15 | return tp.tv_sec+tp.tv_nsec*1e-9;
16 | }
17 | // input: radius (1), nsample (1), xyz1 (b,n,3), xyz2 (b,m,3)
18 | // output: idx (b,m,nsample)
19 | __global__ void query_ball_point_gpu(int b, int n, int m, float radius, int nsample, const float *xyz1, const float *xyz2, int *idx) {
20 | int index = threadIdx.x;
21 | xyz1 += n*3*index;
22 | xyz2 += m*3*index;
23 | idx += m*nsample*index;
24 |
25 | for (int j=0;j>>(b,n,m,radius,nsample,xyz1,xyz2,idx);
113 | cudaDeviceSynchronize();
114 | printf("query_ball_point gpu time %f\n",get_time()-t0);
115 |
116 | t0=get_time();
117 | group_point_gpu<<<1,b>>>(b,n,c,m,nsample,points,idx,out);
118 | cudaDeviceSynchronize();
119 | printf("grou_point gpu time %f\n",get_time()-t0);
120 |
121 | t0=get_time();
122 | group_point_grad_gpu<<<1,b>>>(b,n,c,m,nsample,grad_out,idx,grad_points);
123 | cudaDeviceSynchronize();
124 | printf("grou_point_grad gpu time %f\n",get_time()-t0);
125 |
126 | cudaFree(xyz1);
127 | cudaFree(xyz2);
128 | cudaFree(points);
129 | cudaFree(idx);
130 | cudaFree(out);
131 | cudaFree(grad_out);
132 | cudaFree(grad_points);
133 | return 0;
134 | }
135 |
--------------------------------------------------------------------------------
/tf_ops/grouping/test/query_ball_point_grid.cu:
--------------------------------------------------------------------------------
1 | #include
2 | #include
3 | #include // memset
4 | #include // rand, RAND_MAX
5 | #include // sqrtf
6 | #include
7 | #include
8 | using namespace std;
9 | float randomf(){
10 | return (rand()+0.5)/(RAND_MAX+1.0);
11 | }
12 | static double get_time(){
13 | timespec tp;
14 | clock_gettime(CLOCK_MONOTONIC,&tp);
15 | return tp.tv_sec+tp.tv_nsec*1e-9;
16 | }
17 | // input: radius (1), nsample (1), xyz1 (b,n,3), xyz2 (b,m,3)
18 | // output: idx (b,m,nsample)
19 | __global__ void query_ball_point_gpu(int b, int n, int m, float radius, int nsample, const float *xyz1, const float *xyz2, int *idx) {
20 | int batch_index = blockIdx.x;
21 | xyz1 += n*3*batch_index;
22 | xyz2 += m*3*batch_index;
23 | idx += m*nsample*batch_index;
24 |
25 | int index = threadIdx.x;
26 | int stride = blockDim.x;
27 |
28 | for (int j=index;j>>(b,n,m,radius,nsample,xyz1,xyz2,idx);
123 | cudaDeviceSynchronize();
124 | printf("query_ball_point gpu time %f\n",get_time()-t0);
125 |
126 | t0=get_time();
127 | group_point_gpu<<>>(b,n,c,m,nsample,points,idx,out);
128 | cudaDeviceSynchronize();
129 | printf("grou_point gpu time %f\n",get_time()-t0);
130 |
131 | t0=get_time();
132 | group_point_grad_gpu<<>>(b,n,c,m,nsample,grad_out,idx,grad_points);
133 | cudaDeviceSynchronize();
134 | printf("grou_point_grad gpu time %f\n",get_time()-t0);
135 |
136 | cudaFree(xyz1);
137 | cudaFree(xyz2);
138 | cudaFree(points);
139 | cudaFree(idx);
140 | cudaFree(out);
141 | cudaFree(grad_out);
142 | cudaFree(grad_points);
143 | return 0;
144 | }
145 |
--------------------------------------------------------------------------------
/tf_ops/grouping/test/selection_sort.cpp:
--------------------------------------------------------------------------------
1 | #include
2 | #include
3 | #include // memset
4 | #include // rand, RAND_MAX
5 | #include // sqrtf
6 | #include
7 | #include
8 | using namespace std;
9 | float randomf(){
10 | return (rand()+0.5)/(RAND_MAX+1.0);
11 | }
12 | static double get_time(){
13 | timespec tp;
14 | clock_gettime(CLOCK_MONOTONIC,&tp);
15 | return tp.tv_sec+tp.tv_nsec*1e-9;
16 | }
17 |
18 | // input: k (1), distance matrix dist (b,m,n)
19 | // output: idx (b,m,n), val (b,m,n)
20 | void selection_sort_cpu(int b, int n, int m, int k, const float *dist, int *idx, float *val) {
21 | float *p_dist;
22 | float tmp;
23 | int tmpi;
24 | for (int i=0;i
2 | #include
3 | #include // memset
4 | #include // rand, RAND_MAX
5 | #include // sqrtf
6 | #include
7 | #include
8 | using namespace std;
9 | float randomf(){
10 | return (rand()+0.5)/(RAND_MAX+1.0);
11 | }
12 | static double get_time(){
13 | timespec tp;
14 | clock_gettime(CLOCK_MONOTONIC,&tp);
15 | return tp.tv_sec+tp.tv_nsec*1e-9;
16 | }
17 |
18 | // input: k (1), distance matrix dist (b,m,n)
19 | // output: idx (b,m,k), val (b,m,k)
20 | __global__ void selection_sort_gpu(int b, int n, int m, int k, float *dist, int *idx, float *val) {
21 | int batch_index = blockIdx.x;
22 | dist+=m*n*batch_index;
23 | idx+=m*k*batch_index;
24 | val+=m*k*batch_index;
25 |
26 | int index = threadIdx.x;
27 | int stride = blockDim.x;
28 |
29 | float *p_dist;
30 | for (int j=index;j>>(b,n,m,k,dist,idx,val);
68 | cudaDeviceSynchronize();
69 | printf("selection sort cpu time %f\n",get_time()-t0);
70 |
71 | return 0;
72 | }
73 |
--------------------------------------------------------------------------------
/tf_ops/grouping/test/selection_sort_const.cu:
--------------------------------------------------------------------------------
1 | #include
2 | #include
3 | #include // memset
4 | #include // rand, RAND_MAX
5 | #include // sqrtf
6 | #include
7 | #include
8 | using namespace std;
9 | float randomf(){
10 | return (rand()+0.5)/(RAND_MAX+1.0);
11 | }
12 | static double get_time(){
13 | timespec tp;
14 | clock_gettime(CLOCK_MONOTONIC,&tp);
15 | return tp.tv_sec+tp.tv_nsec*1e-9;
16 | }
17 |
18 | // input: k (1), distance matrix dist (b,m,n)
19 | // output: idx (b,m,n), dist_out (b,m,n)
20 | __global__ void selection_sort_gpu(int b, int n, int m, int k, const float *dist, int *outi, float *out) {
21 | int batch_index = blockIdx.x;
22 | dist+=m*n*batch_index;
23 | outi+=m*n*batch_index;
24 | out+=m*n*batch_index;
25 |
26 | int index = threadIdx.x;
27 | int stride = blockDim.x;
28 |
29 | // copy from dist to dist_out
30 | for (int j=index;j>>(b,n,m,k,dist,idx,dist_out);
84 | cudaDeviceSynchronize();
85 | printf("selection sort cpu time %f\n",get_time()-t0);
86 |
87 | //for (int i=0;i
2 | #include
3 | #include // memset
4 | #include // rand, RAND_MAX
5 | #include // sqrtf
6 | #include "tensorflow/core/framework/op.h"
7 | #include "tensorflow/core/framework/op_kernel.h"
8 | #include "tensorflow/core/framework/shape_inference.h"
9 | #include "tensorflow/core/framework/common_shape_fns.h"
10 | #include
11 | using namespace tensorflow;
12 |
13 | REGISTER_OP("QueryBallPoint")
14 | .Attr("radius: float")
15 | .Attr("nsample: int")
16 | .Input("xyz1: float32")
17 | .Input("xyz2: float32")
18 | .Output("idx: int32")
19 | .Output("pts_cnt: int32")
20 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
21 | ::tensorflow::shape_inference::ShapeHandle dims2; // batch_size * npoint * 3
22 | c->WithRank(c->input(1), 3, &dims2);
23 | int nsample;
24 | TF_RETURN_IF_ERROR(c->GetAttr("nsample", &nsample));
25 | ::tensorflow::shape_inference::ShapeHandle output1 = c->MakeShape({c->Dim(dims2, 0), c->Dim(dims2, 1), nsample});
26 | c->set_output(0, output1);
27 | ::tensorflow::shape_inference::ShapeHandle output2 = c->MakeShape({c->Dim(dims2, 0), c->Dim(dims2, 1)});
28 | c->set_output(1, output2);
29 | return Status::OK();
30 | });
31 | REGISTER_OP("SelectionSort")
32 | .Attr("k: int")
33 | .Input("dist: float32")
34 | .Output("outi: int32")
35 | .Output("out: float32")
36 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
37 | c->set_output(0, c->input(0));
38 | c->set_output(1, c->input(0));
39 | return Status::OK();
40 | });
41 | REGISTER_OP("GroupPoint")
42 | .Input("points: float32")
43 | .Input("idx: int32")
44 | .Output("out: float32")
45 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
46 | ::tensorflow::shape_inference::ShapeHandle dims1; // batch_size * ndataset * channels
47 | c->WithRank(c->input(0), 3, &dims1);
48 | ::tensorflow::shape_inference::ShapeHandle dims2; // batch_size * npoints * nsample
49 | c->WithRank(c->input(1), 3, &dims2);
50 | // batch_size * npoints * nsample * channels
51 | ::tensorflow::shape_inference::ShapeHandle output = c->MakeShape({c->Dim(dims2, 0), c->Dim(dims2, 1), c->Dim(dims2, 2), c->Dim(dims1, 2)});
52 | c->set_output(0, output);
53 | return Status::OK();
54 | });
55 | REGISTER_OP("GroupPointGrad")
56 | .Input("points: float32")
57 | .Input("idx: int32")
58 | .Input("grad_out: float32")
59 | .Output("grad_points: float32")
60 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
61 | c->set_output(0, c->input(0));
62 | return Status::OK();
63 | });
64 |
65 |
66 | void queryBallPointLauncher(int b, int n, int m, float radius, int nsample, const float *xyz1, const float *xyz2, int *idx, int *pts_cnt);
67 | class QueryBallPointGpuOp : public OpKernel {
68 | public:
69 | explicit QueryBallPointGpuOp(OpKernelConstruction* context) : OpKernel(context) {
70 | OP_REQUIRES_OK(context, context->GetAttr("radius", &radius_));
71 | OP_REQUIRES(context, radius_ > 0, errors::InvalidArgument("QueryBallPoint expects positive radius"));
72 |
73 | OP_REQUIRES_OK(context, context->GetAttr("nsample", &nsample_));
74 | OP_REQUIRES(context, nsample_ > 0, errors::InvalidArgument("QueryBallPoint expects positive nsample"));
75 | }
76 |
77 | void Compute(OpKernelContext* context) override {
78 | const Tensor& xyz1_tensor = context->input(0);
79 | OP_REQUIRES(context, xyz1_tensor.dims()==3 && xyz1_tensor.shape().dim_size(2)==3, errors::InvalidArgument("QueryBallPoint expects (batch_size, ndataset, 3) xyz1 shape."));
80 | int b = xyz1_tensor.shape().dim_size(0);
81 | int n = xyz1_tensor.shape().dim_size(1);
82 |
83 | const Tensor& xyz2_tensor = context->input(1);
84 | OP_REQUIRES(context, xyz2_tensor.dims()==3 && xyz2_tensor.shape().dim_size(2)==3, errors::InvalidArgument("QueryBallPoint expects (batch_size, npoint, 3) xyz2 shape."));
85 | int m = xyz2_tensor.shape().dim_size(1);
86 |
87 | Tensor *idx_tensor = nullptr;
88 | OP_REQUIRES_OK(context, context->allocate_output(0, TensorShape{b,m,nsample_}, &idx_tensor));
89 | Tensor *pts_cnt_tensor = nullptr;
90 | OP_REQUIRES_OK(context, context->allocate_output(1, TensorShape{b,m}, &pts_cnt_tensor));
91 |
92 | auto xyz1_flat = xyz1_tensor.flat();
93 | const float *xyz1 = &(xyz1_flat(0));
94 | auto xyz2_flat = xyz2_tensor.flat();
95 | const float *xyz2 = &(xyz2_flat(0));
96 | auto idx_flat = idx_tensor->flat();
97 | int *idx = &(idx_flat(0));
98 | auto pts_cnt_flat = pts_cnt_tensor->flat();
99 | int *pts_cnt = &(pts_cnt_flat(0));
100 | queryBallPointLauncher(b,n,m,radius_,nsample_,xyz1,xyz2,idx,pts_cnt);
101 | }
102 | private:
103 | float radius_;
104 | int nsample_;
105 | };
106 | REGISTER_KERNEL_BUILDER(Name("QueryBallPoint").Device(DEVICE_GPU), QueryBallPointGpuOp);
107 |
108 | void selectionSortLauncher(int b, int n, int m, int k, const float *dist, int *outi, float *out);
109 | class SelectionSortGpuOp : public OpKernel {
110 | public:
111 | explicit SelectionSortGpuOp(OpKernelConstruction* context) : OpKernel(context) {
112 | OP_REQUIRES_OK(context, context->GetAttr("k", &k_));
113 | OP_REQUIRES(context, k_ > 0, errors::InvalidArgument("SelectionSort expects positive k"));
114 | }
115 |
116 | void Compute(OpKernelContext* context) override {
117 | const Tensor& dist_tensor = context->input(0);
118 | OP_REQUIRES(context, dist_tensor.dims()==3, errors::InvalidArgument("SelectionSort expects (b,m,n) dist shape."));
119 | int b = dist_tensor.shape().dim_size(0);
120 | int m = dist_tensor.shape().dim_size(1);
121 | int n = dist_tensor.shape().dim_size(2);
122 |
123 | Tensor *outi_tensor = nullptr;
124 | OP_REQUIRES_OK(context, context->allocate_output(0, TensorShape{b,m,n}, &outi_tensor));
125 | Tensor *out_tensor = nullptr;
126 | OP_REQUIRES_OK(context, context->allocate_output(1, TensorShape{b,m,n}, &out_tensor));
127 |
128 | auto dist_flat = dist_tensor.flat();
129 | const float *dist = &(dist_flat(0));
130 | auto outi_flat = outi_tensor->flat();
131 | int *outi = &(outi_flat(0));
132 | auto out_flat = out_tensor->flat();
133 | float *out = &(out_flat(0));
134 | selectionSortLauncher(b,n,m,k_,dist,outi,out);
135 | }
136 | private:
137 | int k_;
138 | };
139 | REGISTER_KERNEL_BUILDER(Name("SelectionSort").Device(DEVICE_GPU), SelectionSortGpuOp);
140 |
141 |
142 | void groupPointLauncher(int b, int n, int c, int m, int nsample, const float *points, const int *idx, float *out);
143 | class GroupPointGpuOp: public OpKernel{
144 | public:
145 | explicit GroupPointGpuOp(OpKernelConstruction * context):OpKernel(context){}
146 |
147 | void Compute(OpKernelContext * context) override {
148 | const Tensor& points_tensor=context->input(0);
149 | OP_REQUIRES(context, points_tensor.dims()==3, errors::InvalidArgument("GroupPoint expects (batch_size, num_points, channel) points shape"));
150 | int b = points_tensor.shape().dim_size(0);
151 | int n = points_tensor.shape().dim_size(1);
152 | int c = points_tensor.shape().dim_size(2);
153 |
154 | const Tensor& idx_tensor=context->input(1);
155 | OP_REQUIRES(context,idx_tensor.dims()==3 && idx_tensor.shape().dim_size(0)==b, errors::InvalidArgument("GroupPoint expects (batch_size, npoints, nsample) idx shape"));
156 | int m = idx_tensor.shape().dim_size(1);
157 | int nsample = idx_tensor.shape().dim_size(2);
158 |
159 | Tensor * out_tensor = nullptr;
160 | OP_REQUIRES_OK(context, context->allocate_output(0,TensorShape{b,m,nsample,c}, &out_tensor));
161 |
162 | auto points_flat = points_tensor.flat();
163 | const float *points = &(points_flat(0));
164 | auto idx_flat = idx_tensor.flat();
165 | const int *idx = &(idx_flat(0));
166 | auto out_flat = out_tensor->flat();
167 | float *out = &(out_flat(0));
168 | groupPointLauncher(b,n,c,m,nsample,points,idx,out);
169 | }
170 | };
171 | REGISTER_KERNEL_BUILDER(Name("GroupPoint").Device(DEVICE_GPU),GroupPointGpuOp);
172 |
173 | void groupPointGradLauncher(int b, int n, int c, int m, int nsample, const float *grad_out, const int *idx, float *grad_points);
174 | class GroupPointGradGpuOp: public OpKernel{
175 | public:
176 | explicit GroupPointGradGpuOp(OpKernelConstruction * context):OpKernel(context){}
177 |
178 | void Compute(OpKernelContext * context) override {
179 | const Tensor& points_tensor=context->input(0);
180 | OP_REQUIRES(context, points_tensor.dims()==3, errors::InvalidArgument("GroupPointGrad expects (batch_size, num_points, channel) points shape"));
181 | int b = points_tensor.shape().dim_size(0);
182 | int n = points_tensor.shape().dim_size(1);
183 | int c = points_tensor.shape().dim_size(2);
184 |
185 | const Tensor& idx_tensor=context->input(1);
186 | OP_REQUIRES(context,idx_tensor.dims()==3 && idx_tensor.shape().dim_size(0)==b, errors::InvalidArgument("GroupPointGrad expects (batch_size, npoints, nsample) idx shape"));
187 | int m = idx_tensor.shape().dim_size(1);
188 | int nsample = idx_tensor.shape().dim_size(2);
189 |
190 | const Tensor& grad_out_tensor=context->input(2);
191 | OP_REQUIRES(context,grad_out_tensor.dims()==4 && grad_out_tensor.shape().dim_size(0)==b && grad_out_tensor.shape().dim_size(1)==m && grad_out_tensor.shape().dim_size(2)==nsample && grad_out_tensor.shape().dim_size(3)==c, errors::InvalidArgument("GroupPointGrad expects (batch_size, npoints, nsample, channel) grad_out shape"));
192 |
193 | Tensor * grad_points_tensor = nullptr;
194 | OP_REQUIRES_OK(context, context->allocate_output(0,TensorShape{b,n,c}, &grad_points_tensor));
195 |
196 | auto points_flat = points_tensor.flat();
197 | const float *points = &(points_flat(0));
198 | auto idx_flat = idx_tensor.flat();
199 | const int *idx = &(idx_flat(0));
200 | auto grad_out_flat = grad_out_tensor.flat();
201 | const float *grad_out = &(grad_out_flat(0));
202 | auto grad_points_flat = grad_points_tensor->flat();
203 | float *grad_points = &(grad_points_flat(0));
204 | cudaMemset(grad_points, 0, sizeof(float)*b*n*c);
205 | groupPointGradLauncher(b,n,c,m,nsample,grad_out,idx,grad_points);
206 | }
207 | };
208 | REGISTER_KERNEL_BUILDER(Name("GroupPointGrad").Device(DEVICE_GPU),GroupPointGradGpuOp);
209 |
210 |
211 |
--------------------------------------------------------------------------------
/tf_ops/grouping/tf_grouping.py:
--------------------------------------------------------------------------------
1 | import tensorflow as tf
2 | from tensorflow.python.framework import ops
3 | import sys
4 | import os
5 | BASE_DIR = os.path.dirname(os.path.abspath(__file__))
6 | sys.path.append(BASE_DIR)
7 | grouping_module=tf.load_op_library(os.path.join(BASE_DIR, 'tf_grouping_so.so'))
8 | def query_ball_point(radius, nsample, xyz1, xyz2):
9 | '''
10 | Input:
11 | radius: float32, ball search radius
12 | nsample: int32, number of points selected in each ball region
13 | xyz1: (batch_size, ndataset, 3) float32 array, input points
14 | xyz2: (batch_size, npoint, 3) float32 array, query points
15 | Output:
16 | idx: (batch_size, npoint, nsample) int32 array, indices to input points
17 | pts_cnt: (batch_size, npoint) int32 array, number of unique points in each local region
18 | '''
19 | #return grouping_module.query_ball_point(radius, nsample, xyz1, xyz2)
20 | return grouping_module.query_ball_point(xyz1, xyz2, radius, nsample)
21 | ops.NoGradient('QueryBallPoint')
22 | def select_top_k(k, dist):
23 | '''
24 | Input:
25 | k: int32, number of k SMALLEST elements selected
26 | dist: (b,m,n) float32 array, distance matrix, m query points, n dataset points
27 | Output:
28 | idx: (b,m,n) int32 array, first k in n are indices to the top k
29 | dist_out: (b,m,n) float32 array, first k in n are the top k
30 | '''
31 | return grouping_module.selection_sort(dist, k)
32 | ops.NoGradient('SelectionSort')
33 | def group_point(points, idx):
34 | '''
35 | Input:
36 | points: (batch_size, ndataset, channel) float32 array, points to sample from
37 | idx: (batch_size, npoint, nsample) int32 array, indices to points
38 | Output:
39 | out: (batch_size, npoint, nsample, channel) float32 array, values sampled from points
40 | '''
41 | return grouping_module.group_point(points, idx)
42 | @tf.RegisterGradient('GroupPoint')
43 | def _group_point_grad(op, grad_out):
44 | points = op.inputs[0]
45 | idx = op.inputs[1]
46 | return [grouping_module.group_point_grad(points, idx, grad_out), None]
47 |
48 | def knn_point(k, xyz1, xyz2):
49 | '''
50 | Input:
51 | k: int32, number of k in k-nn search
52 | xyz1: (batch_size, ndataset, c) float32 array, input points
53 | xyz2: (batch_size, npoint, c) float32 array, query points
54 | Output:
55 | val: (batch_size, npoint, k) float32 array, L2 distances
56 | idx: (batch_size, npoint, k) int32 array, indices to input points
57 | '''
58 | b = xyz1.get_shape()[0].value
59 | n = xyz1.get_shape()[1].value
60 | c = xyz1.get_shape()[2].value
61 | m = xyz2.get_shape()[1].value
62 | #print(b, n, c, m)
63 | #print(xyz1, (b,1,n,c))
64 | xyz1 = tf.tile(tf.reshape(xyz1, (b,1,n,c)), [1,m,1,1])
65 | xyz2 = tf.tile(tf.reshape(xyz2, (b,m,1,c)), [1,1,n,1])
66 | dist = tf.reduce_sum((xyz1-xyz2)**2, -1)
67 | #print(dist, k)
68 | outi, out = select_top_k(k, dist)
69 | idx = tf.slice(outi, [0,0,0], [-1,-1,k])
70 | val = tf.slice(out, [0,0,0], [-1,-1,k])
71 | #print(idx, val)
72 | #val, idx = tf.nn.top_k(-dist, k=k) # ONLY SUPPORT CPU
73 | return val, idx
74 |
75 | if __name__=='__main__':
76 | knn=True
77 | import numpy as np
78 | import time
79 | np.random.seed(100)
80 | pts = np.random.random((32,512,64)).astype('float32')
81 | tmp1 = np.random.random((32,512,3)).astype('float32')
82 | tmp2 = np.random.random((32,128,3)).astype('float32')
83 | with tf.device('/gpu:1'):
84 | points = tf.constant(pts)
85 | xyz1 = tf.constant(tmp1)
86 | xyz2 = tf.constant(tmp2)
87 | radius = 0.1
88 | nsample = 64
89 | if knn:
90 | _, idx = knn_point(nsample, xyz1, xyz2)
91 | grouped_points = group_point(points, idx)
92 | else:
93 | idx, _ = query_ball_point(radius, nsample, xyz1, xyz2)
94 | grouped_points = group_point(points, idx)
95 | #grouped_points_grad = tf.ones_like(grouped_points)
96 | #points_grad = tf.gradients(grouped_points, points, grouped_points_grad)
97 | with tf.Session('') as sess:
98 | now = time.time()
99 | for _ in range(100):
100 | ret = sess.run(grouped_points)
101 | print(time.time() - now)
102 | print(ret.shape, ret.dtype)
103 | print(ret)
104 |
105 |
106 |
--------------------------------------------------------------------------------
/tf_ops/grouping/tf_grouping_compile.sh:
--------------------------------------------------------------------------------
1 | #/bin/bash
2 | #/usr/local/cuda-9.0/bin/nvcc tf_grouping_g.cu -o tf_grouping_g.cu.o -c -O2 -DGOOGLE_CUDA=1 -x cu -Xcompiler -fPIC
3 |
4 | # TF1.2
5 | #g++ -std=c++11 tf_grouping.cpp tf_grouping_g.cu.o -o tf_grouping_so.so -shared -fPIC -I /usr/local/lib/python2.7/dist-packages/tensorflow/include -I /usr/local/cuda-8.0/include -lcudart -L /usr/local/cuda-8.0/lib64/ -O2 -D_GLIBCXX_USE_CXX11_ABI=0
6 |
7 | # TF1.4
8 | #g++ -std=c++11 tf_grouping.cpp tf_grouping_g.cu.o -o tf_grouping_so.so -shared -fPIC -I /usr/local/lib/python2.7/dist-packages/tensorflow/include -I /usr/local/cuda-8.0/include -I /usr/local/lib/python2.7/dist-packages/tensorflow/include/external/nsync/public -lcudart -L /usr/local/cuda-8.0/lib64/ -L/usr/local/lib/python2.7/dist-packages/tensorflow -ltensorflow_framework -O2 -D_GLIBCXX_USE_CXX11_ABI=0
9 |
10 | # TF1.12 (venv local)
11 | #g++ -std=c++11 tf_grouping.cpp tf_grouping_g.cu.o -o tf_grouping_so.so -shared -fPIC -I /home/ubuntu16/venv/lib/python3.5/site-packages/tensorflow/include -I /usr/local/cuda-9.0/include -I /home/ubuntu16/venv/lib/python3.5/site-packages/tensorflow/include/external/nsync/public -lcudart -L /usr/local/cuda-9.0/lib64/ -L/home/ubuntu16/venv/lib/python3.5/site-packages/tensorflow -ltensorflow_framework -O2 -D_GLIBCXX_USE_CXX11_ABI=0
12 |
13 | # TF1.8 cluster version
14 | CUDA_ROOT=/usr/local/cuda-9.0
15 | TF_INC=$(python -c 'import tensorflow as tf; print(tf.sysconfig.get_include())')
16 | TF_LIB=$(python -c 'import tensorflow as tf; print(tf.sysconfig.get_lib())')
17 |
18 | echo $CUDA_ROOT
19 | echo $TF_INC
20 | echo $TF_LIB
21 |
22 | $CUDA_ROOT/bin/nvcc tf_grouping_g.cu -o tf_grouping_g.cu.o -c -O2 -DGOOGLE_CUDA=1 -x cu -Xcompiler -fPIC
23 |
24 | # TF>=1.4.0
25 | g++ -std=c++11 tf_grouping.cpp tf_grouping_g.cu.o -o tf_grouping_so.so -shared -fPIC -I$TF_INC/ -I$TF_INC/external/nsync/public -L$TF_LIB -ltensorflow_framework -I$CUDA_ROOT/include -lcudart -L$CUDA_ROOT/lib64/ -O2 -D_GLIBCXX_USE_CXX11_ABI=0
--------------------------------------------------------------------------------
/tf_ops/grouping/tf_grouping_g.cu:
--------------------------------------------------------------------------------
1 | // input: radius (1), nsample (1), xyz1 (b,n,3), xyz2 (b,m,3)
2 | // output: idx (b,m,nsample), pts_cnt (b,m)
3 | __global__ void query_ball_point_gpu(int b, int n, int m, float radius, int nsample, const float *xyz1, const float *xyz2, int *idx, int *pts_cnt) {
4 | int batch_index = blockIdx.x;
5 | xyz1 += n*3*batch_index;
6 | xyz2 += m*3*batch_index;
7 | idx += m*nsample*batch_index;
8 | pts_cnt += m*batch_index; // counting how many unique points selected in local region
9 |
10 | int index = threadIdx.x;
11 | int stride = blockDim.x;
12 |
13 | for (int j=index;j>>(b,n,m,radius,nsample,xyz1,xyz2,idx,pts_cnt);
127 | //cudaDeviceSynchronize();
128 | }
129 | void selectionSortLauncher(int b, int n, int m, int k, const float *dist, int *outi, float *out) {
130 | selection_sort_gpu<<>>(b,n,m,k,dist,outi,out);
131 | //cudaDeviceSynchronize();
132 | }
133 | void groupPointLauncher(int b, int n, int c, int m, int nsample, const float *points, const int *idx, float *out){
134 | group_point_gpu<<>>(b,n,c,m,nsample,points,idx,out);
135 | //cudaDeviceSynchronize();
136 | }
137 | void groupPointGradLauncher(int b, int n, int c, int m, int nsample, const float *grad_out, const int *idx, float *grad_points){
138 | group_point_grad_gpu<<>>(b,n,c,m,nsample,grad_out,idx,grad_points);
139 | //group_point_grad_gpu<<<1,1>>>(b,n,c,m,nsample,grad_out,idx,grad_points);
140 | //cudaDeviceSynchronize();
141 | }
142 |
--------------------------------------------------------------------------------
/tf_ops/grouping/tf_grouping_g.cu.o:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/wsustcid/FlowDriveNet/3604495269ae45e5b43964046104f685ec66e383/tf_ops/grouping/tf_grouping_g.cu.o
--------------------------------------------------------------------------------
/tf_ops/grouping/tf_grouping_op_test.py:
--------------------------------------------------------------------------------
1 | import tensorflow as tf
2 | import numpy as np
3 | from tf_grouping import query_ball_point, group_point
4 |
5 | class GroupPointTest(tf.test.TestCase):
6 | def test(self):
7 | pass
8 |
9 | def test_grad(self):
10 | with tf.device('/gpu:0'):
11 | points = tf.constant(np.random.random((1,128,16)).astype('float32'))
12 | print points
13 | xyz1 = tf.constant(np.random.random((1,128,3)).astype('float32'))
14 | xyz2 = tf.constant(np.random.random((1,8,3)).astype('float32'))
15 | radius = 0.3
16 | nsample = 32
17 | idx, pts_cnt = query_ball_point(radius, nsample, xyz1, xyz2)
18 | grouped_points = group_point(points, idx)
19 | print grouped_points
20 |
21 | with self.test_session():
22 | print "---- Going to compute gradient error"
23 | err = tf.test.compute_gradient_error(points, (1,128,16), grouped_points, (1,8,32,16))
24 | print err
25 | self.assertLess(err, 1e-4)
26 |
27 | if __name__=='__main__':
28 | tf.test.main()
29 |
--------------------------------------------------------------------------------
/tf_ops/grouping/tf_grouping_so.so:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/wsustcid/FlowDriveNet/3604495269ae45e5b43964046104f685ec66e383/tf_ops/grouping/tf_grouping_so.so
--------------------------------------------------------------------------------
/tf_ops/sampling/__pycache__/tf_sampling.cpython-35.pyc:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/wsustcid/FlowDriveNet/3604495269ae45e5b43964046104f685ec66e383/tf_ops/sampling/__pycache__/tf_sampling.cpython-35.pyc
--------------------------------------------------------------------------------
/tf_ops/sampling/tf_sampling.cpp:
--------------------------------------------------------------------------------
1 | /* Furthest point sampling
2 | * Original author: Haoqiang Fan
3 | * Modified by Charles R. Qi
4 | * All Rights Reserved. 2017.
5 | */
6 | #include "tensorflow/core/framework/op.h"
7 | #include "tensorflow/core/framework/op_kernel.h"
8 | #include "tensorflow/core/framework/shape_inference.h"
9 | #include "tensorflow/core/framework/common_shape_fns.h"
10 | #include
11 |
12 | using namespace tensorflow;
13 |
14 | REGISTER_OP("ProbSample")
15 | .Input("inp: float32")
16 | .Input("inpr: float32")
17 | .Output("out: int32")
18 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
19 | ::tensorflow::shape_inference::ShapeHandle dims1; // batch_size * ncategory
20 | c->WithRank(c->input(0), 2, &dims1);
21 | ::tensorflow::shape_inference::ShapeHandle dims2; // batch_size * npoints
22 | c->WithRank(c->input(1), 2, &dims2);
23 | // batch_size * npoints
24 | ::tensorflow::shape_inference::ShapeHandle output = c->MakeShape({c->Dim(dims2, 0), c->Dim(dims2, 1)});
25 | c->set_output(0, output);
26 | return Status::OK();
27 | });
28 | REGISTER_OP("FarthestPointSample")
29 | .Attr("npoint: int")
30 | .Input("inp: float32")
31 | .Output("out: int32")
32 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
33 | ::tensorflow::shape_inference::ShapeHandle dims1; // batch_size * npoint * 3
34 | c->WithRank(c->input(0), 3, &dims1);
35 | int npoint;
36 | TF_RETURN_IF_ERROR(c->GetAttr("npoint", &npoint));
37 | ::tensorflow::shape_inference::ShapeHandle output = c->MakeShape({c->Dim(dims1, 0), npoint});
38 | c->set_output(0, output);
39 | return Status::OK();
40 | });
41 | REGISTER_OP("GatherPoint")
42 | .Input("inp: float32")
43 | .Input("idx: int32")
44 | .Output("out: float32")
45 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
46 | ::tensorflow::shape_inference::ShapeHandle dims1; // batch_size * ndataset * 3
47 | c->WithRank(c->input(0), 3, &dims1);
48 | ::tensorflow::shape_inference::ShapeHandle dims2; // batch_size * npoints
49 | c->WithRank(c->input(1), 2, &dims2);
50 | // batch_size * npoints * 3
51 | ::tensorflow::shape_inference::ShapeHandle output = c->MakeShape({c->Dim(dims1, 0), c->Dim(dims2, 1), c->Dim(dims1, 2)});
52 | c->set_output(0, output);
53 | return Status::OK();
54 | });
55 | REGISTER_OP("GatherPointGrad")
56 | .Input("inp: float32")
57 | .Input("idx: int32")
58 | .Input("out_g: float32")
59 | .Output("inp_g: float32")
60 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
61 | c->set_output(0, c->input(0));
62 | return Status::OK();
63 | });
64 |
65 | void probsampleLauncher(int b,int n,int m,const float * inp_p,const float * inp_r,float * temp,int * out);
66 | class ProbSampleGpuOp: public OpKernel{
67 | public:
68 | explicit ProbSampleGpuOp(OpKernelConstruction* context):OpKernel(context){}
69 | void Compute(OpKernelContext * context)override{
70 | const Tensor& inp_tensor=context->input(0);
71 | const Tensor& inpr_tensor=context->input(1);
72 | auto inp_flat=inp_tensor.flat();
73 | auto inpr_flat=inpr_tensor.flat();
74 | const float * inp=&(inp_flat(0));
75 | const float * inpr=&(inpr_flat(0));
76 | OP_REQUIRES(context,inp_tensor.dims()==2,errors::InvalidArgument("ProbSample expects (batch_size,num_choices) inp shape"));
77 | int b=inp_tensor.shape().dim_size(0);
78 | int n=inp_tensor.shape().dim_size(1);
79 | OP_REQUIRES(context,inpr_tensor.dims()==2 && inpr_tensor.shape().dim_size(0)==b,errors::InvalidArgument("ProbSample expects (batch_size,num_points) inpr shape"));
80 | int m=inpr_tensor.shape().dim_size(1);
81 | Tensor * out_tensor=NULL;
82 | OP_REQUIRES_OK(context,context->allocate_output(0,TensorShape{b,m},&out_tensor));
83 | auto out_flat=out_tensor->flat();
84 | int * out=&(out_flat(0));
85 | Tensor temp_tensor;
86 | OP_REQUIRES_OK(context,context->allocate_temp(DataTypeToEnum::value,TensorShape{b,n},&temp_tensor));
87 | auto temp_flat=temp_tensor.flat();
88 | float * temp=&(temp_flat(0));
89 | probsampleLauncher(b,n,m,inp,inpr,temp,out);
90 | }
91 | };
92 | REGISTER_KERNEL_BUILDER(Name("ProbSample").Device(DEVICE_GPU), ProbSampleGpuOp);
93 |
94 | void farthestpointsamplingLauncher(int b,int n,int m,const float * inp,float * temp,int * out);
95 | class FarthestPointSampleGpuOp: public OpKernel{
96 | public:
97 | explicit FarthestPointSampleGpuOp(OpKernelConstruction* context):OpKernel(context) {
98 | OP_REQUIRES_OK(context, context->GetAttr("npoint", &npoint_));
99 | OP_REQUIRES(context, npoint_ > 0, errors::InvalidArgument("FarthestPointSample expects positive npoint"));
100 | }
101 | void Compute(OpKernelContext * context)override{
102 | int m = npoint_;
103 |
104 | const Tensor& inp_tensor=context->input(0);
105 | OP_REQUIRES(context,inp_tensor.dims()==3 && inp_tensor.shape().dim_size(2)==3,errors::InvalidArgument("FarthestPointSample expects (batch_size,num_points,3) inp shape"));
106 | int b=inp_tensor.shape().dim_size(0);
107 | int n=inp_tensor.shape().dim_size(1);
108 | auto inp_flat=inp_tensor.flat();
109 | const float * inp=&(inp_flat(0));
110 | Tensor * out_tensor;
111 | OP_REQUIRES_OK(context,context->allocate_output(0,TensorShape{b,m},&out_tensor));
112 | auto out_flat=out_tensor->flat();
113 | int * out=&(out_flat(0));
114 | Tensor temp_tensor;
115 | OP_REQUIRES_OK(context,context->allocate_temp(DataTypeToEnum::value,TensorShape{32,n},&temp_tensor));
116 | auto temp_flat=temp_tensor.flat();
117 | float * temp=&(temp_flat(0));
118 | farthestpointsamplingLauncher(b,n,m,inp,temp,out);
119 | }
120 | private:
121 | int npoint_;
122 | };
123 | REGISTER_KERNEL_BUILDER(Name("FarthestPointSample").Device(DEVICE_GPU),FarthestPointSampleGpuOp);
124 |
125 | void gatherpointLauncher(int b,int n,int m,const float * inp,const int * idx,float * out);
126 | class GatherPointGpuOp: public OpKernel{
127 | public:
128 | explicit GatherPointGpuOp(OpKernelConstruction * context):OpKernel(context){}
129 | void Compute(OpKernelContext * context)override{
130 | const Tensor& inp_tensor=context->input(0);
131 | OP_REQUIRES(context,inp_tensor.dims()==3 && inp_tensor.shape().dim_size(2)==3,errors::InvalidArgument("GatherPoint expects (batch_size,num_points,3) inp shape"));
132 | int b=inp_tensor.shape().dim_size(0);
133 | int n=inp_tensor.shape().dim_size(1);
134 | const Tensor& idx_tensor=context->input(1);
135 | OP_REQUIRES(context,idx_tensor.dims()==2 && idx_tensor.shape().dim_size(0)==b,errors::InvalidArgument("GatherPoint expects (batch_size,num_result) idx shape"));
136 | int m=idx_tensor.shape().dim_size(1);
137 | auto inp_flat=inp_tensor.flat();
138 | const float * inp=&(inp_flat(0));
139 | auto idx_flat=idx_tensor.flat();
140 | const int * idx=&(idx_flat(0));
141 | Tensor * out_tensor=NULL;
142 | OP_REQUIRES_OK(context,context->allocate_output(0,TensorShape{b,m,3},&out_tensor));
143 | auto out_flat=out_tensor->flat();
144 | float * out=&(out_flat(0));
145 | gatherpointLauncher(b,n,m,inp,idx,out);
146 | }
147 | };
148 | REGISTER_KERNEL_BUILDER(Name("GatherPoint").Device(DEVICE_GPU),GatherPointGpuOp);
149 |
150 | void scatteraddpointLauncher(int b,int n,int m,const float * out_g,const int * idx,float * inp_g);
151 | class GatherPointGradGpuOp: public OpKernel{
152 | public:
153 | explicit GatherPointGradGpuOp(OpKernelConstruction * context):OpKernel(context){}
154 | void Compute(OpKernelContext * context)override{
155 | const Tensor& inp_tensor=context->input(0);
156 | OP_REQUIRES(context,inp_tensor.dims()==3 && inp_tensor.shape().dim_size(2)==3,errors::InvalidArgument("GatherPointGradGpuOp expects (batch_size,num_points,3) inp"));
157 | int b=inp_tensor.shape().dim_size(0);
158 | int n=inp_tensor.shape().dim_size(1);
159 | const Tensor& idx_tensor=context->input(1);
160 | OP_REQUIRES(context,idx_tensor.dims()==2 && idx_tensor.shape().dim_size(0)==b,errors::InvalidArgument("GatherPointGradGpuOp expects (batch_size,num_result) idx shape"));
161 | int m=idx_tensor.shape().dim_size(1);
162 | auto inp_flat=inp_tensor.flat();
163 | const float * inp=&(inp_flat(0));
164 | auto idx_flat=idx_tensor.flat();
165 | const int * idx=&(idx_flat(0));
166 | const Tensor& out_g_tensor=context->input(2);
167 | OP_REQUIRES(context,out_g_tensor.dims()==3 && out_g_tensor.shape().dim_size(0)==b && out_g_tensor.shape().dim_size(1)==m && out_g_tensor.shape().dim_size(2)==3,errors::InvalidArgument("GatherPointGradGpuOp expects (batch_size,num_result,3) out_g shape"));
168 | auto out_g_flat=out_g_tensor.flat();
169 | const float * out_g=&(out_g_flat(0));
170 | Tensor * inp_g_tensor=NULL;
171 | OP_REQUIRES_OK(context,context->allocate_output(0,TensorShape{b,n,3},&inp_g_tensor));
172 | auto inp_g_flat=inp_g_tensor->flat();
173 | float * inp_g=&(inp_g_flat(0));
174 | cudaMemset(inp_g,0,b*n*3*4);
175 | scatteraddpointLauncher(b,n,m,out_g,idx,inp_g);
176 | }
177 | };
178 | REGISTER_KERNEL_BUILDER(Name("GatherPointGrad").Device(DEVICE_GPU),GatherPointGradGpuOp);
179 |
180 |
--------------------------------------------------------------------------------
/tf_ops/sampling/tf_sampling.py:
--------------------------------------------------------------------------------
1 | '''
2 | @Author: Shuai Wang
3 | @Github: https://github.com/wsustcid
4 | @Version: 1.0.0
5 | @Date: 1970-01-01 08:00:00
6 | @LastEditTime: 2020-06-16 22:21:11
7 | @Description:
8 | '''
9 | ''' Furthest point sampling
10 | Original author: Haoqiang Fan
11 | Modified by Charles R. Qi
12 | All Rights Reserved. 2017.
13 | '''
14 | import tensorflow as tf
15 | from tensorflow.python.framework import ops
16 | import sys
17 | import os
18 | BASE_DIR = os.path.dirname(os.path.abspath(__file__))
19 | sys.path.append(BASE_DIR)
20 | sampling_module=tf.load_op_library(os.path.join(BASE_DIR, 'tf_sampling_so.so'))
21 | def prob_sample(inp,inpr):
22 | '''
23 | input:
24 | batch_size * ncategory float32
25 | batch_size * npoints float32
26 | returns:
27 | batch_size * npoints int32
28 | '''
29 | return sampling_module.prob_sample(inp,inpr)
30 | ops.NoGradient('ProbSample')
31 | # TF1.0 API requires set shape in C++
32 | #@tf.RegisterShape('ProbSample')
33 | #def _prob_sample_shape(op):
34 | # shape1=op.inputs[0].get_shape().with_rank(2)
35 | # shape2=op.inputs[1].get_shape().with_rank(2)
36 | # return [tf.TensorShape([shape2.dims[0],shape2.dims[1]])]
37 | def gather_point(inp,idx):
38 | '''
39 | input:
40 | batch_size * ndataset * 3 float32
41 | batch_size * npoints int32
42 | returns:
43 | batch_size * npoints * 3 float32
44 | '''
45 | return sampling_module.gather_point(inp,idx)
46 | #@tf.RegisterShape('GatherPoint')
47 | #def _gather_point_shape(op):
48 | # shape1=op.inputs[0].get_shape().with_rank(3)
49 | # shape2=op.inputs[1].get_shape().with_rank(2)
50 | # return [tf.TensorShape([shape1.dims[0],shape2.dims[1],shape1.dims[2]])]
51 | @tf.RegisterGradient('GatherPoint')
52 | def _gather_point_grad(op,out_g):
53 | inp=op.inputs[0]
54 | idx=op.inputs[1]
55 | return [sampling_module.gather_point_grad(inp,idx,out_g),None]
56 | def farthest_point_sample(npoint,inp):
57 | '''
58 | input:
59 | int32
60 | batch_size * ndataset * 3 float32
61 | returns:
62 | batch_size * npoint int32
63 | '''
64 | return sampling_module.farthest_point_sample(inp, npoint)
65 | ops.NoGradient('FarthestPointSample')
66 |
67 |
68 | if __name__=='__main__':
69 | import numpy as np
70 | np.random.seed(100)
71 | triangles=np.random.rand(1,5,3,3).astype('float32')
72 | with tf.device('/gpu:1'):
73 | inp=tf.constant(triangles)
74 | tria=inp[:,:,0,:]
75 | trib=inp[:,:,1,:]
76 | tric=inp[:,:,2,:]
77 | areas=tf.sqrt(tf.reduce_sum(tf.cross(trib-tria,tric-tria)**2,2)+1e-9)
78 | randomnumbers=tf.random_uniform((1,8192))
79 | triids=prob_sample(areas,randomnumbers)
80 | tria_sample=gather_point(tria,triids)
81 | trib_sample=gather_point(trib,triids)
82 | tric_sample=gather_point(tric,triids)
83 | us=tf.random_uniform((1,8192))
84 | vs=tf.random_uniform((1,8192))
85 | uplusv=1-tf.abs(us+vs-1)
86 | uminusv=us-vs
87 | us=(uplusv+uminusv)*0.5
88 | vs=(uplusv-uminusv)*0.5
89 | pt_sample=tria_sample+(trib_sample-tria_sample)*tf.expand_dims(us,-1)+(tric_sample-tria_sample)*tf.expand_dims(vs,-1)
90 | print('pt_sample: ', pt_sample)
91 | reduced_sample=gather_point(pt_sample,farthest_point_sample(1024,pt_sample))
92 | print(reduced_sample)
93 | with tf.Session('') as sess:
94 | ret=sess.run(reduced_sample)
95 | print(ret.shape,ret.dtype)
96 | import cPickle as pickle
97 | pickle.dump(ret,open('1.pkl','wb'),-1)
98 |
--------------------------------------------------------------------------------
/tf_ops/sampling/tf_sampling_compile.sh:
--------------------------------------------------------------------------------
1 | #/bin/bash
2 | #/usr/local/cuda-9.0/bin/nvcc tf_sampling_g.cu -o tf_sampling_g.cu.o -c -O2 -DGOOGLE_CUDA=1 -x cu -Xcompiler -fPIC
3 |
4 | # TF1.2
5 | #g++ -std=c++11 tf_sampling.cpp tf_sampling_g.cu.o -o tf_sampling_so.so -shared -fPIC -I /usr/local/lib/python2.7/dist-packages/tensorflow/include -I /usr/local/cuda-8.0/include -lcudart -L /usr/local/cuda-8.0/lib64/ -O2 -D_GLIBCXX_USE_CXX11_ABI=0
6 |
7 | # TF1.4
8 | #g++ -std=c++11 tf_sampling.cpp tf_sampling_g.cu.o -o tf_sampling_so.so -shared -fPIC -I /usr/local/lib/python2.7/dist-packages/tensorflow/include -I /usr/local/cuda-8.0/include -I /usr/local/lib/python2.7/dist-packages/tensorflow/include/external/nsync/public -lcudart -L /usr/local/cuda-8.0/lib64/ -L/usr/local/lib/python2.7/dist-packages/tensorflow -ltensorflow_framework -O2 -D_GLIBCXX_USE_CXX11_ABI=0
9 |
10 | # TF1.12 (local venv)
11 | #g++ -std=c++11 tf_sampling.cpp tf_sampling_g.cu.o -o tf_sampling_so.so -shared -fPIC -I /home/ubuntu16/venv/lib/python3.5/site-packages/tensorflow/include -I /usr/local/cuda-9.0/include -I /home/ubuntu16/venv/lib/python3.5/site-packages/tensorflow/include/external/nsync/public -lcudart -L /usr/local/cuda-9.0/lib64/ -L/home/ubuntu16/venv/lib/python3.5/site-packages/tensorflow -ltensorflow_framework -O2 -D_GLIBCXX_USE_CXX11_ABI=0
12 |
13 | # ==== cluster version ====
14 | CUDA_ROOT=/usr/local/cuda-9.0
15 | TF_INC=$(python -c 'import tensorflow as tf; print(tf.sysconfig.get_include())')
16 | TF_LIB=$(python -c 'import tensorflow as tf; print(tf.sysconfig.get_lib())')
17 |
18 | echo $CUDA_ROOT
19 | echo $TF_INC
20 | echo $TF_LIB
21 |
22 | $CUDA_ROOT/bin/nvcc tf_sampling_g.cu -o tf_sampling_g.cu.o -c -O2 -DGOOGLE_CUDA=1 -x cu -Xcompiler -fPIC
23 |
24 | # TF>=1.4.0
25 | g++ -std=c++11 tf_sampling.cpp tf_sampling_g.cu.o -o tf_sampling_so.so -shared -fPIC -I$TF_INC/ -I$TF_INC/external/nsync/public -L$TF_LIB -ltensorflow_framework -I$CUDA_ROOT/include -lcudart -L$CUDA_ROOT/lib64/ -O2 -D_GLIBCXX_USE_CXX11_ABI=0
--------------------------------------------------------------------------------
/tf_ops/sampling/tf_sampling_g.cu:
--------------------------------------------------------------------------------
1 | /* Furthest point sampling GPU implementation
2 | * Original author: Haoqiang Fan
3 | * Modified by Charles R. Qi
4 | * All Rights Reserved. 2017.
5 | */
6 |
7 | __global__ void cumsumKernel(int b,int n,const float * __restrict__ inp,float * __restrict__ out){
8 | const int BlockSize=2048;
9 | const int paddingLevel=5;
10 | __shared__ float buffer4[BlockSize*4];
11 | __shared__ float buffer[BlockSize+(BlockSize>>paddingLevel)];
12 | for (int i=blockIdx.x;i>2;
18 | for (int k=threadIdx.x*4;k>2)+(k>>(2+paddingLevel))]=v4;
33 | }else{
34 | float v=0;
35 | for (int k2=k;k2>2)+(k>>(2+paddingLevel))]=v;
43 | }
44 | }
45 | int u=0;
46 | for (;(2<>(u+1));k+=blockDim.x){
49 | int i1=(((k<<1)+2)<>paddingLevel;
52 | i2+=i2>>paddingLevel;
53 | buffer[i1]+=buffer[i2];
54 | }
55 | }
56 | u--;
57 | for (;u>=0;u--){
58 | __syncthreads();
59 | for (int k=threadIdx.x;k>(u+1));k+=blockDim.x){
60 | int i1=(((k<<1)+3)<>paddingLevel;
63 | i2+=i2>>paddingLevel;
64 | buffer[i1]+=buffer[i2];
65 | }
66 | }
67 | __syncthreads();
68 | for (int k=threadIdx.x*4;k>2)-1)+(((k>>2)-1)>>paddingLevel);
71 | buffer4[k]+=buffer[k2];
72 | buffer4[k+1]+=buffer[k2];
73 | buffer4[k+2]+=buffer[k2];
74 | buffer4[k+3]+=buffer[k2];
75 | }
76 | }
77 | __syncthreads();
78 | for (int k=threadIdx.x;k>paddingLevel)]+runningsum2;
82 | float r2=runningsum+t;
83 | runningsum2=t-(r2-runningsum);
84 | runningsum=r2;
85 | __syncthreads();
86 | }
87 | }
88 | }
89 |
90 | __global__ void binarysearchKernel(int b,int n,int m,const float * __restrict__ dataset,const float * __restrict__ query, int * __restrict__ result){
91 | int base=1;
92 | while (base=1;k>>=1)
99 | if (r>=k && dataset[i*n+r-k]>=q)
100 | r-=k;
101 | result[i*m+j]=r;
102 | }
103 | }
104 | }
105 | __global__ void farthestpointsamplingKernel(int b,int n,int m,const float * __restrict__ dataset,float * __restrict__ temp,int * __restrict__ idxs){
106 | if (m<=0)
107 | return;
108 | const int BlockSize=512;
109 | __shared__ float dists[BlockSize];
110 | __shared__ int dists_i[BlockSize];
111 | const int BufferSize=3072;
112 | __shared__ float buf[BufferSize*3];
113 | for (int i=blockIdx.x;ibest){
147 | best=d2;
148 | besti=k;
149 | }
150 | }
151 | dists[threadIdx.x]=best;
152 | dists_i[threadIdx.x]=besti;
153 | for (int u=0;(1<>(u+1))){
156 | int i1=(threadIdx.x*2)<>>(b,n,inp,out);
196 | }
197 | //require b*n working space
198 | void probsampleLauncher(int b,int n,int m,const float * inp_p,const float * inp_r,float * temp,int * out){
199 | cumsumKernel<<<32,512>>>(b,n,inp_p,temp);
200 | binarysearchKernel<<>>(b,n,m,temp,inp_r,out);
201 | }
202 | //require 32*n working space
203 | void farthestpointsamplingLauncher(int b,int n,int m,const float * inp,float * temp,int * out){
204 | farthestpointsamplingKernel<<<32,512>>>(b,n,m,inp,temp,out);
205 | }
206 | void gatherpointLauncher(int b,int n,int m,const float * inp,const int * idx,float * out){
207 | gatherpointKernel<<>>(b,n,m,inp,idx,out);
208 | }
209 | void scatteraddpointLauncher(int b,int n,int m,const float * out_g,const int * idx,float * inp_g){
210 | scatteraddpointKernel<<>>(b,n,m,out_g,idx,inp_g);
211 | }
212 |
213 |
--------------------------------------------------------------------------------
/tf_ops/sampling/tf_sampling_g.cu.o:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/wsustcid/FlowDriveNet/3604495269ae45e5b43964046104f685ec66e383/tf_ops/sampling/tf_sampling_g.cu.o
--------------------------------------------------------------------------------
/tf_ops/sampling/tf_sampling_so.so:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/wsustcid/FlowDriveNet/3604495269ae45e5b43964046104f685ec66e383/tf_ops/sampling/tf_sampling_so.so
--------------------------------------------------------------------------------
/train.py:
--------------------------------------------------------------------------------
1 | '''
2 | @Author: Shuai Wang
3 | @Github: https://github.com/wsustcid
4 | @Version: 1.0.0
5 | @Date: 2020-06-03 16:56:56
6 | @LastEditTime: 2020-10-07 16:14:17
7 | '''
8 |
9 | import os
10 | import sys
11 | import argparse
12 | from datetime import datetime
13 | import time
14 |
15 | import numpy as np
16 | import tensorflow as tf
17 |
18 | base_dir = os.path.dirname(os.path.abspath(__file__))
19 | sys.path.append(base_dir)
20 | from utils.tf_util import log_string
21 |
22 | from data_gen import DataLoader
23 | from models.flowdrivenet import FlowDriveNet
24 |
25 | parser = argparse.ArgumentParser()
26 | parser.add_argument('--data_root', default='/gdata/wangshuai/Udacity/CH2',
27 | help='data_root path [default: local path]')
28 | parser.add_argument('--input_cfg', default='GRAY',
29 | help='Input type: GRAY, GRAYF, GRAYF-T, XYZ, XYZF, XYZF-T, GRAYF-XYZF-T')
30 | parser.add_argument('--model_cfg', default='VFE',
31 | help='Model type: VFE, VFE-TFP, PFE, PFE-TFP, VFE-PFE-TFP')
32 | parser.add_argument('--loss_cfg', default='MSE',
33 | help='loss type: weighted, step, exp')
34 | parser.add_argument('--height', type=int, default=200, help='img height')
35 | parser.add_argument('--width', type=int, default=200, help='img width')
36 | parser.add_argument('--seq_len', type=int, default=5, help='sel length')
37 | parser.add_argument('--aug_cfg', default='None', help='None, IA, RP, SC, BA, BS')
38 | #parser.add_argument('--use_side_cam', default=False, action='store_true')
39 | parser.add_argument('--num_point', type=int, default=10000, help='Point N')
40 | parser.add_argument('--log_dir', default='test',
41 | help='Log dir [default: test]')
42 | parser.add_argument('--max_epoch', type=int, default=300,
43 | help='Epoch to run [default: 1000]')
44 | parser.add_argument('--early_stop', type=int, default=10,
45 | help='stop training when loss stop decreasing [default: 20]')
46 | parser.add_argument('--batch_size', type=int, default=16,
47 | help='Batch Size during training [default: 16]')
48 | parser.add_argument('--learning_rate', type=float, default=0.0001,
49 | help='Learning rate during training [default: 0.001]')
50 | parser.add_argument('--optimizer', default='adam',
51 | help='adam or momentum [default: adam]')
52 | parser.add_argument('--decay_steps', type=int, default=300000,
53 | help='Decay step for lr decay [default: 200000]') # decay_steps = n_train * epochs
54 | parser.add_argument('--decay_rate', type=float, default=0.7,
55 | help='Decay rate for lr decay [default: 0.7]')
56 |
57 | FLAGS = parser.parse_args()
58 |
59 | BATCH_SIZE = FLAGS.batch_size
60 |
61 | log_dir = os.path.join(base_dir, 'logs', FLAGS.log_dir)
62 | os.makedirs(log_dir, exist_ok=True)
63 | train_log_dir = os.path.join(log_dir, 'log_train.txt')
64 | log_string(train_log_dir, str(FLAGS)+'\n')
65 |
66 | model_file = os.path.join(base_dir, 'models/flowdrivenet.py')
67 | train_file = os.path.join(base_dir,'train.py')
68 | os.system('cp %s %s' % (model_file, log_dir))
69 | os.system('cp %s %s' % (train_file, log_dir))
70 |
71 | #
72 | dataloader = DataLoader(FLAGS.data_root, FLAGS.input_cfg,
73 | FLAGS.height, FLAGS.width,
74 | FLAGS.seq_len,
75 | FLAGS.num_point,
76 | FLAGS.aug_cfg)
77 | model = FlowDriveNet(FLAGS.input_cfg, FLAGS.model_cfg,
78 | FLAGS.height, FLAGS.width,
79 | FLAGS.seq_len, FLAGS.num_point)
80 |
81 | def get_bn_decay(batch):
82 | bn_momentum = tf.train.exponential_decay(
83 | 0.5,
84 | batch*BATCH_SIZE,
85 | float(FLAGS.decay_steps),
86 | 0.5,
87 | staircase=True)
88 | bn_decay = tf.minimum(0.99, 1 - bn_momentum)
89 | return bn_decay
90 |
91 | #def get_lr(batch):
92 | # lr = tf.train.exponential_decay(learning_rate=FLAGS.learning_rate,
93 | # global_step=batch*BATCH_SIZE,
94 | # decay_steps=FLAGS.decay_steps,
95 | # decay_rate=FLAGS.decay_rate,
96 | # staircase=True)
97 | # lr = tf.maximum(lr, 0.00001)
98 | # return lr
99 |
100 | def train():
101 | with tf.Graph().as_default():
102 | image_pl, points_pl, label_pl = model.get_inputs_pl(BATCH_SIZE)
103 | is_training_pl = tf.placeholder(tf.bool, shape=())
104 | # define global_step; optimizer will increase it in every training loop
105 | batch = tf.get_variable('batch', [],
106 | initializer=tf.constant_initializer(0),
107 | trainable=False)
108 | bn_decay = get_bn_decay(batch)
109 | tf.summary.scalar('bn_decay', bn_decay)
110 |
111 | pred = model.get_model(image_pl, points_pl, is_training_pl, bn_decay)
112 | loss = model.get_loss(pred, label_pl, batch, FLAGS.loss_cfg)
113 | rmse_angle, rmse_speed = model.get_rmse(pred, label_pl)
114 |
115 | #learning_rate = get_lr(batch)
116 | if FLAGS.optimizer == 'adam':
117 | optimizer = tf.train.AdamOptimizer(FLAGS.learning_rate)
118 | tf.summary.scalar('learning_rate', FLAGS.learning_rate)
119 |
120 | train_op = optimizer.minimize(loss, global_step=batch)
121 |
122 | # Add ops to save and restore all the variables.
123 | saver = tf.train.Saver()
124 |
125 | # Create a session
126 | config = tf.ConfigProto()
127 | config.gpu_options.allow_growth = True
128 | config.allow_soft_placement = True
129 | config.log_device_placement = False
130 | sess = tf.Session(config=config)
131 |
132 | # Add summary writers
133 | merged = tf.summary.merge_all()
134 | train_writer = tf.summary.FileWriter(os.path.join(log_dir, 'train'), sess.graph)
135 | test_writer = tf.summary.FileWriter(os.path.join(log_dir, 'test'), sess.graph)
136 |
137 | # Init variables
138 | init = tf.global_variables_initializer()
139 | sess.run(init)
140 |
141 | # save all tensor
142 | ops = {'image_pl': image_pl,
143 | 'points_pl': points_pl,
144 | 'label_pl': label_pl,
145 | 'is_training_pl': is_training_pl,
146 | 'train_op': train_op,
147 | 'loss': loss,
148 | 'rmse_angle': rmse_angle,
149 | 'rmse_speed':rmse_speed,
150 | 'merged': merged,
151 | 'batch': batch}
152 |
153 | test_err_min = 100000
154 | for epoch in range(FLAGS.max_epoch):
155 | log_string(train_log_dir, '**** EPOCH %03d ****' % (epoch))
156 |
157 | train_one_epoch(sess, ops, train_writer)
158 | test_err = test_one_epoch(sess, ops, test_writer)
159 | # save best
160 | if test_err < test_err_min:
161 | es_count = 0
162 | test_err_min = test_err
163 | save_path = saver.save(sess, os.path.join(log_dir, "model_best.ckpt"))
164 | log_string(train_log_dir, "Best model saved in : %s" % save_path)
165 | else:
166 | es_count +=1
167 |
168 | #if epoch % 10 == 0:
169 | # save_path = saver.save(sess, os.path.join(log_dir, "model.ckpt"))
170 | # log_string(train_log_dir, "Model saved in file: %s" % save_path)
171 |
172 | # Early Stopping
173 | if es_count >= FLAGS.early_stop:
174 | break
175 |
176 |
177 | def train_one_epoch(sess, ops, train_writer):
178 | """ ops: dict mapping from string to tf ops """
179 | # shuffle data
180 | #data_loader.Xs_train, data_loader.y_train = shuffle(data_loader.Xs_train, data_loader.y_train)
181 | is_training = True
182 | log_string(train_log_dir, datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
183 |
184 | num_batches = dataloader.num_train // BATCH_SIZE
185 | loss_sum = 0.0
186 | rmse_angle_sum = 0.0
187 | rmse_speed_sum = 0.0
188 |
189 | for i in range(num_batches):
190 | #t1 = time.time()
191 | X_image_batch, X_cloud_batch, y_batch = dataloader.load_train_batch(BATCH_SIZE)
192 | #t2 = time.time()
193 |
194 | feed_dict = {ops['image_pl']: X_image_batch,
195 | ops['points_pl']: X_cloud_batch,
196 | ops['label_pl']: y_batch,
197 | ops['is_training_pl']: is_training}
198 |
199 | summary, step, _, loss_batch, rmse_angle_batch, rmse_speed_batch = sess.run([ops['merged'], ops['batch'], ops['train_op'], ops['loss'], ops['rmse_angle'], ops['rmse_speed']], feed_dict=feed_dict)
200 | #t3 = time.time()
201 | #print("data time: {}; train time: {}".format(t2-t1, t3-t2))
202 |
203 | train_writer.add_summary(summary, step)
204 |
205 | loss_sum += loss_batch
206 | rmse_angle_sum += rmse_angle_batch
207 | rmse_speed_sum += rmse_speed_batch
208 |
209 | log_string(train_log_dir, 'Train loss: %f' % (loss_sum / num_batches))
210 | log_string(train_log_dir, 'Train rmse_angle: %f' % (rmse_angle_sum / num_batches))
211 | log_string(train_log_dir, 'Train rmse_speed: %f' % (rmse_speed_sum / num_batches))
212 | log_string(train_log_dir, 'Train rmse_average: %f' % ((rmse_angle_sum+rmse_speed_sum)/num_batches/2))
213 |
214 | def test_one_epoch(sess, ops, test_writer):
215 | """ ops: dict mapping from string to tf ops """
216 |
217 | is_training = False
218 | num_batches = dataloader.num_val // BATCH_SIZE
219 | loss_sum = 0.0
220 | rmse_angle_sum = 0.0
221 | rmse_speed_sum = 0.0
222 |
223 | for i in range(num_batches):
224 | X_image_batch, X_cloud_batch, y_batch = dataloader.load_val_batch(BATCH_SIZE)
225 |
226 | feed_dict = {ops['image_pl']: X_image_batch,
227 | ops['points_pl']: X_cloud_batch,
228 | ops['label_pl']: y_batch,
229 | ops['is_training_pl']: is_training}
230 |
231 | summary, step, loss_batch, rmse_angle_batch, rmse_speed_batch = sess.run([ops['merged'], ops['batch'], ops['loss'], ops['rmse_angle'], ops['rmse_speed']],feed_dict=feed_dict)
232 |
233 | test_writer.add_summary(summary, step)
234 |
235 | loss_sum += loss_batch
236 | rmse_angle_sum += rmse_angle_batch
237 | rmse_speed_sum += rmse_speed_batch
238 |
239 | log_string(train_log_dir, 'Val loss: %f' % (loss_sum / num_batches))
240 | log_string(train_log_dir, 'Val rmse_angle: %f' % (rmse_angle_sum / num_batches))
241 | log_string(train_log_dir, 'Val rmse_speed: %f' % (rmse_speed_sum / num_batches))
242 | log_string(train_log_dir, 'Val rmse_average: %f' % ((rmse_angle_sum+rmse_speed_sum)/num_batches/2))
243 |
244 | return (rmse_angle_sum+rmse_speed_sum)/num_batches/2
245 |
246 | if __name__ == "__main__":
247 | train()
248 |
--------------------------------------------------------------------------------
/train_cmp.py:
--------------------------------------------------------------------------------
1 | '''
2 | @Author: Shuai Wang
3 | @Github: https://github.com/wsustcid
4 | @Version: 1.0.0
5 | @Date: 2020-09-15 21:58:38
6 | @LastEditTime: 2020-09-16 20:57:19
7 | @Description: Training the comparision models
8 | '''
9 |
10 | import os
11 | import sys
12 | import argparse
13 | from datetime import datetime
14 | import time
15 |
16 | import numpy as np
17 | import tensorflow as tf
18 |
19 | base_dir = os.path.dirname(os.path.abspath(__file__))
20 | sys.path.append(base_dir)
21 | from utils.tf_util import log_string
22 |
23 | from data_gen import DataLoader
24 |
25 | parser = argparse.ArgumentParser()
26 | parser.add_argument('--data_root', default='/gdata/wangshuai/Udacity/CH2',
27 | help='data_root path [default: local path]')
28 | parser.add_argument('--input_cfg', default='BGR',
29 | help='Input type: BGR, GRAYF-T, XYZ, GRAY')
30 | parser.add_argument('--model_cfg', default='PilotNet',
31 | help='Model type: PilotNet, BMWNet, PointNet, DroNet')
32 | parser.add_argument('--use_side_cam', default=False, action='store_true')
33 | parser.add_argument('--log_dir', default='test',
34 | help='Log dir [default: test]')
35 | parser.add_argument('--max_epoch', type=int, default=300,
36 | help='Epoch to run [default: 1000]')
37 | parser.add_argument('--early_stop', type=int, default=20,
38 | help='stop training when loss stop decreasing [default: 20]')
39 | parser.add_argument('--batch_size', type=int, default=8,
40 | help='Batch Size during training [default: 16]')
41 | parser.add_argument('--learning_rate', type=float, default=0.0001,
42 | help='Learning rate during training [default: 0.001]')
43 | parser.add_argument('--optimizer', default='adam',
44 | help='adam or momentum [default: adam]')
45 | parser.add_argument('--decay_steps', type=int, default=300000,
46 | help='Decay step for lr decay [default: 200000]') # decay_steps = n_train * epochs
47 | parser.add_argument('--decay_rate', type=float, default=0.7,
48 | help='Decay rate for lr decay [default: 0.7]')
49 |
50 | FLAGS = parser.parse_args()
51 |
52 | BATCH_SIZE = FLAGS.batch_size
53 |
54 | log_dir = os.path.join(base_dir, 'logs', FLAGS.log_dir)
55 | os.makedirs(log_dir, exist_ok=True)
56 | train_log_dir = os.path.join(log_dir, 'log_train.txt')
57 | log_string(train_log_dir, str(FLAGS)+'\n')
58 |
59 |
60 | if FLAGS.model_cfg == 'PilotNet':
61 | from models.pilotnet import PilotNet
62 | dataloader = DataLoader(FLAGS.data_root, "BGR",
63 | height=66, width=200,
64 | seq_len=None,
65 | num_point=None,
66 | use_side_cam=FLAGS.use_side_cam)
67 | model = PilotNet()
68 | elif FLAGS.model_cfg == 'BMWNet':
69 | from models.bmwnet import BMWNet
70 | # TODO add seq_len
71 | dataloader = DataLoader(FLAGS.data_root, 'GRAYF-T',
72 | height=66, width=200,
73 | seq_len=10,
74 | num_point=None,
75 | use_side_cam=FLAGS.use_side_cam)
76 | model = BMWNet()
77 | elif FLAGS.model_cfg == 'PointNet':
78 | from models.pointnet import PointNet
79 | dataloader = DataLoader(FLAGS.data_root, 'XYZ',
80 | height=None, width=None,
81 | seq_len=None,
82 | num_point=10000,
83 | use_side_cam=FLAGS.use_side_cam)
84 | model = PointNet(num_point=10000)
85 | elif FLAGS.model_cfg == 'DroNet':
86 | from models.dronet import DroNet
87 | dataloader = DataLoader(FLAGS.data_root, 'GRAY',
88 | height=200, width=200,
89 | seq_len=None,
90 | num_point=None,
91 | use_side_cam=FLAGS.use_side_cam)
92 | model = DroNet()
93 | else:
94 | raise TypeError
95 |
96 | def get_bn_decay(batch):
97 | bn_momentum = tf.train.exponential_decay(
98 | 0.5,
99 | batch*BATCH_SIZE,
100 | float(FLAGS.decay_steps),
101 | 0.5,
102 | staircase=True)
103 | bn_decay = tf.minimum(0.99, 1 - bn_momentum)
104 | return bn_decay
105 |
106 | #def get_lr(batch):
107 | # lr = tf.train.exponential_decay(learning_rate=FLAGS.learning_rate,
108 | # global_step=batch*BATCH_SIZE,
109 | # decay_steps=FLAGS.decay_steps,
110 | # decay_rate=FLAGS.decay_rate,
111 | # staircase=True)
112 | # lr = tf.maximum(lr, 0.00001)
113 | # return lr
114 |
115 | def train():
116 | with tf.Graph().as_default():
117 | feature_pl, label_pl = model.get_inputs_pl(BATCH_SIZE)
118 | is_training_pl = tf.placeholder(tf.bool, shape=())
119 | # define global_step; optimizer will increase it in every training loop
120 | batch = tf.get_variable('batch', [],
121 | initializer=tf.constant_initializer(0),
122 | trainable=False)
123 | bn_decay = get_bn_decay(batch)
124 |
125 | pred = model.get_model(feature_pl, is_training_pl, bn_decay)
126 | loss = model.get_loss(pred, label_pl)
127 |
128 | #learning_rate = get_lr(batch)
129 | if FLAGS.optimizer == 'adam':
130 | optimizer = tf.train.AdamOptimizer(FLAGS.learning_rate)
131 | tf.summary.scalar('learning_rate', FLAGS.learning_rate)
132 |
133 | train_op = optimizer.minimize(loss, global_step=batch)
134 |
135 | # Add ops to save and restore all the variables.
136 | saver = tf.train.Saver()
137 |
138 | # Create a session
139 | config = tf.ConfigProto()
140 | config.gpu_options.allow_growth = True
141 | config.allow_soft_placement = True
142 | config.log_device_placement = False
143 | sess = tf.Session(config=config)
144 |
145 | # Add summary writers
146 | merged = tf.summary.merge_all()
147 | train_writer = tf.summary.FileWriter(os.path.join(log_dir, 'train'), sess.graph)
148 | test_writer = tf.summary.FileWriter(os.path.join(log_dir, 'test'), sess.graph)
149 |
150 | # Init variables
151 | init = tf.global_variables_initializer()
152 | sess.run(init)
153 |
154 | # save all tensor
155 | ops = {'feature_pl': feature_pl,
156 | 'label_pl': label_pl,
157 | 'is_training_pl': is_training_pl,
158 | 'train_op': train_op,
159 | 'loss': loss,
160 | 'pred': pred,
161 | 'merged': merged,
162 | 'batch': batch}
163 |
164 | test_err_min = 100000
165 | for epoch in range(FLAGS.max_epoch):
166 | log_string(train_log_dir, '**** EPOCH %03d ****' % (epoch))
167 |
168 | train_one_epoch(sess, ops, train_writer)
169 | test_err = test_one_epoch(sess, ops, test_writer)
170 | # save best
171 | if test_err < test_err_min:
172 | es_count = 0
173 | test_err_min = test_err
174 | save_path = saver.save(sess, os.path.join(log_dir, "model_best.ckpt"))
175 | log_string(train_log_dir, "Best model saved in : %s" % save_path)
176 | else:
177 | es_count +=1
178 |
179 | #if epoch % 10 == 0:
180 | # save_path = saver.save(sess, os.path.join(log_dir, "model.ckpt"))
181 | # log_string(train_log_dir, "Model saved in file: %s" % save_path)
182 |
183 | # Early Stopping
184 | if es_count >= FLAGS.early_stop:
185 | break
186 |
187 |
188 | def train_one_epoch(sess, ops, train_writer):
189 | """ ops: dict mapping from string to tf ops """
190 | # shuffle data
191 | #data_loader.Xs_train, data_loader.y_train = shuffle(data_loader.Xs_train, data_loader.y_train)
192 | is_training = True
193 | log_string(train_log_dir, datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
194 |
195 | num_batches = dataloader.num_train // BATCH_SIZE
196 | loss_sum = 0.0
197 | rmse_angle_sum = 0.0
198 | rmse_speed_sum = 0.0
199 |
200 | for i in range(num_batches):
201 | #t1 = time.time()
202 | if FLAGS.model_cfg in ['PilotNet', 'BMWNet']:
203 | X_batch, y = dataloader.load_image_train_batch(BATCH_SIZE)
204 | y_batch = y[:,0:1]
205 | elif FLAGS.model_cfg == 'PointNet':
206 | X_batch, y = dataloader.load_cloud_train_batch(BATCH_SIZE)
207 | y_batch = y[:,1:2]
208 | elif FLAGS.model_cfg == 'DroNet':
209 | X_batch, y_batch = dataloader.load_image_train_batch(BATCH_SIZE)
210 | else:
211 | raise TypeError
212 | #t2 = time.time()
213 |
214 | feed_dict = {ops['feature_pl']: X_batch,
215 | ops['label_pl']: y_batch,
216 | ops['is_training_pl']: is_training}
217 |
218 | summary, step, _, loss_batch, pred_batch = sess.run([ops['merged'], ops['batch'], ops['train_op'], ops['loss'], ops['pred']], feed_dict=feed_dict)
219 | #t3 = time.time()
220 | #print("data time: {}; train time: {}".format(t2-t1, t3-t2))
221 |
222 | train_writer.add_summary(summary, step)
223 |
224 | loss_sum += loss_batch
225 | if FLAGS.model_cfg in ['PilotNet', 'BMWNet']:
226 | rmse_angle_batch = np.sqrt(np.mean(np.square(pred_batch-y_batch)))
227 | rmse_angle_sum += rmse_angle_batch
228 | elif FLAGS.model_cfg == 'PointNet':
229 | rmse_speed_batch = np.sqrt(np.mean(np.square(pred_batch-y_batch)))
230 | rmse_speed_sum += rmse_speed_batch
231 | elif FLAGS.model_cfg == 'DroNet':
232 | rmse_batch = np.sqrt(np.mean(np.square(pred_batch-y_batch), axis=0))
233 | rmse_angle_sum += rmse_batch[0]
234 | rmse_speed_sum += rmse_batch[1]
235 | else:
236 | raise TypeError
237 |
238 | log_string(train_log_dir, 'Train loss: %f' % (loss_sum / num_batches))
239 | if FLAGS.model_cfg in ['PilotNet', 'BMWNet']:
240 | log_string(train_log_dir, 'Train rmse_angle: %f' % (rmse_angle_sum / num_batches))
241 | elif FLAGS.model_cfg == 'PointNet':
242 | log_string(train_log_dir, 'Train rmse_speed: %f' % (rmse_speed_sum / num_batches))
243 | elif FLAGS.model_cfg == 'DroNet':
244 | log_string(train_log_dir, 'Train rmse_angle: %f' % (rmse_angle_sum / num_batches))
245 | log_string(train_log_dir, 'Train rmse_speed: %f' % (rmse_speed_sum / num_batches))
246 | else:
247 | raise TypeError
248 |
249 | def test_one_epoch(sess, ops, test_writer):
250 | """ ops: dict mapping from string to tf ops """
251 |
252 | is_training = True
253 | log_string(train_log_dir, datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
254 |
255 | num_batches = dataloader.num_val // BATCH_SIZE
256 | loss_sum = 0.0
257 | rmse_angle_sum = 0.0
258 | rmse_speed_sum = 0.0
259 |
260 | for i in range(num_batches):
261 | #t1 = time.time()
262 | if FLAGS.model_cfg in ['PilotNet', 'BMWNet']:
263 | X_batch, y = dataloader.load_image_val_batch(BATCH_SIZE)
264 | y_batch = y[:,0:1]
265 | elif FLAGS.model_cfg == 'PointNet':
266 | X_batch, y = dataloader.load_cloud_val_batch(BATCH_SIZE)
267 | y_batch = y[:,1:2]
268 | elif FLAGS.model_cfg == 'DroNet':
269 | X_batch, y_batch = dataloader.load_image_val_batch(BATCH_SIZE)
270 | else:
271 | raise TypeError
272 | #t2 = time.time()
273 |
274 | feed_dict = {ops['feature_pl']: X_batch,
275 | ops['label_pl']: y_batch,
276 | ops['is_training_pl']: is_training}
277 |
278 | summary, step, _, loss_batch, pred_batch = sess.run([ops['merged'], ops['batch'], ops['train_op'], ops['loss'], ops['pred']], feed_dict=feed_dict)
279 | #t3 = time.time()
280 | #print("data time: {}; train time: {}".format(t2-t1, t3-t2))
281 |
282 | test_writer.add_summary(summary, step)
283 |
284 | loss_sum += loss_batch
285 | if FLAGS.model_cfg in ['PilotNet', 'BMWNet']:
286 | rmse_angle_batch = np.sqrt(np.mean(np.square(pred_batch-y_batch)))
287 | rmse_angle_sum += rmse_angle_batch
288 | elif FLAGS.model_cfg == 'PointNet':
289 | rmse_speed_batch = np.sqrt(np.mean(np.square(pred_batch-y_batch)))
290 | rmse_speed_sum += rmse_speed_batch
291 | elif FLAGS.model_cfg == 'DroNet':
292 | rmse_batch = np.sqrt(np.mean(np.square(pred_batch-y_batch), axis=0))
293 | rmse_angle_sum += rmse_batch[0]
294 | rmse_speed_sum += rmse_batch[1]
295 | else:
296 | raise TypeError
297 |
298 | log_string(train_log_dir, 'Val loss: %f' % (loss_sum / num_batches))
299 | if FLAGS.model_cfg in ['PilotNet', 'BMWNet']:
300 | log_string(train_log_dir, 'Val rmse_angle: %f' % (rmse_angle_sum / num_batches))
301 | return rmse_angle_sum/num_batches
302 | elif FLAGS.model_cfg == 'PointNet':
303 | log_string(train_log_dir, 'Val rmse_speed: %f' % (rmse_speed_sum / num_batches))
304 | return rmse_speed_sum/num_batches
305 | elif FLAGS.model_cfg == 'DroNet':
306 | log_string(train_log_dir, 'Val rmse_angle: %f' % (rmse_angle_sum / num_batches))
307 | log_string(train_log_dir, 'Val rmse_speed: %f' % (rmse_speed_sum / num_batches))
308 | return (rmse_angle_sum+rmse_speed_sum)/num_batches/2
309 | else:
310 | raise TypeError
311 |
312 |
313 |
314 | if __name__ == "__main__":
315 | train()
316 |
--------------------------------------------------------------------------------
/train_multi_gpus.py:
--------------------------------------------------------------------------------
1 | '''
2 | @Author: Shuai Wang
3 | @Github: https://github.com/wsustcid
4 | @Version: 1.0.0
5 | @Date: 2020-06-03 16:56:56
6 | @LastEditTime: 2020-09-10 23:38:05
7 | '''
8 |
9 | import os
10 | import sys
11 | import argparse
12 | from datetime import datetime
13 | import time
14 |
15 | import numpy as np
16 | import tensorflow as tf
17 |
18 | base_dir = os.path.dirname(os.path.abspath(__file__))
19 | sys.path.append(base_dir)
20 | from utils.tf_util import log_string
21 |
22 | from data_gen import DataLoader
23 | from models.flowdrivenet import FlowDriveNet
24 |
25 | parser = argparse.ArgumentParser()
26 | parser.add_argument('--num_gpus', type=int, default=2,
27 | help='GPU to use [default: GPU 0]')
28 | parser.add_argument('--data_root', default='/gdata/wangshuai/Udacity/CH2',
29 | help='data_root path [default: local path]')
30 | parser.add_argument('--input_cfg', default='GRAY',
31 | help='Input type: GRAY, GRAYF, GRAYF-T, XYZ, XYZF, XYZF-T, GRAYF-XYZF-T')
32 | parser.add_argument('--model_cfg', default='VFE',
33 | help='Model type: VFE, VFE-TFP, PFE, PFE-TFP, VFE-PFE-TFP')
34 | parser.add_argument('--loss_cfg', default='MSE',
35 | help='loss type: weighted, step, exp')
36 | parser.add_argument('--height', type=int, default=200, help='img height')
37 | parser.add_argument('--width', type=int, default=200, help='img width')
38 | parser.add_argument('--seq_len', type=int, default=5, help='sel length')
39 | parser.add_argument('--use_side_cam', default=False, action='store_true')
40 | parser.add_argument('--num_point', type=int, default=20000, help='Point N')
41 | parser.add_argument('--log_dir', default='test',
42 | help='Log dir [default: test]')
43 | parser.add_argument('--max_epoch', type=int, default=300,
44 | help='Epoch to run [default: 1000]')
45 | parser.add_argument('--early_stop', type=int, default=20,
46 | help='stop training when loss stop decreasing [default: 20]')
47 | parser.add_argument('--batch_size', type=int, default=1,
48 | help='Batch Size during training [default: 1]')
49 | parser.add_argument('--learning_rate', type=float, default=0.0001,
50 | help='Learning rate during training [default: 0.001]')
51 | parser.add_argument('--optimizer', default='adam',
52 | help='adam or momentum [default: adam]')
53 | parser.add_argument('--decay_steps', type=int, default=300000,
54 | help='Decay step for lr decay [default: 200000]') # decay_steps = n_train * epochs
55 | parser.add_argument('--decay_rate', type=float, default=0.7,
56 | help='Decay rate for lr decay [default: 0.7]')
57 |
58 | FLAGS = parser.parse_args()
59 |
60 | BATCH_SIZE = FLAGS.num_gpus * FLAGS.batch_size
61 | DEVICE_BATCH_SIZE = FLAGS.batch_size
62 |
63 | log_dir = os.path.join(base_dir, 'logs', FLAGS.log_dir)
64 | os.makedirs(log_dir, exist_ok=True)
65 | train_log_dir = os.path.join(log_dir, 'log_train.txt')
66 | log_string(train_log_dir, datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
67 | log_string(train_log_dir, str(FLAGS)+'\n')
68 |
69 | model_file = os.path.join(base_dir, 'models/flowdrivenet.py')
70 | train_file = os.path.join(base_dir,'train.py')
71 | os.system('cp %s %s' % (model_file, log_dir))
72 | os.system('cp %s %s' % (train_file, log_dir))
73 |
74 | #
75 | dataloader = DataLoader(FLAGS.data_root, FLAGS.input_cfg,
76 | FLAGS.height, FLAGS.width, FLAGS.use_side_cam)
77 | model = FlowDriveNet(FLAGS.input_cfg, FLAGS.model_cfg,
78 | FLAGS.height, FLAGS.width, FLAGS.seq_len)
79 |
80 | def average_gradients(tower_grads):
81 | """Calculate the average gradient for each shared variable across all towers.
82 | Note that this function provides a synchronization point across all towers.
83 | From tensorflow tutorial: cifar10/cifar10_multi_gpu_train.py
84 | Args:
85 | tower_grads: List of lists of (gradient, variable) tuples. The outer list
86 | is over individual gradients. The inner list is over the gradient
87 | calculation for each tower.
88 | Returns:
89 | List of pairs of (gradient, variable) where the gradient has been averaged
90 | across all towers.
91 | """
92 | average_grads = []
93 | for grad_and_vars in zip(*tower_grads):
94 | # Note that each grad_and_vars looks like the following:
95 | # ((grad0_gpu0, var0_gpu0), ... , (grad0_gpuN, var0_gpuN))
96 | grads = []
97 | #for g, _ in grad_and_vars:
98 | for g, v in grad_and_vars:
99 | # Add 0 dimension to the gradients to represent the tower.
100 | expanded_g = tf.expand_dims(g, 0)
101 |
102 | # Append on a 'tower' dimension which we will average over below.
103 | grads.append(expanded_g)
104 |
105 | # Average over the 'tower' dimension.
106 | grad = tf.concat(axis=0, values=grads)
107 | grad = tf.reduce_mean(grad, 0)
108 |
109 | # Keep in mind that the Variables are redundant because they are shared
110 | # across towers. So .. we will just return the first tower's pointer to
111 | # the Variable.
112 | v = grad_and_vars[0][1]
113 | grad_and_var = (grad, v)
114 | average_grads.append(grad_and_var)
115 | return average_grads
116 |
117 | def get_bn_decay(batch):
118 | bn_momentum = tf.train.exponential_decay(
119 | 0.5,
120 | batch*BATCH_SIZE,
121 | float(FLAGS.decay_steps),
122 | 0.5,
123 | staircase=True)
124 | bn_decay = tf.minimum(0.99, 1 - bn_momentum)
125 | return bn_decay
126 |
127 | #def get_lr(batch):
128 | # lr = tf.train.exponential_decay(learning_rate=FLAGS.learning_rate,
129 | # global_step=batch*BATCH_SIZE,
130 | # decay_steps=FLAGS.decay_steps,
131 | # decay_rate=FLAGS.decay_rate,
132 | # staircase=True)
133 | # lr = tf.maximum(lr, 0.00001)
134 | # return lr
135 |
136 | def train():
137 | with tf.Graph().as_default():
138 | with tf.device('/cpu:0'):
139 | image_pl, points_pl, label_pl = model.get_inputs_pl(BATCH_SIZE)
140 | is_training_pl = tf.placeholder(tf.bool, shape=())
141 | # define global_step; optimizer will increase it in every training loop
142 | batch = tf.get_variable('batch', [],
143 | initializer=tf.constant_initializer(0),
144 | trainable=False)
145 | bn_decay = get_bn_decay(batch)
146 | tf.summary.scalar('bn_decay', bn_decay)
147 |
148 | if FLAGS.optimizer == 'adam':
149 | optimizer = tf.train.AdamOptimizer(FLAGS.learning_rate)
150 | tf.summary.scalar('learning_rate', FLAGS.learning_rate)
151 |
152 | # Allocating variables on CPU first to accelerate multi-gpu training.
153 | model.get_model(image_pl, points_pl, is_training_pl, bn_decay)
154 |
155 | # -------------------------------------------
156 | # Get model and loss on multiple GPU devices
157 | # Data Parallel
158 | # -------------------------------------------
159 | tower_grads = []
160 | pred_gpus = []
161 | loss_gpus = []
162 | for i in range(FLAGS.num_gpus):
163 | with tf.variable_scope(tf.get_variable_scope(), reuse=True):
164 | with tf.device('/gpu:%d'%(i)), tf.name_scope('gpu_%d'%(i)) as scope:
165 | # split data to each gpu
166 | image_pl_gpu = image_pl[i*DEVICE_BATCH_SIZE: (i+1)*DEVICE_BATCH_SIZE]
167 | points_pl_gpu = points_pl[i*DEVICE_BATCH_SIZE: (i+1)*DEVICE_BATCH_SIZE]
168 | label_pl_gpu = label_pl[i*DEVICE_BATCH_SIZE: (i+1)*DEVICE_BATCH_SIZE]
169 |
170 | pred_gpu = model.get_model(image_pl_gpu, points_pl_gpu,
171 | is_training_pl, bn_decay)
172 | loss_gpu = model.get_loss(pred_gpu, label_pl_gpu,
173 | batch, FLAGS.loss_cfg, name=scope)
174 |
175 | grads_gpu = optimizer.compute_gradients(loss_gpu)
176 | tower_grads.append(grads_gpu)
177 | pred_gpus.append(pred_gpu)
178 | loss_gpus.append(loss_gpu)
179 |
180 | # merge pred and losses from multiple gpus
181 | grads = average_gradients(tower_grads)
182 | pred = tf.concat(pred_gpus, 0)
183 | loss = tf.reduce_mean(loss_gpus)
184 |
185 | train_op = optimizer.apply_gradients(grads, global_step=batch)
186 |
187 | rmse_angle, rmse_speed = model.get_rmse(pred, label_pl)
188 |
189 |
190 | # Add ops to save and restore all the variables.
191 | saver = tf.train.Saver()
192 |
193 | # Create a session
194 | config = tf.ConfigProto()
195 | config.gpu_options.allow_growth = True
196 | config.allow_soft_placement = True
197 | config.log_device_placement = False
198 | sess = tf.Session(config=config)
199 |
200 | # Add summary writers
201 | merged = tf.summary.merge_all()
202 | train_writer = tf.summary.FileWriter(os.path.join(log_dir, 'train'), sess.graph)
203 | test_writer = tf.summary.FileWriter(os.path.join(log_dir, 'test'), sess.graph)
204 |
205 | # Init variables
206 | init = tf.global_variables_initializer()
207 | sess.run(init, feed_dict={is_training_pl: True})
208 |
209 | # save all tensor
210 | ops = {'image_pl': image_pl,
211 | 'points_pl': points_pl,
212 | 'label_pl': label_pl,
213 | 'is_training_pl': is_training_pl,
214 | 'train_op': train_op,
215 | 'loss': loss,
216 | 'rmse_angle': rmse_angle,
217 | 'rmse_speed':rmse_speed,
218 | 'merged': merged,
219 | 'batch': batch}
220 |
221 | test_err_min = 100000
222 | for epoch in range(FLAGS.max_epoch):
223 | log_string(train_log_dir, '**** EPOCH %03d ****' % (epoch))
224 |
225 | train_one_epoch(sess, ops, train_writer)
226 | test_err = test_one_epoch(sess, ops, test_writer)
227 | # save best
228 | if test_err < test_err_min:
229 | es_count = 0
230 | test_err_min = test_err
231 | save_path = saver.save(sess, os.path.join(log_dir, "model_best.ckpt"))
232 | log_string(train_log_dir, "Best model saved in : %s" % save_path)
233 | else:
234 | es_count +=1
235 |
236 | #if epoch % 10 == 0:
237 | # save_path = saver.save(sess, os.path.join(log_dir, "model.ckpt"))
238 | # log_string(train_log_dir, "Model saved in file: %s" % save_path)
239 |
240 | # Early Stopping
241 | if es_count >= FLAGS.early_stop:
242 | break
243 |
244 |
245 | def train_one_epoch(sess, ops, train_writer):
246 | """ ops: dict mapping from string to tf ops """
247 | # shuffle data
248 | #data_loader.Xs_train, data_loader.y_train = shuffle(data_loader.Xs_train, data_loader.y_train)
249 | is_training = True
250 | num_batches = dataloader.num_train // BATCH_SIZE
251 | loss_sum = 0.0
252 | rmse_angle_sum = 0.0
253 | rmse_speed_sum = 0.0
254 |
255 | for i in range(num_batches):
256 | #t1 = time.time()
257 | X_image_batch, X_cloud_batch, y_batch = dataloader.load_train_batch(BATCH_SIZE)
258 | #t2 = time.time()
259 |
260 | feed_dict = {ops['image_pl']: X_image_batch,
261 | ops['points_pl']: X_cloud_batch,
262 | ops['label_pl']: y_batch,
263 | ops['is_training_pl']: is_training}
264 |
265 | summary, step, _, loss_batch, rmse_angle_batch, rmse_speed_batch = sess.run([ops['merged'], ops['batch'], ops['train_op'], ops['loss'], ops['rmse_angle'], ops['rmse_speed']], feed_dict=feed_dict)
266 | #t3 = time.time()
267 | #print("data time: {}; train time: {}".format(t2-t1, t3-t2))
268 |
269 | train_writer.add_summary(summary, step)
270 |
271 | loss_sum += loss_batch
272 | rmse_angle_sum += rmse_angle_batch
273 | rmse_speed_sum += rmse_speed_batch
274 |
275 | log_string(train_log_dir, 'Train loss: %f' % (loss_sum / num_batches))
276 | log_string(train_log_dir, 'Train rmse_angle: %f' % (rmse_angle_sum / num_batches))
277 | log_string(train_log_dir, 'Train rmse_speed: %f' % (rmse_speed_sum / num_batches))
278 | log_string(train_log_dir, 'Train rmse_average: %f' % ((rmse_angle_sum+rmse_speed_sum)/ num_batches/2))
279 |
280 | def test_one_epoch(sess, ops, test_writer):
281 | """ ops: dict mapping from string to tf ops """
282 |
283 | is_training = False
284 | num_batches = dataloader.num_val // BATCH_SIZE
285 | loss_sum = 0.0
286 | rmse_angle_sum = 0.0
287 | rmse_speed_sum = 0.0
288 |
289 | for i in range(num_batches):
290 | X_image_batch, X_cloud_batch, y_batch = dataloader.load_val_batch(BATCH_SIZE)
291 |
292 | feed_dict = {ops['image_pl']: X_image_batch,
293 | ops['points_pl']: X_cloud_batch,
294 | ops['label_pl']: y_batch,
295 | ops['is_training_pl']: is_training}
296 |
297 | summary, step, loss_batch, rmse_angle_batch, rmse_speed_batch = sess.run([ops['merged'], ops['batch'], ops['loss'], ops['rmse_angle'], ops['rmse_speed']],feed_dict=feed_dict)
298 |
299 | test_writer.add_summary(summary, step)
300 |
301 | loss_sum += loss_batch
302 | rmse_angle_sum += rmse_angle_batch
303 | rmse_speed_sum += rmse_speed_batch
304 |
305 | log_string(train_log_dir, 'Val loss: %f' % (loss_sum / num_batches))
306 | log_string(train_log_dir, 'Val rmse_angle: %f' % (rmse_angle_sum / num_batches))
307 | log_string(train_log_dir, 'Val rmse_speed: %f' % (rmse_speed_sum / num_batches))
308 | log_string(train_log_dir, 'Val rmse_average: %f' % ((rmse_angle_sum+rmse_speed_sum)/ num_batches/2))
309 |
310 | return (rmse_angle_sum+rmse_speed_sum)/num_batches/2
311 |
312 | if __name__ == "__main__":
313 | train()
314 |
--------------------------------------------------------------------------------
/utils/README.md:
--------------------------------------------------------------------------------
1 | ## Utilility Functions for 3D Point Cloud Deep Learning
2 |
3 | ### visualization tool
4 |
5 | sh compile_render_balls_so.sh
6 | python show3d_balls.py
7 |
--------------------------------------------------------------------------------
/utils/__init__.py:
--------------------------------------------------------------------------------
1 |
2 |
--------------------------------------------------------------------------------
/utils/compile_render_balls_so.sh:
--------------------------------------------------------------------------------
1 | g++ -std=c++11 render_balls_so.cpp -o render_balls_so.so -shared -fPIC -O2 -D_GLIBCXX_USE_CXX11_ABI=0
2 |
3 |
--------------------------------------------------------------------------------
/utils/pc_util.py:
--------------------------------------------------------------------------------
1 | '''
2 | @Author: Shuai Wang
3 | @Github: https://github.com/wsustcid
4 | @Version: 1.0.0
5 | @Date: 2020-06-18 23:32:49
6 | @LastEditTime: 2020-08-07 21:40:49
7 | @Description:
8 | '''
9 |
10 | import os
11 |
12 | BASE_DIR = os.path.dirname(os.path.abspath(__file__))
13 |
14 | import time
15 | import glob
16 | import numpy as np
17 | import open3d as o3d
18 | import pcl
19 | import pcl.pcl_visualization
20 | import pandas as pd
21 |
22 | from pypcd import pypcd
23 | import pprint # for pretty print the meta data of the PCD
24 | import matplotlib.pyplot as plt
25 | from mpl_toolkits.mplot3d import Axes3D
26 | from matplotlib.patches import Circle
27 |
28 |
29 | # ----------------------------------------
30 | # Load pcd file and visualize clouds
31 | # The time used for loadng 1000 cloud:
32 | # - pypcd: 165s
33 | # - open3d: 18s
34 | # - np.fromfile: 5.6s
35 | # ----------------------------------------
36 |
37 | def load_pcd_pypcd(filepath):
38 | """
39 | load point cloud from pcd file use pypcd
40 | show point cloud in scatter
41 |
42 | - The getadata of the cloud saved in pcd file:
43 | {'count': [1, 1, 1, 1, 1, 1],
44 | 'data': 'ascii',
45 | 'fields': ['x', 'y', 'z', 'intensity', 'ring', 'time'],
46 | 'height': 1,
47 | 'points': 23633,
48 | 'size': [4, 4, 4, 4, 2, 4],
49 | 'type': ['F', 'F', 'F', 'F', 'U', 'F'],
50 | 'version': '0.7',
51 | 'viewpoint': [0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0],
52 | 'width': 23633}
53 | """
54 | cloud = pypcd.PointCloud.from_path(filepath)
55 |
56 | pprint.pprint(cloud.get_metadata())
57 |
58 | # access the point cloud as a numpy structured array
59 | print("Raw shape: {}".format(cloud.points))
60 | print(cloud.pc_data[:5])
61 |
62 | new_cloud = np.stack([cloud.pc_data['x'], cloud.pc_data['y'], cloud.pc_data['z']], axis=1)
63 |
64 | print("New shape: {}".format(new_cloud.shape))
65 | print(new_cloud[:5])
66 |
67 |
68 | fig = plt.figure(figsize=(30,30)) # create a figure object
69 | ax = fig.add_subplot(111,projection='3d')
70 | ax.scatter(new_cloud[:,0],new_cloud[:,1], new_cloud[:,2])
71 | ax.axis('scaled')
72 | ax.set_zlabel('Z')
73 | ax.set_ylabel('Y')
74 | ax.set_xlabel('X')
75 |
76 | plt.show()
77 |
78 | def load_pcd_o3d(file_path):
79 | """ load and visualize pcd file use o3d
80 |
81 | Hint:
82 | 1. some useful methods in o3d.geometry (http://www.open3d.org/docs/release/python_api/open3d.geometry.html)
83 | - o3d.geometry.PointCloud.rotate() scale, transform, translate
84 | - o3d.geometry.voxel_down_sample()
85 | - o3d.geometry.compute_point_cloud_mean_and_covariance
86 | - o3d.geometry.crop_point_cloud()
87 | """
88 |
89 | cloud = o3d.io.read_point_cloud(file_path) # return a o3d.geometry.PointCloud
90 | print(cloud.get_max_bound())
91 | print(cloud.get_min_bound())
92 | o3d.visualization.draw_geometries([cloud])
93 |
94 | cloud_np = np.asarray(cloud.points, np.float32)
95 | print(cloud_np.shape)
96 | print(np.max(cloud_np, axis=0))
97 | print(np.min(cloud_np, axis=0))
98 |
99 | def load_pcd_pcl(file_path):
100 | # load
101 | cloud_pcl = pcl.load(file_path)
102 | cloud_np = np.asarray(cloud_pcl, np.float32)
103 | print(cloud_np.shape)
104 | # or convert narray to pcl cloud
105 | #cloud_pcl = pcl.PointCloud()
106 | #cloud_pcl.from_array(cloud_np)
107 |
108 | # show cloud
109 | visual = pcl.pcl_visualization.CloudViewing()
110 | #visual.ShowGrayCloud(cloud_pcl, b'cloud')
111 | #visual.ShowColorACloud()
112 | #visual.ShowColorCloud()
113 | for i in range(1):
114 | visual.ShowMonochromeCloud(cloud_pcl, b'cloud')
115 | time.sleep(100)
116 | #if visual.WasStopped():
117 |
118 | def load_cloud_seq(label_path):
119 | data = pd.read_csv(label_path)
120 | clouds_name = data['point_filename'].values
121 | num_cloud = len(clouds_name)
122 | print("Total cloud: {}".format(num_cloud))
123 |
124 | path_prefix = os.path.dirname(label_path)
125 | visual = pcl.pcl_visualization.CloudViewing()
126 | for i in range(num_cloud):
127 | cloud_path = os.path.join(path_prefix, 'points', clouds_name[i])
128 | cloud_pcl = pcl.load(cloud_path)
129 | visual.ShowMonochromeCloud(cloud_pcl, b'cloud')
130 | print(cloud_path)
131 |
132 | cloud_np = np.asarray(cloud_pcl, np.float32)
133 | #print('Max range: {}; MIn range: {}'.format(np.max(cloud_np, axis=0),
134 | # np.min(cloud_np, axis=0)))
135 |
136 |
137 | time.sleep(0.01)
138 | #if visual.WasStopped():
139 |
140 | def load_cloud_bin_seq(label_path):
141 | data = pd.read_csv(label_path)
142 | clouds_name = data['point_filename'].values
143 | num_cloud = len(clouds_name)
144 | print("Total cloud: {}".format(num_cloud))
145 |
146 | path_prefix = os.path.dirname(label_path)
147 | cloud_pcl = pcl.PointCloud()
148 | visual = pcl.pcl_visualization.CloudViewing()
149 | for i in range(num_cloud):
150 | cloud_path = os.path.join(path_prefix, 'points_bin', clouds_name[i][:-3]+'bin')
151 | cloud_np = np.fromfile(cloud_path, np.float32).reshape(-1,3)
152 | cloud_pcl.from_array(cloud_np)
153 | visual.ShowMonochromeCloud(cloud_pcl, b'cloud')
154 | print(cloud_path)
155 |
156 | time.sleep(0.1)
157 | #if visual.WasStopped():
158 |
159 | def cloud_crop(file_path):
160 | """ Crop the pointcloud and visulaize it
161 |
162 | Hint: o3d.geometry.crop_point_cloud()
163 | Args:
164 | input (open3d.geometry.PointCloud): The input point cloud.
165 | min_bound (numpy.ndarray[float64[3, 1]]): Minimum bound for point coordinate
166 | max_bound (numpy.ndarray[float64[3, 1]]): Maximum bound for point coordinate
167 | Returns:
168 | open3d.geometry.PointCloud
169 | """
170 | # show origin cloud
171 | cloud = o3d.io.read_point_cloud(file_path)
172 | print("Origin Points: ", cloud.dimension)
173 | print("Origin Bound: ", cloud.get_min_bound(), cloud.get_max_bound())
174 | o3d.visualization.draw_geometries([cloud])
175 |
176 | # crop
177 | min_bound = np.array([-70, -40, -5])
178 | max_bound = np.array([70,40,15])
179 | cloud_crop = o3d.geometry.crop_point_cloud(cloud, min_bound,max_bound)
180 |
181 | # show cropped cloud
182 | cloud_np = np.asarray(cloud_crop.points, np.float32)
183 | print("Cropped Points: ", cloud_np.shape)
184 | print("Cropped Bound: ", np.min(cloud_np, axis=0), np.max(cloud_np, axis=0))
185 | o3d.visualization.draw_geometries([cloud_crop])
186 |
187 | def clouds_info(train_path, test_path,
188 | show_distribute=False,
189 | play_clouds=False):
190 | """ plot pointclouds distribution in train set and test set
191 |
192 | 1. plot original point set distribution
193 | 2. get the proper crop bounds
194 | 3. crop the point and visualize the sequential cloud
195 | """
196 | # load train and test clouds file path
197 | train_files = glob.glob(os.path.join(train_path, '*/center.csv'))
198 | test_file = os.path.join(test_path, 'center.csv')
199 |
200 | clouds_path_train = []
201 | for train_file in train_files:
202 | clouds_name = pd.read_csv(train_file)['point_filename'].values
203 | path_prefix = os.path.dirname(train_file)
204 | for cloud_name in clouds_name:
205 | cloud_path = os.path.join(path_prefix, 'points', cloud_name)
206 | clouds_path_train.append(cloud_path)
207 | num_train = len(clouds_path_train)
208 |
209 | clouds_path_test = []
210 | clouds_name = pd.read_csv(test_file)['point_filename'].values
211 | path_prefix = os.path.dirname(test_file)
212 | for cloud_name in clouds_name:
213 | cloud_path = os.path.join(path_prefix, 'points', cloud_name)
214 | clouds_path_test.append(cloud_path)
215 | num_test = len(clouds_path_test)
216 |
217 | print("Total train: ", num_train)
218 | print("Total test: ", num_test)
219 |
220 | # get clouds bound distribution
221 | if show_distribute == True:
222 | x_bound, y_bound, z_bound = [], [], []
223 | for i in range(num_test):
224 | cloud = o3d.io.read_point_cloud(clouds_path_test[i])
225 | x_bound.append(cloud.get_min_bound()[0])
226 | x_bound.append(cloud.get_max_bound()[0])
227 |
228 | y_bound.append(cloud.get_min_bound()[1])
229 | y_bound.append(cloud.get_max_bound()[1])
230 |
231 | z_bound.append(cloud.get_min_bound()[2])
232 | z_bound.append(cloud.get_max_bound()[2])
233 |
234 | fig, axes = plt.subplots(1,3,figsize=(12,2),dpi=300)
235 | axes[0].hist(x_bound, rwidth=0.8)
236 | axes[1].hist(y_bound, rwidth=0.8)
237 | axes[2].hist(z_bound, rwidth=0.8)
238 | plt.tight_layout()
239 | #plt.show()
240 | fig.savefig(os.path.join(BASE_DIR,'../docs/figs/clouds_bound_test.png'))
241 |
242 | x_bound, y_bound, z_bound = [], [], []
243 | for i in range(num_train):
244 | cloud = o3d.io.read_point_cloud(clouds_path_train[i])
245 | x_bound.append(cloud.get_min_bound()[0])
246 | x_bound.append(cloud.get_max_bound()[0])
247 |
248 | y_bound.append(cloud.get_min_bound()[1])
249 | y_bound.append(cloud.get_max_bound()[1])
250 |
251 | z_bound.append(cloud.get_min_bound()[2])
252 | z_bound.append(cloud.get_max_bound()[2])
253 |
254 | fig, axes = plt.subplots(1,3,figsize=(12,2),dpi=300)
255 | axes[0].hist(x_bound, rwidth=0.8)
256 | axes[1].hist(y_bound, rwidth=0.8)
257 | axes[2].hist(z_bound, rwidth=0.8)
258 | plt.tight_layout()
259 | #plt.show()
260 | fig.savefig(os.path.join(BASE_DIR,'../docs/figs/clouds_bound_train.png'))
261 |
262 | # paly clouds
263 | # crop # [-100,-75,-10] [100,75,30]
264 | min_bound = np.array([-70, -40, -10])
265 | max_bound = np.array([70,40,20])
266 |
267 | if play_clouds == True:
268 | visual = pcl.pcl_visualization.CloudViewing()
269 | cloud_pcl = pcl.PointCloud()
270 | for i in range(num_train):
271 | cloud = o3d.io.read_point_cloud(clouds_path_train[i])
272 | cloud_crop = o3d.geometry.crop_point_cloud(cloud, min_bound,max_bound)
273 |
274 | cloud_np = np.asarray(cloud.points, np.float32)
275 |
276 | cloud_pcl.from_array(cloud_np)
277 | visual.ShowMonochromeCloud(cloud_pcl, b'cloud')
278 |
279 | print('{}: Max range: {}; MIn range: {}'.format(cloud_np.shape[0], np.max(cloud_np, axis=0),
280 | np.min(cloud_np, axis=0)))
281 |
282 |
283 | time.sleep(0.5)
284 |
285 | def cloud_show(cloud):
286 | fig, axes = plt.subplots(1,2, figsize=(100,25)) # create a figure object
287 | axes[0].axis('equal')
288 | axes[0].scatter(cloud[:,0],cloud[:,1]) # BEV
289 | axes[0].set_xlabel('X')
290 | axes[0].set_ylabel('Y')
291 |
292 | axes[1].axis('equal')
293 | axes[1].scatter(cloud[:,0],cloud[:,2]) # side view
294 | axes[1].set_xlabel('X')
295 | axes[1].set_ylabel('Z')
296 |
297 | #fig.savefig(path)
298 | plt.show()
299 |
300 | def cloud_show_3d(cloud):
301 | fig = plt.figure(figsize=(100,30)) # create a figure object
302 | ax = fig.add_subplot(111, projection='3d')
303 | ax.set_xlabel('X')
304 | ax.set_ylabel('Y')
305 | ax.set_zlabel('Z')
306 |
307 | ax.scatter(cloud[:,0],cloud[:,1], cloud[:,2])
308 | plt.show()
309 |
310 |
311 | def cloud_plot_circle(cloud, radius=0.5):
312 | fig, axes = plt.subplots(1,2, figsize=(100,25)) # create a figure object
313 | axes[0].axis('equal')
314 | axes[1].axis('equal')
315 |
316 | axes[0].scatter(cloud[:,0],cloud[:,1]) # BEV
317 | axes[1].scatter(cloud[:,0],cloud[:,2]) # side view
318 | for i in range(len(cloud)):
319 | cir1 = Circle((cloud[i][0], cloud[i][1]), radius, color='r', fill=False)
320 | axes[0].add_patch(cir1)
321 |
322 | cir2 = Circle((cloud[i][0], cloud[i][2]), radius, color='r', fill=False)
323 | axes[1].add_patch(cir2)
324 |
325 | axes[0].set_xlabel('X')
326 | axes[0].set_ylabel('Y')
327 | axes[1].set_xlabel('X')
328 | axes[1].set_ylabel('Z')
329 |
330 | #fig.savefig(path)
331 | plt.show()
332 |
333 | def cloud_plot_circle_center(cloud, path, radius=10):
334 | fig, axes = plt.subplots(1,2, figsize=(100,25)) # create a figure object
335 | axes[0].axis('equal')
336 | axes[1].axis('equal')
337 |
338 | axes[0].scatter(cloud[:,0],cloud[:,1]) # BEV
339 | axes[1].scatter(cloud[:,0],cloud[:,2]) # side view
340 |
341 | cir1 = Circle((0,0), radius, color='r', fill=False)
342 | axes[0].add_patch(cir1)
343 |
344 | cir2 = Circle((0,0), radius, color='r', fill=False)
345 | axes[1].add_patch(cir2)
346 |
347 | axes[0].set_xlabel('X')
348 | axes[0].set_ylabel('Y')
349 | axes[1].set_xlabel('X')
350 | axes[1].set_ylabel('Z')
351 |
352 | fig.savefig(path)
353 |
354 | if __name__ == '__main__':
355 | train_path = '/media/ubuntu16/Documents/Datasets/Udacity/CH2/CH2_002'
356 | test_path = '/media/ubuntu16/Documents/Datasets/Udacity/CH2/CH2_001'
357 |
358 | file_name = 'points/1479425440130861.pcd'
359 | #load_pcd_pypcd(os.path.join(test_path,file_name))
360 | #load_pcd_o3d(os.path.join(data_path,file_name))
361 | #load_pcd_pcl(os.path.join(data_path,file_name))
362 | #load_cloud_seq(os.path.join(data_path, 'center.csv'))
363 | #cloud_crop(os.path.join(data_path,file_name))
364 |
365 | #clouds_info(train_path, test_path, show_distribute=False, play_clouds=True)
366 |
367 | load_cloud_bin_seq(os.path.join(test_path, 'center.csv'))
368 |
--------------------------------------------------------------------------------
/utils/provider.py:
--------------------------------------------------------------------------------
1 | import os
2 | import sys
3 | import numpy as np
4 | import h5py
5 | BASE_DIR = os.path.dirname(os.path.abspath(__file__))
6 | sys.path.append(BASE_DIR)
7 |
8 | def shuffle_data(data, labels):
9 | """ Shuffle data and labels.
10 | Input:
11 | data: B,N,... numpy array
12 | label: B,... numpy array
13 | Return:
14 | shuffled data, label and shuffle indices
15 | """
16 | idx = np.arange(len(labels))
17 | np.random.shuffle(idx)
18 | return data[idx, ...], labels[idx], idx
19 |
20 | def shuffle_points(batch_data):
21 | """ Shuffle orders of points in each point cloud -- changes FPS behavior.
22 | Use the same shuffling idx for the entire batch.
23 | Input:
24 | BxNxC array
25 | Output:
26 | BxNxC array
27 | """
28 | idx = np.arange(batch_data.shape[1])
29 | np.random.shuffle(idx)
30 | return batch_data[:,idx,:]
31 |
32 | def rotate_point_cloud(batch_data):
33 | """ Randomly rotate the point clouds to augument the dataset
34 | rotation is per shape based along up direction
35 | Input:
36 | BxNx3 array, original batch of point clouds
37 | Return:
38 | BxNx3 array, rotated batch of point clouds
39 | """
40 | rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
41 | for k in range(batch_data.shape[0]):
42 | rotation_angle = np.random.uniform() * 2 * np.pi
43 | cosval = np.cos(rotation_angle)
44 | sinval = np.sin(rotation_angle)
45 | rotation_matrix = np.array([[cosval, 0, sinval],
46 | [0, 1, 0],
47 | [-sinval, 0, cosval]])
48 | shape_pc = batch_data[k, ...]
49 | rotated_data[k, ...] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix)
50 | return rotated_data
51 |
52 | def rotate_point_cloud_z(batch_data):
53 | """ Randomly rotate the point clouds to augument the dataset
54 | rotation is per shape based along up direction
55 | Input:
56 | BxNx3 array, original batch of point clouds
57 | Return:
58 | BxNx3 array, rotated batch of point clouds
59 | """
60 | rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
61 | for k in range(batch_data.shape[0]):
62 | rotation_angle = np.random.uniform() * 2 * np.pi
63 | cosval = np.cos(rotation_angle)
64 | sinval = np.sin(rotation_angle)
65 | rotation_matrix = np.array([[cosval, sinval, 0],
66 | [-sinval, cosval, 0],
67 | [0, 0, 1]])
68 | shape_pc = batch_data[k, ...]
69 | rotated_data[k, ...] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix)
70 | return rotated_data
71 |
72 | def rotate_point_cloud_with_normal(batch_xyz_normal):
73 | ''' Randomly rotate XYZ, normal point cloud.
74 | Input:
75 | batch_xyz_normal: B,N,6, first three channels are XYZ, last 3 all normal
76 | Output:
77 | B,N,6, rotated XYZ, normal point cloud
78 | '''
79 | for k in range(batch_xyz_normal.shape[0]):
80 | rotation_angle = np.random.uniform() * 2 * np.pi
81 | cosval = np.cos(rotation_angle)
82 | sinval = np.sin(rotation_angle)
83 | rotation_matrix = np.array([[cosval, 0, sinval],
84 | [0, 1, 0],
85 | [-sinval, 0, cosval]])
86 | shape_pc = batch_xyz_normal[k,:,0:3]
87 | shape_normal = batch_xyz_normal[k,:,3:6]
88 | batch_xyz_normal[k,:,0:3] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix)
89 | batch_xyz_normal[k,:,3:6] = np.dot(shape_normal.reshape((-1, 3)), rotation_matrix)
90 | return batch_xyz_normal
91 |
92 | def rotate_perturbation_point_cloud_with_normal(batch_data, angle_sigma=0.06, angle_clip=0.18):
93 | """ Randomly perturb the point clouds by small rotations
94 | Input:
95 | BxNx6 array, original batch of point clouds and point normals
96 | Return:
97 | BxNx3 array, rotated batch of point clouds
98 | """
99 | rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
100 | for k in range(batch_data.shape[0]):
101 | angles = np.clip(angle_sigma*np.random.randn(3), -angle_clip, angle_clip)
102 | Rx = np.array([[1,0,0],
103 | [0,np.cos(angles[0]),-np.sin(angles[0])],
104 | [0,np.sin(angles[0]),np.cos(angles[0])]])
105 | Ry = np.array([[np.cos(angles[1]),0,np.sin(angles[1])],
106 | [0,1,0],
107 | [-np.sin(angles[1]),0,np.cos(angles[1])]])
108 | Rz = np.array([[np.cos(angles[2]),-np.sin(angles[2]),0],
109 | [np.sin(angles[2]),np.cos(angles[2]),0],
110 | [0,0,1]])
111 | R = np.dot(Rz, np.dot(Ry,Rx))
112 | shape_pc = batch_data[k,:,0:3]
113 | shape_normal = batch_data[k,:,3:6]
114 | rotated_data[k,:,0:3] = np.dot(shape_pc.reshape((-1, 3)), R)
115 | rotated_data[k,:,3:6] = np.dot(shape_normal.reshape((-1, 3)), R)
116 | return rotated_data
117 |
118 |
119 | def rotate_point_cloud_by_angle(batch_data, rotation_angle):
120 | """ Rotate the point cloud along up direction with certain angle.
121 | Input:
122 | BxNx3 array, original batch of point clouds
123 | Return:
124 | BxNx3 array, rotated batch of point clouds
125 | """
126 | rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
127 | for k in range(batch_data.shape[0]):
128 | #rotation_angle = np.random.uniform() * 2 * np.pi
129 | cosval = np.cos(rotation_angle)
130 | sinval = np.sin(rotation_angle)
131 | rotation_matrix = np.array([[cosval, 0, sinval],
132 | [0, 1, 0],
133 | [-sinval, 0, cosval]])
134 | shape_pc = batch_data[k,:,0:3]
135 | rotated_data[k,:,0:3] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix)
136 | return rotated_data
137 |
138 | def rotate_point_cloud_by_angle_with_normal(batch_data, rotation_angle):
139 | """ Rotate the point cloud along up direction with certain angle.
140 | Input:
141 | BxNx6 array, original batch of point clouds with normal
142 | scalar, angle of rotation
143 | Return:
144 | BxNx6 array, rotated batch of point clouds iwth normal
145 | """
146 | rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
147 | for k in range(batch_data.shape[0]):
148 | #rotation_angle = np.random.uniform() * 2 * np.pi
149 | cosval = np.cos(rotation_angle)
150 | sinval = np.sin(rotation_angle)
151 | rotation_matrix = np.array([[cosval, 0, sinval],
152 | [0, 1, 0],
153 | [-sinval, 0, cosval]])
154 | shape_pc = batch_data[k,:,0:3]
155 | shape_normal = batch_data[k,:,3:6]
156 | rotated_data[k,:,0:3] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix)
157 | rotated_data[k,:,3:6] = np.dot(shape_normal.reshape((-1,3)), rotation_matrix)
158 | return rotated_data
159 |
160 |
161 |
162 | def rotate_perturbation_point_cloud(batch_data, angle_sigma=0.06, angle_clip=0.18):
163 | """ Randomly perturb the point clouds by small rotations
164 | Input:
165 | BxNx3 array, original batch of point clouds
166 | Return:
167 | BxNx3 array, rotated batch of point clouds
168 | """
169 | rotated_data = np.zeros(batch_data.shape, dtype=np.float32)
170 | for k in range(batch_data.shape[0]):
171 | angles = np.clip(angle_sigma*np.random.randn(3), -angle_clip, angle_clip)
172 | Rx = np.array([[1,0,0],
173 | [0,np.cos(angles[0]),-np.sin(angles[0])],
174 | [0,np.sin(angles[0]),np.cos(angles[0])]])
175 | Ry = np.array([[np.cos(angles[1]),0,np.sin(angles[1])],
176 | [0,1,0],
177 | [-np.sin(angles[1]),0,np.cos(angles[1])]])
178 | Rz = np.array([[np.cos(angles[2]),-np.sin(angles[2]),0],
179 | [np.sin(angles[2]),np.cos(angles[2]),0],
180 | [0,0,1]])
181 | R = np.dot(Rz, np.dot(Ry,Rx))
182 | shape_pc = batch_data[k, ...]
183 | rotated_data[k, ...] = np.dot(shape_pc.reshape((-1, 3)), R)
184 | return rotated_data
185 |
186 |
187 | def jitter_point_cloud(batch_data, sigma=0.01, clip=0.05):
188 | """ Randomly jitter points. jittering is per point.
189 | Input:
190 | BxNx3 array, original batch of point clouds
191 | Return:
192 | BxNx3 array, jittered batch of point clouds
193 | """
194 | B, N, C = batch_data.shape
195 | assert(clip > 0)
196 | jittered_data = np.clip(sigma * np.random.randn(B, N, C), -1*clip, clip)
197 | jittered_data += batch_data
198 | return jittered_data
199 |
200 | def shift_point_cloud(batch_data, shift_range=0.1):
201 | """ Randomly shift point cloud. Shift is per point cloud.
202 | Input:
203 | BxNx3 array, original batch of point clouds
204 | Return:
205 | BxNx3 array, shifted batch of point clouds
206 | """
207 | B, N, C = batch_data.shape
208 | shifts = np.random.uniform(-shift_range, shift_range, (B,3))
209 | for batch_index in range(B):
210 | batch_data[batch_index,:,:] += shifts[batch_index,:]
211 | return batch_data
212 |
213 |
214 | def random_scale_point_cloud(batch_data, scale_low=0.8, scale_high=1.25):
215 | """ Randomly scale the point cloud. Scale is per point cloud.
216 | Input:
217 | BxNx3 array, original batch of point clouds
218 | Return:
219 | BxNx3 array, scaled batch of point clouds
220 | """
221 | B, N, C = batch_data.shape
222 | scales = np.random.uniform(scale_low, scale_high, B)
223 | for batch_index in range(B):
224 | batch_data[batch_index,:,:] *= scales[batch_index]
225 | return batch_data
226 |
227 | def random_point_dropout(batch_pc, max_dropout_ratio=0.875):
228 | ''' batch_pc: BxNx3 '''
229 | for b in range(batch_pc.shape[0]):
230 | dropout_ratio = np.random.random()*max_dropout_ratio # 0~0.875
231 | drop_idx = np.where(np.random.random((batch_pc.shape[1]))<=dropout_ratio)[0]
232 | if len(drop_idx)>0:
233 | batch_pc[b,drop_idx,:] = batch_pc[b,0,:] # set to the first point
234 | return batch_pc
235 |
236 |
237 | def getDataFiles(list_filename):
238 | return [line.rstrip() for line in open(list_filename)]
239 |
240 | def load_h5(h5_filename):
241 | f = h5py.File(h5_filename)
242 | data = f['data'][:]
243 | label = f['label'][:]
244 | return (data, label)
245 |
246 | def loadDataFile(filename):
247 | return load_h5(filename)
248 |
--------------------------------------------------------------------------------
/utils/render_balls_so.cpp:
--------------------------------------------------------------------------------
1 | #include
2 | #include
3 | #include
4 | #include
5 | using namespace std;
6 |
7 | struct PointInfo{
8 | int x,y,z;
9 | float r,g,b;
10 | };
11 |
12 | extern "C"{
13 |
14 | void render_ball(int h,int w,unsigned char * show,int n,int * xyzs,float * c0,float * c1,float * c2,int r){
15 | r=max(r,1);
16 | vector depth(h*w,-2100000000);
17 | vector pattern;
18 | for (int dx=-r;dx<=r;dx++)
19 | for (int dy=-r;dy<=r;dy++)
20 | if (dx*dx+dy*dy=h || y2<0 || y2>=w) && depth[x2*w+y2]0:
89 | show[:,:,0]=np.maximum(show[:,:,0],np.roll(show[:,:,0],1,axis=0))
90 | if magnifyBlue>=2:
91 | show[:,:,0]=np.maximum(show[:,:,0],np.roll(show[:,:,0],-1,axis=0))
92 | show[:,:,0]=np.maximum(show[:,:,0],np.roll(show[:,:,0],1,axis=1))
93 | if magnifyBlue>=2:
94 | show[:,:,0]=np.maximum(show[:,:,0],np.roll(show[:,:,0],-1,axis=1))
95 | if showrot:
96 | cv2.putText(show,'xangle %d'%(int(xangle/np.pi*180)),(30,showsz-30),0,0.5,cv2.cv.CV_RGB(255,0,0))
97 | cv2.putText(show,'yangle %d'%(int(yangle/np.pi*180)),(30,showsz-50),0,0.5,cv2.cv.CV_RGB(255,0,0))
98 | cv2.putText(show,'zoom %d%%'%(int(zoom*100)),(30,showsz-70),0,0.5,cv2.cv.CV_RGB(255,0,0))
99 | changed=True
100 | while True:
101 | if changed:
102 | render()
103 | changed=False
104 | cv2.imshow('show3d',show)
105 | if waittime==0:
106 | cmd=cv2.waitKey(10)%256
107 | else:
108 | cmd=cv2.waitKey(waittime)%256
109 | if cmd==ord('q'):
110 | break
111 | elif cmd==ord('Q'):
112 | sys.exit(0)
113 |
114 | if cmd==ord('t') or cmd == ord('p'):
115 | if cmd == ord('t'):
116 | if c_gt is None:
117 | c0=np.zeros((len(xyz),),dtype='float32')+255
118 | c1=np.zeros((len(xyz),),dtype='float32')+255
119 | c2=np.zeros((len(xyz),),dtype='float32')+255
120 | else:
121 | c0=c_gt[:,0]
122 | c1=c_gt[:,1]
123 | c2=c_gt[:,2]
124 | else:
125 | if c_pred is None:
126 | c0=np.zeros((len(xyz),),dtype='float32')+255
127 | c1=np.zeros((len(xyz),),dtype='float32')+255
128 | c2=np.zeros((len(xyz),),dtype='float32')+255
129 | else:
130 | c0=c_pred[:,0]
131 | c1=c_pred[:,1]
132 | c2=c_pred[:,2]
133 | if normalizecolor:
134 | c0/=(c0.max()+1e-14)/255.0
135 | c1/=(c1.max()+1e-14)/255.0
136 | c2/=(c2.max()+1e-14)/255.0
137 | c0=np.require(c0,'float32','C')
138 | c1=np.require(c1,'float32','C')
139 | c2=np.require(c2,'float32','C')
140 | changed = True
141 |
142 |
143 |
144 | if cmd==ord('n'):
145 | zoom*=1.1
146 | changed=True
147 | elif cmd==ord('m'):
148 | zoom/=1.1
149 | changed=True
150 | elif cmd==ord('r'):
151 | zoom=1.0
152 | changed=True
153 | elif cmd==ord('s'):
154 | cv2.imwrite('show3d.png',show)
155 | if waittime!=0:
156 | break
157 | return cmd
158 |
159 | if __name__=='__main__':
160 | import open3d as o3d
161 | file_path = '/media/ubuntu16/Documents/Datasets/Udacity/CH2/CH2_001/points/1479425440130861.pcd'
162 | cloud = o3d.io.read_point_cloud(file_path)
163 | cloud_np = np.asarray(cloud.points, np.float32)
164 |
165 | showpoints(cloud_np)
166 |
167 |
--------------------------------------------------------------------------------