├── LICENSE ├── README.md ├── SECOND_LICENSE ├── demo.jpg ├── second ├── __init__.py ├── builder │ ├── __init__.py │ ├── anchor_generator_builder.py │ ├── dataset_builder.py │ ├── dbsampler_builder.py │ ├── preprocess_builder.py │ ├── similarity_calculator_builder.py │ ├── target_assigner_builder.py │ └── voxel_builder.py ├── configs │ └── nuscenes │ │ └── all.pp.mhead.vpn.config ├── core │ ├── __init__.py │ ├── anchor_generator.py │ ├── box_coders.py │ ├── box_np_ops.py │ ├── geometry.py │ ├── inference.py │ ├── non_max_suppression │ │ ├── __init__.py │ │ ├── nms_cpu.py │ │ └── nms_gpu.py │ ├── preprocess.py │ ├── region_similarity.py │ ├── sample_ops.py │ ├── target_assigner.py │ └── target_ops.py ├── create_data.py ├── data │ ├── __init__.py │ ├── all_dataset.py │ ├── dataset.py │ ├── kitti_common.py │ ├── kitti_dataset.py │ ├── nusc_eval.py │ ├── nuscenes_dataset.py │ └── preprocess.py ├── protos │ ├── __init__.py │ ├── anchors.proto │ ├── anchors_pb2.py │ ├── box_coder.proto │ ├── box_coder_pb2.py │ ├── input_reader.proto │ ├── input_reader_pb2.py │ ├── losses.proto │ ├── losses_pb2.py │ ├── model.proto │ ├── model_pb2.py │ ├── optimizer.proto │ ├── optimizer_pb2.py │ ├── pipeline.proto │ ├── pipeline_pb2.py │ ├── preprocess.proto │ ├── preprocess_pb2.py │ ├── sampler.proto │ ├── sampler_pb2.py │ ├── second.proto │ ├── second_pb2.py │ ├── similarity.proto │ ├── similarity_pb2.py │ ├── target.proto │ ├── target_pb2.py │ ├── train.proto │ ├── train_pb2.py │ ├── voxel_generator.proto │ └── voxel_generator_pb2.py ├── pytorch │ ├── __init__.py │ ├── builder │ │ ├── __init__.py │ │ ├── box_coder_builder.py │ │ ├── input_reader_builder.py │ │ ├── losses_builder.py │ │ ├── lr_scheduler_builder.py │ │ ├── optimizer_builder.py │ │ └── second_builder.py │ ├── core │ │ ├── __init__.py │ │ ├── box_coders.py │ │ ├── box_torch_ops.py │ │ ├── ghm_loss.py │ │ └── losses.py │ ├── inference.py │ ├── models │ │ ├── __init__.py │ │ ├── middle.py │ │ ├── net_multi_head.py │ │ ├── pointpillars.py │ │ ├── resnet.py │ │ ├── rpn.py │ │ ├── voxel_encoder.py │ │ └── voxelnet.py │ ├── train.py │ └── utils │ │ └── __init__.py ├── script.py └── utils │ ├── __init__.py │ ├── bbox_plot.py │ ├── check.py │ ├── compare_runtime.py │ ├── config_tool.py │ ├── config_tool │ ├── __init__.py │ └── train.py │ ├── dump_kitti_files.py │ ├── dump_nusc_json.py │ ├── eval.py │ ├── find.py │ ├── loader.py │ ├── log_tool.py │ ├── mapping │ ├── CMakeLists.txt │ └── mapping.cpp │ ├── merge_kitti_results.py │ ├── merge_result.py │ ├── model_tool.py │ ├── plot_acc.py │ ├── plot_kitti_maps.py │ ├── plot_kitti_results.py │ ├── print_acc.py │ ├── print_attr.py │ ├── print_kitti_maps.py │ ├── print_test.py │ ├── print_tp_errors.py │ ├── print_vis.py │ ├── progress_bar.py │ ├── simplevis.py │ └── timer.py └── torchplus ├── __init__.py ├── metrics.py ├── nn ├── __init__.py ├── functional.py └── modules │ ├── __init__.py │ ├── common.py │ └── normalization.py ├── ops ├── __init__.py └── array_ops.py ├── tools.py └── train ├── __init__.py ├── checkpoint.py ├── common.py ├── fastai_optim.py ├── learning_schedules.py ├── learning_schedules_fastai.py └── optim.py /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright <2020> 2 | 3 | Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met: 4 | 5 | 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer. 6 | 7 | 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution. 8 | 9 | 3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission. 10 | 11 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 12 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | ![Demo result](https://raw.githubusercontent.com/peiyunh/wysiwyg/master/demo.jpg) 2 | 3 | # What You See Is What You Get:
Exploiting Visibility for 3D Object Detection 4 | By Peiyun Hu, Jason Ziglar, David Held, and Deva Ramanan 5 | 6 | ## Citing us 7 | You can find our paper on [CVF Open Access](http://openaccess.thecvf.com/content_CVPR_2020/papers/Hu_What_You_See_is_What_You_Get_Exploiting_Visibility_for_CVPR_2020_paper.pdf). If you find our work useful, please consider citing: 8 | ``` 9 | @inproceedings{hu20wysiwyg, 10 | title={What You See Is What You Get: Exploiting Visibility for 3d Object Detection}, 11 | author={Hu, Peiyun and Ziglar, Jason and Held, David and Ramanan, Deva}, 12 | booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition}, 13 | pages={11001--11009}, 14 | year={2020} 15 | } 16 | ``` 17 | 18 | ## Setup 19 | The code is developed based on [SECOND](https://github.com/traveller59/second.pytorch), a well-known deep network based 3D object detector. Please refer to this [README](https://github.com/traveller59/second.pytorch/blob/master/README.md) and [NUSCENES-GUIDE](https://github.com/traveller59/second.pytorch/blob/master/NUSCENES-GUIDE.md) on how to set up a working environment for the SECOND detector. 20 | 21 | Perhaps most importantly, you have to install [spconv](https://github.com/traveller59/spconv). 22 | 23 | ## Pre-trained models 24 | Download the pre-trained weights of the proposed model (early fusion) from [Google Drive](https://drive.google.com/file/d/1PeS6KCwi9JJlWG55MgEmNUHlgpkHeBHy/view?usp=sharing). 25 | 26 | The model configuration is located at [second/configs/nuscenes/all.pp.mhead.vpn.config](https://github.com/peiyunh/wysiwyg/blob/master/second/configs/nuscenes/all.pp.mhead.vpn.config). 27 | 28 | ## Training 29 | Run the following command to start training our model. 30 | ``` 31 | python script.py train_nuscenes --base_cfg="all.pp.mhead.vpn.config" --sample_factor=1 --epochs=20 --eval_epoch=2 --sweep_db=True --label=vp_pp_oa_ta_learn --resume=True 32 | ``` 33 | 34 | ## Raycasting 35 | You can find the code for raycasting under [second/utils/mapping](https://github.com/peiyunh/wysiwyg/blob/master/second/utils/mapping/). The default setting for raycasting `drilling`. 36 | 37 | We call raycasting functions in [second/data/preprocess.py](https://github.com/peiyunh/wysiwyg/blob/master/second/data/preprocess.py#L341-L359). If you are interested in integrating visibility into your own work, you can use it as an example of how to extract a visibility volume out of a LiDAR point cloud. 38 | 39 | You will need to compile the code using CMake. 40 | -------------------------------------------------------------------------------- /SECOND_LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2018 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /demo.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/peiyunh/wysiwyg/bb0f6ecbf3f15082fe4ec14a701ccbb14198babd/demo.jpg -------------------------------------------------------------------------------- /second/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/peiyunh/wysiwyg/bb0f6ecbf3f15082fe4ec14a701ccbb14198babd/second/__init__.py -------------------------------------------------------------------------------- /second/builder/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/peiyunh/wysiwyg/bb0f6ecbf3f15082fe4ec14a701ccbb14198babd/second/builder/__init__.py -------------------------------------------------------------------------------- /second/builder/anchor_generator_builder.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from second.protos import box_coder_pb2 4 | from second.core.anchor_generator import (AnchorGeneratorStride, 5 | AnchorGeneratorRange) 6 | 7 | 8 | def build(class_cfg): 9 | """Create optimizer based on config. 10 | 11 | Args: 12 | optimizer_config: A Optimizer proto message. 13 | 14 | Returns: 15 | An optimizer and a list of variables for summary. 16 | 17 | Raises: 18 | ValueError: when using an unsupported input data type. 19 | """ 20 | ag_type = class_cfg.WhichOneof('anchor_generator') 21 | 22 | if ag_type == 'anchor_generator_stride': 23 | config = class_cfg.anchor_generator_stride 24 | ag = AnchorGeneratorStride( 25 | sizes=list(config.sizes), 26 | anchor_strides=list(config.strides), 27 | anchor_offsets=list(config.offsets), 28 | rotations=list(config.rotations), 29 | match_threshold=class_cfg.matched_threshold, 30 | unmatch_threshold=class_cfg.unmatched_threshold, 31 | class_name=class_cfg.class_name, 32 | custom_values=list(config.custom_values)) 33 | return ag 34 | elif ag_type == 'anchor_generator_range': 35 | config = class_cfg.anchor_generator_range 36 | ag = AnchorGeneratorRange( 37 | sizes=list(config.sizes), 38 | anchor_ranges=list(config.anchor_ranges), 39 | rotations=list(config.rotations), 40 | match_threshold=class_cfg.matched_threshold, 41 | unmatch_threshold=class_cfg.unmatched_threshold, 42 | class_name=class_cfg.class_name, 43 | custom_values=list(config.custom_values)) 44 | return ag 45 | elif ag_type == 'no_anchor': 46 | return None 47 | else: 48 | raise ValueError(" unknown anchor generator type") -------------------------------------------------------------------------------- /second/builder/dataset_builder.py: -------------------------------------------------------------------------------- 1 | # Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Input reader builder. 16 | 17 | Creates data sources for DetectionModels from an InputReader config. See 18 | input_reader.proto for options. 19 | 20 | Note: If users wishes to also use their own InputReaders with the Object 21 | Detection configuration framework, they should define their own builder function 22 | that wraps the build function. 23 | """ 24 | 25 | from second.protos import input_reader_pb2 26 | from second.data.dataset import get_dataset_class 27 | from second.data.preprocess import prep_pointcloud 28 | from second.core import box_np_ops 29 | import numpy as np 30 | from second.builder import dbsampler_builder 31 | from functools import partial 32 | from second.utils.config_tool import get_downsample_factor 33 | 34 | def build(input_reader_config, 35 | model_config, 36 | training, 37 | voxel_generator, 38 | target_assigner, 39 | multi_gpu=False): 40 | """Builds a tensor dictionary based on the InputReader config. 41 | 42 | Args: 43 | input_reader_config: A input_reader_pb2.InputReader object. 44 | 45 | Returns: 46 | A tensor dict based on the input_reader_config. 47 | 48 | Raises: 49 | ValueError: On invalid input reader proto. 50 | ValueError: If no input paths are specified. 51 | """ 52 | if not isinstance(input_reader_config, input_reader_pb2.InputReader): 53 | raise ValueError('input_reader_config not of type ' 54 | 'input_reader_pb2.InputReader.') 55 | prep_cfg = input_reader_config.preprocess 56 | dataset_cfg = input_reader_config.dataset 57 | num_point_features = model_config.num_point_features 58 | out_size_factor = get_downsample_factor(model_config) 59 | assert out_size_factor > 0 60 | cfg = input_reader_config 61 | db_sampler_cfg = prep_cfg.database_sampler 62 | db_sampler = None 63 | if len(db_sampler_cfg.sample_groups) > 0 or db_sampler_cfg.database_info_path != "": # enable sample 64 | db_sampler = dbsampler_builder.build(db_sampler_cfg) 65 | grid_size = voxel_generator.grid_size 66 | print(f'grid_size: {grid_size}') 67 | feature_map_size = grid_size[:2] // out_size_factor 68 | feature_map_size = [*feature_map_size, 1][::-1] 69 | print(f'feature_map_size: {feature_map_size}') 70 | assert all([n != '' for n in target_assigner.classes]), "you must specify class_name in anchor_generators." 71 | dataset_cls = get_dataset_class(dataset_cfg.dataset_class_name) 72 | assert dataset_cls.NumPointFeatures >= 3, "you must set this to correct value" 73 | assert dataset_cls.NumPointFeatures == num_point_features, "currently you need keep them same" 74 | prep_func = partial( 75 | prep_pointcloud, 76 | root_path=dataset_cfg.kitti_root_path, 77 | voxel_generator=voxel_generator, 78 | target_assigner=target_assigner, 79 | training=training, 80 | max_voxels=prep_cfg.max_number_of_voxels, 81 | remove_outside_points=False, 82 | remove_unknown=prep_cfg.remove_unknown_examples, 83 | create_targets=training, 84 | shuffle_points=prep_cfg.shuffle_points, 85 | gt_rotation_noise=list(prep_cfg.groundtruth_rotation_uniform_noise), 86 | gt_loc_noise_std=list(prep_cfg.groundtruth_localization_noise_std), 87 | global_rotation_noise=list(prep_cfg.global_rotation_uniform_noise), 88 | global_scaling_noise=list(prep_cfg.global_scaling_uniform_noise), 89 | global_random_rot_range=list( 90 | prep_cfg.global_random_rotation_range_per_object), 91 | global_translate_noise_std=list(prep_cfg.global_translate_noise_std), 92 | db_sampler=db_sampler, 93 | num_point_features=dataset_cls.NumPointFeatures, 94 | anchor_area_threshold=prep_cfg.anchor_area_threshold, 95 | gt_points_drop=prep_cfg.groundtruth_points_drop_percentage, 96 | gt_drop_max_keep=prep_cfg.groundtruth_drop_max_keep_points, 97 | remove_points_after_sample=prep_cfg.remove_points_after_sample, 98 | remove_environment=prep_cfg.remove_environment, 99 | use_group_id=prep_cfg.use_group_id, 100 | out_size_factor=out_size_factor, 101 | multi_gpu=multi_gpu, 102 | min_points_in_gt=prep_cfg.min_num_of_points_in_gt, 103 | random_flip_x=prep_cfg.random_flip_x, 104 | random_flip_y=prep_cfg.random_flip_y, 105 | sample_importance=prep_cfg.sample_importance) 106 | 107 | ret = target_assigner.generate_anchors(feature_map_size) 108 | class_names = target_assigner.classes 109 | print(f'feature_map_size in dataset_builder.build(): {feature_map_size}') 110 | 111 | anchors_dict = target_assigner.generate_anchors_dict(feature_map_size) 112 | anchors_list = [] 113 | for k, v in anchors_dict.items(): 114 | print(f'anchors_dict[{k}]["anchors"].shape: {anchors_dict[k]["anchors"].shape}') 115 | anchors_list.append(v["anchors"]) 116 | 117 | # anchors = ret["anchors"] 118 | anchors = np.concatenate(anchors_list, axis=0) 119 | anchors = anchors.reshape([-1, target_assigner.box_ndim]) 120 | assert np.allclose(anchors, ret["anchors"].reshape(-1, target_assigner.box_ndim)) 121 | matched_thresholds = ret["matched_thresholds"] 122 | unmatched_thresholds = ret["unmatched_thresholds"] 123 | anchors_bv = box_np_ops.rbbox2d_to_near_bbox( 124 | anchors[:, [0, 1, 3, 4, 6]]) 125 | anchor_cache = { 126 | "anchors": anchors, 127 | "anchors_bv": anchors_bv, 128 | "matched_thresholds": matched_thresholds, 129 | "unmatched_thresholds": unmatched_thresholds, 130 | "anchors_dict": anchors_dict, 131 | } 132 | prep_func = partial(prep_func, anchor_cache=anchor_cache) 133 | dataset = dataset_cls( 134 | info_path=dataset_cfg.kitti_info_path, 135 | root_path=dataset_cfg.kitti_root_path, 136 | class_names=class_names, 137 | prep_func=prep_func) 138 | 139 | return dataset 140 | -------------------------------------------------------------------------------- /second/builder/dbsampler_builder.py: -------------------------------------------------------------------------------- 1 | import pickle 2 | 3 | import second.core.preprocess as prep 4 | from second.builder import preprocess_builder 5 | from second.core.preprocess import DataBasePreprocessor 6 | from second.core.sample_ops import DataBaseSamplerV2 7 | 8 | 9 | def build(sampler_config): 10 | cfg = sampler_config 11 | groups = list(cfg.sample_groups) 12 | prepors = [ 13 | preprocess_builder.build_db_preprocess(c) 14 | for c in cfg.database_prep_steps 15 | ] 16 | db_prepor = DataBasePreprocessor(prepors) 17 | rate = cfg.rate 18 | grot_range = cfg.global_random_rotation_range_per_object 19 | groups = [dict(g.name_to_max_num) for g in groups] 20 | info_path = cfg.database_info_path 21 | with open(info_path, 'rb') as f: 22 | db_infos = pickle.load(f) 23 | grot_range = list(grot_range) 24 | if len(grot_range) == 0: 25 | grot_range = None 26 | sampler = DataBaseSamplerV2(db_infos, groups, db_prepor, rate, grot_range) 27 | return sampler 28 | -------------------------------------------------------------------------------- /second/builder/preprocess_builder.py: -------------------------------------------------------------------------------- 1 | import second.core.preprocess as prep 2 | 3 | def build_db_preprocess(db_prep_config): 4 | prep_type = db_prep_config.WhichOneof('database_preprocessing_step') 5 | 6 | if prep_type == 'filter_by_difficulty': 7 | cfg = db_prep_config.filter_by_difficulty 8 | return prep.DBFilterByDifficulty(list(cfg.removed_difficulties)) 9 | elif prep_type == 'filter_by_min_num_points': 10 | cfg = db_prep_config.filter_by_min_num_points 11 | return prep.DBFilterByMinNumPoint(dict(cfg.min_num_point_pairs)) 12 | else: 13 | raise ValueError("unknown database prep type") 14 | 15 | -------------------------------------------------------------------------------- /second/builder/similarity_calculator_builder.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from second.core import region_similarity 4 | from second.protos import similarity_pb2 5 | 6 | 7 | def build(similarity_config): 8 | """Create optimizer based on config. 9 | 10 | Args: 11 | optimizer_config: A Optimizer proto message. 12 | 13 | Returns: 14 | An optimizer and a list of variables for summary. 15 | 16 | Raises: 17 | ValueError: when using an unsupported input data type. 18 | """ 19 | similarity_type = similarity_config.WhichOneof('region_similarity') 20 | if similarity_type == 'rotate_iou_similarity': 21 | return region_similarity.RotateIouSimilarity() 22 | elif similarity_type == 'nearest_iou_similarity': 23 | return region_similarity.NearestIouSimilarity() 24 | elif similarity_type == 'distance_similarity': 25 | cfg = similarity_config.distance_similarity 26 | return region_similarity.DistanceSimilarity( 27 | distance_norm=cfg.distance_norm, 28 | with_rotation=cfg.with_rotation, 29 | rotation_alpha=cfg.rotation_alpha) 30 | else: 31 | raise ValueError("unknown similarity type") 32 | -------------------------------------------------------------------------------- /second/builder/target_assigner_builder.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from second.core.target_assigner import TargetAssigner 4 | from second.protos import target_pb2, anchors_pb2 5 | from second.builder import similarity_calculator_builder 6 | from second.builder import anchor_generator_builder 7 | 8 | def build(target_assigner_config, bv_range, box_coder): 9 | """Builds a tensor dictionary based on the InputReader config. 10 | 11 | Args: 12 | input_reader_config: A input_reader_pb2.InputReader object. 13 | 14 | Returns: 15 | A tensor dict based on the input_reader_config. 16 | 17 | Raises: 18 | ValueError: On invalid input reader proto. 19 | ValueError: If no input paths are specified. 20 | """ 21 | if not isinstance(target_assigner_config, (target_pb2.TargetAssigner)): 22 | raise ValueError('input_reader_config not of type ' 23 | 'input_reader_pb2.InputReader.') 24 | classes_cfg = target_assigner_config.class_settings 25 | anchor_generators = [] 26 | classes = [] 27 | feature_map_sizes = [] 28 | for class_setting in classes_cfg: 29 | anchor_generator = anchor_generator_builder.build(class_setting) 30 | if anchor_generator is not None: 31 | anchor_generators.append(anchor_generator) 32 | else: 33 | assert target_assigner_config.assign_per_class is False 34 | classes.append(class_setting.class_name) 35 | feature_map_sizes.append(class_setting.feature_map_size) 36 | similarity_calcs = [] 37 | for class_setting in classes_cfg: 38 | similarity_calcs.append(similarity_calculator_builder.build( 39 | class_setting.region_similarity_calculator)) 40 | 41 | positive_fraction = target_assigner_config.sample_positive_fraction 42 | if positive_fraction < 0: 43 | positive_fraction = None 44 | target_assigner = TargetAssigner( 45 | box_coder=box_coder, 46 | anchor_generators=anchor_generators, 47 | feature_map_sizes=feature_map_sizes, 48 | positive_fraction=positive_fraction, 49 | sample_size=target_assigner_config.sample_size, 50 | region_similarity_calculators=similarity_calcs, 51 | classes=classes, 52 | assign_per_class=target_assigner_config.assign_per_class) 53 | return target_assigner 54 | -------------------------------------------------------------------------------- /second/builder/voxel_builder.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from spconv.utils import VoxelGeneratorV2 4 | from second.protos import voxel_generator_pb2 5 | 6 | 7 | def build(voxel_config): 8 | """Builds a tensor dictionary based on the InputReader config. 9 | 10 | Args: 11 | input_reader_config: A input_reader_pb2.InputReader object. 12 | 13 | Returns: 14 | A tensor dict based on the input_reader_config. 15 | 16 | Raises: 17 | ValueError: On invalid input reader proto. 18 | ValueError: If no input paths are specified. 19 | """ 20 | if not isinstance(voxel_config, (voxel_generator_pb2.VoxelGenerator)): 21 | raise ValueError('input_reader_config not of type ' 22 | 'input_reader_pb2.InputReader.') 23 | voxel_generator = VoxelGeneratorV2( 24 | voxel_size=list(voxel_config.voxel_size), 25 | point_cloud_range=list(voxel_config.point_cloud_range), 26 | max_num_points=voxel_config.max_number_of_points_per_voxel, 27 | max_voxels=20000, 28 | full_mean=voxel_config.full_empty_part_with_mean, 29 | block_filtering=voxel_config.block_filtering, 30 | block_factor=voxel_config.block_factor, 31 | block_size=voxel_config.block_size, 32 | height_threshold=voxel_config.height_threshold) 33 | return voxel_generator 34 | -------------------------------------------------------------------------------- /second/core/__init__.py: -------------------------------------------------------------------------------- 1 | # from . import box_np_ops, box_tf_ops, geometry, preprocess, non_max_suppression -------------------------------------------------------------------------------- /second/core/anchor_generator.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from second.core import box_np_ops 3 | 4 | class AnchorGenerator: 5 | @property 6 | def class_name(self): 7 | raise NotImplementedError 8 | 9 | @property 10 | def num_anchors_per_localization(self): 11 | raise NotImplementedError 12 | 13 | def generate(self, feature_map_size): 14 | raise NotImplementedError 15 | 16 | @property 17 | def ndim(self): 18 | raise NotImplementedError 19 | 20 | 21 | class AnchorGeneratorStride(AnchorGenerator): 22 | def __init__(self, 23 | sizes=[1.6, 3.9, 1.56], 24 | anchor_strides=[0.4, 0.4, 1.0], 25 | anchor_offsets=[0.2, -39.8, -1.78], 26 | rotations=[0, np.pi / 2], 27 | class_name=None, 28 | match_threshold=-1, 29 | unmatch_threshold=-1, 30 | custom_values=(), 31 | dtype=np.float32): 32 | super().__init__() 33 | self._sizes = sizes 34 | self._anchor_strides = anchor_strides 35 | self._anchor_offsets = anchor_offsets 36 | self._rotations = rotations 37 | self._dtype = dtype 38 | self._class_name = class_name 39 | self.match_threshold = match_threshold 40 | self.unmatch_threshold = unmatch_threshold 41 | self._custom_values = custom_values 42 | 43 | @property 44 | def class_name(self): 45 | return self._class_name 46 | 47 | @property 48 | def num_anchors_per_localization(self): 49 | num_rot = len(self._rotations) 50 | num_size = np.array(self._sizes).reshape([-1, 3]).shape[0] 51 | return num_rot * num_size 52 | 53 | def generate(self, feature_map_size): 54 | res = box_np_ops.create_anchors_3d_stride( 55 | feature_map_size, self._sizes, self._anchor_strides, 56 | self._anchor_offsets, self._rotations, self._dtype) 57 | if len(self._custom_values) > 0: 58 | custom_ndim = len(self._custom_values) 59 | custom = np.zeros([*res.shape[:-1], custom_ndim]) 60 | custom[:] = self._custom_values 61 | res = np.concatenate([res, custom], axis=-1) 62 | return res 63 | 64 | @property 65 | def ndim(self): 66 | return 7 + len(self._custom_values) 67 | 68 | @property 69 | def custom_ndim(self): 70 | return len(self._custom_values) 71 | 72 | class AnchorGeneratorRange(AnchorGenerator): 73 | def __init__(self, 74 | anchor_ranges, 75 | sizes=[1.6, 3.9, 1.56], 76 | rotations=[0, np.pi / 2], 77 | class_name=None, 78 | match_threshold=-1, 79 | unmatch_threshold=-1, 80 | custom_values=(), 81 | dtype=np.float32): 82 | super().__init__() 83 | self._sizes = sizes 84 | self._anchor_ranges = anchor_ranges 85 | self._rotations = rotations 86 | self._dtype = dtype 87 | self._class_name = class_name 88 | self.match_threshold = match_threshold 89 | self.unmatch_threshold = unmatch_threshold 90 | self._custom_values = custom_values 91 | 92 | @property 93 | def class_name(self): 94 | return self._class_name 95 | 96 | @property 97 | def num_anchors_per_localization(self): 98 | num_rot = len(self._rotations) 99 | num_size = np.array(self._sizes).reshape([-1, 3]).shape[0] 100 | return num_rot * num_size 101 | 102 | def generate(self, feature_map_size): 103 | res = box_np_ops.create_anchors_3d_range( 104 | feature_map_size, self._anchor_ranges, self._sizes, 105 | self._rotations, self._dtype) 106 | 107 | if len(self._custom_values) > 0: 108 | custom_ndim = len(self._custom_values) 109 | custom = np.zeros([*res.shape[:-1], custom_ndim]) 110 | custom[:] = self._custom_values 111 | res = np.concatenate([res, custom], axis=-1) 112 | return res 113 | 114 | @property 115 | def ndim(self): 116 | return 7 + len(self._custom_values) 117 | 118 | @property 119 | def custom_ndim(self): 120 | return len(self._custom_values) -------------------------------------------------------------------------------- /second/core/box_coders.py: -------------------------------------------------------------------------------- 1 | from abc import ABCMeta 2 | from abc import abstractmethod 3 | from abc import abstractproperty 4 | from second.core import box_np_ops 5 | import numpy as np 6 | 7 | class BoxCoder(object): 8 | """Abstract base class for box coder.""" 9 | __metaclass__ = ABCMeta 10 | 11 | @abstractproperty 12 | def code_size(self): 13 | pass 14 | 15 | def encode(self, boxes, anchors): 16 | return self._encode(boxes, anchors) 17 | 18 | def decode(self, rel_codes, anchors): 19 | return self._decode(rel_codes, anchors) 20 | 21 | @abstractmethod 22 | def _encode(self, boxes, anchors): 23 | pass 24 | 25 | @abstractmethod 26 | def _decode(self, rel_codes, anchors): 27 | pass 28 | 29 | 30 | class GroundBox3dCoder(BoxCoder): 31 | def __init__(self, linear_dim=False, vec_encode=False, custom_ndim=0): 32 | super().__init__() 33 | self.linear_dim = linear_dim 34 | self.vec_encode = vec_encode 35 | self.custom_ndim = custom_ndim 36 | 37 | @property 38 | def code_size(self): 39 | res = 8 if self.vec_encode else 7 40 | return self.custom_ndim + res 41 | 42 | def _encode(self, boxes, anchors): 43 | return box_np_ops.second_box_encode(boxes, anchors, self.vec_encode, self.linear_dim) 44 | 45 | def _decode(self, encodings, anchors): 46 | return box_np_ops.second_box_decode(encodings, anchors, self.vec_encode, self.linear_dim) 47 | 48 | 49 | class BevBoxCoder(BoxCoder): 50 | """WARNING: this coder will return encoding with size=5, but 51 | takes size=7 boxes, anchors 52 | """ 53 | def __init__(self, linear_dim=False, vec_encode=False, z_fixed=-1.0, h_fixed=2.0, custom_ndim=0): 54 | super().__init__() 55 | self.linear_dim = linear_dim 56 | self.z_fixed = z_fixed 57 | self.h_fixed = h_fixed 58 | self.vec_encode = vec_encode 59 | self.custom_ndim = custom_ndim 60 | assert self.custom_ndim == 0 61 | 62 | @property 63 | def code_size(self): 64 | res = 6 if self.vec_encode else 5 65 | return self.custom_ndim + res 66 | 67 | def _encode(self, boxes, anchors): 68 | anchors = anchors[..., [0, 1, 3, 4, 6]] 69 | boxes = boxes[..., [0, 1, 3, 4, 6]] 70 | return box_np_ops.bev_box_encode(boxes, anchors, self.vec_encode, self.linear_dim) 71 | 72 | def _decode(self, encodings, anchors): 73 | anchors = anchors[..., [0, 1, 3, 4, 6]] 74 | ret = box_np_ops.bev_box_decode(encodings, anchors, self.vec_encode, self.linear_dim) 75 | z_fixed = np.full([*ret.shape[:-1], 1], self.z_fixed, dtype=ret.dtype) 76 | h_fixed = np.full([*ret.shape[:-1], 1], self.h_fixed, dtype=ret.dtype) 77 | return np.concatenate([ret[..., :2], z_fixed, ret[..., 2:4], h_fixed, ret[..., 4:]], axis=-1) 78 | 79 | 80 | 81 | -------------------------------------------------------------------------------- /second/core/inference.py: -------------------------------------------------------------------------------- 1 | import abc 2 | import contextlib 3 | 4 | import numpy as np 5 | from google.protobuf import text_format 6 | 7 | from second.data.preprocess import merge_second_batch, prep_pointcloud 8 | from second.protos import pipeline_pb2 9 | import second.data.kitti_common as kitti 10 | 11 | class InferenceContext: 12 | def __init__(self): 13 | self.config = None 14 | self.root_path = None 15 | self.target_assigner = None 16 | self.voxel_generator = None 17 | self.anchor_cache = None 18 | self.built = False 19 | 20 | def get_inference_input_dict(self, info, points): 21 | assert self.anchor_cache is not None 22 | assert self.target_assigner is not None 23 | assert self.voxel_generator is not None 24 | assert self.config is not None 25 | assert self.built is True 26 | kitti.convert_to_kitti_info_version2(info) 27 | pc_info = info["point_cloud"] 28 | image_info = info["image"] 29 | calib = info["calib"] 30 | 31 | rect = calib['R0_rect'] 32 | Trv2c = calib['Tr_velo_to_cam'] 33 | P2 = calib['P2'] 34 | 35 | input_cfg = self.config.eval_input_reader 36 | model_cfg = self.config.model.second 37 | 38 | input_dict = { 39 | 'points': points, 40 | "calib": { 41 | 'rect': rect, 42 | 'Trv2c': Trv2c, 43 | 'P2': P2, 44 | }, 45 | "image": { 46 | 'image_shape': np.array(image_info["image_shape"], dtype=np.int32), 47 | 'image_idx': image_info['image_idx'], 48 | 'image_path': image_info['image_path'], 49 | }, 50 | } 51 | out_size_factor = np.prod(model_cfg.rpn.layer_strides) 52 | if len(model_cfg.rpn.upsample_strides) > 0: 53 | out_size_factor /= model_cfg.rpn.upsample_strides[-1] 54 | out_size_factor *= model_cfg.middle_feature_extractor.downsample_factor 55 | out_size_factor = int(out_size_factor) 56 | example = prep_pointcloud( 57 | input_dict=input_dict, 58 | root_path=str(self.root_path), 59 | voxel_generator=self.voxel_generator, 60 | target_assigner=self.target_assigner, 61 | max_voxels=input_cfg.max_number_of_voxels, 62 | class_names=self.target_assigner.classes, 63 | training=False, 64 | create_targets=False, 65 | shuffle_points=input_cfg.shuffle_points, 66 | generate_bev=False, 67 | without_reflectivity=model_cfg.without_reflectivity, 68 | num_point_features=model_cfg.num_point_features, 69 | anchor_area_threshold=input_cfg.anchor_area_threshold, 70 | anchor_cache=self.anchor_cache, 71 | out_size_factor=out_size_factor, 72 | out_dtype=np.float32) 73 | example["metadata"] = {} 74 | if "image" in info: 75 | example["metadata"]["image"] = input_dict["image"] 76 | 77 | if "anchors_mask" in example: 78 | example["anchors_mask"] = example["anchors_mask"].astype(np.uint8) 79 | ############# 80 | # convert example to batched example 81 | ############# 82 | example = merge_second_batch([example]) 83 | return example 84 | 85 | def get_config(self, path): 86 | config = pipeline_pb2.TrainEvalPipelineConfig() 87 | with open(path, "r") as f: 88 | proto_str = f.read() 89 | text_format.Merge(proto_str, config) 90 | return config 91 | 92 | @abc.abstractclassmethod 93 | def _build(self): 94 | raise NotImplementedError() 95 | 96 | def build(self, config_path): 97 | self.config = self.get_config(config_path) 98 | ret = self._build() 99 | self.built = True 100 | return ret 101 | 102 | @abc.abstractclassmethod 103 | def _inference(self, example): 104 | raise NotImplementedError() 105 | 106 | def inference(self, example): 107 | return self._inference(example) 108 | 109 | @abc.abstractclassmethod 110 | def _restore(self, ckpt_path): 111 | raise NotImplementedError() 112 | 113 | def restore(self, ckpt_path): 114 | return self._restore(ckpt_path) 115 | 116 | @abc.abstractclassmethod 117 | def _ctx(self): 118 | raise NotImplementedError() 119 | 120 | @contextlib.contextmanager 121 | def ctx(self): 122 | yield self._ctx() 123 | -------------------------------------------------------------------------------- /second/core/non_max_suppression/__init__.py: -------------------------------------------------------------------------------- 1 | from second.core.non_max_suppression.nms_cpu import nms_jit, soft_nms_jit 2 | from second.core.non_max_suppression.nms_gpu import (nms_gpu, rotate_iou_gpu, 3 | rotate_nms_gpu) 4 | -------------------------------------------------------------------------------- /second/core/non_max_suppression/nms_cpu.py: -------------------------------------------------------------------------------- 1 | import math 2 | from pathlib import Path 3 | import numba 4 | import numpy as np 5 | from spconv.utils import ( 6 | non_max_suppression_cpu, rotate_non_max_suppression_cpu) 7 | from second.core import box_np_ops 8 | from second.core.non_max_suppression.nms_gpu import rotate_iou_gpu 9 | 10 | 11 | def nms_cc(dets, thresh): 12 | scores = dets[:, 4] 13 | order = scores.argsort()[::-1].astype(np.int32) # highest->lowest 14 | return non_max_suppression_cpu(dets, order, thresh, 1.0) 15 | 16 | 17 | def rotate_nms_cc(dets, thresh): 18 | scores = dets[:, 5] 19 | order = scores.argsort()[::-1].astype(np.int32) # highest->lowest 20 | dets_corners = box_np_ops.center_to_corner_box2d(dets[:, :2], dets[:, 2:4], 21 | dets[:, 4]) 22 | 23 | dets_standup = box_np_ops.corner_to_standup_nd(dets_corners) 24 | 25 | standup_iou = box_np_ops.iou_jit(dets_standup, dets_standup, eps=0.0) 26 | # print(dets_corners.shape, order.shape, standup_iou.shape) 27 | return rotate_non_max_suppression_cpu(dets_corners, order, standup_iou, 28 | thresh) 29 | 30 | @numba.jit(nopython=True) 31 | def nms_jit(dets, thresh, eps=0.0): 32 | x1 = dets[:, 0] 33 | y1 = dets[:, 1] 34 | x2 = dets[:, 2] 35 | y2 = dets[:, 3] 36 | scores = dets[:, 4] 37 | areas = (x2 - x1 + eps) * (y2 - y1 + eps) 38 | order = scores.argsort()[::-1].astype(np.int32) # highest->lowest 39 | ndets = dets.shape[0] 40 | suppressed = np.zeros((ndets), dtype=np.int32) 41 | keep = [] 42 | for _i in range(ndets): 43 | i = order[_i] # start with highest score box 44 | if suppressed[ 45 | i] == 1: # if any box have enough iou with this, remove it 46 | continue 47 | keep.append(i) 48 | for _j in range(_i + 1, ndets): 49 | j = order[_j] 50 | if suppressed[j] == 1: 51 | continue 52 | # calculate iou between i and j box 53 | w = max(min(x2[i], x2[j]) - max(x1[i], x1[j]) + eps, 0.0) 54 | h = max(min(y2[i], y2[j]) - max(y1[i], y1[j]) + eps, 0.0) 55 | inter = w * h 56 | ovr = inter / (areas[i] + areas[j] - inter) 57 | # ovr = inter / areas[j] 58 | if ovr >= thresh: 59 | suppressed[j] = 1 60 | return keep 61 | 62 | 63 | @numba.jit('float32[:, :], float32, float32, float32, uint32', nopython=True) 64 | def soft_nms_jit(boxes, sigma=0.5, Nt=0.3, threshold=0.001, method=0): 65 | N = boxes.shape[0] 66 | pos = 0 67 | maxscore = 0 68 | maxpos = 0 69 | for i in range(N): 70 | maxscore = boxes[i, 4] 71 | maxpos = i 72 | 73 | tx1 = boxes[i, 0] 74 | ty1 = boxes[i, 1] 75 | tx2 = boxes[i, 2] 76 | ty2 = boxes[i, 3] 77 | ts = boxes[i, 4] 78 | pos = i + 1 79 | # get max box 80 | while pos < N: 81 | if maxscore < boxes[pos, 4]: 82 | maxscore = boxes[pos, 4] 83 | maxpos = pos 84 | pos = pos + 1 85 | 86 | # add max box as a detection 87 | boxes[i, 0] = boxes[maxpos, 0] 88 | boxes[i, 1] = boxes[maxpos, 1] 89 | boxes[i, 2] = boxes[maxpos, 2] 90 | boxes[i, 3] = boxes[maxpos, 3] 91 | boxes[i, 4] = boxes[maxpos, 4] 92 | 93 | # swap ith box with position of max box 94 | boxes[maxpos, 0] = tx1 95 | boxes[maxpos, 1] = ty1 96 | boxes[maxpos, 2] = tx2 97 | boxes[maxpos, 3] = ty2 98 | boxes[maxpos, 4] = ts 99 | 100 | tx1 = boxes[i, 0] 101 | ty1 = boxes[i, 1] 102 | tx2 = boxes[i, 2] 103 | ty2 = boxes[i, 3] 104 | ts = boxes[i, 4] 105 | 106 | pos = i + 1 107 | # NMS iterations, note that N changes if detection boxes fall below threshold 108 | while pos < N: 109 | x1 = boxes[pos, 0] 110 | y1 = boxes[pos, 1] 111 | x2 = boxes[pos, 2] 112 | y2 = boxes[pos, 3] 113 | s = boxes[pos, 4] 114 | 115 | area = (x2 - x1 + 1) * (y2 - y1 + 1) 116 | iw = (min(tx2, x2) - max(tx1, x1) + 1) 117 | if iw > 0: 118 | ih = (min(ty2, y2) - max(ty1, y1) + 1) 119 | if ih > 0: 120 | ua = float((tx2 - tx1 + 1) * (ty2 - ty1 + 1) + area - 121 | iw * ih) 122 | ov = iw * ih / ua #iou between max box and detection box 123 | 124 | if method == 1: # linear 125 | if ov > Nt: 126 | weight = 1 - ov 127 | else: 128 | weight = 1 129 | elif method == 2: # gaussian 130 | weight = np.exp(-(ov * ov) / sigma) 131 | else: # original NMS 132 | if ov > Nt: 133 | weight = 0 134 | else: 135 | weight = 1 136 | 137 | boxes[pos, 4] = weight * boxes[pos, 4] 138 | 139 | # if box score falls below threshold, discard the box by swapping with last box 140 | # update N 141 | if boxes[pos, 4] < threshold: 142 | boxes[pos, 0] = boxes[N - 1, 0] 143 | boxes[pos, 1] = boxes[N - 1, 1] 144 | boxes[pos, 2] = boxes[N - 1, 2] 145 | boxes[pos, 3] = boxes[N - 1, 3] 146 | boxes[pos, 4] = boxes[N - 1, 4] 147 | N = N - 1 148 | pos = pos - 1 149 | 150 | pos = pos + 1 151 | 152 | keep = [i for i in range(N)] 153 | return keep 154 | 155 | -------------------------------------------------------------------------------- /second/core/region_similarity.py: -------------------------------------------------------------------------------- 1 | # Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | 16 | """Region Similarity Calculators for BoxLists. 17 | 18 | Region Similarity Calculators compare a pairwise measure of similarity 19 | between the boxes in two BoxLists. 20 | """ 21 | from abc import ABCMeta 22 | from abc import abstractmethod 23 | 24 | from second.core import box_np_ops 25 | 26 | class RegionSimilarityCalculator(object): 27 | """Abstract base class for 2d region similarity calculator.""" 28 | __metaclass__ = ABCMeta 29 | 30 | def compare(self, boxes1, boxes2): 31 | """Computes matrix of pairwise similarity between BoxLists. 32 | 33 | This op (to be overriden) computes a measure of pairwise similarity between 34 | the boxes in the given BoxLists. Higher values indicate more similarity. 35 | 36 | Note that this method simply measures similarity and does not explicitly 37 | perform a matching. 38 | 39 | Args: 40 | boxes1: [N, 5] [x,y,w,l,r] tensor. 41 | boxes2: [M, 5] [x,y,w,l,r] tensor. 42 | 43 | Returns: 44 | a (float32) tensor of shape [N, M] with pairwise similarity score. 45 | """ 46 | return self._compare(boxes1, boxes2) 47 | 48 | @abstractmethod 49 | def _compare(self, boxes1, boxes2): 50 | pass 51 | 52 | 53 | class RotateIouSimilarity(RegionSimilarityCalculator): 54 | """Class to compute similarity based on Intersection over Union (IOU) metric. 55 | 56 | This class computes pairwise similarity between two BoxLists based on IOU. 57 | """ 58 | 59 | def _compare(self, boxes1, boxes2): 60 | """Compute pairwise IOU similarity between the two BoxLists. 61 | 62 | Args: 63 | boxlist1: BoxList holding N boxes. 64 | boxlist2: BoxList holding M boxes. 65 | 66 | Returns: 67 | A tensor with shape [N, M] representing pairwise iou scores. 68 | """ 69 | 70 | return box_np_ops.riou_cc(boxes1, boxes2) 71 | 72 | 73 | class NearestIouSimilarity(RegionSimilarityCalculator): 74 | """Class to compute similarity based on the squared distance metric. 75 | 76 | This class computes pairwise similarity between two BoxLists based on the 77 | negative squared distance metric. 78 | """ 79 | 80 | def _compare(self, boxes1, boxes2): 81 | """Compute matrix of (negated) sq distances. 82 | 83 | Args: 84 | boxlist1: BoxList holding N boxes. 85 | boxlist2: BoxList holding M boxes. 86 | 87 | Returns: 88 | A tensor with shape [N, M] representing negated pairwise squared distance. 89 | """ 90 | boxes1_bv = box_np_ops.rbbox2d_to_near_bbox(boxes1) 91 | boxes2_bv = box_np_ops.rbbox2d_to_near_bbox(boxes2) 92 | ret = box_np_ops.iou_jit(boxes1_bv, boxes2_bv, eps=0.0) 93 | return ret 94 | 95 | 96 | class DistanceSimilarity(RegionSimilarityCalculator): 97 | """Class to compute similarity based on Intersection over Area (IOA) metric. 98 | 99 | This class computes pairwise similarity between two BoxLists based on their 100 | pairwise intersections divided by the areas of second BoxLists. 101 | """ 102 | def __init__(self, distance_norm, with_rotation=False, rotation_alpha=0.5): 103 | self._distance_norm = distance_norm 104 | self._with_rotation = with_rotation 105 | self._rotation_alpha = rotation_alpha 106 | 107 | def _compare(self, boxes1, boxes2): 108 | """Compute matrix of (negated) sq distances. 109 | 110 | Args: 111 | boxlist1: BoxList holding N boxes. 112 | boxlist2: BoxList holding M boxes. 113 | 114 | Returns: 115 | A tensor with shape [N, M] representing negated pairwise squared distance. 116 | """ 117 | return box_np_ops.distance_similarity( 118 | boxes1[..., [0, 1, -1]], 119 | boxes2[..., [0, 1, -1]], 120 | dist_norm=self._distance_norm, 121 | with_rotation=self._with_rotation, 122 | rot_alpha=self._rotation_alpha) 123 | -------------------------------------------------------------------------------- /second/create_data.py: -------------------------------------------------------------------------------- 1 | import copy 2 | from pathlib import Path 3 | import pickle 4 | 5 | import fire 6 | 7 | import second.data.kitti_dataset as kitti_ds 8 | import second.data.nuscenes_dataset as nu_ds 9 | from second.data.all_dataset import create_groundtruth_database, create_groundtruth_database_with_sweep_info 10 | 11 | import numpy as np 12 | from tqdm import trange 13 | 14 | def kitti_data_prep(root_path): 15 | kitti_ds.create_kitti_info_file(root_path) 16 | kitti_ds.create_reduced_point_cloud(root_path) 17 | create_groundtruth_database("KittiDataset", root_path, Path(root_path) / "kitti_infos_train.pkl") 18 | 19 | def kitti_gt_fgm_data_prep( 20 | old_root_path, old_trainval_info_path, new_root_path, new_train_info_path 21 | ): 22 | from second.core import box_np_ops 23 | from second.data.dataset import get_dataset_class 24 | dataset = get_dataset_class('KittiDataset')( 25 | root_path = old_root_path, 26 | info_path = old_trainval_info_path 27 | ) 28 | for i in trange(len(dataset)): 29 | image_idx = i 30 | sensor_data = dataset.get_sensor_data(i) 31 | if 'image_idx' in sensor_data['metadata']: 32 | image_idx = sensor_data['metadata']['image_idx'] 33 | points = sensor_data['lidar']['points'] 34 | annos = sensor_data['lidar']['annotations'] 35 | gt_boxes = annos['boxes'] 36 | gt_mask = box_np_ops.points_in_rbbox(points, gt_boxes) 37 | points_aug = np.concatenate((points, gt_mask.max(axis=1, keepdims=True)), axis=1) 38 | points_aug = points_aug.astype(np.float32) 39 | velo_file = 'training/velodyne_reduced/%06d.bin' % (image_idx) 40 | with open(f'{new_root_path}/{velo_file}', 'w') as f: 41 | points_aug.tofile(f) 42 | create_groundtruth_database( 43 | dataset_class_name='KittiFGMDataset', 44 | data_path=new_root_path, 45 | info_path=new_train_info_path 46 | ) 47 | 48 | def nuscenes_data_prep(root_path, version, dataset_name, max_sweeps=10): 49 | nu_ds.create_nuscenes_infos(root_path, version=version, max_sweeps=max_sweeps) 50 | name = "infos_train.pkl" 51 | if version == "v1.0-test": 52 | name = "infos_test.pkl" 53 | # create_groundtruth_database(dataset_name, root_path, Path(root_path) / name) 54 | 55 | def nuscenes_data_prep_with_sweep_info(root_path, version, dataset_name, max_sweeps=10): 56 | # nu_ds.create_nuscenes_infos(root_path, version=version, max_sweeps=max_sweeps, with_sweep_info=True) 57 | name = "infos_train_with_sweep_info.pkl" 58 | if version == "v1.0-test": 59 | name = "infos_test.pkl" # because we will not have annotations 60 | create_groundtruth_database_with_sweep_info(dataset_name, root_path, Path(root_path) / name) 61 | 62 | 63 | if __name__ == '__main__': 64 | fire.Fire() 65 | -------------------------------------------------------------------------------- /second/data/__init__.py: -------------------------------------------------------------------------------- 1 | from . import kitti_dataset 2 | from . import nuscenes_dataset -------------------------------------------------------------------------------- /second/data/dataset.py: -------------------------------------------------------------------------------- 1 | import pathlib 2 | import pickle 3 | import time 4 | from functools import partial 5 | 6 | import numpy as np 7 | 8 | from second.core import box_np_ops 9 | from second.core import preprocess as prep 10 | from second.data import kitti_common as kitti 11 | 12 | REGISTERED_DATASET_CLASSES = {} 13 | 14 | def register_dataset(cls, name=None): 15 | global REGISTERED_DATASET_CLASSES 16 | if name is None: 17 | name = cls.__name__ 18 | assert name not in REGISTERED_DATASET_CLASSES, f"exist class: {REGISTERED_DATASET_CLASSES}" 19 | REGISTERED_DATASET_CLASSES[name] = cls 20 | return cls 21 | 22 | def get_dataset_class(name): 23 | global REGISTERED_DATASET_CLASSES 24 | assert name in REGISTERED_DATASET_CLASSES, f"available class: {REGISTERED_DATASET_CLASSES}" 25 | return REGISTERED_DATASET_CLASSES[name] 26 | 27 | 28 | class Dataset(object): 29 | """An abstract class representing a pytorch-like Dataset. 30 | All other datasets should subclass it. All subclasses should override 31 | ``__len__``, that provides the size of the dataset, and ``__getitem__``, 32 | supporting integer indexing in range from 0 to len(self) exclusive. 33 | """ 34 | NumPointFeatures = -1 35 | def __getitem__(self, index): 36 | """This function is used for preprocess. 37 | you need to create a input dict in this function for network inference. 38 | format: { 39 | anchors 40 | voxels 41 | num_points 42 | coordinates 43 | if training: 44 | labels 45 | reg_targets 46 | [optional]anchors_mask, slow in SECOND v1.5, don't use this. 47 | [optional]metadata, in kitti, image index is saved in metadata 48 | } 49 | """ 50 | raise NotImplementedError 51 | 52 | def __len__(self): 53 | raise NotImplementedError 54 | 55 | def get_sensor_data(self, query): 56 | """Dataset must provide a unified function to get data. 57 | Args: 58 | query: int or dict. this param must support int for training. 59 | if dict, should have this format (no example yet): 60 | { 61 | sensor_name: { 62 | sensor_meta 63 | } 64 | } 65 | if int, will return all sensor data. 66 | (TODO: how to deal with unsynchronized data?) 67 | Returns: 68 | sensor_data: dict. 69 | if query is int (return all), return a dict with all sensors: 70 | { 71 | sensor_name: sensor_data 72 | ... 73 | metadata: ... (for kitti, contains image_idx) 74 | } 75 | 76 | if sensor is lidar (all lidar point cloud must be concatenated to one array): 77 | e.g. If your dataset have two lidar sensor, you need to return a single dict: 78 | { 79 | "lidar": { 80 | "points": ... 81 | ... 82 | } 83 | } 84 | sensor_data: { 85 | points: [N, 3+] 86 | [optional]annotations: { 87 | "boxes": [N, 7] locs, dims, yaw, in lidar coord system. must tested 88 | in provided visualization tools such as second.utils.simplevis 89 | or web tool. 90 | "names": array of string. 91 | } 92 | } 93 | if sensor is camera (not used yet): 94 | sensor_data: { 95 | data: image string (array is too large) 96 | [optional]annotations: { 97 | "boxes": [N, 4] 2d bbox 98 | "names": array of string. 99 | } 100 | } 101 | metadata: { 102 | # dataset-specific information. 103 | # for kitti, must have image_idx for label file generation. 104 | image_idx: ... 105 | } 106 | [optional]calib # only used for kitti 107 | """ 108 | raise NotImplementedError 109 | 110 | def evaluation(self, dt_annos, output_dir): 111 | """Dataset must provide a evaluation function to evaluate model.""" 112 | raise NotImplementedError 113 | -------------------------------------------------------------------------------- /second/data/nusc_eval.py: -------------------------------------------------------------------------------- 1 | import fire 2 | 3 | from nuscenes import NuScenes 4 | from nuscenes.eval.detection.config import config_factory 5 | from nuscenes.eval.detection.evaluate import NuScenesEval 6 | 7 | def eval_main(root_path, version, eval_version, res_path, eval_set, output_dir): 8 | nusc = NuScenes( 9 | version=version, dataroot=str(root_path), verbose=False) 10 | 11 | cfg = config_factory(eval_version) 12 | nusc_eval = NuScenesEval(nusc, config=cfg, result_path=res_path, eval_set=eval_set, 13 | output_dir=output_dir, 14 | verbose=False) 15 | nusc_eval.main(render_curves=False) 16 | 17 | if __name__ == "__main__": 18 | fire.Fire(eval_main) -------------------------------------------------------------------------------- /second/protos/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/peiyunh/wysiwyg/bb0f6ecbf3f15082fe4ec14a701ccbb14198babd/second/protos/__init__.py -------------------------------------------------------------------------------- /second/protos/anchors.proto: -------------------------------------------------------------------------------- 1 | syntax = "proto3"; 2 | 3 | package second.protos; 4 | 5 | message AnchorGeneratorStride { 6 | repeated float sizes = 1; 7 | repeated float strides = 2; 8 | repeated float offsets = 3; 9 | repeated float rotations = 4; 10 | repeated float custom_values = 5; 11 | } 12 | 13 | message AnchorGeneratorRange { 14 | repeated float sizes = 1; 15 | repeated float anchor_ranges = 2; 16 | repeated float rotations = 3; 17 | repeated float custom_values = 4; 18 | } 19 | 20 | message NoAnchor { 21 | } 22 | 23 | -------------------------------------------------------------------------------- /second/protos/box_coder.proto: -------------------------------------------------------------------------------- 1 | syntax = "proto3"; 2 | 3 | package second.protos; 4 | 5 | // Configuration proto for the box coder to be used in the object detection 6 | // pipeline. See core/box_coder.py for details. 7 | message BoxCoder { 8 | oneof box_coder { 9 | GroundBox3dCoder ground_box3d_coder = 1; 10 | BevBoxCoder bev_box_coder = 2; 11 | } 12 | } 13 | 14 | message GroundBox3dCoder { 15 | bool linear_dim = 1; 16 | bool encode_angle_vector = 2; 17 | } 18 | 19 | message BevBoxCoder { 20 | bool linear_dim = 1; 21 | bool encode_angle_vector = 2; 22 | float z_fixed = 3; 23 | float h_fixed = 4; 24 | } 25 | -------------------------------------------------------------------------------- /second/protos/input_reader.proto: -------------------------------------------------------------------------------- 1 | syntax = "proto3"; 2 | 3 | package second.protos; 4 | import "second/protos/preprocess.proto"; 5 | import "second/protos/sampler.proto"; 6 | 7 | message InputReader { 8 | uint32 batch_size = 1; 9 | message Dataset { 10 | string kitti_info_path = 1; 11 | string kitti_root_path = 2; 12 | string dataset_class_name = 3; // support KittiDataset and NuScenesDataset 13 | } 14 | Dataset dataset = 2; 15 | message Preprocess { 16 | bool shuffle_points = 1; 17 | uint32 max_number_of_voxels = 2; 18 | repeated float groundtruth_localization_noise_std = 3; 19 | repeated float groundtruth_rotation_uniform_noise = 4; 20 | repeated float global_rotation_uniform_noise = 5; 21 | repeated float global_scaling_uniform_noise = 6; 22 | repeated float global_translate_noise_std = 7; 23 | bool remove_unknown_examples = 8; 24 | uint32 num_workers = 9; 25 | float anchor_area_threshold = 10; 26 | bool remove_points_after_sample = 11; 27 | float groundtruth_points_drop_percentage = 12; 28 | uint32 groundtruth_drop_max_keep_points = 13; 29 | bool remove_environment = 14; 30 | repeated float global_random_rotation_range_per_object = 15; 31 | repeated DatabasePreprocessingStep database_prep_steps = 16; 32 | Sampler database_sampler = 17; 33 | bool use_group_id = 18; // this will enable group sample and noise 34 | int64 min_num_of_points_in_gt = 19; // gt boxes contains less than this will be ignored. 35 | bool random_flip_x = 20; 36 | bool random_flip_y = 21; 37 | float sample_importance = 22; 38 | } 39 | Preprocess preprocess = 3; 40 | uint32 max_num_epochs = 4; // deprecated 41 | uint32 prefetch_size = 5; // deprecated 42 | } 43 | -------------------------------------------------------------------------------- /second/protos/losses.proto: -------------------------------------------------------------------------------- 1 | syntax = "proto3"; 2 | 3 | package second.protos; 4 | 5 | // Message for configuring the localization loss, classification loss and hard 6 | // example miner used for training object detection models. See core/losses.py 7 | // for details 8 | message Loss { 9 | // Localization loss to use. 10 | LocalizationLoss localization_loss = 1; 11 | 12 | // Classification loss to use. 13 | ClassificationLoss classification_loss = 2; 14 | 15 | // If not left to default, applies hard example mining. 16 | HardExampleMiner hard_example_miner = 3; 17 | 18 | // Classification loss weight. 19 | float classification_weight = 4; 20 | 21 | // Localization loss weight. 22 | float localization_weight = 5; 23 | 24 | } 25 | 26 | // Configuration for bounding box localization loss function. 27 | message LocalizationLoss { 28 | oneof localization_loss { 29 | WeightedL2LocalizationLoss weighted_l2 = 1; 30 | WeightedSmoothL1LocalizationLoss weighted_smooth_l1 = 2; 31 | WeightedGHMLocalizationLoss weighted_ghm = 3; 32 | } 33 | bool encode_rad_error_by_sin = 4; 34 | } 35 | 36 | // L2 location loss: 0.5 * ||weight * (a - b)|| ^ 2 37 | message WeightedL2LocalizationLoss { 38 | // DEPRECATED, do not use. 39 | // Output loss per anchor. 40 | bool anchorwise_output = 1; 41 | repeated float code_weight = 2; 42 | } 43 | 44 | // SmoothL1 (Huber) location loss: .5 * x ^ 2 if |x| < 1 else |x| - .5 45 | message WeightedSmoothL1LocalizationLoss { 46 | // DEPRECATED, do not use. 47 | // Output loss per anchor. 48 | bool anchorwise_output = 1; 49 | float sigma = 2; 50 | repeated float code_weight = 3; 51 | } 52 | message WeightedGHMLocalizationLoss { 53 | // DEPRECATED, do not use. 54 | // Output loss per anchor. 55 | bool anchorwise_output = 1; 56 | float mu = 2; 57 | int32 bins = 3; 58 | float momentum = 4; 59 | repeated float code_weight = 5; 60 | } 61 | 62 | 63 | // Configuration for class prediction loss function. 64 | message ClassificationLoss { 65 | oneof classification_loss { 66 | WeightedSigmoidClassificationLoss weighted_sigmoid = 1; 67 | WeightedSoftmaxClassificationLoss weighted_softmax = 2; 68 | BootstrappedSigmoidClassificationLoss bootstrapped_sigmoid = 3; 69 | SigmoidFocalClassificationLoss weighted_sigmoid_focal = 4; 70 | SoftmaxFocalClassificationLoss weighted_softmax_focal = 5; 71 | GHMClassificationLoss weighted_ghm = 6; 72 | } 73 | } 74 | 75 | // Classification loss using a sigmoid function over class predictions. 76 | message WeightedSigmoidClassificationLoss { 77 | // DEPRECATED, do not use. 78 | // Output loss per anchor. 79 | bool anchorwise_output = 1; 80 | } 81 | 82 | // Sigmoid Focal cross entropy loss as described in 83 | // https://arxiv.org/abs/1708.02002 84 | message SigmoidFocalClassificationLoss { 85 | // DEPRECATED, do not use. 86 | bool anchorwise_output = 1; 87 | // modulating factor for the loss. 88 | float gamma = 2; 89 | // alpha weighting factor for the loss. 90 | float alpha = 3; 91 | } 92 | // Sigmoid Focal cross entropy loss as described in 93 | // https://arxiv.org/abs/1708.02002 94 | message SoftmaxFocalClassificationLoss { 95 | // DEPRECATED, do not use. 96 | bool anchorwise_output = 1; 97 | // modulating factor for the loss. 98 | float gamma = 2; 99 | // alpha weighting factor for the loss. 100 | float alpha = 3; 101 | } 102 | message GHMClassificationLoss { 103 | bool anchorwise_output = 1; 104 | int32 bins = 2; 105 | float momentum = 3; 106 | } 107 | // Classification loss using a softmax function over class predictions. 108 | message WeightedSoftmaxClassificationLoss { 109 | // DEPRECATED, do not use. 110 | // Output loss per anchor. 111 | bool anchorwise_output = 1; 112 | // Scale logit (input) value before calculating softmax classification loss. 113 | // Typically used for softmax distillation. 114 | float logit_scale = 2; 115 | } 116 | 117 | // Classification loss using a sigmoid function over the class prediction with 118 | // the highest prediction score. 119 | message BootstrappedSigmoidClassificationLoss { 120 | // Interpolation weight between 0 and 1. 121 | float alpha = 1; 122 | 123 | // Whether hard boot strapping should be used or not. If true, will only use 124 | // one class favored by model. Othewise, will use all predicted class 125 | // probabilities. 126 | bool hard_bootstrap = 2; 127 | 128 | // DEPRECATED, do not use. 129 | // Output loss per anchor. 130 | bool anchorwise_output = 3; 131 | } 132 | 133 | // Configuation for hard example miner. 134 | message HardExampleMiner { 135 | // Maximum number of hard examples to be selected per image (prior to 136 | // enforcing max negative to positive ratio constraint). If set to 0, 137 | // all examples obtained after NMS are considered. 138 | int32 num_hard_examples = 1; 139 | 140 | // Minimum intersection over union for an example to be discarded during NMS. 141 | float iou_threshold = 2; 142 | 143 | // Whether to use classification losses ('cls', default), localization losses 144 | // ('loc') or both losses ('both'). In the case of 'both', cls_loss_weight and 145 | // loc_loss_weight are used to compute weighted sum of the two losses. 146 | enum LossType { 147 | BOTH = 0; 148 | CLASSIFICATION = 1; 149 | LOCALIZATION = 2; 150 | } 151 | LossType loss_type = 3; 152 | 153 | // Maximum number of negatives to retain for each positive anchor. If 154 | // num_negatives_per_positive is 0 no prespecified negative:positive ratio is 155 | // enforced. 156 | int32 max_negatives_per_positive = 4; 157 | 158 | // Minimum number of negative anchors to sample for a given image. Setting 159 | // this to a positive number samples negatives in an image without any 160 | // positive anchors and thus not bias the model towards having at least one 161 | // detection per image. 162 | int32 min_negatives_per_image = 5; 163 | } 164 | -------------------------------------------------------------------------------- /second/protos/model.proto: -------------------------------------------------------------------------------- 1 | syntax = "proto3"; 2 | 3 | package second.protos; 4 | import "second/protos/second.proto"; 5 | message DetectionModel{ 6 | oneof model { 7 | VoxelNet second = 1; 8 | } 9 | } 10 | -------------------------------------------------------------------------------- /second/protos/model_pb2.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # Generated by the protocol buffer compiler. DO NOT EDIT! 3 | # source: second/protos/model.proto 4 | 5 | import sys 6 | _b=sys.version_info[0]<3 and (lambda x:x) or (lambda x:x.encode('latin1')) 7 | from google.protobuf import descriptor as _descriptor 8 | from google.protobuf import message as _message 9 | from google.protobuf import reflection as _reflection 10 | from google.protobuf import symbol_database as _symbol_database 11 | # @@protoc_insertion_point(imports) 12 | 13 | _sym_db = _symbol_database.Default() 14 | 15 | 16 | from second.protos import second_pb2 as second_dot_protos_dot_second__pb2 17 | 18 | 19 | DESCRIPTOR = _descriptor.FileDescriptor( 20 | name='second/protos/model.proto', 21 | package='second.protos', 22 | syntax='proto3', 23 | serialized_options=None, 24 | serialized_pb=_b('\n\x19second/protos/model.proto\x12\rsecond.protos\x1a\x1asecond/protos/second.proto\"D\n\x0e\x44\x65tectionModel\x12)\n\x06second\x18\x01 \x01(\x0b\x32\x17.second.protos.VoxelNetH\x00\x42\x07\n\x05modelb\x06proto3') 25 | , 26 | dependencies=[second_dot_protos_dot_second__pb2.DESCRIPTOR,]) 27 | 28 | 29 | 30 | 31 | _DETECTIONMODEL = _descriptor.Descriptor( 32 | name='DetectionModel', 33 | full_name='second.protos.DetectionModel', 34 | filename=None, 35 | file=DESCRIPTOR, 36 | containing_type=None, 37 | fields=[ 38 | _descriptor.FieldDescriptor( 39 | name='second', full_name='second.protos.DetectionModel.second', index=0, 40 | number=1, type=11, cpp_type=10, label=1, 41 | has_default_value=False, default_value=None, 42 | message_type=None, enum_type=None, containing_type=None, 43 | is_extension=False, extension_scope=None, 44 | serialized_options=None, file=DESCRIPTOR), 45 | ], 46 | extensions=[ 47 | ], 48 | nested_types=[], 49 | enum_types=[ 50 | ], 51 | serialized_options=None, 52 | is_extendable=False, 53 | syntax='proto3', 54 | extension_ranges=[], 55 | oneofs=[ 56 | _descriptor.OneofDescriptor( 57 | name='model', full_name='second.protos.DetectionModel.model', 58 | index=0, containing_type=None, fields=[]), 59 | ], 60 | serialized_start=72, 61 | serialized_end=140, 62 | ) 63 | 64 | _DETECTIONMODEL.fields_by_name['second'].message_type = second_dot_protos_dot_second__pb2._VOXELNET 65 | _DETECTIONMODEL.oneofs_by_name['model'].fields.append( 66 | _DETECTIONMODEL.fields_by_name['second']) 67 | _DETECTIONMODEL.fields_by_name['second'].containing_oneof = _DETECTIONMODEL.oneofs_by_name['model'] 68 | DESCRIPTOR.message_types_by_name['DetectionModel'] = _DETECTIONMODEL 69 | _sym_db.RegisterFileDescriptor(DESCRIPTOR) 70 | 71 | DetectionModel = _reflection.GeneratedProtocolMessageType('DetectionModel', (_message.Message,), dict( 72 | DESCRIPTOR = _DETECTIONMODEL, 73 | __module__ = 'second.protos.model_pb2' 74 | # @@protoc_insertion_point(class_scope:second.protos.DetectionModel) 75 | )) 76 | _sym_db.RegisterMessage(DetectionModel) 77 | 78 | 79 | # @@protoc_insertion_point(module_scope) 80 | -------------------------------------------------------------------------------- /second/protos/optimizer.proto: -------------------------------------------------------------------------------- 1 | syntax = "proto3"; 2 | 3 | package second.protos; 4 | 5 | // Messages for configuring the optimizing strategy for training object 6 | // detection models. 7 | 8 | // Top level optimizer message. 9 | message Optimizer { 10 | oneof optimizer { 11 | RMSPropOptimizer rms_prop_optimizer = 1; 12 | MomentumOptimizer momentum_optimizer = 2; 13 | AdamOptimizer adam_optimizer = 3; 14 | } 15 | bool use_moving_average = 4; 16 | float moving_average_decay = 5; 17 | bool fixed_weight_decay = 6; // i.e. AdamW 18 | } 19 | 20 | // Configuration message for the RMSPropOptimizer 21 | // See: https://www.tensorflow.org/api_docs/python/tf/train/RMSPropOptimizer 22 | message RMSPropOptimizer { 23 | LearningRate learning_rate = 1; 24 | float momentum_optimizer_value = 2; 25 | float decay = 3; 26 | float epsilon = 4; 27 | float weight_decay = 5; 28 | } 29 | 30 | // Configuration message for the MomentumOptimizer 31 | // See: https://www.tensorflow.org/api_docs/python/tf/train/MomentumOptimizer 32 | message MomentumOptimizer { 33 | LearningRate learning_rate = 1; 34 | float momentum_optimizer_value = 2; 35 | float weight_decay = 3; 36 | } 37 | 38 | // Configuration message for the AdamOptimizer 39 | // See: https://www.tensorflow.org/api_docs/python/tf/train/AdamOptimizer 40 | message AdamOptimizer { 41 | LearningRate learning_rate = 1; 42 | float weight_decay = 2; 43 | bool amsgrad = 3; 44 | } 45 | 46 | message LearningRate { 47 | oneof learning_rate { 48 | MultiPhase multi_phase = 1; 49 | OneCycle one_cycle = 2; 50 | ExponentialDecay exponential_decay = 3; 51 | ManualStepping manual_stepping = 4; 52 | } 53 | } 54 | 55 | message LearningRatePhase { 56 | float start = 1; 57 | string lambda_func = 2; 58 | string momentum_lambda_func = 3; 59 | } 60 | 61 | message MultiPhase { 62 | repeated LearningRatePhase phases = 1; 63 | } 64 | 65 | message OneCycle { 66 | float lr_max = 1; 67 | repeated float moms = 2; 68 | float div_factor = 3; 69 | float pct_start = 4; 70 | } 71 | 72 | /* 73 | ManualStepping example: 74 | initial_learning_rate = 0.001 75 | decay_length = 0.1 76 | decay_factor = 0.8 77 | staircase = True 78 | detail: 79 | progress 0%~10%, lr=0.001 80 | progress 10%~20%, lr=0.001 * 0.8 81 | progress 20%~30%, lr=0.001 * 0.8 * 0.8 82 | ...... 83 | */ 84 | 85 | 86 | message ExponentialDecay { 87 | float initial_learning_rate = 1; 88 | float decay_length = 2; // must in range (0, 1) 89 | float decay_factor = 3; 90 | bool staircase = 4; 91 | } 92 | 93 | /* 94 | ManualStepping example: 95 | boundaries = [0.8, 0.9] 96 | rates = [0.001, 0.002, 0.003] 97 | detail: 98 | progress 0%~80%, lr=0.001 99 | progress 80%~90%, lr=0.002 100 | progress 90%~100%, lr=0.003 101 | */ 102 | 103 | message ManualStepping { 104 | repeated float boundaries = 1; // must in range (0, 1) 105 | repeated float rates = 2; 106 | } -------------------------------------------------------------------------------- /second/protos/pipeline.proto: -------------------------------------------------------------------------------- 1 | syntax = "proto3"; 2 | 3 | package second.protos; 4 | 5 | import "second/protos/input_reader.proto"; 6 | import "second/protos/model.proto"; 7 | import "second/protos/train.proto"; 8 | // Convenience message for configuring a training and eval pipeline. Allows all 9 | // of the pipeline parameters to be configured from one file. 10 | message TrainEvalPipelineConfig { 11 | DetectionModel model = 1; 12 | InputReader train_input_reader = 2; 13 | TrainConfig train_config = 3; 14 | InputReader eval_input_reader = 4; 15 | } 16 | 17 | -------------------------------------------------------------------------------- /second/protos/pipeline_pb2.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # Generated by the protocol buffer compiler. DO NOT EDIT! 3 | # source: second/protos/pipeline.proto 4 | 5 | import sys 6 | _b=sys.version_info[0]<3 and (lambda x:x) or (lambda x:x.encode('latin1')) 7 | from google.protobuf import descriptor as _descriptor 8 | from google.protobuf import message as _message 9 | from google.protobuf import reflection as _reflection 10 | from google.protobuf import symbol_database as _symbol_database 11 | # @@protoc_insertion_point(imports) 12 | 13 | _sym_db = _symbol_database.Default() 14 | 15 | 16 | from second.protos import input_reader_pb2 as second_dot_protos_dot_input__reader__pb2 17 | from second.protos import model_pb2 as second_dot_protos_dot_model__pb2 18 | from second.protos import train_pb2 as second_dot_protos_dot_train__pb2 19 | 20 | 21 | DESCRIPTOR = _descriptor.FileDescriptor( 22 | name='second/protos/pipeline.proto', 23 | package='second.protos', 24 | syntax='proto3', 25 | serialized_options=None, 26 | serialized_pb=_b('\n\x1csecond/protos/pipeline.proto\x12\rsecond.protos\x1a second/protos/input_reader.proto\x1a\x19second/protos/model.proto\x1a\x19second/protos/train.proto\"\xe8\x01\n\x17TrainEvalPipelineConfig\x12,\n\x05model\x18\x01 \x01(\x0b\x32\x1d.second.protos.DetectionModel\x12\x36\n\x12train_input_reader\x18\x02 \x01(\x0b\x32\x1a.second.protos.InputReader\x12\x30\n\x0ctrain_config\x18\x03 \x01(\x0b\x32\x1a.second.protos.TrainConfig\x12\x35\n\x11\x65val_input_reader\x18\x04 \x01(\x0b\x32\x1a.second.protos.InputReaderb\x06proto3') 27 | , 28 | dependencies=[second_dot_protos_dot_input__reader__pb2.DESCRIPTOR,second_dot_protos_dot_model__pb2.DESCRIPTOR,second_dot_protos_dot_train__pb2.DESCRIPTOR,]) 29 | 30 | 31 | 32 | 33 | _TRAINEVALPIPELINECONFIG = _descriptor.Descriptor( 34 | name='TrainEvalPipelineConfig', 35 | full_name='second.protos.TrainEvalPipelineConfig', 36 | filename=None, 37 | file=DESCRIPTOR, 38 | containing_type=None, 39 | fields=[ 40 | _descriptor.FieldDescriptor( 41 | name='model', full_name='second.protos.TrainEvalPipelineConfig.model', index=0, 42 | number=1, type=11, cpp_type=10, label=1, 43 | has_default_value=False, default_value=None, 44 | message_type=None, enum_type=None, containing_type=None, 45 | is_extension=False, extension_scope=None, 46 | serialized_options=None, file=DESCRIPTOR), 47 | _descriptor.FieldDescriptor( 48 | name='train_input_reader', full_name='second.protos.TrainEvalPipelineConfig.train_input_reader', index=1, 49 | number=2, type=11, cpp_type=10, label=1, 50 | has_default_value=False, default_value=None, 51 | message_type=None, enum_type=None, containing_type=None, 52 | is_extension=False, extension_scope=None, 53 | serialized_options=None, file=DESCRIPTOR), 54 | _descriptor.FieldDescriptor( 55 | name='train_config', full_name='second.protos.TrainEvalPipelineConfig.train_config', index=2, 56 | number=3, type=11, cpp_type=10, label=1, 57 | has_default_value=False, default_value=None, 58 | message_type=None, enum_type=None, containing_type=None, 59 | is_extension=False, extension_scope=None, 60 | serialized_options=None, file=DESCRIPTOR), 61 | _descriptor.FieldDescriptor( 62 | name='eval_input_reader', full_name='second.protos.TrainEvalPipelineConfig.eval_input_reader', index=3, 63 | number=4, type=11, cpp_type=10, label=1, 64 | has_default_value=False, default_value=None, 65 | message_type=None, enum_type=None, containing_type=None, 66 | is_extension=False, extension_scope=None, 67 | serialized_options=None, file=DESCRIPTOR), 68 | ], 69 | extensions=[ 70 | ], 71 | nested_types=[], 72 | enum_types=[ 73 | ], 74 | serialized_options=None, 75 | is_extendable=False, 76 | syntax='proto3', 77 | extension_ranges=[], 78 | oneofs=[ 79 | ], 80 | serialized_start=136, 81 | serialized_end=368, 82 | ) 83 | 84 | _TRAINEVALPIPELINECONFIG.fields_by_name['model'].message_type = second_dot_protos_dot_model__pb2._DETECTIONMODEL 85 | _TRAINEVALPIPELINECONFIG.fields_by_name['train_input_reader'].message_type = second_dot_protos_dot_input__reader__pb2._INPUTREADER 86 | _TRAINEVALPIPELINECONFIG.fields_by_name['train_config'].message_type = second_dot_protos_dot_train__pb2._TRAINCONFIG 87 | _TRAINEVALPIPELINECONFIG.fields_by_name['eval_input_reader'].message_type = second_dot_protos_dot_input__reader__pb2._INPUTREADER 88 | DESCRIPTOR.message_types_by_name['TrainEvalPipelineConfig'] = _TRAINEVALPIPELINECONFIG 89 | _sym_db.RegisterFileDescriptor(DESCRIPTOR) 90 | 91 | TrainEvalPipelineConfig = _reflection.GeneratedProtocolMessageType('TrainEvalPipelineConfig', (_message.Message,), dict( 92 | DESCRIPTOR = _TRAINEVALPIPELINECONFIG, 93 | __module__ = 'second.protos.pipeline_pb2' 94 | # @@protoc_insertion_point(class_scope:second.protos.TrainEvalPipelineConfig) 95 | )) 96 | _sym_db.RegisterMessage(TrainEvalPipelineConfig) 97 | 98 | 99 | # @@protoc_insertion_point(module_scope) 100 | -------------------------------------------------------------------------------- /second/protos/preprocess.proto: -------------------------------------------------------------------------------- 1 | syntax = "proto3"; 2 | 3 | package second.protos; 4 | 5 | message DatabasePreprocessingStep { 6 | oneof database_preprocessing_step { 7 | DBFilterByDifficulty filter_by_difficulty = 1; 8 | DBFilterByMinNumPointInGroundTruth filter_by_min_num_points = 2; 9 | } 10 | } 11 | 12 | message DBFilterByDifficulty{ 13 | repeated int32 removed_difficulties = 1; 14 | } 15 | 16 | message DBFilterByMinNumPointInGroundTruth{ 17 | map min_num_point_pairs = 1; 18 | } 19 | -------------------------------------------------------------------------------- /second/protos/sampler.proto: -------------------------------------------------------------------------------- 1 | syntax = "proto3"; 2 | 3 | package second.protos; 4 | import "second/protos/preprocess.proto"; 5 | 6 | message Group{ 7 | map name_to_max_num = 1; 8 | } 9 | 10 | message Sampler{ 11 | string database_info_path = 1; 12 | repeated Group sample_groups = 2; 13 | repeated DatabasePreprocessingStep database_prep_steps = 3; 14 | repeated float global_random_rotation_range_per_object = 4; 15 | float rate = 5; 16 | } 17 | -------------------------------------------------------------------------------- /second/protos/second.proto: -------------------------------------------------------------------------------- 1 | syntax = "proto3"; 2 | 3 | package second.protos; 4 | import "second/protos/losses.proto"; 5 | import "second/protos/box_coder.proto"; 6 | import "second/protos/target.proto"; 7 | import "second/protos/voxel_generator.proto"; 8 | 9 | message VoxelNet { 10 | string network_class_name = 1; 11 | VoxelGenerator voxel_generator = 2; 12 | message VoxelFeatureExtractor { 13 | string module_class_name = 1; 14 | repeated int32 num_filters = 2; 15 | bool with_distance = 3; 16 | int32 num_input_features = 4; 17 | } 18 | VoxelFeatureExtractor voxel_feature_extractor = 3; 19 | message MiddleFeatureExtractor { 20 | string module_class_name = 1; 21 | repeated int32 num_filters_down1 = 2; 22 | repeated int32 num_filters_down2 = 3; 23 | int32 num_input_features = 4; 24 | int32 downsample_factor = 5; 25 | } 26 | MiddleFeatureExtractor middle_feature_extractor = 4; 27 | message RPN { 28 | string module_class_name = 1; 29 | repeated int32 layer_nums = 2; 30 | repeated int32 layer_strides = 3; 31 | repeated int32 num_filters = 4; 32 | repeated double upsample_strides = 5; 33 | repeated int32 num_upsample_filters = 6; 34 | bool use_groupnorm = 7; 35 | int32 num_groups = 8; 36 | int32 num_input_features = 9; 37 | } 38 | RPN rpn = 5; 39 | RPN vpn = 25; 40 | uint32 num_point_features = 6; 41 | bool use_sigmoid_score = 7; 42 | Loss loss = 8; 43 | bool encode_rad_error_by_sin = 9; 44 | bool encode_background_as_zeros = 10; 45 | bool use_direction_classifier = 11; 46 | float direction_loss_weight = 12; 47 | float pos_class_weight = 13; 48 | float neg_class_weight = 14; 49 | enum LossNormType { 50 | NormByNumExamples = 0; 51 | NormByNumPositives = 1; 52 | NormByNumPosNeg = 2; 53 | DontNorm = 3; 54 | } 55 | LossNormType loss_norm_type = 15; 56 | BoxCoder box_coder = 16; 57 | TargetAssigner target_assigner = 17; 58 | repeated float post_center_limit_range = 18; 59 | float direction_offset = 19; 60 | float sin_error_factor = 20; 61 | bool nms_class_agnostic = 21; 62 | int64 num_direction_bins = 22; 63 | float direction_limit_offset = 23; 64 | 65 | // deprecated in future 66 | bool lidar_input = 24; 67 | 68 | } 69 | -------------------------------------------------------------------------------- /second/protos/similarity.proto: -------------------------------------------------------------------------------- 1 | syntax = "proto3"; 2 | 3 | package second.protos; 4 | 5 | // Configuration proto for region similarity calculators. See 6 | // core/region_similarity_calculator.py for details. 7 | message RegionSimilarityCalculator { 8 | oneof region_similarity { 9 | RotateIouSimilarity rotate_iou_similarity = 1; 10 | NearestIouSimilarity nearest_iou_similarity = 2; 11 | DistanceSimilarity distance_similarity = 3; 12 | } 13 | } 14 | 15 | // Configuration for intersection-over-union (IOU) similarity calculator. 16 | message RotateIouSimilarity { 17 | } 18 | 19 | // Configuration for intersection-over-union (IOU) similarity calculator. 20 | message NearestIouSimilarity { 21 | } 22 | 23 | // Configuration for intersection-over-union (IOU) similarity calculator. 24 | message DistanceSimilarity { 25 | float distance_norm = 1; 26 | bool with_rotation = 2; 27 | float rotation_alpha = 3; 28 | } -------------------------------------------------------------------------------- /second/protos/target.proto: -------------------------------------------------------------------------------- 1 | syntax = "proto3"; 2 | 3 | package second.protos; 4 | import "second/protos/anchors.proto"; 5 | import "second/protos/similarity.proto"; 6 | 7 | message ClassSetting { 8 | oneof anchor_generator { 9 | AnchorGeneratorStride anchor_generator_stride = 1; 10 | AnchorGeneratorRange anchor_generator_range = 2; 11 | NoAnchor no_anchor = 3; 12 | } 13 | RegionSimilarityCalculator region_similarity_calculator = 4; 14 | bool use_multi_class_nms = 5; 15 | bool use_rotate_nms = 6; 16 | int32 nms_pre_max_size = 7; 17 | int32 nms_post_max_size = 8; 18 | float nms_score_threshold = 9; 19 | float nms_iou_threshold = 10; 20 | float matched_threshold = 11; 21 | float unmatched_threshold = 12; 22 | string class_name = 13; 23 | repeated int64 feature_map_size = 14; // 3D zyx (DHW) size 24 | } 25 | 26 | message TargetAssigner { 27 | repeated ClassSetting class_settings = 1; 28 | float sample_positive_fraction = 2; 29 | uint32 sample_size = 3; 30 | bool assign_per_class = 4; 31 | repeated int64 nms_pre_max_sizes = 5; // this will override setting in ClassSettings if provide. 32 | repeated int64 nms_post_max_sizes = 6; // this will override setting in ClassSettings if provide. 33 | repeated int64 nms_score_thresholds = 7; // this will override setting in ClassSettings if provide. 34 | repeated int64 nms_iou_thresholds = 8; // this will override setting in ClassSettings if provide. 35 | } -------------------------------------------------------------------------------- /second/protos/train.proto: -------------------------------------------------------------------------------- 1 | syntax = "proto3"; 2 | 3 | package second.protos; 4 | 5 | import "second/protos/optimizer.proto"; 6 | import "second/protos/preprocess.proto"; 7 | 8 | message TrainConfig{ 9 | Optimizer optimizer = 1; 10 | uint32 steps = 2; 11 | uint32 steps_per_eval = 3; 12 | uint32 save_checkpoints_secs = 4; 13 | uint32 save_summary_steps = 5; 14 | bool enable_mixed_precision = 6; 15 | float loss_scale_factor = 7; 16 | bool clear_metrics_every_epoch = 8; 17 | } -------------------------------------------------------------------------------- /second/protos/train_pb2.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # Generated by the protocol buffer compiler. DO NOT EDIT! 3 | # source: second/protos/train.proto 4 | 5 | import sys 6 | _b=sys.version_info[0]<3 and (lambda x:x) or (lambda x:x.encode('latin1')) 7 | from google.protobuf import descriptor as _descriptor 8 | from google.protobuf import message as _message 9 | from google.protobuf import reflection as _reflection 10 | from google.protobuf import symbol_database as _symbol_database 11 | # @@protoc_insertion_point(imports) 12 | 13 | _sym_db = _symbol_database.Default() 14 | 15 | 16 | from second.protos import optimizer_pb2 as second_dot_protos_dot_optimizer__pb2 17 | from second.protos import preprocess_pb2 as second_dot_protos_dot_preprocess__pb2 18 | 19 | 20 | DESCRIPTOR = _descriptor.FileDescriptor( 21 | name='second/protos/train.proto', 22 | package='second.protos', 23 | syntax='proto3', 24 | serialized_options=None, 25 | serialized_pb=_b('\n\x19second/protos/train.proto\x12\rsecond.protos\x1a\x1dsecond/protos/optimizer.proto\x1a\x1esecond/protos/preprocess.proto\"\xfa\x01\n\x0bTrainConfig\x12+\n\toptimizer\x18\x01 \x01(\x0b\x32\x18.second.protos.Optimizer\x12\r\n\x05steps\x18\x02 \x01(\r\x12\x16\n\x0esteps_per_eval\x18\x03 \x01(\r\x12\x1d\n\x15save_checkpoints_secs\x18\x04 \x01(\r\x12\x1a\n\x12save_summary_steps\x18\x05 \x01(\r\x12\x1e\n\x16\x65nable_mixed_precision\x18\x06 \x01(\x08\x12\x19\n\x11loss_scale_factor\x18\x07 \x01(\x02\x12!\n\x19\x63lear_metrics_every_epoch\x18\x08 \x01(\x08\x62\x06proto3') 26 | , 27 | dependencies=[second_dot_protos_dot_optimizer__pb2.DESCRIPTOR,second_dot_protos_dot_preprocess__pb2.DESCRIPTOR,]) 28 | 29 | 30 | 31 | 32 | _TRAINCONFIG = _descriptor.Descriptor( 33 | name='TrainConfig', 34 | full_name='second.protos.TrainConfig', 35 | filename=None, 36 | file=DESCRIPTOR, 37 | containing_type=None, 38 | fields=[ 39 | _descriptor.FieldDescriptor( 40 | name='optimizer', full_name='second.protos.TrainConfig.optimizer', index=0, 41 | number=1, type=11, cpp_type=10, label=1, 42 | has_default_value=False, default_value=None, 43 | message_type=None, enum_type=None, containing_type=None, 44 | is_extension=False, extension_scope=None, 45 | serialized_options=None, file=DESCRIPTOR), 46 | _descriptor.FieldDescriptor( 47 | name='steps', full_name='second.protos.TrainConfig.steps', index=1, 48 | number=2, type=13, cpp_type=3, label=1, 49 | has_default_value=False, default_value=0, 50 | message_type=None, enum_type=None, containing_type=None, 51 | is_extension=False, extension_scope=None, 52 | serialized_options=None, file=DESCRIPTOR), 53 | _descriptor.FieldDescriptor( 54 | name='steps_per_eval', full_name='second.protos.TrainConfig.steps_per_eval', index=2, 55 | number=3, type=13, cpp_type=3, label=1, 56 | has_default_value=False, default_value=0, 57 | message_type=None, enum_type=None, containing_type=None, 58 | is_extension=False, extension_scope=None, 59 | serialized_options=None, file=DESCRIPTOR), 60 | _descriptor.FieldDescriptor( 61 | name='save_checkpoints_secs', full_name='second.protos.TrainConfig.save_checkpoints_secs', index=3, 62 | number=4, type=13, cpp_type=3, label=1, 63 | has_default_value=False, default_value=0, 64 | message_type=None, enum_type=None, containing_type=None, 65 | is_extension=False, extension_scope=None, 66 | serialized_options=None, file=DESCRIPTOR), 67 | _descriptor.FieldDescriptor( 68 | name='save_summary_steps', full_name='second.protos.TrainConfig.save_summary_steps', index=4, 69 | number=5, type=13, cpp_type=3, label=1, 70 | has_default_value=False, default_value=0, 71 | message_type=None, enum_type=None, containing_type=None, 72 | is_extension=False, extension_scope=None, 73 | serialized_options=None, file=DESCRIPTOR), 74 | _descriptor.FieldDescriptor( 75 | name='enable_mixed_precision', full_name='second.protos.TrainConfig.enable_mixed_precision', index=5, 76 | number=6, type=8, cpp_type=7, label=1, 77 | has_default_value=False, default_value=False, 78 | message_type=None, enum_type=None, containing_type=None, 79 | is_extension=False, extension_scope=None, 80 | serialized_options=None, file=DESCRIPTOR), 81 | _descriptor.FieldDescriptor( 82 | name='loss_scale_factor', full_name='second.protos.TrainConfig.loss_scale_factor', index=6, 83 | number=7, type=2, cpp_type=6, label=1, 84 | has_default_value=False, default_value=float(0), 85 | message_type=None, enum_type=None, containing_type=None, 86 | is_extension=False, extension_scope=None, 87 | serialized_options=None, file=DESCRIPTOR), 88 | _descriptor.FieldDescriptor( 89 | name='clear_metrics_every_epoch', full_name='second.protos.TrainConfig.clear_metrics_every_epoch', index=7, 90 | number=8, type=8, cpp_type=7, label=1, 91 | has_default_value=False, default_value=False, 92 | message_type=None, enum_type=None, containing_type=None, 93 | is_extension=False, extension_scope=None, 94 | serialized_options=None, file=DESCRIPTOR), 95 | ], 96 | extensions=[ 97 | ], 98 | nested_types=[], 99 | enum_types=[ 100 | ], 101 | serialized_options=None, 102 | is_extendable=False, 103 | syntax='proto3', 104 | extension_ranges=[], 105 | oneofs=[ 106 | ], 107 | serialized_start=108, 108 | serialized_end=358, 109 | ) 110 | 111 | _TRAINCONFIG.fields_by_name['optimizer'].message_type = second_dot_protos_dot_optimizer__pb2._OPTIMIZER 112 | DESCRIPTOR.message_types_by_name['TrainConfig'] = _TRAINCONFIG 113 | _sym_db.RegisterFileDescriptor(DESCRIPTOR) 114 | 115 | TrainConfig = _reflection.GeneratedProtocolMessageType('TrainConfig', (_message.Message,), dict( 116 | DESCRIPTOR = _TRAINCONFIG, 117 | __module__ = 'second.protos.train_pb2' 118 | # @@protoc_insertion_point(class_scope:second.protos.TrainConfig) 119 | )) 120 | _sym_db.RegisterMessage(TrainConfig) 121 | 122 | 123 | # @@protoc_insertion_point(module_scope) 124 | -------------------------------------------------------------------------------- /second/protos/voxel_generator.proto: -------------------------------------------------------------------------------- 1 | syntax = "proto3"; 2 | 3 | package second.protos; 4 | 5 | message VoxelGenerator{ 6 | repeated float voxel_size = 1; 7 | repeated float point_cloud_range = 2; 8 | uint32 max_number_of_points_per_voxel = 3; 9 | bool full_empty_part_with_mean = 4; 10 | bool block_filtering = 5; 11 | int64 block_factor = 6; 12 | int64 block_size = 7; 13 | float height_threshold = 8; 14 | } 15 | -------------------------------------------------------------------------------- /second/protos/voxel_generator_pb2.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # Generated by the protocol buffer compiler. DO NOT EDIT! 3 | # source: second/protos/voxel_generator.proto 4 | 5 | import sys 6 | _b=sys.version_info[0]<3 and (lambda x:x) or (lambda x:x.encode('latin1')) 7 | from google.protobuf import descriptor as _descriptor 8 | from google.protobuf import message as _message 9 | from google.protobuf import reflection as _reflection 10 | from google.protobuf import symbol_database as _symbol_database 11 | # @@protoc_insertion_point(imports) 12 | 13 | _sym_db = _symbol_database.Default() 14 | 15 | 16 | 17 | 18 | DESCRIPTOR = _descriptor.FileDescriptor( 19 | name='second/protos/voxel_generator.proto', 20 | package='second.protos', 21 | syntax='proto3', 22 | serialized_options=None, 23 | serialized_pb=_b('\n#second/protos/voxel_generator.proto\x12\rsecond.protos\"\xe7\x01\n\x0eVoxelGenerator\x12\x12\n\nvoxel_size\x18\x01 \x03(\x02\x12\x19\n\x11point_cloud_range\x18\x02 \x03(\x02\x12&\n\x1emax_number_of_points_per_voxel\x18\x03 \x01(\r\x12!\n\x19\x66ull_empty_part_with_mean\x18\x04 \x01(\x08\x12\x17\n\x0f\x62lock_filtering\x18\x05 \x01(\x08\x12\x14\n\x0c\x62lock_factor\x18\x06 \x01(\x03\x12\x12\n\nblock_size\x18\x07 \x01(\x03\x12\x18\n\x10height_threshold\x18\x08 \x01(\x02\x62\x06proto3') 24 | ) 25 | 26 | 27 | 28 | 29 | _VOXELGENERATOR = _descriptor.Descriptor( 30 | name='VoxelGenerator', 31 | full_name='second.protos.VoxelGenerator', 32 | filename=None, 33 | file=DESCRIPTOR, 34 | containing_type=None, 35 | fields=[ 36 | _descriptor.FieldDescriptor( 37 | name='voxel_size', full_name='second.protos.VoxelGenerator.voxel_size', index=0, 38 | number=1, type=2, cpp_type=6, label=3, 39 | has_default_value=False, default_value=[], 40 | message_type=None, enum_type=None, containing_type=None, 41 | is_extension=False, extension_scope=None, 42 | serialized_options=None, file=DESCRIPTOR), 43 | _descriptor.FieldDescriptor( 44 | name='point_cloud_range', full_name='second.protos.VoxelGenerator.point_cloud_range', index=1, 45 | number=2, type=2, cpp_type=6, label=3, 46 | has_default_value=False, default_value=[], 47 | message_type=None, enum_type=None, containing_type=None, 48 | is_extension=False, extension_scope=None, 49 | serialized_options=None, file=DESCRIPTOR), 50 | _descriptor.FieldDescriptor( 51 | name='max_number_of_points_per_voxel', full_name='second.protos.VoxelGenerator.max_number_of_points_per_voxel', index=2, 52 | number=3, type=13, cpp_type=3, label=1, 53 | has_default_value=False, default_value=0, 54 | message_type=None, enum_type=None, containing_type=None, 55 | is_extension=False, extension_scope=None, 56 | serialized_options=None, file=DESCRIPTOR), 57 | _descriptor.FieldDescriptor( 58 | name='full_empty_part_with_mean', full_name='second.protos.VoxelGenerator.full_empty_part_with_mean', index=3, 59 | number=4, type=8, cpp_type=7, label=1, 60 | has_default_value=False, default_value=False, 61 | message_type=None, enum_type=None, containing_type=None, 62 | is_extension=False, extension_scope=None, 63 | serialized_options=None, file=DESCRIPTOR), 64 | _descriptor.FieldDescriptor( 65 | name='block_filtering', full_name='second.protos.VoxelGenerator.block_filtering', index=4, 66 | number=5, type=8, cpp_type=7, label=1, 67 | has_default_value=False, default_value=False, 68 | message_type=None, enum_type=None, containing_type=None, 69 | is_extension=False, extension_scope=None, 70 | serialized_options=None, file=DESCRIPTOR), 71 | _descriptor.FieldDescriptor( 72 | name='block_factor', full_name='second.protos.VoxelGenerator.block_factor', index=5, 73 | number=6, type=3, cpp_type=2, label=1, 74 | has_default_value=False, default_value=0, 75 | message_type=None, enum_type=None, containing_type=None, 76 | is_extension=False, extension_scope=None, 77 | serialized_options=None, file=DESCRIPTOR), 78 | _descriptor.FieldDescriptor( 79 | name='block_size', full_name='second.protos.VoxelGenerator.block_size', index=6, 80 | number=7, type=3, cpp_type=2, label=1, 81 | has_default_value=False, default_value=0, 82 | message_type=None, enum_type=None, containing_type=None, 83 | is_extension=False, extension_scope=None, 84 | serialized_options=None, file=DESCRIPTOR), 85 | _descriptor.FieldDescriptor( 86 | name='height_threshold', full_name='second.protos.VoxelGenerator.height_threshold', index=7, 87 | number=8, type=2, cpp_type=6, label=1, 88 | has_default_value=False, default_value=float(0), 89 | message_type=None, enum_type=None, containing_type=None, 90 | is_extension=False, extension_scope=None, 91 | serialized_options=None, file=DESCRIPTOR), 92 | ], 93 | extensions=[ 94 | ], 95 | nested_types=[], 96 | enum_types=[ 97 | ], 98 | serialized_options=None, 99 | is_extendable=False, 100 | syntax='proto3', 101 | extension_ranges=[], 102 | oneofs=[ 103 | ], 104 | serialized_start=55, 105 | serialized_end=286, 106 | ) 107 | 108 | DESCRIPTOR.message_types_by_name['VoxelGenerator'] = _VOXELGENERATOR 109 | _sym_db.RegisterFileDescriptor(DESCRIPTOR) 110 | 111 | VoxelGenerator = _reflection.GeneratedProtocolMessageType('VoxelGenerator', (_message.Message,), dict( 112 | DESCRIPTOR = _VOXELGENERATOR, 113 | __module__ = 'second.protos.voxel_generator_pb2' 114 | # @@protoc_insertion_point(class_scope:second.protos.VoxelGenerator) 115 | )) 116 | _sym_db.RegisterMessage(VoxelGenerator) 117 | 118 | 119 | # @@protoc_insertion_point(module_scope) 120 | -------------------------------------------------------------------------------- /second/pytorch/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/peiyunh/wysiwyg/bb0f6ecbf3f15082fe4ec14a701ccbb14198babd/second/pytorch/__init__.py -------------------------------------------------------------------------------- /second/pytorch/builder/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/peiyunh/wysiwyg/bb0f6ecbf3f15082fe4ec14a701ccbb14198babd/second/pytorch/builder/__init__.py -------------------------------------------------------------------------------- /second/pytorch/builder/box_coder_builder.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from second.protos import box_coder_pb2 4 | from second.pytorch.core.box_coders import (BevBoxCoderTorch, 5 | GroundBox3dCoderTorch) 6 | 7 | 8 | def build(box_coder_config): 9 | """Create optimizer based on config. 10 | 11 | Args: 12 | optimizer_config: A Optimizer proto message. 13 | 14 | Returns: 15 | An optimizer and a list of variables for summary. 16 | 17 | Raises: 18 | ValueError: when using an unsupported input data type. 19 | """ 20 | box_coder_type = box_coder_config.WhichOneof('box_coder') 21 | if box_coder_type == 'ground_box3d_coder': 22 | cfg = box_coder_config.ground_box3d_coder 23 | return GroundBox3dCoderTorch(cfg.linear_dim, cfg.encode_angle_vector) 24 | elif box_coder_type == 'bev_box_coder': 25 | cfg = box_coder_config.bev_box_coder 26 | return BevBoxCoderTorch(cfg.linear_dim, cfg.encode_angle_vector, cfg.z_fixed, cfg.h_fixed) 27 | else: 28 | raise ValueError("unknown box_coder type") 29 | -------------------------------------------------------------------------------- /second/pytorch/builder/input_reader_builder.py: -------------------------------------------------------------------------------- 1 | # Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Input reader builder. 16 | 17 | Creates data sources for DetectionModels from an InputReader config. See 18 | input_reader.proto for options. 19 | 20 | Note: If users wishes to also use their own InputReaders with the Object 21 | Detection configuration framework, they should define their own builder function 22 | that wraps the build function. 23 | """ 24 | 25 | from torch.utils.data import Dataset 26 | 27 | from second.builder import dataset_builder 28 | from second.protos import input_reader_pb2 29 | 30 | 31 | class DatasetWrapper(Dataset): 32 | """ convert our dataset to Dataset class in pytorch. 33 | """ 34 | 35 | def __init__(self, dataset): 36 | self._dataset = dataset 37 | 38 | def __len__(self): 39 | return len(self._dataset) 40 | 41 | def __getitem__(self, idx): 42 | return self._dataset[idx] 43 | 44 | @property 45 | def dataset(self): 46 | return self._dataset 47 | 48 | 49 | def build(input_reader_config, 50 | model_config, 51 | training, 52 | voxel_generator, 53 | target_assigner=None, 54 | multi_gpu=False) -> DatasetWrapper: 55 | """Builds a tensor dictionary based on the InputReader config. 56 | 57 | Args: 58 | input_reader_config: A input_reader_pb2.InputReader object. 59 | 60 | Returns: 61 | A tensor dict based on the input_reader_config. 62 | 63 | Raises: 64 | ValueError: On invalid input reader proto. 65 | ValueError: If no input paths are specified. 66 | """ 67 | if not isinstance(input_reader_config, input_reader_pb2.InputReader): 68 | raise ValueError('input_reader_config not of type ' 69 | 'input_reader_pb2.InputReader.') 70 | dataset = dataset_builder.build( 71 | input_reader_config, 72 | model_config, 73 | training, 74 | voxel_generator, 75 | target_assigner, 76 | multi_gpu=multi_gpu) 77 | dataset = DatasetWrapper(dataset) 78 | return dataset 79 | -------------------------------------------------------------------------------- /second/pytorch/builder/losses_builder.py: -------------------------------------------------------------------------------- 1 | # Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | 16 | """A function to build localization and classification losses from config.""" 17 | 18 | from second.pytorch.core import losses 19 | from second.pytorch.core.ghm_loss import GHMCLoss, GHMRLoss 20 | from second.protos import losses_pb2 21 | 22 | 23 | def build(loss_config): 24 | """Build losses based on the config. 25 | 26 | Builds classification, localization losses and optionally a hard example miner 27 | based on the config. 28 | 29 | Args: 30 | loss_config: A losses_pb2.Loss object. 31 | 32 | Returns: 33 | classification_loss: Classification loss object. 34 | localization_loss: Localization loss object. 35 | classification_weight: Classification loss weight. 36 | localization_weight: Localization loss weight. 37 | hard_example_miner: Hard example miner object. 38 | 39 | Raises: 40 | ValueError: If hard_example_miner is used with sigmoid_focal_loss. 41 | """ 42 | classification_loss = _build_classification_loss( 43 | loss_config.classification_loss) 44 | localization_loss = _build_localization_loss( 45 | loss_config.localization_loss) 46 | classification_weight = loss_config.classification_weight 47 | localization_weight = loss_config.localization_weight 48 | hard_example_miner = None 49 | if loss_config.HasField('hard_example_miner'): 50 | raise ValueError('Pytorch don\'t support HardExampleMiner') 51 | return (classification_loss, localization_loss, 52 | classification_weight, 53 | localization_weight, hard_example_miner) 54 | 55 | def build_faster_rcnn_classification_loss(loss_config): 56 | """Builds a classification loss for Faster RCNN based on the loss config. 57 | 58 | Args: 59 | loss_config: A losses_pb2.ClassificationLoss object. 60 | 61 | Returns: 62 | Loss based on the config. 63 | 64 | Raises: 65 | ValueError: On invalid loss_config. 66 | """ 67 | if not isinstance(loss_config, losses_pb2.ClassificationLoss): 68 | raise ValueError('loss_config not of type losses_pb2.ClassificationLoss.') 69 | 70 | loss_type = loss_config.WhichOneof('classification_loss') 71 | 72 | if loss_type == 'weighted_sigmoid': 73 | return losses.WeightedSigmoidClassificationLoss() 74 | if loss_type == 'weighted_softmax': 75 | config = loss_config.weighted_softmax 76 | return losses.WeightedSoftmaxClassificationLoss( 77 | logit_scale=config.logit_scale) 78 | 79 | # By default, Faster RCNN second stage classifier uses Softmax loss 80 | # with anchor-wise outputs. 81 | config = loss_config.weighted_softmax 82 | return losses.WeightedSoftmaxClassificationLoss( 83 | logit_scale=config.logit_scale) 84 | 85 | 86 | def _build_localization_loss(loss_config): 87 | """Builds a localization loss based on the loss config. 88 | 89 | Args: 90 | loss_config: A losses_pb2.LocalizationLoss object. 91 | 92 | Returns: 93 | Loss based on the config. 94 | 95 | Raises: 96 | ValueError: On invalid loss_config. 97 | """ 98 | if not isinstance(loss_config, losses_pb2.LocalizationLoss): 99 | raise ValueError('loss_config not of type losses_pb2.LocalizationLoss.') 100 | 101 | loss_type = loss_config.WhichOneof('localization_loss') 102 | 103 | if loss_type == 'weighted_l2': 104 | config = loss_config.weighted_l2 105 | if len(config.code_weight) == 0: 106 | code_weight = None 107 | else: 108 | code_weight = config.code_weight 109 | return losses.WeightedL2LocalizationLoss(code_weight) 110 | 111 | if loss_type == 'weighted_smooth_l1': 112 | config = loss_config.weighted_smooth_l1 113 | if len(config.code_weight) == 0: 114 | code_weight = None 115 | else: 116 | code_weight = config.code_weight 117 | return losses.WeightedSmoothL1LocalizationLoss(config.sigma, code_weight) 118 | if loss_type == 'weighted_ghm': 119 | config = loss_config.weighted_ghm 120 | if len(config.code_weight) == 0: 121 | code_weight = None 122 | else: 123 | code_weight = config.code_weight 124 | return GHMRLoss(config.mu, config.bins, config.momentum, code_weight) 125 | 126 | raise ValueError('Empty loss config.') 127 | 128 | 129 | def _build_classification_loss(loss_config): 130 | """Builds a classification loss based on the loss config. 131 | 132 | Args: 133 | loss_config: A losses_pb2.ClassificationLoss object. 134 | 135 | Returns: 136 | Loss based on the config. 137 | 138 | Raises: 139 | ValueError: On invalid loss_config. 140 | """ 141 | if not isinstance(loss_config, losses_pb2.ClassificationLoss): 142 | raise ValueError('loss_config not of type losses_pb2.ClassificationLoss.') 143 | 144 | loss_type = loss_config.WhichOneof('classification_loss') 145 | 146 | if loss_type == 'weighted_sigmoid': 147 | return losses.WeightedSigmoidClassificationLoss() 148 | 149 | if loss_type == 'weighted_sigmoid_focal': 150 | config = loss_config.weighted_sigmoid_focal 151 | # alpha = None 152 | # if config.HasField('alpha'): 153 | # alpha = config.alpha 154 | if config.alpha > 0: 155 | alpha = config.alpha 156 | else: 157 | alpha = None 158 | return losses.SigmoidFocalClassificationLoss( 159 | gamma=config.gamma, 160 | alpha=alpha) 161 | if loss_type == 'weighted_softmax_focal': 162 | config = loss_config.weighted_softmax_focal 163 | # alpha = None 164 | # if config.HasField('alpha'): 165 | # alpha = config.alpha 166 | if config.alpha > 0: 167 | alpha = config.alpha 168 | else: 169 | alpha = None 170 | return losses.SoftmaxFocalClassificationLoss( 171 | gamma=config.gamma, 172 | alpha=alpha) 173 | if loss_type == 'weighted_ghm': 174 | config = loss_config.weighted_ghm 175 | return GHMCLoss( 176 | bins=config.bins, 177 | momentum=config.momentum) 178 | 179 | if loss_type == 'weighted_softmax': 180 | config = loss_config.weighted_softmax 181 | return losses.WeightedSoftmaxClassificationLoss( 182 | logit_scale=config.logit_scale) 183 | 184 | if loss_type == 'bootstrapped_sigmoid': 185 | config = loss_config.bootstrapped_sigmoid 186 | return losses.BootstrappedSigmoidClassificationLoss( 187 | alpha=config.alpha, 188 | bootstrap_type=('hard' if config.hard_bootstrap else 'soft')) 189 | 190 | raise ValueError('Empty loss config.') 191 | -------------------------------------------------------------------------------- /second/pytorch/builder/lr_scheduler_builder.py: -------------------------------------------------------------------------------- 1 | # Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | 16 | """Functions to build DetectionModel training optimizers.""" 17 | 18 | from torchplus.train import learning_schedules_fastai as lsf 19 | import torch 20 | 21 | def build(optimizer_config, optimizer, total_step): 22 | """Create lr scheduler based on config. note that 23 | lr_scheduler must accept a optimizer that has been restored. 24 | 25 | Args: 26 | optimizer_config: A Optimizer proto message. 27 | 28 | Returns: 29 | An optimizer and a list of variables for summary. 30 | 31 | Raises: 32 | ValueError: when using an unsupported input data type. 33 | """ 34 | optimizer_type = optimizer_config.WhichOneof('optimizer') 35 | 36 | if optimizer_type == 'rms_prop_optimizer': 37 | config = optimizer_config.rms_prop_optimizer 38 | lr_scheduler = _create_learning_rate_scheduler( 39 | config.learning_rate, optimizer, total_step=total_step) 40 | 41 | if optimizer_type == 'momentum_optimizer': 42 | config = optimizer_config.momentum_optimizer 43 | lr_scheduler = _create_learning_rate_scheduler( 44 | config.learning_rate, optimizer, total_step=total_step) 45 | 46 | if optimizer_type == 'adam_optimizer': 47 | config = optimizer_config.adam_optimizer 48 | lr_scheduler = _create_learning_rate_scheduler( 49 | config.learning_rate, optimizer, total_step=total_step) 50 | 51 | return lr_scheduler 52 | 53 | def _create_learning_rate_scheduler(learning_rate_config, optimizer, total_step): 54 | """Create optimizer learning rate scheduler based on config. 55 | 56 | Args: 57 | learning_rate_config: A LearningRate proto message. 58 | 59 | Returns: 60 | A learning rate. 61 | 62 | Raises: 63 | ValueError: when using an unsupported input data type. 64 | """ 65 | lr_scheduler = None 66 | learning_rate_type = learning_rate_config.WhichOneof('learning_rate') 67 | if learning_rate_type == 'multi_phase': 68 | config = learning_rate_config.multi_phase 69 | lr_phases = [] 70 | mom_phases = [] 71 | for phase_cfg in config.phases: 72 | lr_phases.append((phase_cfg.start, phase_cfg.lambda_func)) 73 | mom_phases.append((phase_cfg.start, phase_cfg.momentum_lambda_func)) 74 | lr_scheduler = lsf.LRSchedulerStep( 75 | optimizer,total_step, lr_phases, mom_phases) 76 | 77 | if learning_rate_type == 'one_cycle': 78 | config = learning_rate_config.one_cycle 79 | lr_scheduler = lsf.OneCycle( 80 | optimizer, total_step, config.lr_max, list(config.moms), config.div_factor, config.pct_start) 81 | if learning_rate_type == 'exponential_decay': 82 | config = learning_rate_config.exponential_decay 83 | lr_scheduler = lsf.ExponentialDecay( 84 | optimizer, total_step, config.initial_learning_rate, config.decay_length, config.decay_factor, config.staircase) 85 | if learning_rate_type == 'manual_stepping': 86 | config = learning_rate_config.manual_stepping 87 | lr_scheduler = lsf.ManualStepping( 88 | optimizer, total_step, list(config.boundaries), list(config.rates)) 89 | 90 | if lr_scheduler is None: 91 | raise ValueError('Learning_rate %s not supported.' % learning_rate_type) 92 | 93 | return lr_scheduler -------------------------------------------------------------------------------- /second/pytorch/builder/optimizer_builder.py: -------------------------------------------------------------------------------- 1 | # Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Functions to build DetectionModel training optimizers.""" 16 | 17 | from torchplus.train import learning_schedules 18 | from torchplus.train import optim 19 | import torch 20 | from torch import nn 21 | from torchplus.train.fastai_optim import OptimWrapper, FastAIMixedOptim 22 | from functools import partial 23 | 24 | def children(m: nn.Module): 25 | "Get children of `m`." 26 | return list(m.children()) 27 | 28 | 29 | def num_children(m: nn.Module) -> int: 30 | "Get number of children modules in `m`." 31 | return len(children(m)) 32 | 33 | flatten_model = lambda m: sum(map(flatten_model,m.children()),[]) if num_children(m) else [m] 34 | 35 | get_layer_groups = lambda m: [nn.Sequential(*flatten_model(m))] 36 | 37 | 38 | def build(optimizer_config, net, name=None, mixed=False, loss_scale=512.0): 39 | """Create optimizer based on config. 40 | 41 | Args: 42 | optimizer_config: A Optimizer proto message. 43 | 44 | Returns: 45 | An optimizer and a list of variables for summary. 46 | 47 | Raises: 48 | ValueError: when using an unsupported input data type. 49 | """ 50 | optimizer_type = optimizer_config.WhichOneof('optimizer') 51 | optimizer = None 52 | 53 | if optimizer_type == 'rms_prop_optimizer': 54 | config = optimizer_config.rms_prop_optimizer 55 | optimizer_func = partial( 56 | torch.optim.RMSprop, 57 | alpha=config.decay, 58 | momentum=config.momentum_optimizer_value, 59 | eps=config.epsilon) 60 | 61 | if optimizer_type == 'momentum_optimizer': 62 | config = optimizer_config.momentum_optimizer 63 | optimizer_func = partial( 64 | torch.optim.SGD, 65 | momentum=config.momentum_optimizer_value, 66 | eps=config.epsilon) 67 | 68 | if optimizer_type == 'adam_optimizer': 69 | config = optimizer_config.adam_optimizer 70 | if optimizer_config.fixed_weight_decay: 71 | optimizer_func = partial( 72 | torch.optim.Adam, betas=(0.9, 0.99), amsgrad=config.amsgrad) 73 | else: 74 | # regular adam 75 | optimizer_func = partial( 76 | torch.optim.Adam, amsgrad=config.amsgrad) 77 | 78 | 79 | 80 | # optimizer = OptimWrapper(optimizer, true_wd=optimizer_config.fixed_weight_decay, wd=config.weight_decay) 81 | optimizer = OptimWrapper.create( 82 | optimizer_func, 83 | 3e-3, 84 | get_layer_groups(net), 85 | wd=config.weight_decay, 86 | true_wd=optimizer_config.fixed_weight_decay, 87 | bn_wd=True) 88 | print(hasattr(optimizer, "_amp_stash"), '_amp_stash') 89 | if optimizer is None: 90 | raise ValueError('Optimizer %s not supported.' % optimizer_type) 91 | 92 | if optimizer_config.use_moving_average: 93 | raise ValueError('torch don\'t support moving average') 94 | if name is None: 95 | # assign a name to optimizer for checkpoint system 96 | optimizer.name = optimizer_type 97 | else: 98 | optimizer.name = name 99 | return optimizer 100 | -------------------------------------------------------------------------------- /second/pytorch/core/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/peiyunh/wysiwyg/bb0f6ecbf3f15082fe4ec14a701ccbb14198babd/second/pytorch/core/__init__.py -------------------------------------------------------------------------------- /second/pytorch/core/box_coders.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from second.core.box_coders import BevBoxCoder, GroundBox3dCoder 4 | from second.pytorch.core import box_torch_ops 5 | 6 | 7 | class GroundBox3dCoderTorch(GroundBox3dCoder): 8 | def encode_torch(self, boxes, anchors): 9 | return box_torch_ops.second_box_encode(boxes, anchors, self.vec_encode, 10 | self.linear_dim) 11 | 12 | def decode_torch(self, boxes, anchors): 13 | return box_torch_ops.second_box_decode(boxes, anchors, self.vec_encode, 14 | self.linear_dim) 15 | 16 | 17 | class BevBoxCoderTorch(BevBoxCoder): 18 | def encode_torch(self, boxes, anchors): 19 | anchors = anchors[..., [0, 1, 3, 4, 6]] 20 | boxes = boxes[..., [0, 1, 3, 4, 6]] 21 | return box_torch_ops.bev_box_encode(boxes, anchors, self.vec_encode, 22 | self.linear_dim) 23 | 24 | def decode_torch(self, encodings, anchors): 25 | anchors = anchors[..., [0, 1, 3, 4, 6]] 26 | ret = box_torch_ops.bev_box_decode(encodings, anchors, self.vec_encode, 27 | self.linear_dim) 28 | z_fixed = torch.full([*ret.shape[:-1], 1], 29 | self.z_fixed, 30 | dtype=ret.dtype, 31 | device=ret.device) 32 | h_fixed = torch.full([*ret.shape[:-1], 1], 33 | self.h_fixed, 34 | dtype=ret.dtype, 35 | device=ret.device) 36 | return torch.cat( 37 | [ret[..., :2], z_fixed, ret[..., 2:4], h_fixed, ret[..., 4:]], 38 | dim=-1) 39 | -------------------------------------------------------------------------------- /second/pytorch/core/ghm_loss.py: -------------------------------------------------------------------------------- 1 | ##################### 2 | # THIS LOSS IS NOT WORKING!!!! 3 | ##################### 4 | 5 | """ 6 | The implementation of GHM-C and GHM-R losses. 7 | Details can be found in the paper `Gradient Harmonized Single-stage Detector`: 8 | https://arxiv.org/abs/1811.05181 9 | Copyright (c) 2018 Multimedia Laboratory, CUHK. 10 | Licensed under the MIT License (see LICENSE for details) 11 | Written by Buyu Li 12 | """ 13 | 14 | from second.pytorch.core.losses import Loss, _sigmoid_cross_entropy_with_logits 15 | import torch 16 | 17 | class GHMCLoss(Loss): 18 | def __init__(self, bins=10, momentum=0): 19 | self.bins = bins 20 | self.momentum = momentum 21 | self.edges = [float(x) / bins for x in range(bins+1)] 22 | self.edges[-1] += 1e-6 23 | if momentum > 0: 24 | self.acc_sum = [0.0 for _ in range(bins)] 25 | self.count = 50 26 | 27 | def _compute_loss(self, 28 | prediction_tensor, 29 | target_tensor, 30 | weights, 31 | class_indices=None): 32 | """ Args: 33 | input [batch_num, class_num]: 34 | The direct prediction of classification fc layer. 35 | target [batch_num, class_num]: 36 | Binary target (0 or 1) for each sample each class. The value is -1 37 | when the sample is ignored. 38 | """ 39 | input = prediction_tensor 40 | target = target_tensor 41 | batch_size = prediction_tensor.shape[0] 42 | num_anchors = prediction_tensor.shape[1] 43 | num_class = prediction_tensor.shape[2] 44 | edges = self.edges 45 | weights_ghm = torch.zeros_like(input).view(-1, num_class) 46 | per_entry_cross_ent = (_sigmoid_cross_entropy_with_logits( 47 | labels=target_tensor, logits=prediction_tensor)) 48 | # gradient length 49 | g = torch.abs(input.sigmoid().detach() - target).view(-1, num_class) 50 | valid = weights.view(-1, 1).expand(-1, num_class) > 0 51 | num_examples = max(valid.float().sum().item(), 1.0) 52 | num_valid_bins = 0 # n valid bins 53 | self.count -= 1 54 | num_bins = [] 55 | for i in range(self.bins): 56 | inds = (g >= edges[i]) & (g < edges[i+1]) & valid 57 | num_in_bin = inds.sum().item() 58 | num_bins.append(num_in_bin) 59 | if num_in_bin > 0: 60 | if self.momentum > 0: 61 | self.acc_sum[i] = self.momentum * self.acc_sum[i] \ 62 | + (1 - self.momentum) * num_in_bin 63 | weights_ghm[inds] = num_examples / self.acc_sum[i] 64 | else: 65 | weights_ghm[inds] = num_examples / num_in_bin 66 | num_valid_bins += 1 67 | if self.count <= 0: 68 | print("GHMC loss bins:", num_bins) 69 | self.count = 50 70 | if num_valid_bins > 0: 71 | weights_ghm = weights_ghm / num_valid_bins 72 | return per_entry_cross_ent * weights_ghm.view(batch_size, num_anchors, num_class) / num_examples 73 | 74 | 75 | class GHMRLoss(Loss): 76 | def __init__(self, mu=0.02, bins=10, momentum=0, code_weights=None): 77 | self.mu = mu 78 | self.bins = bins 79 | self.edges = [float(x) / bins for x in range(bins+1)] 80 | self.edges[-1] = 1e3 81 | self.momentum = momentum 82 | if momentum > 0: 83 | self.acc_sum = [0.0 for _ in range(bins)] 84 | self._codewise = True 85 | 86 | def _compute_loss(self, 87 | prediction_tensor, 88 | target_tensor, 89 | weights): 90 | """ Args: 91 | input [batch_num, class_num]: 92 | The direct prediction of classification fc layer. 93 | target [batch_num, class_num]: 94 | Binary target (0 or 1) for each sample each class. The value is -1 95 | when the sample is ignored. 96 | """ 97 | # ASL1 loss 98 | diff = prediction_tensor - target_tensor 99 | loss = torch.sqrt(diff * diff + self.mu * self.mu) - self.mu 100 | batch_size = prediction_tensor.shape[0] 101 | num_anchors = prediction_tensor.shape[1] 102 | num_codes = prediction_tensor.shape[2] 103 | 104 | # gradient length 105 | g = torch.abs(diff / torch.sqrt(self.mu * self.mu + diff * diff)).detach().view(-1, num_codes) 106 | weights_ghm = torch.zeros_like(g) 107 | 108 | valid = weights.view(-1, 1).expand(-1, num_codes) > 0 109 | # print(g.shape, prediction_tensor.shape, valid.shape) 110 | num_examples = max(valid.float().sum().item() / num_codes, 1.0) 111 | num_valid_bins = 0 # n: valid bins 112 | for i in range(self.bins): 113 | inds = (g >= self.edges[i]) & (g < self.edges[i+1]) & valid 114 | num_in_bin = inds.sum().item() 115 | if num_in_bin > 0: 116 | num_valid_bins += 1 117 | if self.momentum > 0: 118 | self.acc_sum[i] = self.momentum * self.acc_sum[i] \ 119 | + (1 - self.momentum) * num_in_bin 120 | weights_ghm[inds] = num_examples / self.acc_sum[i] 121 | else: 122 | weights_ghm[inds] = num_examples / num_in_bin 123 | if num_valid_bins > 0: 124 | weights_ghm /= num_valid_bins 125 | weights_ghm = weights_ghm.view(batch_size, num_anchors, num_codes) 126 | loss = loss * weights_ghm / num_examples 127 | return loss 128 | -------------------------------------------------------------------------------- /second/pytorch/inference.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | 3 | import numpy as np 4 | import torch 5 | 6 | import torchplus 7 | from second.core import box_np_ops 8 | from second.core.inference import InferenceContext 9 | from second.builder import target_assigner_builder, voxel_builder 10 | from second.pytorch.builder import box_coder_builder, second_builder 11 | from second.pytorch.models.voxelnet import VoxelNet 12 | from second.pytorch.train import predict_to_kitti_label, example_convert_to_torch 13 | 14 | 15 | class TorchInferenceContext(InferenceContext): 16 | def __init__(self): 17 | super().__init__() 18 | self.net = None 19 | self.anchor_cache = None 20 | 21 | def _build(self): 22 | config = self.config 23 | input_cfg = config.eval_input_reader 24 | model_cfg = config.model.second 25 | train_cfg = config.train_config 26 | batch_size = 1 27 | voxel_generator = voxel_builder.build(model_cfg.voxel_generator) 28 | bv_range = voxel_generator.point_cloud_range[[0, 1, 3, 4]] 29 | grid_size = voxel_generator.grid_size 30 | self.voxel_generator = voxel_generator 31 | vfe_num_filters = list(model_cfg.voxel_feature_extractor.num_filters) 32 | 33 | box_coder = box_coder_builder.build(model_cfg.box_coder) 34 | target_assigner_cfg = model_cfg.target_assigner 35 | target_assigner = target_assigner_builder.build( 36 | target_assigner_cfg, bv_range, box_coder) 37 | self.target_assigner = target_assigner 38 | out_size_factor = model_cfg.rpn.layer_strides[0] / model_cfg.rpn.upsample_strides[0] 39 | out_size_factor *= model_cfg.middle_feature_extractor.downsample_factor 40 | out_size_factor = int(out_size_factor) 41 | self.net = second_builder.build(model_cfg, voxel_generator, 42 | target_assigner) 43 | self.net.cuda().eval() 44 | if train_cfg.enable_mixed_precision: 45 | self.net.half() 46 | self.net.metrics_to_float() 47 | self.net.convert_norm_to_float(self.net) 48 | feature_map_size = grid_size[:2] // out_size_factor 49 | feature_map_size = [*feature_map_size, 1][::-1] 50 | ret = target_assigner.generate_anchors(feature_map_size) 51 | anchors_dict = target_assigner.generate_anchors_dict(feature_map_size) 52 | anchors = ret["anchors"] 53 | anchors = anchors.reshape([-1, 7]) 54 | matched_thresholds = ret["matched_thresholds"] 55 | unmatched_thresholds = ret["unmatched_thresholds"] 56 | anchors_bv = box_np_ops.rbbox2d_to_near_bbox( 57 | anchors[:, [0, 1, 3, 4, 6]]) 58 | anchor_cache = { 59 | "anchors": anchors, 60 | "anchors_bv": anchors_bv, 61 | "matched_thresholds": matched_thresholds, 62 | "unmatched_thresholds": unmatched_thresholds, 63 | "anchors_dict": anchors_dict, 64 | } 65 | self.anchor_cache = anchor_cache 66 | 67 | def _restore(self, ckpt_path): 68 | ckpt_path = Path(ckpt_path) 69 | assert ckpt_path.suffix == ".tckpt" 70 | torchplus.train.restore(str(ckpt_path), self.net) 71 | 72 | def _inference(self, example): 73 | train_cfg = self.config.train_config 74 | input_cfg = self.config.eval_input_reader 75 | model_cfg = self.config.model.second 76 | example_torch = example_convert_to_torch(example) 77 | result_annos = predict_to_kitti_label( 78 | self.net, example_torch, list( 79 | self.target_assigner.classes), 80 | model_cfg.post_center_limit_range, model_cfg.lidar_input) 81 | return result_annos 82 | 83 | def _ctx(self): 84 | return None 85 | -------------------------------------------------------------------------------- /second/pytorch/models/__init__.py: -------------------------------------------------------------------------------- 1 | from . import net_multi_head -------------------------------------------------------------------------------- /second/pytorch/models/resnet.py: -------------------------------------------------------------------------------- 1 | import spconv 2 | from torch import nn 3 | from torch.nn import functional as F 4 | 5 | from torchplus.nn import Empty, GroupNorm, Sequential 6 | 7 | 8 | def conv3x3(in_planes, out_planes, stride=1, indice_key=None): 9 | """3x3 convolution with padding""" 10 | return spconv.SubMConv3d( 11 | in_planes, 12 | out_planes, 13 | kernel_size=3, 14 | stride=stride, 15 | padding=1, 16 | bias=False, 17 | indice_key=indice_key) 18 | 19 | 20 | def conv1x1(in_planes, out_planes, stride=1, indice_key=None): 21 | """1x1 convolution""" 22 | return spconv.SubMConv3d( 23 | in_planes, 24 | out_planes, 25 | kernel_size=1, 26 | stride=stride, 27 | padding=1, 28 | bias=False, 29 | indice_key=indice_key) 30 | 31 | 32 | class SparseBasicBlock(spconv.SparseModule): 33 | expansion = 1 34 | 35 | def __init__(self, 36 | inplanes, 37 | planes, 38 | stride=1, 39 | downsample=None, 40 | indice_key=None): 41 | super(SparseBasicBlock, self).__init__() 42 | self.conv1 = conv3x3(inplanes, planes, stride, indice_key=indice_key) 43 | self.bn1 = nn.BatchNorm1d(planes) 44 | self.relu = nn.ReLU() 45 | self.conv2 = conv3x3(planes, planes, indice_key=indice_key) 46 | self.bn2 = nn.BatchNorm1d(planes) 47 | self.downsample = downsample 48 | self.stride = stride 49 | 50 | def forward(self, x): 51 | identity = x 52 | 53 | out = self.conv1(x) 54 | out.features = self.bn1(out.features) 55 | out.features = self.relu(out.features) 56 | 57 | out = self.conv2(out) 58 | out.features = self.bn2(out.features) 59 | 60 | if self.downsample is not None: 61 | identity = self.downsample(x) 62 | 63 | out.features += identity.features 64 | out.features = self.relu(out.features) 65 | 66 | return out 67 | 68 | 69 | class SparseBottleneck(spconv.SparseModule): 70 | expansion = 4 71 | 72 | def __init__(self, 73 | inplanes, 74 | planes, 75 | stride=1, 76 | downsample=None, 77 | indice_key=None): 78 | super(SparseBottleneck, self).__init__() 79 | self.conv1 = conv1x1(inplanes, planes, indice_key=indice_key) 80 | self.bn1 = nn.BatchNorm1d(planes) 81 | self.conv2 = conv3x3(planes, planes, stride, indice_key=indice_key) 82 | self.bn2 = nn.BatchNorm1d(planes) 83 | self.conv3 = conv1x1( 84 | planes, planes * self.expansion, indice_key=indice_key) 85 | self.bn3 = nn.BatchNorm1d(planes * self.expansion) 86 | self.relu = nn.ReLU() 87 | self.downsample = downsample 88 | self.stride = stride 89 | 90 | def forward(self, x): 91 | identity = x 92 | 93 | out = self.conv1(x) 94 | out.features = self.bn1(out.features) 95 | out.features = self.relu(out.features) 96 | 97 | out = self.conv2(out) 98 | out.features = self.bn2(out.features) 99 | out.features = self.relu(out.features) 100 | 101 | out = self.conv3(out) 102 | out.features = self.bn3(out.features) 103 | 104 | if self.downsample is not None: 105 | identity = self.downsample(x) 106 | 107 | out += identity 108 | out.features = self.relu(out.features) 109 | 110 | return out 111 | -------------------------------------------------------------------------------- /second/pytorch/utils/__init__.py: -------------------------------------------------------------------------------- 1 | import time 2 | import contextlib 3 | import torch 4 | 5 | @contextlib.contextmanager 6 | def torch_timer(name=''): 7 | torch.cuda.synchronize() 8 | t = time.time() 9 | yield 10 | torch.cuda.synchronize() 11 | print(name, "time:", time.time() - t) -------------------------------------------------------------------------------- /second/script.py: -------------------------------------------------------------------------------- 1 | import os 2 | import fire 3 | from second.pytorch.train import train, evaluate 4 | from google.protobuf import text_format 5 | from second.protos import pipeline_pb2 6 | from pathlib import Path 7 | from second.utils import config_tool 8 | from second.data.all_dataset import get_dataset_class 9 | 10 | def _div_up(a, b): 11 | return (a + b - 1) // b 12 | 13 | def _get_cfg(cfg_path): 14 | cfg = pipeline_pb2.TrainEvalPipelineConfig() 15 | with open(cfg_path, 'r') as f: 16 | proto_str = f.read() 17 | text_format.Merge(proto_str, cfg) 18 | return cfg 19 | 20 | def _fix_kitti_dir(cfg, kitti_dir): 21 | cfg.train_input_reader.dataset.kitti_info_path = f'{kitti_dir}/kitti_infos_train.pkl' 22 | cfg.train_input_reader.dataset.kitti_root_path = f'{kitti_dir}' 23 | cfg.train_input_reader.preprocess.database_sampler.database_info_path = f'{kitti_dir}/dbinfos_train.pkl' 24 | cfg.eval_input_reader.dataset.kitti_info_path = f'{kitti_dir}/kitti_infos_val.pkl' 25 | cfg.eval_input_reader.dataset.kitti_root_path = f'{kitti_dir}' 26 | 27 | def _modify_kitti_dataset(cfg, new_kitti_dir, new_dataset_class, new_input_dim): 28 | # overwrite dataset directories 29 | cfg.train_input_reader.dataset.kitti_info_path = f'{new_kitti_dir}/kitti_infos_train.pkl' 30 | cfg.train_input_reader.dataset.kitti_root_path = f'{new_kitti_dir}' 31 | cfg.train_input_reader.preprocess.database_sampler.database_info_path = f'{new_kitti_dir}/dbinfos_train.pkl' 32 | cfg.eval_input_reader.dataset.kitti_info_path = f'{new_kitti_dir}/kitti_infos_val.pkl' 33 | cfg.eval_input_reader.dataset.kitti_root_path = f'{new_kitti_dir}' 34 | # 35 | cfg.train_input_reader.dataset.dataset_class_name = new_dataset_class 36 | cfg.eval_input_reader.dataset.dataset_class_name = new_dataset_class 37 | # 38 | cfg.model.second.num_point_features = new_input_dim 39 | cfg.model.second.voxel_feature_extractor.num_input_features = new_input_dim # VFE 40 | cfg.model.second.middle_feature_extractor.num_input_features = new_input_dim-1 # MFE 41 | 42 | def _fix_nusc(cfg, nusc_dir, sweep_db): 43 | train_info_name = 'infos_train_with_sweep_info' if sweep_db else 'infos_train' 44 | cfg.train_input_reader.dataset.kitti_info_path = f'{nusc_dir}/{train_info_name}.pkl' 45 | cfg.train_input_reader.dataset.kitti_root_path = f'{nusc_dir}' 46 | db_name = 'dbinfos_train_with_sweep_info' if sweep_db else 'dbinfos_train' 47 | cfg.train_input_reader.preprocess.database_sampler.database_info_path = f'{nusc_dir}/{db_name}.pkl' 48 | val_info_name = 'infos_val_with_sweep_info' if sweep_db else 'infos_val' 49 | cfg.eval_input_reader.dataset.kitti_info_path = f'{nusc_dir}/{val_info_name}.pkl' 50 | cfg.eval_input_reader.dataset.kitti_root_path = f'{nusc_dir}' 51 | 52 | def _fix_nusc_step(cfg, 53 | epochs, 54 | eval_epoch, 55 | data_sample_factor, 56 | gpus, 57 | num_examples=28130): 58 | input_cfg = cfg.train_input_reader 59 | train_cfg = cfg.train_config 60 | batch_size = input_cfg.batch_size 61 | data_sample_factor_to_name = { 62 | 1: "NuScenesDataset", 63 | 2: "NuScenesDatasetD2", 64 | 3: "NuScenesDatasetD3", 65 | 4: "NuScenesDatasetD4", 66 | 5: "NuScenesDatasetD5", 67 | 6: "NuScenesDatasetD6", 68 | 7: "NuScenesDatasetD7", 69 | 8: "NuScenesDatasetD8", 70 | } 71 | dataset_name = data_sample_factor_to_name[data_sample_factor] 72 | input_cfg.dataset.dataset_class_name = dataset_name 73 | ds = get_dataset_class(dataset_name)( 74 | root_path=input_cfg.dataset.kitti_root_path, 75 | info_path=input_cfg.dataset.kitti_info_path, 76 | ) 77 | num_examples_after_sample = len(ds) 78 | step_per_epoch = _div_up(num_examples_after_sample, batch_size) 79 | step_per_eval = _div_up(step_per_epoch * eval_epoch, len(gpus)) 80 | total_step = _div_up(step_per_epoch * epochs, len(gpus)) 81 | train_cfg.steps = total_step 82 | train_cfg.steps_per_eval = step_per_eval 83 | 84 | 85 | def train_kitti( 86 | kitti_dir = '/data/kitti/object', 87 | base_cfg = 'all.fhd.config', 88 | ext_input = '', 89 | resume = False, 90 | ): 91 | cfg_path = Path(f'configs/{base_cfg}') 92 | cfg_name = base_cfg.replace('.config', '').replace('.', '_') 93 | cfg = _get_cfg(cfg_path) 94 | if ext_input == '': 95 | _fix_kitti_dir(cfg, kitti_dir) 96 | else: 97 | if ext_input == 'gt_fgm': 98 | _modify_kitti_dataset(cfg, '/data/kitti/object_gt_fgm', 'KittiFGMDataset', 5) 99 | cfg_name = f'{cfg_name}_{ext_input}' 100 | model_dir = Path('models/kitti') / cfg_name 101 | train(cfg, model_dir, resume=resume) 102 | 103 | def train_nuscenes( 104 | nusc_dir = '/data/nuscenes', 105 | base_cfg = 'all.fhd.config', 106 | sample_factor = 8, 107 | epochs = 50, 108 | eval_epoch = 5, 109 | sweep_db = False, 110 | label = "", 111 | resume = False, 112 | ): 113 | cfg_path = Path(f'configs/nuscenes/{base_cfg}') 114 | cfg_name = base_cfg.replace('.config', '').replace('.', '_') 115 | cfg_name += "_swpdb" if sweep_db else "" 116 | cfg_name += f"_d{sample_factor}_ep{epochs}_ev{eval_epoch}" 117 | cfg_name += "" if label == "" else f"_{label}" 118 | cfg = _get_cfg(cfg_path) 119 | gpus = [int(i) for i in os.environ['CUDA_VISIBLE_DEVICES'].split(',')] 120 | _fix_nusc(cfg, nusc_dir, sweep_db) 121 | _fix_nusc_step(cfg, epochs, eval_epoch, sample_factor, gpus) 122 | model_dir = Path('models/nuscenes') / cfg_name 123 | train(cfg, model_dir, multi_gpu=len(gpus)>1, resume=resume) 124 | 125 | def test_nuscenes( 126 | nusc_dir = '/data/nuscenes', 127 | model_name = 'all_pp_mhead_vpn_swpdb_d1_ep20_ev2_vp_pp_oa_ta_ots_early_fusion', 128 | eval_step = 187540, 129 | ): 130 | model_dir = f'models/nuscenes/{model_name}' 131 | ckpt_path = f'{model_dir}/voxelnet-{eval_step}.tckpt' 132 | 133 | cfg_path = f'{model_dir}/pipeline.config' 134 | cfg = _get_cfg(cfg_path) 135 | cfg.eval_input_reader.dataset.kitti_info_path = f'{nusc_dir}/infos_test.pkl' 136 | cfg.eval_input_reader.dataset.kitti_root_path = f'{nusc_dir}' 137 | 138 | evaluate(cfg, model_dir, ckpt_path=ckpt_path, clean_after=False) 139 | 140 | if __name__ == "__main__": 141 | fire.Fire() 142 | -------------------------------------------------------------------------------- /second/utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/peiyunh/wysiwyg/bb0f6ecbf3f15082fe4ec14a701ccbb14198babd/second/utils/__init__.py -------------------------------------------------------------------------------- /second/utils/check.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | def is_array_like(x): 4 | return isinstance(x, (list, tuple, np.ndarray)) 5 | 6 | def shape_mergeable(x, expected_shape): 7 | mergeable = True 8 | if is_array_like(x) and is_array_like(expected_shape): 9 | x = np.array(x) 10 | if len(x.shape) == len(expected_shape): 11 | for s, s_ex in zip(x.shape, expected_shape): 12 | if s_ex is not None and s != s_ex: 13 | mergeable = False 14 | break 15 | return mergeable -------------------------------------------------------------------------------- /second/utils/compare_runtime.py: -------------------------------------------------------------------------------- 1 | import ast 2 | import numpy as np 3 | 4 | # rootdir = "/data/nuscenes_models/second" 5 | rootdir = "models/nuscenes" 6 | 7 | models = [ 8 | "all_pp_mhead_d2", 9 | "all_pp_mhead_vpn_swpdb_d1_ep20_ev2_vp_pp_oa_ta_ots_early_fusion", 10 | ] 11 | 12 | for model in models: 13 | path = f"{rootdir}/{model}/log.json.lst" 14 | print(path) 15 | 16 | step_times = [] 17 | voxel_times = [] 18 | prep_times = [] 19 | with open(path) as f: 20 | for line in f: 21 | if line.startswith(r'{"runtime":'): 22 | data = ast.literal_eval(line) 23 | step_times.append(data["runtime"]["steptime"]) 24 | voxel_times.append(data["runtime"]["voxel_gene_time"]) 25 | prep_times.append(data["runtime"]["prep_time"]) 26 | 27 | step_times = np.array(step_times) 28 | voxel_times = np.array(voxel_times) 29 | prep_times = np.array(prep_times) 30 | 31 | # I = (step_times <= 1) 32 | I = np.arange(len(step_times)) 33 | print( 34 | "total:", step_times[I].mean(), 35 | "voxel:", voxel_times[I].mean(), 36 | "prep:", prep_times[I].mean() 37 | ) -------------------------------------------------------------------------------- /second/utils/config_tool.py: -------------------------------------------------------------------------------- 1 | # This file contains some config modification function. 2 | # some functions should be only used for KITTI dataset. 3 | 4 | from google.protobuf import text_format 5 | from second.protos import pipeline_pb2, second_pb2 6 | from pathlib import Path 7 | import numpy as np 8 | 9 | 10 | def change_detection_range(model_config, new_range): 11 | assert len(new_range) == 4, "you must provide a list such as [-50, -50, 50, 50]" 12 | old_pc_range = list(model_config.voxel_generator.point_cloud_range) 13 | old_pc_range[:2] = new_range[:2] 14 | old_pc_range[3:5] = new_range[2:] 15 | model_config.voxel_generator.point_cloud_range[:] = old_pc_range 16 | for anchor_generator in model_config.target_assigner.anchor_generators: 17 | a_type = anchor_generator.WhichOneof('anchor_generator') 18 | if a_type == "anchor_generator_range": 19 | a_cfg = anchor_generator.anchor_generator_range 20 | old_a_range = list(a_cfg.anchor_ranges) 21 | old_a_range[:2] = new_range[:2] 22 | old_a_range[3:5] = new_range[2:] 23 | a_cfg.anchor_ranges[:] = old_a_range 24 | elif a_type == "anchor_generator_stride": 25 | a_cfg = anchor_generator.anchor_generator_stride 26 | old_offset = list(a_cfg.offsets) 27 | stride = list(a_cfg.strides) 28 | old_offset[0] = new_range[0] + stride[0] / 2 29 | old_offset[1] = new_range[1] + stride[1] / 2 30 | a_cfg.offsets[:] = old_offset 31 | else: 32 | raise ValueError("unknown") 33 | old_post_range = list(model_config.post_center_limit_range) 34 | old_post_range[:2] = new_range[:2] 35 | old_post_range[3:5] = new_range[2:] 36 | model_config.post_center_limit_range[:] = old_post_range 37 | 38 | def get_downsample_factor(model_config): 39 | downsample_factor = np.prod(model_config.rpn.layer_strides) 40 | if len(model_config.rpn.upsample_strides) > 0: 41 | downsample_factor /= model_config.rpn.upsample_strides[-1] 42 | downsample_factor *= model_config.middle_feature_extractor.downsample_factor 43 | downsample_factor = int(downsample_factor) 44 | assert downsample_factor > 0 45 | return downsample_factor 46 | 47 | 48 | if __name__ == "__main__": 49 | config_path = "/home/yy/deeplearning/deeplearning/mypackages/second/configs/car.lite.1.config" 50 | config = pipeline_pb2.TrainEvalPipelineConfig() 51 | 52 | with open(config_path, "r") as f: 53 | proto_str = f.read() 54 | text_format.Merge(proto_str, config) 55 | 56 | change_detection_range(config, [-50, -50, 50, 50]) 57 | proto_str = text_format.MessageToString(config, indent=2) 58 | print(proto_str) 59 | 60 | -------------------------------------------------------------------------------- /second/utils/config_tool/__init__.py: -------------------------------------------------------------------------------- 1 | # This file contains some config modification function. 2 | # some functions should be only used for KITTI dataset. 3 | 4 | from google.protobuf import text_format 5 | from second.protos import pipeline_pb2, second_pb2 6 | from pathlib import Path 7 | import numpy as np 8 | 9 | def read_config(path): 10 | config = pipeline_pb2.TrainEvalPipelineConfig() 11 | 12 | with open(path, "r") as f: 13 | proto_str = f.read() 14 | text_format.Merge(proto_str, config) 15 | return config 16 | 17 | def change_detection_range(model_config, new_range): 18 | assert len(new_range) == 4, "you must provide a list such as [-50, -50, 50, 50]" 19 | old_pc_range = list(model_config.voxel_generator.point_cloud_range) 20 | old_pc_range[:2] = new_range[:2] 21 | old_pc_range[3:5] = new_range[2:] 22 | model_config.voxel_generator.point_cloud_range[:] = old_pc_range 23 | for anchor_generator in model_config.target_assigner.anchor_generators: 24 | a_type = anchor_generator.WhichOneof('anchor_generator') 25 | if a_type == "anchor_generator_range": 26 | a_cfg = anchor_generator.anchor_generator_range 27 | old_a_range = list(a_cfg.anchor_ranges) 28 | old_a_range[:2] = new_range[:2] 29 | old_a_range[3:5] = new_range[2:] 30 | a_cfg.anchor_ranges[:] = old_a_range 31 | elif a_type == "anchor_generator_stride": 32 | a_cfg = anchor_generator.anchor_generator_stride 33 | old_offset = list(a_cfg.offsets) 34 | stride = list(a_cfg.strides) 35 | old_offset[0] = new_range[0] + stride[0] / 2 36 | old_offset[1] = new_range[1] + stride[1] / 2 37 | a_cfg.offsets[:] = old_offset 38 | else: 39 | raise ValueError("unknown") 40 | old_post_range = list(model_config.post_center_limit_range) 41 | old_post_range[:2] = new_range[:2] 42 | old_post_range[3:5] = new_range[2:] 43 | model_config.post_center_limit_range[:] = old_post_range 44 | 45 | def get_downsample_factor(model_config): 46 | downsample_factor = np.prod(model_config.rpn.layer_strides) 47 | if len(model_config.rpn.upsample_strides) > 0: 48 | downsample_factor /= model_config.rpn.upsample_strides[-1] 49 | downsample_factor *= model_config.middle_feature_extractor.downsample_factor 50 | downsample_factor = np.round(downsample_factor).astype(np.int64) 51 | assert downsample_factor > 0 52 | return downsample_factor 53 | 54 | 55 | if __name__ == "__main__": 56 | config_path = "/home/yy/deeplearning/deeplearning/mypackages/second/configs/car.lite.1.config" 57 | config = pipeline_pb2.TrainEvalPipelineConfig() 58 | 59 | with open(config_path, "r") as f: 60 | proto_str = f.read() 61 | text_format.Merge(proto_str, config) 62 | 63 | change_detection_range(config, [-50, -50, 50, 50]) 64 | proto_str = text_format.MessageToString(config, indent=2) 65 | print(proto_str) 66 | 67 | -------------------------------------------------------------------------------- /second/utils/config_tool/train.py: -------------------------------------------------------------------------------- 1 | from second.protos.optimizer_pb2 import Optimizer, LearningRate, OneCycle, ManualStepping, ExponentialDecay 2 | from second.protos.sampler_pb2 import Sampler 3 | from second.utils.config_tool import read_config 4 | from pathlib import Path 5 | from google.protobuf import text_format 6 | from second.data.all_dataset import get_dataset_class 7 | 8 | def _get_optim_cfg(train_config, optim): 9 | if optim == "adam_optimizer": 10 | return train_config.optimizer.adam_optimizer 11 | elif optim == "rms_prop_optimizer": 12 | return train_config.optimizer.rms_prop_optimizer 13 | elif optim == "momentum_optimizer": 14 | return train_config.optimizer.momentum_optimizer 15 | else: 16 | raise NotImplementedError 17 | 18 | 19 | def manual_stepping(train_config, boundaries, rates, optim="adam_optimizer"): 20 | optim_cfg = _get_optim_cfg(train_config, optim) 21 | optim_cfg.learning_rate.manual_stepping.CopyFrom( 22 | ManualStepping(boundaries=boundaries, rates=rates)) 23 | 24 | 25 | def exp_decay(train_config, 26 | init_lr, 27 | decay_length, 28 | decay_factor, 29 | staircase=True, 30 | optim="adam_optimizer"): 31 | optim_cfg = _get_optim_cfg(train_config, optim) 32 | optim_cfg.learning_rate.exponential_decay.CopyFrom( 33 | ExponentialDecay( 34 | initial_learning_rate=init_lr, 35 | decay_length=decay_length, 36 | decay_factor=decay_factor, 37 | staircase=staircase)) 38 | 39 | 40 | def one_cycle(train_config, 41 | lr_max, 42 | moms, 43 | div_factor, 44 | pct_start, 45 | optim="adam_optimizer"): 46 | optim_cfg = _get_optim_cfg(train_config, optim) 47 | optim_cfg.learning_rate.one_cycle.CopyFrom( 48 | OneCycle( 49 | lr_max=lr_max, 50 | moms=moms, 51 | div_factor=div_factor, 52 | pct_start=pct_start)) 53 | 54 | def _div_up(a, b): 55 | return (a + b - 1) // b 56 | 57 | def set_train_step(config, 58 | epochs, 59 | eval_epoch): 60 | input_cfg = config.train_input_reader 61 | train_cfg = config.train_config 62 | batch_size = input_cfg.batch_size 63 | dataset_name = input_cfg.dataset.dataset_class_name 64 | ds = get_dataset_class(dataset_name)( 65 | root_path=input_cfg.dataset.kitti_root_path, 66 | info_path=input_cfg.dataset.kitti_info_path, 67 | ) 68 | num_examples_after_sample = len(ds) 69 | step_per_epoch = _div_up(num_examples_after_sample, batch_size) 70 | step_per_eval = step_per_epoch * eval_epoch 71 | total_step = step_per_epoch * epochs 72 | train_cfg.steps = total_step 73 | train_cfg.steps_per_eval = step_per_eval 74 | 75 | def disable_sample(config): 76 | input_cfg = config.train_input_reader 77 | input_cfg.database_sampler.CopyFrom(Sampler()) 78 | 79 | def disable_per_gt_aug(config): 80 | prep_cfg = config.train_input_reader.preprocess 81 | prep_cfg.groundtruth_localization_noise_std[:] = [0, 0, 0] 82 | prep_cfg.groundtruth_rotation_uniform_noise[:] = [0, 0] 83 | 84 | def disable_global_aug(config): 85 | prep_cfg = config.train_input_reader.preprocess 86 | prep_cfg.global_rotation_uniform_noise[:] = [0, 0] 87 | prep_cfg.global_scaling_uniform_noise[:] = [0, 0] 88 | prep_cfg.global_random_rotation_range_per_object[:] = [0, 0] 89 | prep_cfg.global_translate_noise_std[:] = [0, 0, 0] 90 | 91 | if __name__ == "__main__": 92 | path = Path(__file__).resolve().parents[2] / "configs/car.lite.config" 93 | config = read_config(path) 94 | manual_stepping(config.train_config, [0.8, 0.9], [1e-4, 1e-5, 1e-6]) 95 | 96 | print(text_format.MessageToString(config, indent=2)) -------------------------------------------------------------------------------- /second/utils/dump_kitti_files.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import pickle 3 | from second.data.kitti_dataset import KittiDataset 4 | 5 | split = "val" 6 | 7 | root_path = '/data/kitti/object' 8 | info_path = f'{root_path}/kitti_infos_{split}.pkl' 9 | 10 | model_name = "pointpillars_car_vfn_oc_xyres_16" 11 | eval_step = 296960 12 | 13 | dataset = KittiDataset(root_path, info_path, class_names=["Car"]) 14 | model_dir = f'models/kitti/{model_name}' 15 | result_dir = f'{model_dir}/{split}_results/step_{eval_step}' 16 | result_path = f'{result_dir}/result.pkl' 17 | 18 | with open(result_path, 'rb') as f: 19 | detections = pickle.load(f) 20 | 21 | dataset.evaluation(detections, result_dir) -------------------------------------------------------------------------------- /second/utils/dump_nusc_json.py: -------------------------------------------------------------------------------- 1 | # This function needs to be called in SECOND's environment 2 | 3 | import sys 4 | import pickle 5 | from second.data.nuscenes_dataset import NuScenesDataset 6 | 7 | split = 'val' if len(sys.argv) <= 3 else sys.argv[3] 8 | 9 | model_name = sys.argv[1] 10 | if len(sys.argv) > 2: 11 | eval_step = int(sys.argv[2]) 12 | else: 13 | eval_step = 58650 14 | 15 | root_path = '/data/nuscenes' 16 | if split == 'val': 17 | info_path = f'{root_path}/infos_val.pkl' 18 | else: 19 | info_path = f'{root_path}/infos_test.pkl' 20 | 21 | # THIS ORDER IS IMPORTANT AS IT TRANSLATES ANCHOR INDEX TO CLASS LABELS 22 | # IT IS TIED WITH THE ORDER CLASSES ARE DEFINED INSIDE THE CONFIG FILE 23 | class_names = ['bus', 'car', 'construction_vehicle', 'trailer', 'truck', 24 | 'barrier', 'bicycle', 'motorcycle', 'pedestrian', 'traffic_cone'] 25 | 26 | dataset = NuScenesDataset(root_path, info_path, class_names=class_names) 27 | model_dir = f'models/nuscenes/{model_name}' 28 | if split == 'val': 29 | result_dir = f'{model_dir}/results/step_{eval_step}' 30 | else: 31 | result_dir = f'{model_dir}/eval_results/step_{eval_step}' 32 | 33 | result_path = f'{result_dir}/result.pkl' 34 | print(f'loading results from {result_path}') 35 | with open(result_path, 'rb') as f: 36 | detections = pickle.load(f) 37 | 38 | dataset.evaluation_nusc(detections, result_dir, clean_after=False) 39 | -------------------------------------------------------------------------------- /second/utils/loader.py: -------------------------------------------------------------------------------- 1 | import importlib 2 | from pathlib import Path 3 | import sys 4 | import os 5 | import logging 6 | logger = logging.getLogger('second.utils.loader') 7 | 8 | CUSTOM_LOADED_MODULES = {} 9 | 10 | 11 | def _get_possible_module_path(paths): 12 | ret = [] 13 | for p in paths: 14 | p = Path(p) 15 | for path in p.glob("*"): 16 | if path.suffix in ["py", ".so"] or (path.is_dir()): 17 | if path.stem.isidentifier(): 18 | ret.append(path) 19 | return ret 20 | 21 | 22 | def _get_regular_import_name(path, module_paths): 23 | path = Path(path) 24 | for mp in module_paths: 25 | mp = Path(mp) 26 | if mp == path: 27 | return path.stem 28 | try: 29 | relative_path = path.relative_to(Path(mp)) 30 | parts = list((relative_path.parent / relative_path.stem).parts) 31 | module_name = '.'.join([mp.stem] + parts) 32 | return module_name 33 | except: 34 | pass 35 | return None 36 | 37 | 38 | def import_file(path, name: str = None, add_to_sys=True, 39 | disable_warning=False): 40 | global CUSTOM_LOADED_MODULES 41 | path = Path(path) 42 | module_name = path.stem 43 | try: 44 | user_paths = os.environ['PYTHONPATH'].split(os.pathsep) 45 | except KeyError: 46 | user_paths = [] 47 | possible_paths = _get_possible_module_path(user_paths) 48 | model_import_name = _get_regular_import_name(path, possible_paths) 49 | if model_import_name is not None: 50 | return import_name(model_import_name) 51 | if name is not None: 52 | module_name = name 53 | spec = importlib.util.spec_from_file_location(module_name, path) 54 | module = importlib.util.module_from_spec(spec) 55 | spec.loader.exec_module(module) 56 | if not disable_warning: 57 | logger.warning(( 58 | f"Failed to perform regular import for file {path}. " 59 | "this means this file isn't in any folder in PYTHONPATH " 60 | "or don't have __init__.py in that project. " 61 | "directly file import may fail and some reflecting features are " 62 | "disabled even if import succeed. please add your project to PYTHONPATH " 63 | "or add __init__.py to ensure this file can be regularly imported. " 64 | )) 65 | 66 | if add_to_sys: # this will enable find objects defined in a file. 67 | # avoid replace system modules. 68 | if module_name in sys.modules and module_name not in CUSTOM_LOADED_MODULES: 69 | raise ValueError(f"{module_name} exists in system.") 70 | CUSTOM_LOADED_MODULES[module_name] = module 71 | sys.modules[module_name] = module 72 | return module 73 | 74 | 75 | def import_name(name, package=None): 76 | module = importlib.import_module(name, package) 77 | return module -------------------------------------------------------------------------------- /second/utils/log_tool.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from tensorboardX import SummaryWriter 3 | import json 4 | from pathlib import Path 5 | 6 | def _flat_nested_json_dict(json_dict, flatted, sep=".", start=""): 7 | for k, v in json_dict.items(): 8 | if isinstance(v, dict): 9 | _flat_nested_json_dict(v, flatted, sep, start + sep + str(k)) 10 | else: 11 | flatted[start + sep + str(k)] = v 12 | 13 | def flat_nested_json_dict(json_dict, sep=".") -> dict: 14 | """flat a nested json-like dict. this function make shadow copy. 15 | """ 16 | flatted = {} 17 | for k, v in json_dict.items(): 18 | if isinstance(v, dict): 19 | _flat_nested_json_dict(v, flatted, sep, str(k)) 20 | else: 21 | flatted[str(k)] = v 22 | return flatted 23 | 24 | def metric_to_str(metrics, sep='.'): 25 | flatted_metrics = flat_nested_json_dict(metrics, sep) 26 | metrics_str_list = [] 27 | for k, v in flatted_metrics.items(): 28 | if isinstance(v, float): 29 | metrics_str_list.append(f"{k}={v:.4}") 30 | elif isinstance(v, (list, tuple)): 31 | if v and isinstance(v[0], float): 32 | v_str = ', '.join([f"{e:.4}" for e in v]) 33 | metrics_str_list.append(f"{k}=[{v_str}]") 34 | else: 35 | metrics_str_list.append(f"{k}={v}") 36 | else: 37 | metrics_str_list.append(f"{k}={v}") 38 | return ', '.join(metrics_str_list) 39 | 40 | class SimpleModelLog: 41 | """For simple log. 42 | generate 4 kinds of log: 43 | 1. simple log.txt, all metric dicts are flattened to produce 44 | readable results. 45 | 2. TensorBoard scalars and texts 46 | 3. multi-line json file log.json.lst 47 | 4. tensorboard_scalars.json, all scalars are stored in this file 48 | in tensorboard json format. 49 | """ 50 | def __init__(self, model_dir): 51 | self.model_dir = Path(model_dir) 52 | self.log_file = None 53 | self.log_mjson_file = None 54 | self.summary_writter = None 55 | self.metrics = [] 56 | self._text_current_gstep = -1 57 | self._tb_texts = [] 58 | 59 | def open(self): 60 | model_dir = self.model_dir 61 | assert model_dir.exists() 62 | summary_dir = model_dir / 'summary' 63 | summary_dir.mkdir(parents=True, exist_ok=True) 64 | 65 | log_mjson_file_path = model_dir / f'log.json.lst' 66 | if log_mjson_file_path.exists(): 67 | with open(log_mjson_file_path, 'r') as f: 68 | for line in f.readlines(): 69 | self.metrics.append(json.loads(line)) 70 | log_file_path = model_dir / f'log.txt' 71 | self.log_mjson_file = open(log_mjson_file_path, 'a') 72 | self.log_file = open(log_file_path, 'a') 73 | self.summary_writter = SummaryWriter(str(summary_dir)) 74 | return self 75 | 76 | def close(self): 77 | assert self.summary_writter is not None 78 | self.log_mjson_file.close() 79 | self.log_file.close() 80 | tb_json_path = str(self.model_dir / "tensorboard_scalars.json") 81 | self.summary_writter.export_scalars_to_json(tb_json_path) 82 | self.summary_writter.close() 83 | self.log_mjson_file = None 84 | self.log_file = None 85 | self.summary_writter = None 86 | 87 | def log_text(self, text, step, tag="regular log"): 88 | """This function only add text to log.txt and tensorboard texts 89 | """ 90 | print(text) 91 | print(text, file=self.log_file) 92 | if step > self._text_current_gstep and self._text_current_gstep != -1: 93 | total_text = '\n'.join(self._tb_texts) 94 | self.summary_writter.add_text(tag, total_text, global_step=step) 95 | self._tb_texts = [] 96 | self._text_current_gstep = step 97 | else: 98 | self._tb_texts.append(text) 99 | if self._text_current_gstep == -1: 100 | self._text_current_gstep = step 101 | 102 | 103 | def log_metrics(self, metrics: dict, step): 104 | flatted_summarys = flat_nested_json_dict(metrics, "/") 105 | for k, v in flatted_summarys.items(): 106 | if isinstance(v, (list, tuple)): 107 | if any([isinstance(e, str) for e in v]): 108 | continue 109 | v_dict = {str(i): e for i, e in enumerate(v)} 110 | for k1, v1 in v_dict.items(): 111 | self.summary_writter.add_scalar(k + "/" + k1, v1, step) 112 | else: 113 | if isinstance(v, str): 114 | continue 115 | self.summary_writter.add_scalar(k, v, step) 116 | log_str = metric_to_str(metrics) 117 | print(log_str) 118 | print(log_str, file=self.log_file) 119 | print(json.dumps(metrics), file=self.log_mjson_file) 120 | 121 | -------------------------------------------------------------------------------- /second/utils/mapping/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.13.3) 2 | project(mapping) 3 | 4 | set (CMAKE_CXX_STANDARD 14) 5 | 6 | set(CMAKE_BUILD_TYPE Release) 7 | 8 | set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/..) 9 | set(CMAKE_CXX_FLAGS "-Wall -Wextra") 10 | set(CMAKE_CXX_FLAGS_RELEASE "-O3") 11 | 12 | find_package(Eigen3) 13 | include_directories(${EIGEN3_INCLUDE_DIRS}) 14 | 15 | # add_subdirectory(pybind11) 16 | find_package(pybind11) 17 | pybind11_add_module(mapping mapping.cpp) 18 | 19 | # find_package(OpenMP) 20 | # target_link_libraries(mapping PRIVATE OpenMP::OpenMP_CXX) -------------------------------------------------------------------------------- /second/utils/merge_kitti_results.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | 3 | """ 4 | # branch: pps-kitti 5 | # car 6 | CUDA_VISIBLE_DEVICES=0 python script.py test_kitti --model_name="pointpillars_car_oc_xyres_16" --eval_step=296960 & 7 | # ped 8 | CUDA_VISIBLE_DEVICES=1 python script.py test_kitti --model_name="pointpillars_ped_cycle_oc_lb_xyres_16" --eval_step=278400 & 9 | # cyc 10 | CUDA_VISIBLE_DEVICES=2 python script.py test_kitti --model_name="pointpillars_ped_cycle_oc_lb_xyres_16" --eval_step=139200 & 11 | 12 | # branch: pps-kitti-vp-oa-early-fusion 13 | # car 14 | CUDA_VISIBLE_DEVICES=0 python script.py test_kitti --model_name="pointpillars_car_vfn_oc_xyres_16" --eval_step=296960 & 15 | # ped 16 | CUDA_VISIBLE_DEVICES=1 python script.py test_kitti --model_name="pointpillars_ped_cycle_vfn_oc_lb_xyres_16" --eval_step=204160 & 17 | # cyc 18 | CUDA_VISIBLE_DEVICES=2 python script.py test_kitti --model_name="pointpillars_ped_cycle_vfn_oc_lb_xyres_16" --eval_step=83520 & 19 | """ 20 | # split = "test" 21 | split = "val" 22 | kitti_dir = "/data/kitti/object" 23 | model_dir = "models/kitti" 24 | 25 | """ 26 | submission = { 27 | "name": "WYSIWYG", 28 | "dirs": { 29 | "Car": f"pointpillars_car_vfn_oc_xyres_16/{split}_results/step_296960/kitti", 30 | "Pedestrian": f"pointpillars_ped_cycle_vfn_oc_lb_xyres_16/{split}_results/step_204160/kitti", 31 | "Cyclist": f"pointpillars_ped_cycle_vfn_oc_lb_xyres_16/{split}_results/step_83520/kitti", 32 | } 33 | } 34 | """ 35 | submission = { 36 | "name": "WYSIWYG", 37 | "dirs": { 38 | # "Car": f"pointpillars_car_vfn_oc_xyres_16/{split}_results/step_296960/kitti", 39 | "Car": f"trainval_pointpillars_car_vpn_xyres_16/{split}_results/step_296960/kitti", 40 | "Pedestrian": f"trainval_pointpillars_ped_cycle_vpn_lb_xyres_16/{split}_results/step_296960/kitti", 41 | "Cyclist": f"trainval_pointpillars_ped_cycle_vpn_lb_xyres_16/{split}_results/step_296960/kitti", 42 | } 43 | } 44 | 45 | # score_thr = 0.5 46 | score_thr = 0.0 47 | # score_thr = 0.3 48 | if score_thr > 0: 49 | output_dir = Path(f"results/kitti/{split}/{submission['name']}_above_{score_thr}/data") 50 | else: 51 | output_dir = Path(f"results/kitti/{split}/{submission['name']}/data") 52 | 53 | output_dir.mkdir(parents=True, exist_ok=True) 54 | 55 | with open(f"{kitti_dir}/{split}.txt") as f: 56 | I = [int(line.strip()) for line in f] 57 | 58 | for i in I: 59 | output_file = output_dir / ("%06d.txt" % i) 60 | output_handle = open(output_file, "a") 61 | for cls in submission["dirs"]: 62 | input_file = Path(model_dir) / submission["dirs"][cls] / ("%06d.txt" % i) 63 | input_handle = open(input_file, "r") 64 | for line in input_handle: 65 | if line.startswith(cls): 66 | elements = line.split() 67 | label, trunc, occl, alpha, x1, y1, x2, y2, h, w, l, x, y, z, r, score = elements 68 | if float(score) < score_thr: continue 69 | if not line.endswith('\n'): line += '\n' 70 | output_handle.write(line) 71 | input_handle.close() 72 | output_handle.close() 73 | -------------------------------------------------------------------------------- /second/utils/merge_result.py: -------------------------------------------------------------------------------- 1 | import fire 2 | from pathlib import Path 3 | import re 4 | 5 | def merge(path1, path2, output_path): 6 | filepaths1 = Path(path1).glob('*.txt') 7 | prog = re.compile(r'^\d{6}.txt$') 8 | filepaths1 = filter(lambda f: prog.match(f.name), filepaths1) 9 | for fp1 in list(filepaths1): 10 | with open(fp1) as f1: 11 | contents = f1.readlines() 12 | if len(contents) != 0: 13 | contents += "\n" 14 | with open(Path(path2) / f"{fp1.stem}.txt", 'r') as f2: 15 | contents += f2.readlines() 16 | with open(Path(output_path) / f"{fp1.stem}.txt", 'w') as f: 17 | f.writelines(contents) 18 | 19 | if __name__ == '__main__': 20 | fire.Fire() 21 | 22 | -------------------------------------------------------------------------------- /second/utils/model_tool.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | import re 3 | import shutil 4 | 5 | def rm_invalid_model_dir(directory, step_thresh=200): 6 | directory = Path(directory).resolve() 7 | ckpt_re = r"[a-zA-Z0-9_]+\-([0-9]+)\.tckpt" 8 | ckpt_re = re.compile(ckpt_re) 9 | for path in directory.rglob("*"): 10 | if path.is_dir(): 11 | pipeline_path = (path / "pipeline.config") 12 | log_path = (path / "log.txt") 13 | summary_path = (path / "summary") 14 | must_exists = [pipeline_path, log_path, summary_path] 15 | if not all([e.exists() for e in must_exists]): 16 | continue 17 | ckpts = [] 18 | for subpath in path.iterdir(): 19 | match = ckpt_re.search(subpath.name) 20 | if match is not None: 21 | ckpts.append(int(match.group(1))) 22 | if len(ckpts) == 0 or all([e < step_thresh for e in ckpts]): 23 | shutil.rmtree(str(path)) 24 | -------------------------------------------------------------------------------- /second/utils/plot_kitti_maps.py: -------------------------------------------------------------------------------- 1 | import os 2 | import ast 3 | import json 4 | import numpy as np 5 | import seaborn as sns 6 | sns.set() 7 | import matplotlib.pyplot as plt 8 | 9 | root_dir = "models/kitti/" 10 | 11 | ignored_names = [ 12 | "backup", 13 | "all8_fhd(run0)", 14 | "all8_fhd_gt_fgm", 15 | # "all_fhd", 16 | # "all8_fhd", 17 | # "car_fhd", 18 | # "car_fhd_onestage", 19 | # "people_fhd", 20 | "pointpillars_pp_pretrain", 21 | # "pointpillars_car_xyres_16", 22 | # "pointpillars_ped_cycle_xyres_16", 23 | # "pointpillars_car_vfn_xyres_16", 24 | # "pointpillars_ped_cycle_vfn_xyres_16", 25 | ] 26 | 27 | ignored_labels = [ 28 | # "oc", 29 | ] 30 | 31 | metrics = { 32 | "Car": "3d@0.70", 33 | "Pedestrian": "3d@0.50", 34 | "Cyclist": "3d@0.50", 35 | } 36 | 37 | numbers = { 38 | "Car": {}, 39 | "Pedestrian": {}, 40 | "Cyclist": {}, 41 | } 42 | 43 | difficulties = { 44 | 0: "Easy", 45 | 1: "Medium", 46 | 2: "Hard" 47 | } 48 | 49 | for name in os.listdir(root_dir): 50 | if name in ignored_names: 51 | continue 52 | 53 | if any(label in name for label in ignored_labels): 54 | continue 55 | 56 | last_step = -1 57 | stats_list = [] 58 | path = os.path.join(root_dir, name, "log.json.lst") 59 | with open(path) as f: 60 | for line in f: 61 | data = ast.literal_eval(line) 62 | if "runtime" in data: 63 | last_step = data["runtime"]["step"] 64 | elif "eval.kitti" in data: 65 | mAPs_dict = data["eval.kitti"]["official"] 66 | stats_list.append((last_step, mAPs_dict)) 67 | 68 | for last_step, mAPs_dict in stats_list: 69 | for cls in metrics: 70 | metric = metrics[cls] 71 | if cls in mAPs_dict: 72 | mAPs = (mAPs_dict[cls][metric]) 73 | if name not in numbers[cls]: 74 | numbers[cls][name] = [] 75 | numbers[cls][name].append([last_step]+mAPs) 76 | 77 | plt.figure(figsize=(32,18)) 78 | i = 0 79 | for cls in numbers: 80 | for name in numbers[cls]: 81 | data = np.array(numbers[cls][name]) 82 | steps, mAPs = data[:,0], data[:, 1:] 83 | # pick the best checkpoint based on the performance on the medium subset 84 | k = np.argmax(mAPs[:, 1]) 85 | for j in difficulties: 86 | plt.subplot(len(numbers), 3, i*3+j+1) 87 | plt.plot(steps, mAPs[:,j], '-.') 88 | plt.plot(steps, np.full_like(steps, mAPs[k,j]), label=name) 89 | for j in difficulties: 90 | plt.subplot(len(numbers), 3, i*3+j+1) 91 | plt.title(cls+" - "+difficulties[j]) 92 | plt.legend() 93 | i += 1 94 | plt.savefig(f"results/kitti.png", dpi=250, bbox_inches="tight") 95 | -------------------------------------------------------------------------------- /second/utils/plot_kitti_results.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from pathlib import Path 3 | from skimage.io import imread 4 | import matplotlib.pyplot as plt 5 | import matplotlib.patches as patches 6 | 7 | split = "test" 8 | kitti_dir = Path("/data/kitti/object") 9 | image_dir = kitti_dir / ("testing" if split=="test" else "training") / "image_2" 10 | 11 | model_names = [ 12 | # "PointPillars", 13 | "WYSIWYG", 14 | ] 15 | 16 | res_dirs = [ 17 | Path(f"results/kitti/{split}/{model_name}/data") for model_name in model_names 18 | ] 19 | if split == "val": 20 | model_names = ["GT"] + model_names 21 | res_dirs = [Path("/data/kitti/object/training/label_2")] + res_dirs 22 | 23 | 24 | with open(f"{kitti_dir}/{split}.txt") as f: 25 | I = [int(line.strip()) for line in f] 26 | 27 | CLASSES = ["Car", "Pedestrian", "Cyclist", "Van", "Truck", "Misc", "Tram", "Person_sitting"] 28 | colors = { 29 | cls : (0.7 * np.random.rand(3) + 0.3).tolist() for cls in CLASSES 30 | } 31 | 32 | def sigmoid(x): 33 | return 5 / (1 + np.exp(-x)) 34 | 35 | # np.random.shuffle(I) 36 | for i in I: 37 | plt.clf() 38 | 39 | image_file = image_dir / ("%06d.png" % i) 40 | image = imread(image_file) 41 | for m, res_dir in enumerate(res_dirs): 42 | plt.subplot(len(res_dirs), 1, m+1) 43 | plt.title(model_names[m]) 44 | 45 | plt.imshow(image) 46 | res_file = res_dir / ("%06d.txt" % i) 47 | with open(res_file) as f: 48 | for line in f: 49 | elements = line.split() 50 | if len(elements) == 15: 51 | label, trunc, occl, alpha, x1, y1, x2, y2, h, w, l, x, y, z, r = elements 52 | score = 1 53 | else: 54 | label, trunc, occl, alpha, x1, y1, x2, y2, h, w, l, x, y, z, r, score = elements 55 | 56 | x1, x2, y1, y2, score = float(x1), float(x2), float(y1), float(y2), float(score) 57 | if label=="DontCare" or score < 0.5: continue 58 | 59 | print(model_names[m], ":", line.strip()) 60 | rect = patches.Rectangle((x1, y1), x2-x1+1, y2-y1+1, linewidth=sigmoid(score), edgecolor=colors[label], facecolor="none") 61 | ax = plt.gca() 62 | ax.add_patch(rect) 63 | plt.show() -------------------------------------------------------------------------------- /second/utils/print_attr.py: -------------------------------------------------------------------------------- 1 | import os 2 | import json 3 | import argparse 4 | 5 | parser = argparse.ArgumentParser() 6 | parser.add_argument('--dataset', type=str, default='nuscenes') 7 | parser.add_argument('--step', type=int, default='-1') 8 | parser.add_argument('--metric', type=str, default='label_aps') 9 | parser.add_argument('--thresh', type=str, default="") 10 | args = parser.parse_args() 11 | 12 | dataset = args.dataset 13 | base_dir = f'models/{dataset}' 14 | 15 | classes = [ 16 | # # official class order 17 | # 'car', 'truck', 'bus', 'trailer', 'construction_vehicle', 'pedestrian', 'motorcycle', 'bicycle', 'traffic_cone', 'barrier' 18 | # # sorted by percentage 19 | # 'car', 'pedestrian', 'barrier', 'traffic_cone', 'truck', 'bus', 'trailer', 'construction_vehicle', 'motorcycle', 'bicycle' 20 | 'car', 'car[vehicle.moving]', 'car[vehicle.stopped]', 'car[vehicle.parked]', 21 | 'pedestrian', 'pedestrian[pedestrian.moving]', 'pedestrian[pedestrian.sitting_lying_down]', 'pedestrian[pedestrian.standing]', 22 | 'barrier', 23 | 'traffic_cone', 24 | 'truck', 'truck[vehicle.moving]', 'truck[vehicle.stopped]', 'truck[vehicle.parked]', 25 | 'bus', 'bus[vehicle.moving]', 'bus[vehicle.stopped]', 'bus[vehicle.parked]', 26 | 'trailer', 'trailer[vehicle.moving]', 'trailer[vehicle.stopped]', 'trailer[vehicle.parked]', 27 | 'construction_vehicle', 'construction_vehicle[vehicle.moving]', 'construction_vehicle[vehicle.stopped]', 'construction_vehicle[vehicle.parked]', 28 | 'motorcycle', 'motorcycle[cycle.with_rider]', 'motorcycle[cycle.without_rider]', 29 | 'bicycle', 'bicycle[cycle.with_rider]', 'bicycle[cycle.without_rider]' 30 | ] 31 | 32 | if args.dataset == 'nuscenes': 33 | methods = [ 34 | # "pointpillars", 35 | # "all_pp_mhead_vpn_swpdb_d1_ep20_ev2_vp_pp_oa_ta_learn", 36 | # "all_pp_mhead_vpn_swpdb_d1_ep20_ev2_vp_pp_oa_ta_ots", 37 | "all_pp_mhead_d1_ep20_ev2_pp_oa_ta", 38 | "all_pp_mhead_vpn_swpdb_d1_ep20_ev2_vp_pp_oa_ta_ots_early_fusion", 39 | ] 40 | names = [ 41 | # "pp(official)", 42 | # "vp_pp_oa_ta_learn", 43 | # "vp_pp_oa_ta_ots", 44 | "pp_oa_ta", 45 | "vp_pp_oa_ta_ots_EF", 46 | ] 47 | 48 | cache = {} 49 | 50 | assert(len(methods) == len(names)) 51 | 52 | for method in methods: 53 | res_dir = f'{base_dir}/{method}/results' 54 | if not os.path.exists(res_dir): 55 | continue 56 | 57 | if args.step == -1: # use the final checkpoint 58 | all_steps = [] 59 | for d in os.listdir(res_dir): 60 | if os.path.exists(os.path.join(res_dir, d, "metrics_summary.json")): 61 | all_steps.append(int(d.split("_")[1])) 62 | if len(all_steps) == 0: 63 | continue 64 | step = max(all_steps) 65 | else: 66 | step = args.step 67 | print(method, step) 68 | 69 | res_file = f'{res_dir}/step_{step}/metrics_summary.json' 70 | if os.path.exists(res_file): 71 | with open(res_file, 'r') as f: 72 | summary = json.load(f) 73 | 74 | cache[method] = summary 75 | 76 | 77 | delim = '\t' 78 | # delim = ' & ' 79 | metric = args.metric 80 | 81 | print('{:24}'.format(f'mAP[{args.thresh}]'), end=delim) 82 | for cls in classes: 83 | print('{:5}'.format(cls[:5]), end=delim) 84 | print('{:5}'.format('avg')) 85 | 86 | for name, method in zip(names, methods): 87 | if method not in cache: 88 | continue 89 | 90 | print('{:24}'.format(name), end=delim) 91 | APs = [] 92 | for cls in classes: 93 | n = cache[method][metric][cls] 94 | if args.thresh in n: 95 | AP = n[args.thresh] 96 | else: 97 | AP = sum(n.values())/len(n) 98 | APs.append(AP) 99 | print('{:.3f}'.format(AP), end=delim) 100 | mAP = sum(APs)/len(APs) 101 | print('{:.3f}'.format(mAP)) 102 | -------------------------------------------------------------------------------- /second/utils/print_kitti_maps.py: -------------------------------------------------------------------------------- 1 | import os 2 | import ast 3 | import json 4 | import numpy as np 5 | 6 | root_dir = "models/kitti/" 7 | 8 | ignored_names = [ 9 | "backup", 10 | "all8_fhd(run0)", 11 | "all8_fhd_gt_fgm" 12 | ] 13 | 14 | metrics = { 15 | "Car": "3d@0.70", 16 | "Pedestrian": "3d@0.50", 17 | "Cyclist": "3d@0.50", 18 | } 19 | 20 | numbers = { 21 | "Car": {}, 22 | "Pedestrian": {}, 23 | "Cyclist": {}, 24 | } 25 | 26 | difficulties = { 27 | 0: "Easy", 28 | 1: "Medium", 29 | 2: "Hard" 30 | } 31 | 32 | step_set = [ 33 | 296950, 34 | ] 35 | 36 | for name in os.listdir(root_dir): 37 | if name in ignored_names: 38 | continue 39 | 40 | last_step = -1 41 | stats_list = [] 42 | path = os.path.join(root_dir, name, "log.json.lst") 43 | with open(path) as f: 44 | for line in f: 45 | data = ast.literal_eval(line) 46 | if "runtime" in data: 47 | last_step = data["runtime"]["step"] 48 | elif "eval.kitti" in data: 49 | mAPs_dict = data["eval.kitti"]["official"] 50 | stats_list.append((last_step, mAPs_dict)) 51 | 52 | for last_step, mAPs_dict in stats_list: 53 | if len(step_set) > 0 and last_step not in step_set: 54 | continue 55 | for cls in metrics: 56 | metric = metrics[cls] 57 | if cls in mAPs_dict: 58 | mAPs = (mAPs_dict[cls][metric]) 59 | if name not in numbers[cls]: 60 | numbers[cls][name] = [] 61 | numbers[cls][name].append([last_step]+mAPs) 62 | 63 | 64 | # model selection across checkpoints 65 | best_steps = {} 66 | best_numbers = {} 67 | for cls in numbers: 68 | if cls not in best_numbers: best_numbers[cls] = {} 69 | if cls not in best_steps: best_steps[cls] = {} 70 | 71 | for name in numbers[cls]: 72 | if name not in best_numbers[cls]: best_numbers[cls][name] = [] 73 | 74 | data = np.array(numbers[cls][name]) 75 | steps, mAPs = data[:,0], data[:, 1:] 76 | # pick the best checkpoint based on the performance on the medium subset 77 | k = np.argmax(mAPs[:, 1]) 78 | best_numbers[cls][name] = mAPs[k, :] 79 | best_steps[cls][name] = int(steps[k]) 80 | 81 | 82 | # diff_order = [1, 0, 2] 83 | diff_order = [0, 1, 2] 84 | for cls in best_numbers: 85 | print(f"{cls:>48}", end='\t') 86 | print(f"{'step':>8}", end='\t') 87 | for i in diff_order: 88 | print(f"{difficulties[i]:<4}", end='\t') 89 | print(f"{'Avg':<4}") 90 | 91 | ordered_names = sorted(best_numbers[cls].keys(), key=lambda x: -best_numbers[cls][x][1]) 92 | for name in ordered_names: 93 | print(f"{name:>48}", end='\t') 94 | print(f"{best_steps[cls][name]:>8}", end='\t') 95 | for i in diff_order: 96 | print(f"{best_numbers[cls][name][i]:.2f}", end='\t') 97 | print(f"{np.mean(best_numbers[cls][name]):.2f}") 98 | print() 99 | -------------------------------------------------------------------------------- /second/utils/print_test.py: -------------------------------------------------------------------------------- 1 | import os 2 | import json 3 | import argparse 4 | 5 | parser = argparse.ArgumentParser() 6 | parser.add_argument('--dataset', type=str, default='nuscenes') 7 | parser.add_argument('--step', type=int, default='-1') 8 | parser.add_argument('--metric', type=str, default='mean_dist_aps') 9 | parser.add_argument('--thresh', type=str, default="") 10 | args = parser.parse_args() 11 | 12 | classes = [ 13 | 'car', 'pedestrian', 'barrier', 'traffic_cone', 'truck', 'bus', 'trailer', 'construction_vehicle', 'motorcycle', 'bicycle' 14 | ] 15 | 16 | name = "freespace" 17 | res_file = f"utils/test_results.json" 18 | if os.path.exists(res_file): 19 | with open(res_file, 'r') as f: 20 | summary = json.load(f) 21 | 22 | print(summary) 23 | 24 | # delim = '\t' 25 | delim = ' & ' 26 | metric = args.metric 27 | 28 | print('{:24}'.format(f'mAP[{args.thresh}]'), end=delim) 29 | for cls in classes: 30 | print('{:5}'.format(cls[:5]), end=delim) 31 | print('{:5}'.format('avg')) 32 | 33 | print('{:24}'.format(name), end=delim) 34 | APs = [] 35 | for cls in classes: 36 | n = summary[metric][cls] 37 | if args.thresh in n: 38 | AP = n[args.thresh] 39 | else: 40 | AP = sum(n.values())/len(n) 41 | APs.append(AP) 42 | print('{:.3f}'.format(AP), end=delim) 43 | mAP = sum(APs)/len(APs) 44 | print('{:.3f}'.format(mAP)) 45 | -------------------------------------------------------------------------------- /second/utils/print_tp_errors.py: -------------------------------------------------------------------------------- 1 | import os 2 | import json 3 | import argparse 4 | 5 | parser = argparse.ArgumentParser() 6 | parser.add_argument('--dataset', type=str, default='nuscenes') 7 | parser.add_argument('--step', type=int, default='-1') 8 | parser.add_argument('--metric', type=str, default='trans_err') 9 | parser.add_argument('--thresh', type=str, default="") 10 | args = parser.parse_args() 11 | 12 | dataset = args.dataset 13 | base_dir = f'models/{dataset}' 14 | 15 | classes = [ 16 | # # official class order 17 | # 'car', 'truck', 'bus', 'trailer', 'construction_vehicle', 'pedestrian', 'motorcycle', 'bicycle', 'traffic_cone', 'barrier' 18 | # # sorted by percentage 19 | 'car', 'pedestrian', 'barrier', 'traffic_cone', 'truck', 'bus', 'trailer', 'construction_vehicle', 'motorcycle', 'bicycle' 20 | # 'car', 'car[vehicle.moving]', 'car[vehicle.stopped]', 'car[vehicle.parked]', 21 | # 'pedestrian', 'pedestrian[pedestrian.moving]', 'pedestrian[pedestrian.sitting_lying_down]', 'pedestrian[pedestrian.standing]', 22 | # 'barrier', 23 | # 'traffic_cone', 24 | # 'truck', 'truck[vehicle.moving]', 'truck[vehicle.stopped]', 'truck[vehicle.parked]', 25 | # 'bus', 'bus[vehicle.moving]', 'bus[vehicle.stopped]', 'bus[vehicle.parked]', 26 | # 'trailer', 'trailer[vehicle.moving]', 'trailer[vehicle.stopped]', 'trailer[vehicle.parked]', 27 | # 'construction_vehicle', 'construction_vehicle[vehicle.moving]', 'construction_vehicle[vehicle.stopped]', 'construction_vehicle[vehicle.parked]', 28 | # 'motorcycle', 'motorcycle[cycle.with_rider]', 'motorcycle[cycle.without_rider]', 29 | # 'bicycle', 'bicycle[cycle.with_rider]', 'bicycle[cycle.without_rider]' 30 | ] 31 | 32 | if args.dataset == 'nuscenes': 33 | methods = [ 34 | "pointpillars", 35 | # "all_pp_mhead_vpn_swpdb_d1_ep50_ev5_vp_pp_oa_ta_learn", 36 | "all_pp_mhead_vpn_swpdb_d1_ep20_ev2_vp_pp_oa_ta_learn", 37 | "all_pp_mhead_vpn_swpdb_d1_ep50_ev5_vp_pp_oa_ta_ots", 38 | "all_pp_mhead_vpn_swpdb_d1_ep20_ev2_vp_pp_oa_ta_ots", 39 | "all_pp_mhead_vpn_nodbs_d1_ep20_ev2_vp_pp_ta_ots", 40 | "all_pp_mhead_vpn_swpdb_d1_ep20_ev2_vp_pp_oa_naive_nomask_ta_ots", 41 | "all_pp_mhead_vpn_swpdb_d1_ep20_ev2_vp_pp_oa_cull_nomask_ta_ots", 42 | "all_pp_mhead_vpn_swpdb_d1_ep20_ev2_vp_pp_oa_drill_nomask_ta_ots", 43 | "all_pp_mhead_vpn_vpnonly_swp0_swpdb_d1_ep20_ev2_vp_oa_drill", 44 | "all_pp_mhead_vpn_vpnonly_swp0_swpdb_d1_ep20_ev2_vp_oa_cull", 45 | "all_pp_mhead_d1_ep20_ev2_pp_oa_ta", 46 | "all_pp_mhead_nodbs_d1_ep20_ev2_pp_ta", 47 | "all_pp_mhead_swp0_d1_ep20_ev2_pp_oa", 48 | "all_pp_mhead_vpn_d1_ep20_ev2_vp_pp_ta_ots", 49 | "all_pp_mhead_vpn_rpnonly_swp0_nodbs_d1_ep20_ev2_pp", 50 | "all_pp_mhead_vpn_swp0_nodbs_d1_ep20_ev2_vp_pp", 51 | "all_pp_mhead_vpn_swp0_swpdb_d1_ep20_ev2_vp_pp_oa_cull", 52 | "all_pp_mhead_vpn_swp0_swpdb_d1_ep20_ev2_vp_pp_oa_double_drill", 53 | "all_pp_mhead_vpn_swp0_swpdb_d1_ep20_ev2_vp_pp_oa_drill", 54 | "all_pp_mhead_vpn_swp0_swpdb_d1_ep20_ev2_vp_pp_oa_naive", 55 | "all_pp_mhead_vpn_swpdb_d1_ep20_ev2_vp_pp_oa_cull_ta_ots", 56 | "all_pp_mhead_vpn_swpdb_d1_ep20_ev2_vp_pp_oa_drill_ta_ots", 57 | # "all_pp_mhead_vpn_swpdb_d1_ep20_ev2_vp_pp_oa_ta_learn", 58 | # "all_pp_mhead_vpn_swpdb_d1_ep20_ev2_vp_pp_oa_ta_ots", 59 | "all_pp_mhead_vpn_swpdb_d1_ep50_ev5_vp_pp_oa_ots", 60 | "all_pp_mhead_vpn_vpnonly_nodbs_d1_ep20_ev2_vp_ta_ots", 61 | "all_pp_mhead_vpn_vpnonly_swp0_nodbs_d1_ep20_ev2_vp", 62 | "all_pp_mhead_vpn_vpnonly_swp0_swpdb_d1_ep20_ev2_vp_oa_cull_nomask", 63 | "all_pp_mhead_vpn_vpnonly_swp0_swpdb_d1_ep20_ev2_vp_oa_drill_nomask", 64 | "all_pp_mhead_vpn_vpnonly_swp0_swpdb_d1_ep20_ev2_vp_oa_naive", 65 | "all_pp_mhead_vpn_vpnonly_swpdb_d1_ep20_ev2_vp_ta_ots_oa_cull", 66 | "all_pp_mhead_vpn_vpnonly_swpdb_d1_ep20_ev2_vp_ta_ots_oa_drill", 67 | "all_pp_mhead_vpn_vpnonly_swpdb_d1_ep20_ev2_vp_ta_ots_oa_naive", 68 | "all_pp_mhead_vpn_swpdb_d1_ep20_ev2_vp_pp_oa_ta_ots_sfvis", 69 | "all_pp_mhead_vpn_swpdb_d1_ep20_ev2_vp_pp_oa_ta_ots_early_fusion", 70 | ] 71 | names = [ 72 | "pp(official)", 73 | # "vp_pp_oa_ta_learn", 74 | "vp_pp_oa_ta_learn", 75 | "vp_pp_oa_ta_ots(50ep)", 76 | "vp_pp_oa_ta_ots", 77 | "vp_pp_ta_ots", 78 | "vp_pp_oa_naive_nomask_ta_ots", 79 | "vp_pp_oa_cull_nomask_ta_ots", 80 | "vp_pp_oa_drill_nomask_ta_ots", 81 | "vp_oa_drill", 82 | "vp_oa_cull", 83 | "pp_oa_ta", 84 | "pp_ta", 85 | "pp_oa", 86 | "vp_pp_ta_ots", 87 | "pp", 88 | "vp_pp", 89 | "vp_pp_oa_cull", 90 | "vp_pp_oa_double_drill", 91 | "vp_pp_oa_drill", 92 | "vp_pp_oa_naive", 93 | "vp_pp_oa_cull_ta_ots", 94 | "vp_pp_oa_drill_ta_ots", 95 | # "vp_pp_oa_ta_learn", 96 | # "vp_pp_oa_ta_ots", 97 | "vp_pp_oa_ta_ots(50ep)", 98 | "vp_ta_ots", 99 | "vp", 100 | "vp_oa_cull_nomask", 101 | "vp_oa_drill_nomask", 102 | "vp_oa_naive", 103 | "vp_ta_ots_oa_cull", 104 | "vp_ta_ots_oa_drill", 105 | "vp_ta_ots_oa_naive", 106 | "vp_pp_oa_drill_ta_ots_sfvis", 107 | "vp_pp_oa_ta_ots_early_fusion", 108 | ] 109 | 110 | cache = {} 111 | 112 | assert(len(methods) == len(names)) 113 | 114 | for method in methods: 115 | res_dir = f'{base_dir}/{method}/results' 116 | if not os.path.exists(res_dir): 117 | continue 118 | 119 | if args.step == -1: # use the final checkpoint 120 | all_steps = [] 121 | for d in os.listdir(res_dir): 122 | if os.path.exists(os.path.join(res_dir, d, "metrics_summary.json")): 123 | all_steps.append(int(d.split("_")[1])) 124 | if len(all_steps) == 0: 125 | continue 126 | step = max(all_steps) 127 | else: 128 | step = args.step 129 | print(method, step) 130 | 131 | res_file = f'{res_dir}/step_{step}/metrics_summary.json' 132 | if os.path.exists(res_file): 133 | with open(res_file, 'r') as f: 134 | summary = json.load(f) 135 | 136 | cache[method] = summary 137 | 138 | 139 | # delim = '\t' 140 | delim = ' & ' 141 | 142 | print('{:24}'.format(f'mAP[{args.thresh}]'), end=delim) 143 | for cls in classes: 144 | print('{:5}'.format(cls[:5]), end=delim) 145 | print('{:5}'.format('avg')) 146 | 147 | for name, method in zip(names, methods): 148 | if method not in cache: 149 | continue 150 | 151 | print('{:24}'.format(name), end=delim) 152 | # APs = [] 153 | errs = [] 154 | for cls in classes: 155 | # n = cache[method][metric][cls] 156 | # if args.thresh in n: 157 | # AP = n[args.thresh] 158 | # else: 159 | # AP = sum(n.values())/len(n) 160 | # APs.append(AP) 161 | # print('{:.3f}'.format(AP), end=delim) 162 | err = cache[method]["label_tp_errors"][cls][args.metric] 163 | errs.append(err) 164 | print('{:.3f}'.format(err), end=delim) 165 | # mAP = sum(APs)/len(APs) 166 | mErr = sum(errs) / len(errs) 167 | print('{:.3f}'.format(mErr)) 168 | -------------------------------------------------------------------------------- /second/utils/print_vis.py: -------------------------------------------------------------------------------- 1 | import os 2 | import json 3 | import argparse 4 | 5 | parser = argparse.ArgumentParser() 6 | parser.add_argument('--dataset', type=str, default='nuscenes') 7 | parser.add_argument('--step', type=int, default='-1') 8 | parser.add_argument('--metric', type=str, default='label_aps') 9 | parser.add_argument('--thresh', type=str, default="") 10 | args = parser.parse_args() 11 | 12 | dataset = args.dataset 13 | base_dir = f'models/{dataset}' 14 | 15 | classes = [ 16 | # # official class order 17 | # 'car', 'truck', 'bus', 'trailer', 'construction_vehicle', 'pedestrian', 'motorcycle', 'bicycle', 'traffic_cone', 'barrier' 18 | # # sorted by percentage 19 | # 'car', 'pedestrian', 'barrier', 'traffic_cone', 'truck', 'bus', 'trailer', 'construction_vehicle', 'motorcycle', 'bicycle' 20 | # 'car', 'car[vehicle.moving]', 'car[vehicle.stopped]', 'car[vehicle.parked]', 21 | # 'pedestrian', 'pedestrian[pedestrian.moving]', 'pedestrian[pedestrian.sitting_lying_down]', 'pedestrian[pedestrian.standing]', 22 | # 'barrier', 23 | # 'traffic_cone', 24 | # 'truck', 'truck[vehicle.moving]', 'truck[vehicle.stopped]', 'truck[vehicle.parked]', 25 | # 'bus', 'bus[vehicle.moving]', 'bus[vehicle.stopped]', 'bus[vehicle.parked]', 26 | # 'trailer', 'trailer[vehicle.moving]', 'trailer[vehicle.stopped]', 'trailer[vehicle.parked]', 27 | # 'construction_vehicle', 'construction_vehicle[vehicle.moving]', 'construction_vehicle[vehicle.stopped]', 'construction_vehicle[vehicle.parked]', 28 | # 'motorcycle', 'motorcycle[cycle.with_rider]', 'motorcycle[cycle.without_rider]', 29 | # 'bicycle', 'bicycle[cycle.with_rider]', 'bicycle[cycle.without_rider]' 30 | 'car[v0-40]', 31 | # 'car[v40-60]', 32 | # 'car[v60-80]', 33 | 'car[v80-100]', 34 | 'truck[v0-40]', 35 | # 'truck[v40-60]', 36 | # 'truck[v60-80]', 37 | 'truck[v80-100]', 38 | 'bus[v0-40]', 39 | # 'bus[v40-60]', 40 | # 'bus[v60-80]', 41 | 'bus[v80-100]', 42 | 'trailer[v0-40]', 43 | # 'trailer[v40-60]', 44 | # 'trailer[v60-80]', 45 | 'trailer[v80-100]', 46 | 'construction_vehicle[v0-40]', 47 | # 'construction_vehicle[v40-60]', 48 | # 'construction_vehicle[v60-80]', 49 | 'construction_vehicle[v80-100]', 50 | 'pedestrian[v0-40]', 51 | # 'pedestrian[v40-60]', 52 | # 'pedestrian[v60-80]', 53 | 'pedestrian[v80-100]', 54 | 'motorcycle[v0-40]', 55 | # 'motorcycle[v40-60]', 56 | # 'motorcycle[v60-80]', 57 | 'motorcycle[v80-100]', 58 | 'bicycle[v0-40]', 59 | # 'bicycle[v40-60]', 60 | # 'bicycle[v60-80]', 61 | 'bicycle[v80-100]', 62 | 'traffic_cone[v0-40]', 63 | # 'traffic_cone[v40-60]', 64 | # 'traffic_cone[v60-80]', 65 | 'traffic_cone[v80-100]', 66 | 'barrier[v0-40]', 67 | # 'barrier[v40-60]', 68 | # 'barrier[v60-80]', 69 | 'barrier[v80-100]', 70 | ] 71 | 72 | if args.dataset == 'nuscenes': 73 | methods = [ 74 | "pointpillars", 75 | "all_pp_mhead_vpn_swpdb_d1_ep20_ev2_vp_pp_oa_ta_learn", 76 | "all_pp_mhead_vpn_swpdb_d1_ep20_ev2_vp_pp_oa_ta_ots", 77 | "all_pp_mhead_d1_ep20_ev2_pp_oa_ta", 78 | "all_pp_mhead_vpn_swpdb_d1_ep20_ev2_vp_pp_oa_ta_ots_early_fusion", 79 | ] 80 | names = [ 81 | "pp(official)", 82 | "vp_pp_oa_ta_learn", 83 | "vp_pp_oa_ta_ots", 84 | "pp_oa_ta", 85 | "vp_pp_oa_ta_ots_EF", 86 | ] 87 | 88 | cache = {} 89 | 90 | assert(len(methods) == len(names)) 91 | 92 | for method in methods: 93 | res_dir = f'{base_dir}/{method}/results' 94 | if not os.path.exists(res_dir): 95 | continue 96 | 97 | if args.step == -1: # use the final checkpoint 98 | all_steps = [] 99 | for d in os.listdir(res_dir): 100 | if os.path.exists(os.path.join(res_dir, d, "metrics_summary.json")): 101 | all_steps.append(int(d.split("_")[1])) 102 | if len(all_steps) == 0: 103 | continue 104 | step = max(all_steps) 105 | else: 106 | step = args.step 107 | print(method, step) 108 | 109 | res_file = f'{res_dir}/step_{step}/metrics_summary.json' 110 | if os.path.exists(res_file): 111 | with open(res_file, 'r') as f: 112 | summary = json.load(f) 113 | 114 | cache[method] = summary 115 | 116 | 117 | delim = '\t' 118 | # delim = ' & ' 119 | metric = args.metric 120 | 121 | print('{:24}'.format(f'mAP[{args.thresh}]'), end=delim) 122 | for cls in classes: 123 | print('{:5}'.format(cls[:5]), end=delim) 124 | print('{:5}'.format('avg')) 125 | 126 | for name, method in zip(names, methods): 127 | if method not in cache: 128 | continue 129 | 130 | print('{:24}'.format(name), end=delim) 131 | APs = [] 132 | for cls in classes: 133 | n = cache[method][metric][cls] 134 | if args.thresh in n: 135 | AP = n[args.thresh] 136 | else: 137 | AP = sum(n.values())/len(n) 138 | APs.append(AP) 139 | print('{:.3f}'.format(AP), end=delim) 140 | mAP = sum(APs)/len(APs) 141 | print('{:.3f}'.format(mAP)) 142 | -------------------------------------------------------------------------------- /second/utils/progress_bar.py: -------------------------------------------------------------------------------- 1 | import contextlib 2 | import enum 3 | import math 4 | import time 5 | 6 | import numpy as np 7 | 8 | 9 | def progress_str(val, *args, width=20, with_ptg=True): 10 | val = max(0., min(val, 1.)) 11 | assert width > 1 12 | pos = round(width * val) - 1 13 | if with_ptg is True: 14 | log = '[{}%]'.format(max_point_str(val * 100.0, 4)) 15 | log += '[' 16 | for i in range(width): 17 | if i < pos: 18 | log += '=' 19 | elif i == pos: 20 | log += '>' 21 | else: 22 | log += '.' 23 | log += ']' 24 | for arg in args: 25 | log += '[{}]'.format(arg) 26 | return log 27 | 28 | 29 | def second_to_time_str(second, omit_hours_if_possible=True): 30 | second = int(second) 31 | m, s = divmod(second, 60) 32 | h, m = divmod(m, 60) 33 | if omit_hours_if_possible: 34 | if h == 0: 35 | return '{:02d}:{:02d}'.format(m, s) 36 | return '{:02d}:{:02d}:{:02d}'.format(h, m, s) 37 | 38 | 39 | def progress_bar_iter(task_list, width=20, with_ptg=True, step_time_average=50, name=None): 40 | total_step = len(task_list) 41 | step_times = [] 42 | start_time = 0.0 43 | name = '' if name is None else f"[{name}]" 44 | for i, task in enumerate(task_list): 45 | t = time.time() 46 | yield task 47 | step_times.append(time.time() - t) 48 | start_time += step_times[-1] 49 | start_time_str = second_to_time_str(start_time) 50 | average_step_time = np.mean(step_times[-step_time_average:]) + 1e-6 51 | speed_str = "{:.2f}it/s".format(1 / average_step_time) 52 | remain_time = (total_step - i) * average_step_time 53 | remain_time_str = second_to_time_str(remain_time) 54 | time_str = start_time_str + '>' + remain_time_str 55 | prog_str = progress_str( 56 | (i + 1) / total_step, 57 | speed_str, 58 | time_str, 59 | width=width, 60 | with_ptg=with_ptg) 61 | print(name + prog_str + ' ', end='\r') 62 | print("") 63 | 64 | 65 | list_bar = progress_bar_iter 66 | 67 | def enumerate_bar(task_list, width=20, with_ptg=True, step_time_average=50, name=None): 68 | total_step = len(task_list) 69 | step_times = [] 70 | start_time = 0.0 71 | name = '' if name is None else f"[{name}]" 72 | for i, task in enumerate(task_list): 73 | t = time.time() 74 | yield i, task 75 | step_times.append(time.time() - t) 76 | start_time += step_times[-1] 77 | start_time_str = second_to_time_str(start_time) 78 | average_step_time = np.mean(step_times[-step_time_average:]) + 1e-6 79 | speed_str = "{:.2f}it/s".format(1 / average_step_time) 80 | remain_time = (total_step - i) * average_step_time 81 | remain_time_str = second_to_time_str(remain_time) 82 | time_str = start_time_str + '>' + remain_time_str 83 | prog_str = progress_str( 84 | (i + 1) / total_step, 85 | speed_str, 86 | time_str, 87 | width=width, 88 | with_ptg=with_ptg) 89 | print(name + prog_str + ' ', end='\r') 90 | print("") 91 | 92 | 93 | def max_point_str(val, max_point): 94 | positive = bool(val >= 0.0) 95 | val = np.abs(val) 96 | if val == 0: 97 | point = 1 98 | else: 99 | point = max(int(np.log10(val)), 0) + 1 100 | fmt = "{:." + str(max(max_point - point, 0)) + "f}" 101 | if positive is True: 102 | return fmt.format(val) 103 | else: 104 | return fmt.format(-val) 105 | 106 | 107 | class Unit(enum.Enum): 108 | Iter = 'iter' 109 | Byte = 'byte' 110 | 111 | 112 | def convert_size(size_bytes): 113 | # from https://stackoverflow.com/questions/5194057/better-way-to-convert-file-sizes-in-python 114 | if size_bytes == 0: 115 | return "0B" 116 | size_name = ("B", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB") 117 | i = int(math.floor(math.log(size_bytes, 1024))) 118 | p = math.pow(1024, i) 119 | s = round(size_bytes / p, 2) 120 | return s, size_name[i] 121 | 122 | 123 | class ProgressBar: 124 | def __init__(self, 125 | width=20, 126 | with_ptg=True, 127 | step_time_average=50, 128 | speed_unit=Unit.Iter): 129 | self._width = width 130 | self._with_ptg = with_ptg 131 | self._step_time_average = step_time_average 132 | self._step_times = [] 133 | self._start_time = 0.0 134 | self._total_size = None 135 | self._speed_unit = speed_unit 136 | 137 | def start(self, total_size): 138 | self._start = True 139 | self._step_times = [] 140 | self._finished_sizes = [] 141 | self._time_elapsed = 0.0 142 | self._current_time = time.time() 143 | self._total_size = total_size 144 | self._progress = 0 145 | 146 | def print_bar(self, finished_size=1, pre_string=None, post_string=None): 147 | self._step_times.append(time.time() - self._current_time) 148 | self._finished_sizes.append(finished_size) 149 | self._time_elapsed += self._step_times[-1] 150 | start_time_str = second_to_time_str(self._time_elapsed) 151 | time_per_size = np.array(self._step_times[-self._step_time_average:]) 152 | time_per_size /= np.array( 153 | self._finished_sizes[-self._step_time_average:]) 154 | average_step_time = np.mean(time_per_size) + 1e-6 155 | if self._speed_unit == Unit.Iter: 156 | speed_str = "{:.2f}it/s".format(1 / average_step_time) 157 | elif self._speed_unit == Unit.Byte: 158 | size, size_unit = convert_size(1 / average_step_time) 159 | speed_str = "{:.2f}{}/s".format(size, size_unit) 160 | else: 161 | raise ValueError("unknown speed unit") 162 | remain_time = (self._total_size - self._progress) * average_step_time 163 | remain_time_str = second_to_time_str(remain_time) 164 | time_str = start_time_str + '>' + remain_time_str 165 | prog_str = progress_str( 166 | (self._progress + 1) / self._total_size, 167 | speed_str, 168 | time_str, 169 | width=self._width, 170 | with_ptg=self._with_ptg) 171 | self._progress += finished_size 172 | if pre_string is not None: 173 | prog_str = pre_string + prog_str 174 | if post_string is not None: 175 | prog_str += post_string 176 | if self._progress >= self._total_size: 177 | print(prog_str + ' ') 178 | else: 179 | print(prog_str + ' ', end='\r') 180 | self._current_time = time.time() 181 | -------------------------------------------------------------------------------- /second/utils/timer.py: -------------------------------------------------------------------------------- 1 | import time 2 | from contextlib import contextmanager 3 | 4 | @contextmanager 5 | def simple_timer(name=''): 6 | t = time.time() 7 | yield 8 | print(f"{name} exec time: {time.time() - t}") -------------------------------------------------------------------------------- /torchplus/__init__.py: -------------------------------------------------------------------------------- 1 | from . import train 2 | from . import nn 3 | from . import metrics 4 | from . import tools 5 | 6 | from .tools import change_default_args 7 | from torchplus.ops.array_ops import scatter_nd, gather_nd 8 | -------------------------------------------------------------------------------- /torchplus/nn/__init__.py: -------------------------------------------------------------------------------- 1 | from torchplus.nn.functional import one_hot 2 | from torchplus.nn.modules.common import Empty, Sequential 3 | from torchplus.nn.modules.normalization import GroupNorm 4 | -------------------------------------------------------------------------------- /torchplus/nn/functional.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | def one_hot(tensor, depth, dim=-1, on_value=1.0, dtype=torch.float32): 4 | tensor_onehot = torch.zeros( 5 | *list(tensor.shape), depth, dtype=dtype, device=tensor.device) 6 | tensor_onehot.scatter_(dim, tensor.unsqueeze(dim).long(), on_value) 7 | return tensor_onehot 8 | -------------------------------------------------------------------------------- /torchplus/nn/modules/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/peiyunh/wysiwyg/bb0f6ecbf3f15082fe4ec14a701ccbb14198babd/torchplus/nn/modules/__init__.py -------------------------------------------------------------------------------- /torchplus/nn/modules/common.py: -------------------------------------------------------------------------------- 1 | import sys 2 | from collections import OrderedDict 3 | 4 | import torch 5 | from torch.nn import functional as F 6 | 7 | class Empty(torch.nn.Module): 8 | def __init__(self, *args, **kwargs): 9 | super(Empty, self).__init__() 10 | 11 | def forward(self, *args, **kwargs): 12 | if len(args) == 1: 13 | return args[0] 14 | elif len(args) == 0: 15 | return None 16 | return args 17 | 18 | class Sequential(torch.nn.Module): 19 | r"""A sequential container. 20 | Modules will be added to it in the order they are passed in the constructor. 21 | Alternatively, an ordered dict of modules can also be passed in. 22 | 23 | To make it easier to understand, given is a small example:: 24 | 25 | # Example of using Sequential 26 | model = Sequential( 27 | nn.Conv2d(1,20,5), 28 | nn.ReLU(), 29 | nn.Conv2d(20,64,5), 30 | nn.ReLU() 31 | ) 32 | 33 | # Example of using Sequential with OrderedDict 34 | model = Sequential(OrderedDict([ 35 | ('conv1', nn.Conv2d(1,20,5)), 36 | ('relu1', nn.ReLU()), 37 | ('conv2', nn.Conv2d(20,64,5)), 38 | ('relu2', nn.ReLU()) 39 | ])) 40 | 41 | # Example of using Sequential with kwargs(python 3.6+) 42 | model = Sequential( 43 | conv1=nn.Conv2d(1,20,5), 44 | relu1=nn.ReLU(), 45 | conv2=nn.Conv2d(20,64,5), 46 | relu2=nn.ReLU() 47 | ) 48 | """ 49 | 50 | def __init__(self, *args, **kwargs): 51 | super(Sequential, self).__init__() 52 | if len(args) == 1 and isinstance(args[0], OrderedDict): 53 | for key, module in args[0].items(): 54 | self.add_module(key, module) 55 | else: 56 | for idx, module in enumerate(args): 57 | self.add_module(str(idx), module) 58 | for name, module in kwargs.items(): 59 | if sys.version_info < (3, 6): 60 | raise ValueError("kwargs only supported in py36+") 61 | if name in self._modules: 62 | raise ValueError("name exists.") 63 | self.add_module(name, module) 64 | 65 | def __getitem__(self, idx): 66 | if not (-len(self) <= idx < len(self)): 67 | raise IndexError('index {} is out of range'.format(idx)) 68 | if idx < 0: 69 | idx += len(self) 70 | it = iter(self._modules.values()) 71 | for i in range(idx): 72 | next(it) 73 | return next(it) 74 | 75 | def __len__(self): 76 | return len(self._modules) 77 | 78 | def add(self, module, name=None): 79 | if name is None: 80 | name = str(len(self._modules)) 81 | if name in self._modules: 82 | raise KeyError("name exists") 83 | self.add_module(name, module) 84 | 85 | def forward(self, input): 86 | # i = 0 87 | for module in self._modules.values(): 88 | # print(i) 89 | input = module(input) 90 | # i += 1 91 | return input -------------------------------------------------------------------------------- /torchplus/nn/modules/normalization.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | class GroupNorm(torch.nn.GroupNorm): 5 | def __init__(self, num_channels, num_groups, eps=1e-5, affine=True): 6 | super().__init__( 7 | num_groups=num_groups, 8 | num_channels=num_channels, 9 | eps=eps, 10 | affine=affine) 11 | -------------------------------------------------------------------------------- /torchplus/ops/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/peiyunh/wysiwyg/bb0f6ecbf3f15082fe4ec14a701ccbb14198babd/torchplus/ops/__init__.py -------------------------------------------------------------------------------- /torchplus/ops/array_ops.py: -------------------------------------------------------------------------------- 1 | import ctypes 2 | import math 3 | import time 4 | import torch 5 | 6 | 7 | def scatter_nd(indices, updates, shape): 8 | """pytorch edition of tensorflow scatter_nd. 9 | this function don't contain except handle code. so use this carefully 10 | when indice repeats, don't support repeat add which is supported 11 | in tensorflow. 12 | """ 13 | ret = torch.zeros(*shape, dtype=updates.dtype, device=updates.device) 14 | ndim = indices.shape[-1] 15 | output_shape = list(indices.shape[:-1]) + shape[indices.shape[-1]:] 16 | flatted_indices = indices.view(-1, ndim) 17 | slices = [flatted_indices[:, i] for i in range(ndim)] 18 | slices += [Ellipsis] 19 | ret[slices] = updates.view(*output_shape) 20 | return ret 21 | 22 | 23 | def gather_nd(params, indices): 24 | # this function has a limit that MAX_ADVINDEX_CALC_DIMS=5 25 | ndim = indices.shape[-1] 26 | output_shape = list(indices.shape[:-1]) + list(params.shape[indices.shape[-1]:]) 27 | flatted_indices = indices.view(-1, ndim) 28 | slices = [flatted_indices[:, i] for i in range(ndim)] 29 | slices += [Ellipsis] 30 | return params[slices].view(*output_shape) 31 | -------------------------------------------------------------------------------- /torchplus/tools.py: -------------------------------------------------------------------------------- 1 | import functools 2 | import inspect 3 | import sys 4 | from collections import OrderedDict 5 | 6 | import numba 7 | import numpy as np 8 | import torch 9 | 10 | 11 | def get_pos_to_kw_map(func): 12 | pos_to_kw = {} 13 | fsig = inspect.signature(func) 14 | pos = 0 15 | for name, info in fsig.parameters.items(): 16 | if info.kind is info.POSITIONAL_OR_KEYWORD: 17 | pos_to_kw[pos] = name 18 | pos += 1 19 | return pos_to_kw 20 | 21 | 22 | def get_kw_to_default_map(func): 23 | kw_to_default = {} 24 | fsig = inspect.signature(func) 25 | for name, info in fsig.parameters.items(): 26 | if info.kind is info.POSITIONAL_OR_KEYWORD: 27 | if info.default is not info.empty: 28 | kw_to_default[name] = info.default 29 | return kw_to_default 30 | 31 | 32 | def change_default_args(**kwargs): 33 | def layer_wrapper(layer_class): 34 | class DefaultArgLayer(layer_class): 35 | def __init__(self, *args, **kw): 36 | pos_to_kw = get_pos_to_kw_map(layer_class.__init__) 37 | kw_to_pos = {kw: pos for pos, kw in pos_to_kw.items()} 38 | for key, val in kwargs.items(): 39 | if key not in kw and kw_to_pos[key] > len(args): 40 | kw[key] = val 41 | super().__init__(*args, **kw) 42 | 43 | return DefaultArgLayer 44 | 45 | return layer_wrapper 46 | 47 | def torch_to_np_dtype(ttype): 48 | type_map = { 49 | torch.float16: np.dtype(np.float16), 50 | torch.float32: np.dtype(np.float32), 51 | torch.float16: np.dtype(np.float64), 52 | torch.int32: np.dtype(np.int32), 53 | torch.int64: np.dtype(np.int64), 54 | torch.uint8: np.dtype(np.uint8), 55 | } 56 | return type_map[ttype] 57 | -------------------------------------------------------------------------------- /torchplus/train/__init__.py: -------------------------------------------------------------------------------- 1 | from torchplus.train.checkpoint import (latest_checkpoint, restore, 2 | restore_latest_checkpoints, 3 | restore_models, save, save_models, 4 | try_restore_latest_checkpoints) 5 | from torchplus.train.common import create_folder 6 | from torchplus.train.optim import MixedPrecisionWrapper 7 | -------------------------------------------------------------------------------- /torchplus/train/checkpoint.py: -------------------------------------------------------------------------------- 1 | import json 2 | import logging 3 | import os 4 | import signal 5 | from pathlib import Path 6 | 7 | import torch 8 | 9 | 10 | class DelayedKeyboardInterrupt(object): 11 | def __enter__(self): 12 | self.signal_received = False 13 | self.old_handler = signal.signal(signal.SIGINT, self.handler) 14 | 15 | def handler(self, sig, frame): 16 | self.signal_received = (sig, frame) 17 | logging.debug('SIGINT received. Delaying KeyboardInterrupt.') 18 | 19 | def __exit__(self, type, value, traceback): 20 | signal.signal(signal.SIGINT, self.old_handler) 21 | if self.signal_received: 22 | self.old_handler(*self.signal_received) 23 | 24 | 25 | def latest_checkpoint(model_dir, model_name): 26 | """return path of latest checkpoint in a model_dir 27 | Args: 28 | model_dir: string, indicate your model dir(save ckpts, summarys, 29 | logs, etc). 30 | model_name: name of your model. we find ckpts by name 31 | Returns: 32 | path: None if isn't exist or latest checkpoint path. 33 | """ 34 | ckpt_info_path = Path(model_dir) / "checkpoints.json" 35 | if not ckpt_info_path.is_file(): 36 | return None 37 | with open(ckpt_info_path, 'r') as f: 38 | ckpt_dict = json.loads(f.read()) 39 | if model_name not in ckpt_dict['latest_ckpt']: 40 | return None 41 | latest_ckpt = ckpt_dict['latest_ckpt'][model_name] 42 | ckpt_file_name = Path(model_dir) / latest_ckpt 43 | if not ckpt_file_name.is_file(): 44 | return None 45 | 46 | return str(ckpt_file_name) 47 | 48 | def _ordered_unique(seq): 49 | seen = set() 50 | return [x for x in seq if not (x in seen or seen.add(x))] 51 | 52 | def save(model_dir, 53 | model, 54 | model_name, 55 | global_step, 56 | max_to_keep=8, 57 | keep_latest=True): 58 | """save a model into model_dir. 59 | Args: 60 | model_dir: string, indicate your model dir(save ckpts, summarys, 61 | logs, etc). 62 | model: torch.nn.Module instance. 63 | model_name: name of your model. we find ckpts by name 64 | global_step: int, indicate current global step. 65 | max_to_keep: int, maximum checkpoints to keep. 66 | keep_latest: bool, if True and there are too much ckpts, 67 | will delete oldest ckpt. else will delete ckpt which has 68 | smallest global step. 69 | Returns: 70 | path: None if isn't exist or latest checkpoint path. 71 | """ 72 | 73 | # prevent save incomplete checkpoint due to key interrupt 74 | with DelayedKeyboardInterrupt(): 75 | ckpt_info_path = Path(model_dir) / "checkpoints.json" 76 | ckpt_filename = "{}-{}.tckpt".format(model_name, global_step) 77 | ckpt_path = Path(model_dir) / ckpt_filename 78 | if not ckpt_info_path.is_file(): 79 | ckpt_info_dict = {'latest_ckpt': {}, 'all_ckpts': {}} 80 | else: 81 | with open(ckpt_info_path, 'r') as f: 82 | ckpt_info_dict = json.loads(f.read()) 83 | ckpt_info_dict['latest_ckpt'][model_name] = ckpt_filename 84 | if model_name in ckpt_info_dict['all_ckpts']: 85 | ckpt_info_dict['all_ckpts'][model_name].append(ckpt_filename) 86 | else: 87 | ckpt_info_dict['all_ckpts'][model_name] = [ckpt_filename] 88 | all_ckpts = ckpt_info_dict['all_ckpts'][model_name] 89 | 90 | torch.save(model.state_dict(), ckpt_path) 91 | # check ckpt in all_ckpts is exist, if not, delete it from all_ckpts 92 | all_ckpts_checked = [] 93 | for ckpt in all_ckpts: 94 | ckpt_path_uncheck = Path(model_dir) / ckpt 95 | if ckpt_path_uncheck.is_file(): 96 | all_ckpts_checked.append(str(ckpt_path_uncheck)) 97 | all_ckpts = all_ckpts_checked 98 | if len(all_ckpts) > max_to_keep: 99 | if keep_latest: 100 | ckpt_to_delete = all_ckpts.pop(0) 101 | else: 102 | # delete smallest step 103 | get_step = lambda name: int(name.split('.')[0].split('-')[1]) 104 | min_step = min([get_step(name) for name in all_ckpts]) 105 | ckpt_to_delete = "{}-{}.tckpt".format(model_name, min_step) 106 | all_ckpts.remove(ckpt_to_delete) 107 | os.remove(str(Path(model_dir) / ckpt_to_delete)) 108 | all_ckpts_filename = _ordered_unique([Path(f).name for f in all_ckpts]) 109 | ckpt_info_dict['all_ckpts'][model_name] = all_ckpts_filename 110 | with open(ckpt_info_path, 'w') as f: 111 | f.write(json.dumps(ckpt_info_dict, indent=2)) 112 | 113 | 114 | def restore(ckpt_path, model, map_func=None): 115 | if not Path(ckpt_path).is_file(): 116 | raise ValueError("checkpoint {} not exist.".format(ckpt_path)) 117 | state_dict = torch.load(ckpt_path) 118 | if map_func is not None: 119 | map_func(state_dict) 120 | model.load_state_dict(state_dict) 121 | print("Restoring parameters from {}".format(ckpt_path)) 122 | 123 | 124 | def _check_model_names(models): 125 | model_names = [] 126 | for model in models: 127 | if not hasattr(model, "name"): 128 | raise ValueError("models must have name attr") 129 | model_names.append(model.name) 130 | if len(model_names) != len(set(model_names)): 131 | raise ValueError("models must have unique name: {}".format( 132 | ", ".join(model_names))) 133 | 134 | 135 | def _get_name_to_model_map(models): 136 | if isinstance(models, dict): 137 | name_to_model = {name: m for name, m in models.items()} 138 | else: 139 | _check_model_names(models) 140 | name_to_model = {m.name: m for m in models} 141 | return name_to_model 142 | 143 | 144 | def try_restore_latest_checkpoints(model_dir, models, map_func=None): 145 | name_to_model = _get_name_to_model_map(models) 146 | for name, model in name_to_model.items(): 147 | latest_ckpt = latest_checkpoint(model_dir, name) 148 | if latest_ckpt is not None: 149 | restore(latest_ckpt, model, map_func) 150 | 151 | def restore_latest_checkpoints(model_dir, models, map_func=None): 152 | name_to_model = _get_name_to_model_map(models) 153 | for name, model in name_to_model.items(): 154 | latest_ckpt = latest_checkpoint(model_dir, name) 155 | if latest_ckpt is not None: 156 | restore(latest_ckpt, model, map_func) 157 | else: 158 | raise ValueError("model {}\'s ckpt isn't exist".format(name)) 159 | 160 | def restore_models(model_dir, models, global_step, map_func=None): 161 | name_to_model = _get_name_to_model_map(models) 162 | for name, model in name_to_model.items(): 163 | ckpt_filename = "{}-{}.tckpt".format(name, global_step) 164 | ckpt_path = model_dir + "/" + ckpt_filename 165 | restore(ckpt_path, model, map_func) 166 | 167 | 168 | def save_models(model_dir, 169 | models, 170 | global_step, 171 | max_to_keep=15, 172 | keep_latest=True): 173 | with DelayedKeyboardInterrupt(): 174 | name_to_model = _get_name_to_model_map(models) 175 | for name, model in name_to_model.items(): 176 | save(model_dir, model, name, global_step, max_to_keep, keep_latest) 177 | -------------------------------------------------------------------------------- /torchplus/train/common.py: -------------------------------------------------------------------------------- 1 | import datetime 2 | import os 3 | import shutil 4 | 5 | def create_folder(prefix, add_time=True, add_str=None, delete=False): 6 | additional_str = '' 7 | if delete is True: 8 | if os.path.exists(prefix): 9 | shutil.rmtree(prefix) 10 | os.makedirs(prefix) 11 | folder = prefix 12 | if add_time is True: 13 | # additional_str has a form such as '170903_220351' 14 | additional_str += datetime.datetime.now().strftime("%y%m%d_%H%M%S") 15 | if add_str is not None: 16 | folder += '/' + additional_str + '_' + add_str 17 | else: 18 | folder += '/' + additional_str 19 | if delete is True: 20 | if os.path.exists(folder): 21 | shutil.rmtree(folder) 22 | os.makedirs(folder) 23 | return folder -------------------------------------------------------------------------------- /torchplus/train/learning_schedules_fastai.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import math 3 | from functools import partial 4 | import torch 5 | 6 | 7 | class LRSchedulerStep(object): 8 | def __init__(self, fai_optimizer, total_step, lr_phases, mom_phases): 9 | self.optimizer = fai_optimizer 10 | self.total_step = total_step 11 | self.lr_phases = [] 12 | 13 | for i, (start, lambda_func) in enumerate(lr_phases): 14 | if len(self.lr_phases) != 0: 15 | assert self.lr_phases[-1][0] < int(start * total_step) 16 | if isinstance(lambda_func, str): 17 | lambda_func = eval(lambda_func) 18 | if i < len(lr_phases) - 1: 19 | self.lr_phases.append((int(start * total_step), 20 | int(lr_phases[i + 1][0] * total_step), 21 | lambda_func)) 22 | else: 23 | self.lr_phases.append((int(start * total_step), total_step, 24 | lambda_func)) 25 | assert self.lr_phases[0][0] == 0 26 | self.mom_phases = [] 27 | 28 | for i, (start, lambda_func) in enumerate(mom_phases): 29 | if len(self.mom_phases) != 0: 30 | assert self.mom_phases[-1][0] < int(start * total_step) 31 | if isinstance(lambda_func, str): 32 | lambda_func = eval(lambda_func) 33 | if i < len(mom_phases) - 1: 34 | self.mom_phases.append((int(start * total_step), 35 | int(mom_phases[i + 1][0] * total_step), 36 | lambda_func)) 37 | else: 38 | self.mom_phases.append((int(start * total_step), total_step, 39 | lambda_func)) 40 | if len(mom_phases) > 0: 41 | assert self.mom_phases[0][0] == 0 42 | 43 | def step(self, step): 44 | lrs = [] 45 | moms = [] 46 | for start, end, func in self.lr_phases: 47 | if step >= start: 48 | lrs.append(func((step - start) / (end - start))) 49 | if len(lrs) > 0: 50 | self.optimizer.lr = lrs[-1] 51 | for start, end, func in self.mom_phases: 52 | if step >= start: 53 | moms.append(func((step - start) / (end - start))) 54 | self.optimizer.mom = func((step - start) / (end - start)) 55 | if len(moms) > 0: 56 | self.optimizer.mom = moms[-1] 57 | 58 | @property 59 | def learning_rate(self): 60 | return self.optimizer.lr 61 | 62 | def annealing_cos(start, end, pct): 63 | # print(pct, start, end) 64 | "Cosine anneal from `start` to `end` as pct goes from 0.0 to 1.0." 65 | cos_out = np.cos(np.pi * pct) + 1 66 | return end + (start - end) / 2 * cos_out 67 | 68 | 69 | class OneCycle(LRSchedulerStep): 70 | def __init__(self, fai_optimizer, total_step, lr_max, moms, div_factor, 71 | pct_start): 72 | self.lr_max = lr_max 73 | self.moms = moms 74 | self.div_factor = div_factor 75 | self.pct_start = pct_start 76 | a1 = int(total_step * self.pct_start) 77 | a2 = total_step - a1 78 | low_lr = self.lr_max / self.div_factor 79 | lr_phases = ((0, partial(annealing_cos, low_lr, self.lr_max)), 80 | (self.pct_start, 81 | partial(annealing_cos, self.lr_max, low_lr / 1e4))) 82 | mom_phases = ((0, partial(annealing_cos, *self.moms)), 83 | (self.pct_start, partial(annealing_cos, 84 | *self.moms[::-1]))) 85 | fai_optimizer.lr, fai_optimizer.mom = low_lr, self.moms[0] 86 | super().__init__(fai_optimizer, total_step, lr_phases, mom_phases) 87 | 88 | 89 | class ExponentialDecay(LRSchedulerStep): 90 | def __init__(self, 91 | fai_optimizer, 92 | total_step, 93 | initial_learning_rate, 94 | decay_length, 95 | decay_factor, 96 | staircase=True): 97 | """ 98 | Args: 99 | decay_length: must in (0, 1) 100 | """ 101 | assert decay_length > 0 102 | assert decay_length < 1 103 | self._decay_steps_unified = decay_length 104 | self._decay_factor = decay_factor 105 | self._staircase = staircase 106 | step = 0 107 | stage = 1 108 | lr_phases = [] 109 | if staircase: 110 | while step <= total_step: 111 | func = lambda p, _d=initial_learning_rate * stage: _d 112 | lr_phases.append((step / total_step, func)) 113 | stage *= decay_factor 114 | step += int(decay_length * total_step) 115 | else: 116 | func = lambda p: pow(decay_factor, (p / decay_length)) 117 | lr_phases.append((0, func)) 118 | super().__init__(fai_optimizer, total_step, lr_phases, []) 119 | 120 | 121 | class ManualStepping(LRSchedulerStep): 122 | def __init__(self, fai_optimizer, total_step, boundaries, rates): 123 | assert all([b > 0 and b < 1 for b in boundaries]) 124 | assert len(boundaries) + 1 == len(rates) 125 | boundaries.insert(0, 0.0) 126 | lr_phases = [] 127 | for start, rate in zip(boundaries, rates): 128 | func = lambda p, _d=rate: _d 129 | lr_phases.append((start, func)) 130 | super().__init__(fai_optimizer, total_step, lr_phases, []) 131 | 132 | 133 | class FakeOptim: 134 | def __init__(self): 135 | self.lr = 0 136 | self.mom = 0 137 | 138 | 139 | if __name__ == "__main__": 140 | import matplotlib.pyplot as plt 141 | opt = FakeOptim() # 3e-3, wd=0.4, div_factor=10 142 | # schd = OneCycle(opt, 100, 3e-3, (0.95, 0.85), 10.0, 0.4) 143 | schd = ExponentialDecay(opt, 100, 3e-4, 0.1, 0.8, staircase=True) 144 | schd = ManualStepping(opt, 100, [0.8, 0.9], [0.001, 0.0001, 0.00005]) 145 | lrs = [] 146 | moms = [] 147 | for i in range(100): 148 | schd.step(i) 149 | lrs.append(opt.lr) 150 | moms.append(opt.mom) 151 | 152 | plt.plot(lrs) 153 | # plt.plot(moms) 154 | # plt.show() 155 | # plt.plot(moms) 156 | plt.show() 157 | -------------------------------------------------------------------------------- /torchplus/train/optim.py: -------------------------------------------------------------------------------- 1 | from collections import defaultdict, Iterable 2 | 3 | import torch 4 | from copy import deepcopy 5 | from itertools import chain 6 | from torch.autograd import Variable 7 | 8 | required = object() 9 | 10 | def param_fp32_copy(params): 11 | param_copy = [ 12 | param.clone().type(torch.cuda.FloatTensor).detach() for param in params 13 | ] 14 | for param in param_copy: 15 | param.requires_grad = True 16 | return param_copy 17 | 18 | def set_grad(params, params_with_grad, scale=1.0): 19 | for param, param_w_grad in zip(params, params_with_grad): 20 | if param.grad is None: 21 | param.grad = torch.nn.Parameter( 22 | param.data.new().resize_(*param.data.size())) 23 | grad = param_w_grad.grad.data 24 | if scale is not None: 25 | grad /= scale 26 | if torch.isnan(grad).any() or torch.isinf(grad).any(): 27 | return True # invalid grad 28 | param.grad.data.copy_(grad) 29 | return False 30 | 31 | class MixedPrecisionWrapper(object): 32 | """mixed precision optimizer wrapper. 33 | Arguments: 34 | optimizer (torch.optim.Optimizer): an instance of 35 | :class:`torch.optim.Optimizer` 36 | scale: (float): a scalar for grad scale. 37 | auto_scale: (bool): whether enable auto scale. 38 | The algorihm of auto scale is discribled in 39 | http://docs.nvidia.com/deeplearning/sdk/mixed-precision-training/index.html 40 | """ 41 | 42 | def __init__(self, 43 | optimizer, 44 | scale=None, 45 | auto_scale=True, 46 | inc_factor=2.0, 47 | dec_factor=0.5, 48 | num_iters_be_stable=500): 49 | # if not isinstance(optimizer, torch.optim.Optimizer): 50 | # raise ValueError("must provide a torch.optim.Optimizer") 51 | self.optimizer = optimizer 52 | if hasattr(self.optimizer, 'name'): 53 | self.name = self.optimizer.name # for ckpt system 54 | param_groups_copy = [] 55 | for i, group in enumerate(optimizer.param_groups): 56 | group_copy = {n: v for n, v in group.items() if n != 'params'} 57 | group_copy['params'] = param_fp32_copy(group['params']) 58 | param_groups_copy.append(group_copy) 59 | 60 | # switch param_groups, may be dangerous 61 | self.param_groups = optimizer.param_groups 62 | optimizer.param_groups = param_groups_copy 63 | self.grad_scale = scale 64 | self.auto_scale = auto_scale 65 | self.inc_factor = inc_factor 66 | self.dec_factor = dec_factor 67 | self.stable_iter_count = 0 68 | self.num_iters_be_stable = num_iters_be_stable 69 | 70 | def __getstate__(self): 71 | return self.optimizer.__getstate__() 72 | 73 | def __setstate__(self, state): 74 | return self.optimizer.__setstate__(state) 75 | 76 | def __repr__(self): 77 | return self.optimizer.__repr__() 78 | 79 | def state_dict(self): 80 | return self.optimizer.state_dict() 81 | 82 | def load_state_dict(self, state_dict): 83 | return self.optimizer.load_state_dict(state_dict) 84 | 85 | def zero_grad(self): 86 | return self.optimizer.zero_grad() 87 | 88 | def step(self, closure=None): 89 | for g, g_copy in zip(self.param_groups, self.optimizer.param_groups): 90 | invalid = set_grad(g_copy['params'], g['params'], self.grad_scale) 91 | if invalid: 92 | if self.grad_scale is None or self.auto_scale is False: 93 | raise ValueError("nan/inf detected but auto_scale disabled.") 94 | self.grad_scale *= self.dec_factor 95 | print('scale decay to {}'.format(self.grad_scale)) 96 | return 97 | if self.auto_scale is True: 98 | self.stable_iter_count += 1 99 | if self.stable_iter_count > self.num_iters_be_stable: 100 | if self.grad_scale is not None: 101 | self.grad_scale *= self.inc_factor 102 | self.stable_iter_count = 0 103 | 104 | if closure is None: 105 | self.optimizer.step() 106 | else: 107 | self.optimizer.step(closure) 108 | for g, g_copy in zip(self.param_groups, self.optimizer.param_groups): 109 | for p_copy, p in zip(g_copy['params'], g['params']): 110 | p.data.copy_(p_copy.data) 111 | --------------------------------------------------------------------------------