├── .gitignore ├── LICENSE ├── README.md ├── benchmark ├── __init__.py ├── __pycache__ │ ├── __init__.cpython-310.pyc │ ├── __init__.cpython-38.pyc │ ├── evaluate_semantic_instance.cpython-310.pyc │ ├── evaluate_semantic_instance.cpython-38.pyc │ ├── util.cpython-310.pyc │ ├── util.cpython-38.pyc │ ├── util_3d.cpython-310.pyc │ └── util_3d.cpython-38.pyc ├── evaluate_semantic_instance.py ├── util.py └── util_3d.py ├── conf ├── __init__.py ├── augmentation │ ├── albumentations_aug.yaml │ └── volumentations_aug.yaml ├── callbacks │ └── callbacks_instance_segmentation.yaml ├── config_base_instance_segmentation.yaml ├── data │ ├── collation_functions │ │ ├── voxelize_collate.yaml │ │ └── voxelize_collate_merge.yaml │ ├── data_loaders │ │ ├── simple_loader.yaml │ │ └── simple_loader_save_memory.yaml │ ├── datasets │ │ ├── matterport.yaml │ │ ├── matterport_scannet.yaml │ │ ├── rio.yaml │ │ ├── s3dis.yaml │ │ ├── scannet.yaml │ │ ├── scannet200.yaml │ │ ├── semantic_kitti.yaml │ │ └── stpls3d.yaml │ ├── indoor.yaml │ └── outdoor.yaml ├── logging │ ├── base.yaml │ ├── full.yaml │ ├── minimal.yaml │ └── offline.yaml ├── loss │ ├── cross_entropy.yaml │ ├── set_criterion.yaml │ └── set_criterion_custom_weights_1.yaml ├── matcher │ └── hungarian_matcher.yaml ├── metrics │ └── miou.yaml ├── model │ └── mask3d.yaml ├── optimizer │ ├── adamw.yaml │ └── adamw_lower.yaml ├── scheduler │ ├── exponentiallr.yaml │ ├── lambdalr.yaml │ └── onecyclelr.yaml └── trainer │ ├── trainer.yaml │ └── trainer600.yaml ├── datasets ├── __pycache__ │ ├── random_cuboid.cpython-310.pyc │ ├── random_cuboid.cpython-38.pyc │ ├── semseg.cpython-310.pyc │ ├── semseg.cpython-38.pyc │ ├── utils.cpython-310.pyc │ └── utils.cpython-38.pyc ├── outdoor_semseg.py ├── preprocessing │ ├── __pycache__ │ │ ├── base_preprocessing.cpython-310.pyc │ │ └── base_preprocessing.cpython-38.pyc │ ├── base_preprocessing.py │ ├── matterport_preprocessing.py │ ├── rio_preprocessing.py │ ├── s3dis_preprocessing.py │ ├── scannet_preprocessing.py │ ├── semantic_kitti_preprocessing.py │ └── stpls3d_preprocessing.py ├── random_cuboid.py ├── scannet200 │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-310.pyc │ │ ├── __init__.cpython-38.pyc │ │ ├── scannet200_constants.cpython-310.pyc │ │ ├── scannet200_constants.cpython-38.pyc │ │ ├── scannet200_splits.cpython-310.pyc │ │ └── scannet200_splits.cpython-38.pyc │ ├── scannet200_constants.py │ └── scannet200_splits.py ├── semseg.py └── utils.py ├── df_list ├── 10_database.yaml ├── 1_database.yaml ├── 20_database.yaml └── 5_database.yaml ├── environment.yml ├── main_instance_segmentation.py ├── models ├── __init__.py ├── criterion.py ├── mask3d.py ├── matcher.py ├── metrics │ ├── __init__.py │ ├── confusionmatrix.py │ └── metrics.py ├── misc.py ├── model.py ├── modules │ ├── 3detr_helpers.py │ ├── __init__.py │ ├── common.py │ ├── helpers_3detr.py │ ├── resnet_block.py │ ├── resnet_block.py.tmp │ └── senet_block.py ├── position_embedding.py ├── res16unet.py ├── resnet.py ├── resnet.py.tmp ├── resunet.py └── wrapper.py ├── pseudo_mask_gen ├── README.md ├── get bbox prior │ ├── correspondences.py │ ├── extractor.py │ ├── project.py │ ├── projection.py │ ├── scene_frame_merge.py │ └── scene_frame_ncut.py ├── get_dis_matrix.py ├── post_pro.py ├── project_feature.py ├── projection.py ├── topk_merge.py ├── util.py └── vccs │ ├── data_prepare_ScanNet.py │ └── initialSP_prepare_ScanNet.py ├── scenes ├── 1.txt ├── 10.txt ├── 20.txt └── 5.txt ├── scripts ├── s3dis │ ├── s3dis_from_scratch.sh │ └── s3dis_pretrained.sh ├── scannet │ ├── scannet_benchmark.sh │ ├── scannet_df.sh │ ├── scannet_pretrain_for_s3dis.sh │ └── scannet_val.sh ├── scannet200 │ ├── scannet200_benchmark.sh │ └── scannet200_val.sh └── stpls3d │ ├── merge_exports.py │ ├── stpls3d_benchmark.sh │ └── stpls3d_val.sh ├── third_party └── pointnet2 │ ├── __pycache__ │ ├── pointnet2_utils.cpython-310.pyc │ ├── pointnet2_utils.cpython-38.pyc │ ├── pytorch_utils.cpython-310.pyc │ └── pytorch_utils.cpython-38.pyc │ ├── _ext_src │ ├── include │ │ ├── ball_query.h │ │ ├── cuda_utils.h │ │ ├── group_points.h │ │ ├── interpolate.h │ │ ├── sampling.h │ │ └── utils.h │ └── src │ │ ├── ball_query.cpp │ │ ├── ball_query_gpu.cu │ │ ├── bindings.cpp │ │ ├── group_points.cpp │ │ ├── group_points_gpu.cu │ │ ├── interpolate.cpp │ │ ├── interpolate_gpu.cu │ │ ├── sampling.cpp │ │ └── sampling_gpu.cu │ ├── pointnet2_modules.py │ ├── pointnet2_test.py │ ├── pointnet2_utils.py │ ├── pytorch_utils.py │ └── setup.py ├── trainer ├── __init__.py └── trainer.py └── utils ├── __init__.py ├── gradflow_check.py ├── kfold.py ├── pc_visualizations.py ├── point_cloud_utils.py ├── pointops2 ├── __init__.py ├── functions │ ├── __init__.py │ ├── pointops.py │ ├── pointops2.py │ ├── pointops_ablation.py │ ├── test_attention_op_step1.py │ ├── test_attention_op_step1_v2.py │ ├── test_attention_op_step2.py │ ├── test_relative_pos_encoding_op_step1.py │ ├── test_relative_pos_encoding_op_step1_v2.py │ ├── test_relative_pos_encoding_op_step1_v3.py │ ├── test_relative_pos_encoding_op_step2.py │ └── test_relative_pos_encoding_op_step2_v2.py ├── setup.py └── src │ ├── __init__.py │ ├── aggregation │ ├── aggregation_cuda.cpp │ ├── aggregation_cuda_kernel.cu │ └── aggregation_cuda_kernel.h │ ├── attention │ ├── attention_cuda.cpp │ ├── attention_cuda_kernel.cu │ └── attention_cuda_kernel.h │ ├── attention_v2 │ ├── attention_cuda_kernel_v2.cu │ ├── attention_cuda_kernel_v2.h │ └── attention_cuda_v2.cpp │ ├── cuda_utils.h │ ├── grouping │ ├── grouping_cuda.cpp │ ├── grouping_cuda_kernel.cu │ └── grouping_cuda_kernel.h │ ├── interpolation │ ├── interpolation_cuda.cpp │ ├── interpolation_cuda_kernel.cu │ └── interpolation_cuda_kernel.h │ ├── knnquery │ ├── knnquery_cuda.cpp │ ├── knnquery_cuda_kernel.cu │ └── knnquery_cuda_kernel.h │ ├── pointops_api.cpp │ ├── rpe │ ├── relative_pos_encoding_cuda.cpp │ ├── relative_pos_encoding_cuda_kernel.cu │ └── relative_pos_encoding_cuda_kernel.h │ ├── rpe_v2 │ ├── relative_pos_encoding_cuda_kernel_v2.cu │ ├── relative_pos_encoding_cuda_kernel_v2.h │ └── relative_pos_encoding_cuda_v2.cpp │ ├── sampling │ ├── sampling_cuda.cpp │ ├── sampling_cuda_kernel.cu │ └── sampling_cuda_kernel.h │ └── subtraction │ ├── subtraction_cuda.cpp │ ├── subtraction_cuda_kernel.cu │ └── subtraction_cuda_kernel.h ├── utils.py └── votenet_utils ├── __pycache__ ├── box_util.cpython-310.pyc ├── box_util.cpython-38.pyc ├── eval_det.cpython-310.pyc ├── eval_det.cpython-38.pyc ├── metric_util.cpython-310.pyc └── metric_util.cpython-38.pyc ├── box_util.py ├── eval_det.py ├── metric_util.py ├── nms.py ├── nn_distance.py ├── pc_util.py ├── tf_logger.py └── tf_visualizer.py /.gitignore: -------------------------------------------------------------------------------- 1 | exp 2 | exp_pro -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2024 Cheng Shi 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 | -------------------------------------------------------------------------------- /benchmark/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/benchmark/__init__.py -------------------------------------------------------------------------------- /benchmark/__pycache__/__init__.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/benchmark/__pycache__/__init__.cpython-310.pyc -------------------------------------------------------------------------------- /benchmark/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/benchmark/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /benchmark/__pycache__/evaluate_semantic_instance.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/benchmark/__pycache__/evaluate_semantic_instance.cpython-310.pyc -------------------------------------------------------------------------------- /benchmark/__pycache__/evaluate_semantic_instance.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/benchmark/__pycache__/evaluate_semantic_instance.cpython-38.pyc -------------------------------------------------------------------------------- /benchmark/__pycache__/util.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/benchmark/__pycache__/util.cpython-310.pyc -------------------------------------------------------------------------------- /benchmark/__pycache__/util.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/benchmark/__pycache__/util.cpython-38.pyc -------------------------------------------------------------------------------- /benchmark/__pycache__/util_3d.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/benchmark/__pycache__/util_3d.cpython-310.pyc -------------------------------------------------------------------------------- /benchmark/__pycache__/util_3d.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/benchmark/__pycache__/util_3d.cpython-38.pyc -------------------------------------------------------------------------------- /benchmark/util.py: -------------------------------------------------------------------------------- 1 | import os, sys 2 | import csv 3 | try: 4 | import numpy as np 5 | except: 6 | print("Failed to import numpy package.") 7 | sys.exit(-1) 8 | try: 9 | import imageio 10 | except: 11 | print("Please install the module 'imageio' for image processing, e.g.") 12 | print("pip install imageio") 13 | sys.exit(-1) 14 | 15 | # print an error message and quit 16 | def print_error(message, user_fault=False): 17 | sys.stderr.write('ERROR: ' + str(message) + '\n') 18 | if user_fault: 19 | sys.exit(2) 20 | sys.exit(-1) 21 | 22 | 23 | # if string s represents an int 24 | def represents_int(s): 25 | try: 26 | int(s) 27 | return True 28 | except ValueError: 29 | return False 30 | 31 | 32 | def read_label_mapping(filename, label_from='raw_category', label_to='nyu40id'): 33 | assert os.path.isfile(filename) 34 | mapping = dict() 35 | with open(filename) as csvfile: 36 | reader = csv.DictReader(csvfile, delimiter='\t') 37 | for row in reader: 38 | mapping[row[label_from]] = int(row[label_to]) 39 | # if ints convert 40 | if represents_int(list(mapping.keys())[0]): 41 | mapping = {int(k):v for k,v in mapping.items()} 42 | return mapping 43 | 44 | 45 | # input: scene_types.txt or scene_types_all.txt 46 | def read_scene_types_mapping(filename, remove_spaces=True): 47 | assert os.path.isfile(filename) 48 | mapping = dict() 49 | lines = open(filename).read().splitlines() 50 | lines = [line.split('\t') for line in lines] 51 | if remove_spaces: 52 | mapping = { x[1].strip():int(x[0]) for x in lines } 53 | else: 54 | mapping = { x[1]:int(x[0]) for x in lines } 55 | return mapping 56 | 57 | 58 | # color by label 59 | def visualize_label_image(filename, image): 60 | height = image.shape[0] 61 | width = image.shape[1] 62 | vis_image = np.zeros([height, width, 3], dtype=np.uint8) 63 | color_palette = create_color_palette() 64 | for idx, color in enumerate(color_palette): 65 | vis_image[image==idx] = color 66 | imageio.imwrite(filename, vis_image) 67 | 68 | 69 | # color by different instances (mod length of color palette) 70 | def visualize_instance_image(filename, image): 71 | height = image.shape[0] 72 | width = image.shape[1] 73 | vis_image = np.zeros([height, width, 3], dtype=np.uint8) 74 | color_palette = create_color_palette() 75 | instances = np.unique(image) 76 | for idx, inst in enumerate(instances): 77 | vis_image[image==inst] = color_palette[inst%len(color_palette)] 78 | imageio.imwrite(filename, vis_image) 79 | 80 | 81 | # color palette for nyu40 labels 82 | def create_color_palette(): 83 | return [ 84 | (0, 0, 0), 85 | (174, 199, 232), # wall 86 | (152, 223, 138), # floor 87 | (31, 119, 180), # cabinet 88 | (255, 187, 120), # bed 89 | (188, 189, 34), # chair 90 | (140, 86, 75), # sofa 91 | (255, 152, 150), # table 92 | (214, 39, 40), # door 93 | (197, 176, 213), # window 94 | (148, 103, 189), # bookshelf 95 | (196, 156, 148), # picture 96 | (23, 190, 207), # counter 97 | (178, 76, 76), 98 | (247, 182, 210), # desk 99 | (66, 188, 102), 100 | (219, 219, 141), # curtain 101 | (140, 57, 197), 102 | (202, 185, 52), 103 | (51, 176, 203), 104 | (200, 54, 131), 105 | (92, 193, 61), 106 | (78, 71, 183), 107 | (172, 114, 82), 108 | (255, 127, 14), # refrigerator 109 | (91, 163, 138), 110 | (153, 98, 156), 111 | (140, 153, 101), 112 | (158, 218, 229), # shower curtain 113 | (100, 125, 154), 114 | (178, 127, 135), 115 | (120, 185, 128), 116 | (146, 111, 194), 117 | (44, 160, 44), # toilet 118 | (112, 128, 144), # sink 119 | (96, 207, 209), 120 | (227, 119, 194), # bathtub 121 | (213, 92, 176), 122 | (94, 106, 211), 123 | (82, 84, 163), # otherfurn 124 | (100, 85, 144) 125 | ] 126 | -------------------------------------------------------------------------------- /conf/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/conf/__init__.py -------------------------------------------------------------------------------- /conf/augmentation/albumentations_aug.yaml: -------------------------------------------------------------------------------- 1 | __version__: 0.4.5 2 | transform: 3 | __class_fullname__: albumentations.core.composition.Compose 4 | additional_targets: {} 5 | bbox_params: null 6 | keypoint_params: null 7 | p: 1.0 8 | transforms: 9 | - __class_fullname__: albumentations.augmentations.transforms.RandomBrightnessContrast 10 | always_apply: true 11 | brightness_by_max: true 12 | brightness_limit: 13 | - -0.2 14 | - 0.2 15 | contrast_limit: 16 | - -0.2 17 | - 0.2 18 | p: 0.5 19 | - __class_fullname__: albumentations.augmentations.transforms.RGBShift 20 | always_apply: true 21 | b_shift_limit: 22 | - -20 23 | - 20 24 | g_shift_limit: 25 | - -20 26 | - 20 27 | p: 0.5 28 | r_shift_limit: 29 | - -20 30 | - 20 31 | -------------------------------------------------------------------------------- /conf/augmentation/volumentations_aug.yaml: -------------------------------------------------------------------------------- 1 | # pi = 3.14159265358979 2 | # pi/2 = 1.57079632679489 3 | # pi/3 = 1.04719755119659 4 | # pi/6 = 0.52359877559829 5 | # pi/12 = 0.26179938779914 6 | # pi/24 = 0.13089969389957 7 | # 8 | __version__: 0.1.6 9 | transform: 10 | __class_fullname__: volumentations.core.composition.Compose 11 | additional_targets: {} 12 | p: 1.0 13 | transforms: 14 | - __class_fullname__: volumentations.augmentations.transforms.Scale3d 15 | always_apply: true 16 | p: 0.5 17 | scale_limit: 18 | - - -0.1 19 | - 0.1 20 | - - -0.1 21 | - 0.1 22 | - - -0.1 23 | - 0.1 24 | - __class_fullname__: volumentations.augmentations.transforms.RotateAroundAxis3d 25 | always_apply: true 26 | axis: 27 | - 0 28 | - 0 29 | - 1 30 | p: 0.5 31 | rotation_limit: 32 | - -3.141592653589793 33 | - 3.141592653589793 34 | - __class_fullname__: volumentations.augmentations.transforms.RotateAroundAxis3d 35 | always_apply: true 36 | axis: 37 | - 0 38 | - 1 39 | - 0 40 | p: 0.5 41 | rotation_limit: 42 | - -0.13089969389957 43 | - 0.13089969389957 44 | - __class_fullname__: volumentations.augmentations.transforms.RotateAroundAxis3d 45 | always_apply: true 46 | axis: 47 | - 1 48 | - 0 49 | - 0 50 | p: 0.5 51 | rotation_limit: 52 | - -0.13089969389957 53 | - 0.13089969389957 54 | -------------------------------------------------------------------------------- /conf/callbacks/callbacks_instance_segmentation.yaml: -------------------------------------------------------------------------------- 1 | # @package _group_ 2 | - _target_: pytorch_lightning.callbacks.ModelCheckpoint 3 | monitor: val_bed_val_ap 4 | save_last: true 5 | save_top_k: 60 6 | mode: max 7 | dirpath: ${general.save_dir} 8 | filename: "{epoch}-{val_ap:.3f}" 9 | every_n_epochs: 1 10 | 11 | - _target_: pytorch_lightning.callbacks.LearningRateMonitor 12 | -------------------------------------------------------------------------------- /conf/config_base_instance_segmentation.yaml: -------------------------------------------------------------------------------- 1 | general: 2 | train_mode: true 3 | task: "instance_segmentation" 4 | seed: null 5 | checkpoint: null 6 | backbone_checkpoint: null 7 | freeze_backbone: false # train only last layer 8 | linear_probing_backbone: false 9 | train_on_segments: true 10 | eval_on_segments: true 11 | filter_out_instances: false 12 | save_visualizations: false 13 | visualization_point_size: 20 14 | decoder_id: -1 15 | export: false 16 | use_dbscan: false 17 | ignore_class_threshold: 100 18 | project_name: scannet 19 | workspace: user_name 20 | experiment_name: validation 21 | num_targets: 19 22 | add_instance: true 23 | dbscan_eps: 0.95 24 | dbscan_min_points: 1 25 | 26 | 27 | export_threshold: 0.75 28 | 29 | reps_per_epoch: 1 30 | 31 | on_crops: false 32 | 33 | scores_threshold: 0.1 34 | iou_threshold: 1.0 35 | 36 | area: 5 37 | 38 | eval_inner_core: -1 # disabled 39 | 40 | topk_per_image: -1 41 | 42 | ignore_mask_idx: [] 43 | 44 | max_batch_size: 99999999 45 | 46 | save_dir: save_path 47 | # time/commit/md5(config)_uuid 48 | # time/experiment_id/version_uuid 49 | experiment_id: "1" # commit[:8], or unique from logger 50 | version: 1 # md5[:8] of config 51 | 52 | gpus: 1 53 | 54 | resume: false 55 | 56 | defaults: 57 | - data: indoor 58 | - data/data_loaders: simple_loader 59 | - data/datasets: scannet 60 | - data/collation_functions: voxelize_collate 61 | - logging: offline 62 | - model: mask3d 63 | - metrics: miou 64 | - optimizer: adamw 65 | - scheduler: onecyclelr 66 | - trainer: trainer600 67 | - callbacks: callbacks_instance_segmentation 68 | - matcher: hungarian_matcher 69 | - loss: set_criterion 70 | 71 | hydra: 72 | run: 73 | dir: hydra_logs/${now:%Y-%m-%d}/${now:%H-%M-%S} 74 | sweep: 75 | dir: hydra_logs/${now:%Y-%m-%d}/${now:%H-%M-%S} 76 | # dir: ${general.save_dir} 77 | subdir: ${hydra.job.num}_${hydra.job.id} 78 | -------------------------------------------------------------------------------- /conf/data/collation_functions/voxelize_collate.yaml: -------------------------------------------------------------------------------- 1 | # @package data 2 | 3 | train_collation: 4 | _target_: datasets.utils.VoxelizeCollate 5 | ignore_label: ${data.ignore_label} 6 | voxel_size: ${data.voxel_size} 7 | mode: ${data.train_mode} 8 | small_crops: false 9 | very_small_crops: false 10 | batch_instance: false 11 | probing: ${general.linear_probing_backbone} 12 | task: ${general.task} 13 | ignore_class_threshold: ${general.ignore_class_threshold} 14 | filter_out_classes: ${data.train_dataset.filter_out_classes} 15 | label_offset: ${data.train_dataset.label_offset} 16 | num_queries: ${model.num_queries} 17 | 18 | validation_collation: 19 | _target_: datasets.utils.VoxelizeCollate 20 | ignore_label: ${data.ignore_label} 21 | voxel_size: ${data.voxel_size} 22 | mode: ${data.validation_mode} 23 | batch_instance: false 24 | probing: ${general.linear_probing_backbone} 25 | task: ${general.task} 26 | ignore_class_threshold: ${general.ignore_class_threshold} 27 | filter_out_classes: ${data.validation_dataset.filter_out_classes} 28 | label_offset: ${data.validation_dataset.label_offset} 29 | num_queries: ${model.num_queries} 30 | 31 | test_collation: 32 | _target_: datasets.utils.VoxelizeCollate 33 | ignore_label: ${data.ignore_label} 34 | voxel_size: ${data.voxel_size} 35 | mode: ${data.test_mode} 36 | batch_instance: false 37 | probing: ${general.linear_probing_backbone} 38 | task: ${general.task} 39 | ignore_class_threshold: ${general.ignore_class_threshold} 40 | filter_out_classes: ${data.test_dataset.filter_out_classes} 41 | label_offset: ${data.test_dataset.label_offset} 42 | num_queries: ${model.num_queries} -------------------------------------------------------------------------------- /conf/data/collation_functions/voxelize_collate_merge.yaml: -------------------------------------------------------------------------------- 1 | # @package data 2 | 3 | train_collation: 4 | _target_: datasets.utils.VoxelizeCollateMerge 5 | ignore_label: ${data.ignore_label} 6 | voxel_size: ${data.voxel_size} 7 | mode: ${data.train_mode} 8 | small_crops: false 9 | very_small_crops: false 10 | scenes: 2 11 | batch_instance: false 12 | make_one_pc_noise: false 13 | place_nearby: false 14 | place_far: false 15 | proba: 1 16 | probing: ${general.linear_probing_backbone} 17 | include_ignore: ${general.include_ignore} 18 | task: ${general.task} 19 | 20 | validation_collation: 21 | _target_: datasets.utils.VoxelizeCollate 22 | ignore_label: ${data.ignore_label} 23 | voxel_size: ${data.voxel_size} 24 | mode: ${data.validation_mode} 25 | probing: ${general.linear_probing_backbone} 26 | include_ignore: ${general.include_ignore} 27 | task: ${general.task} 28 | 29 | test_collation: 30 | _target_: datasets.utils.VoxelizeCollate 31 | ignore_label: ${data.ignore_label} 32 | voxel_size: ${data.voxel_size} 33 | mode: ${data.test_mode} 34 | probing: ${general.linear_probing_backbone} 35 | include_ignore: ${general.include_ignore} 36 | task: ${general.task} 37 | -------------------------------------------------------------------------------- /conf/data/data_loaders/simple_loader.yaml: -------------------------------------------------------------------------------- 1 | # @package data 2 | 3 | train_dataloader: 4 | _target_: torch.utils.data.DataLoader 5 | shuffle: False 6 | pin_memory: ${data.pin_memory} 7 | num_workers: ${data.num_workers} 8 | batch_size: ${data.batch_size} 9 | # persistent_workers: True 10 | 11 | validation_dataloader: 12 | _target_: torch.utils.data.DataLoader 13 | shuffle: false 14 | pin_memory: ${data.pin_memory} 15 | num_workers: ${data.num_workers} 16 | batch_size: ${data.test_batch_size} 17 | # persistent_workers: True 18 | test_dataloader: 19 | _target_: torch.utils.data.DataLoader 20 | shuffle: false 21 | pin_memory: ${data.pin_memory} 22 | num_workers: ${data.num_workers} 23 | batch_size: ${data.test_batch_size} 24 | # persistent_workers: True -------------------------------------------------------------------------------- /conf/data/data_loaders/simple_loader_save_memory.yaml: -------------------------------------------------------------------------------- 1 | # @package data 2 | 3 | train_dataloader: 4 | _target_: torch.utils.data.DataLoader 5 | shuffle: true 6 | pin_memory: ${data.pin_memory} 7 | num_workers: ${data.num_workers} 8 | batch_size: ${data.batch_size} 9 | 10 | validation_dataloader: 11 | _target_: torch.utils.data.DataLoader 12 | shuffle: false 13 | pin_memory: ${data.pin_memory} 14 | num_workers: 1 15 | batch_size: ${data.test_batch_size} 16 | 17 | test_dataloader: 18 | _target_: torch.utils.data.DataLoader 19 | shuffle: false 20 | pin_memory: ${data.pin_memory} 21 | num_workers: 1 22 | batch_size: ${data.test_batch_size} 23 | -------------------------------------------------------------------------------- /conf/data/datasets/matterport.yaml: -------------------------------------------------------------------------------- 1 | # @package data 2 | train_dataset: 3 | _target_: mix3d.datasets.semseg.SemanticSegmentationDataset 4 | data_dir: /inspurfs/group/yangsb/yangbin/data/processed/matterport 5 | image_augmentations_path: mix3d/conf/augmentation/albumentations_aug.yaml 6 | volume_augmentations_path: mix3d/conf/augmentation/volumentations_aug.yaml 7 | label_db_filepath: /inspurfs/group/yangsb/yangbin/data/processed/matterport/label_database.yaml 8 | color_mean_std: /inspurfs/group/yangsb/yangbin/data/processed/scannet/color_mean_std.yaml 9 | data_percent: 1.0 10 | mode: ${data.train_mode} 11 | ignore_label: ${data.ignore_label} 12 | num_labels: ${data.num_labels} 13 | add_raw_coordinates: ${data.add_raw_coordinates} 14 | add_colors: ${data.add_colors} 15 | add_normals: ${data.add_normals} 16 | add_instance: ${data.add_instance} 17 | 18 | validation_dataset: 19 | _target_: mix3d.datasets.semseg.SemanticSegmentationDataset 20 | data_dir: /inspurfs/group/yangsb/yangbin/data/processed/scannet 21 | image_augmentations_path: null 22 | volume_augmentations_path: null 23 | label_db_filepath: /inspurfs/group/yangsb/yangbin/data/processed/matterport/label_database.yaml 24 | color_mean_std: /inspurfs/group/yangsb/yangbin/data/processed/scannet/color_mean_std.yaml 25 | data_percent: 1.0 26 | mode: ${data.validation_mode} 27 | ignore_label: ${data.ignore_label} 28 | num_labels: ${data.num_labels} 29 | add_raw_coordinates: ${data.add_raw_coordinates} 30 | add_colors: ${data.add_colors} 31 | add_normals: ${data.add_normals} 32 | add_instance: ${data.add_instance} 33 | 34 | test_dataset: 35 | _target_: mix3d.datasets.semseg.SemanticSegmentationDataset 36 | data_dir: /inspurfs/group/yangsb/yangbin/data/processed/matterport 37 | image_augmentations_path: null 38 | volume_augmentations_path: null 39 | label_db_filepath: /inspurfs/group/yangsb/yangbin/data/processed/matterport/label_database.yaml 40 | color_mean_std: /inspurfs/group/yangsb/yangbin/data/processed/scannet/color_mean_std.yaml 41 | data_percent: 1.0 42 | mode: ${data.test_mode} 43 | ignore_label: ${data.ignore_label} 44 | num_labels: ${data.num_labels} 45 | add_raw_coordinates: ${data.add_raw_coordinates} 46 | add_colors: ${data.add_colors} 47 | add_normals: ${data.add_normals} 48 | add_instance: ${data.add_instance} 49 | -------------------------------------------------------------------------------- /conf/data/datasets/matterport_scannet.yaml: -------------------------------------------------------------------------------- 1 | # @package data 2 | train_dataset: 3 | _target_: mix3d.datasets.semseg.SemanticSegmentationDataset 4 | data_dir: 5 | - /inspurfs/group/yangsb/yangbin/data/processed/scannet 6 | - /inspurfs/group/yangsb/yangbin/data/processed/matterport 7 | image_augmentations_path: mix3d/conf/augmentation/albumentations_aug.yaml 8 | volume_augmentations_path: mix3d/conf/augmentation/volumentations_aug.yaml 9 | label_db_filepath: /inspurfs/group/yangsb/yangbin/data/processed/scannet/label_database.yaml 10 | color_mean_std: /inspurfs/group/yangsb/yangbin/data/processed/scannet/color_mean_std.yaml 11 | data_percent: 1.0 12 | mode: ${data.train_mode} 13 | ignore_label: ${data.ignore_label} 14 | num_labels: ${data.num_labels} 15 | add_raw_coordinates: ${data.add_raw_coordinates} 16 | add_colors: ${data.add_colors} 17 | add_normals: ${data.add_normals} 18 | add_instance: ${data.add_instance} 19 | 20 | validation_dataset: 21 | _target_: mix3d.datasets.semseg.SemanticSegmentationDataset 22 | data_dir: /inspurfs/group/yangsb/yangbin/data/processed/scannet 23 | image_augmentations_path: null 24 | volume_augmentations_path: null 25 | label_db_filepath: /inspurfs/group/yangsb/yangbin/data/processed/scannet/label_database.yaml 26 | color_mean_std: /inspurfs/group/yangsb/yangbin/data/processed/scannet/color_mean_std.yaml 27 | data_percent: 1.0 28 | mode: ${data.validation_mode} 29 | ignore_label: ${data.ignore_label} 30 | num_labels: ${data.num_labels} 31 | add_raw_coordinates: ${data.add_raw_coordinates} 32 | add_colors: ${data.add_colors} 33 | add_normals: ${data.add_normals} 34 | add_instance: ${data.add_instance} 35 | 36 | test_dataset: 37 | _target_: mix3d.datasets.semseg.SemanticSegmentationDataset 38 | data_dir: /inspurfs/group/yangsb/yangbin/data/processed/scannet 39 | image_augmentations_path: null 40 | volume_augmentations_path: null 41 | label_db_filepath: /inspurfs/group/yangsb/yangbin/data/processed/scannet/label_database.yaml 42 | color_mean_std: /inspurfs/group/yangsb/yangbin/data/processed/scannet/color_mean_std.yaml 43 | data_percent: 1.0 44 | mode: ${data.test_mode} 45 | ignore_label: ${data.ignore_label} 46 | num_labels: ${data.num_labels} 47 | add_raw_coordinates: ${data.add_raw_coordinates} 48 | add_colors: ${data.add_colors} 49 | add_normals: ${data.add_normals} 50 | add_instance: ${data.add_instance} 51 | -------------------------------------------------------------------------------- /conf/data/datasets/rio.yaml: -------------------------------------------------------------------------------- 1 | # @package data 2 | train_dataset: 3 | _target_: mix3d.datasets.semseg.SemanticSegmentationDataset 4 | data_dir: data/processed/rio 5 | image_augmentations_path: mix3d/conf/augmentation/albumentations_aug.yaml 6 | volume_augmentations_path: mix3d/conf/augmentation/volumentations_aug.yaml 7 | label_db_filepath: data/processed/scannet/label_database.yaml 8 | color_mean_std: data/processed/scannet/color_mean_std.yaml 9 | data_percent: 1.0 10 | mode: ${data.train_mode} 11 | ignore_label: ${data.ignore_label} 12 | num_labels: ${data.num_labels} 13 | add_raw_coordinates: ${data.add_raw_coordinates} 14 | add_colors: ${data.add_colors} 15 | add_normals: ${data.add_normals} 16 | add_instance: ${data.add_instance} 17 | 18 | validation_dataset: 19 | _target_: mix3d.datasets.semseg.SemanticSegmentationDataset 20 | data_dir: data/processed/rio 21 | image_augmentations_path: null 22 | volume_augmentations_path: null 23 | label_db_filepath: data/processed/scannet/label_database.yaml 24 | color_mean_std: data/processed/scannet/color_mean_std.yaml 25 | data_percent: 1.0 26 | mode: ${data.validation_mode} 27 | ignore_label: ${data.ignore_label} 28 | num_labels: ${data.num_labels} 29 | add_raw_coordinates: ${data.add_raw_coordinates} 30 | add_colors: ${data.add_colors} 31 | add_normals: ${data.add_normals} 32 | add_instance: ${data.add_instance} 33 | 34 | test_dataset: 35 | _target_: mix3d.datasets.semseg.SemanticSegmentationDataset 36 | data_dir: data/processed/rio 37 | image_augmentations_path: null 38 | volume_augmentations_path: null 39 | label_db_filepath: data/processed/scannet/label_database.yaml 40 | color_mean_std: data/processed/scannet/color_mean_std.yaml 41 | data_percent: 1.0 42 | mode: ${data.test_mode} 43 | ignore_label: ${data.ignore_label} 44 | num_labels: ${data.num_labels} 45 | add_raw_coordinates: ${data.add_raw_coordinates} 46 | add_colors: ${data.add_colors} 47 | add_normals: ${data.add_normals} 48 | add_instance: ${data.add_instance} 49 | -------------------------------------------------------------------------------- /conf/data/datasets/s3dis.yaml: -------------------------------------------------------------------------------- 1 | # @package data 2 | train_dataset: 3 | _target_: datasets.semseg.SemanticSegmentationDataset 4 | dataset_name: "s3dis" 5 | data_dir: data/processed/s3dis 6 | image_augmentations_path: conf/augmentation/albumentations_aug.yaml 7 | volume_augmentations_path: conf/augmentation/volumentations_aug.yaml 8 | label_db_filepath: data/processed/s3dis/label_database.yaml 9 | color_mean_std: data/processed/s3dis/color_mean_std.yaml 10 | data_percent: 1.0 11 | mode: ${data.train_mode} 12 | ignore_label: ${data.ignore_label} 13 | num_labels: ${data.num_labels} 14 | add_raw_coordinates: ${data.add_raw_coordinates} 15 | add_colors: ${data.add_colors} 16 | add_normals: ${data.add_normals} 17 | add_instance: ${data.add_instance} 18 | cache_data: ${data.cache_data} 19 | # different augs experiments 20 | instance_oversampling: 0.0 21 | place_around_existing: False 22 | point_per_cut: 0 23 | max_cut_region: 0 24 | flip_in_center: false 25 | noise_rate: 0 26 | resample_points: 0 27 | cropping: ${data.cropping} 28 | cropping_args: ${data.cropping_args} 29 | is_tta: false 30 | crop_min_size: ${data.crop_min_size} 31 | crop_length: ${data.crop_length} 32 | cropping_v1: ${data.cropping_v1} 33 | area: ${general.area} 34 | filter_out_classes: [] 35 | label_offset: 0 36 | 37 | validation_dataset: 38 | _target_: datasets.semseg.SemanticSegmentationDataset 39 | dataset_name: "s3dis" 40 | data_dir: data/processed/s3dis 41 | image_augmentations_path: null 42 | volume_augmentations_path: null 43 | label_db_filepath: data/processed/s3dis/label_database.yaml 44 | color_mean_std: data/processed/s3dis/color_mean_std.yaml 45 | data_percent: 1.0 46 | mode: ${data.validation_mode} 47 | ignore_label: ${data.ignore_label} 48 | num_labels: ${data.num_labels} 49 | add_raw_coordinates: ${data.add_raw_coordinates} 50 | add_colors: ${data.add_colors} 51 | add_normals: ${data.add_normals} 52 | add_instance: ${data.add_instance} 53 | cache_data: ${data.cache_data} 54 | cropping: false 55 | is_tta: false 56 | crop_min_size: ${data.crop_min_size} 57 | crop_length: ${data.crop_length} 58 | cropping_v1: ${data.cropping_v1} 59 | area: ${general.area} 60 | filter_out_classes: [] 61 | label_offset: 0 62 | 63 | test_dataset: 64 | _target_: datasets.semseg.SemanticSegmentationDataset 65 | dataset_name: "s3dis" 66 | data_dir: data/processed/s3dis 67 | image_augmentations_path: null 68 | volume_augmentations_path: null 69 | label_db_filepath: data/processed/s3dis/label_database.yaml 70 | color_mean_std: data/processed/s3dis/color_mean_std.yaml 71 | data_percent: 1.0 72 | mode: ${data.test_mode} 73 | ignore_label: ${data.ignore_label} 74 | num_labels: ${data.num_labels} 75 | add_raw_coordinates: ${data.add_raw_coordinates} 76 | add_colors: ${data.add_colors} 77 | add_normals: ${data.add_normals} 78 | add_instance: ${data.add_instance} 79 | cache_data: ${data.cache_data} 80 | cropping: false 81 | is_tta: false 82 | crop_min_size: ${data.crop_min_size} 83 | crop_length: ${data.crop_length} 84 | cropping_v1: ${data.cropping_v1} 85 | area: ${general.area} 86 | filter_out_classes: [] 87 | label_offset: 0 88 | -------------------------------------------------------------------------------- /conf/data/datasets/scannet.yaml: -------------------------------------------------------------------------------- 1 | # @package data 2 | train_dataset: 3 | _target_: datasets.semseg.SemanticSegmentationDataset 4 | dataset_name: "scannet" 5 | data_dir: data/processed/scannet 6 | image_augmentations_path: conf/augmentation/albumentations_aug.yaml 7 | volume_augmentations_path: conf/augmentation/volumentations_aug.yaml 8 | label_db_filepath: data/processed/scannet/label_database.yaml 9 | color_mean_std: data/processed/scannet/color_mean_std.yaml 10 | data_percent: 1.0 11 | mode: ${data.train_mode} 12 | ignore_label: ${data.ignore_label} 13 | num_labels: ${data.num_labels} 14 | add_raw_coordinates: ${data.add_raw_coordinates} 15 | add_colors: ${data.add_colors} 16 | add_normals: ${data.add_normals} 17 | add_instance: ${data.add_instance} 18 | # different augs experiments 19 | instance_oversampling: 0.0 20 | place_around_existing: false 21 | point_per_cut: 0 22 | max_cut_region: 0 23 | flip_in_center: false 24 | noise_rate: 0 25 | resample_points: 0 26 | add_unlabeled_pc: false 27 | cropping: ${data.cropping} 28 | cropping_args: ${data.cropping_args} 29 | is_tta: false 30 | crop_min_size: ${data.crop_min_size} 31 | crop_length: ${data.crop_length} 32 | filter_out_classes: [0, 1] 33 | label_offset: 2 34 | layer0_path: ${data.layer0_path} 35 | part_path: ${data.part_path} 36 | objrct_path: ${data.objrct_path} 37 | df: ${data.df} 38 | 39 | validation_dataset: 40 | _target_: datasets.semseg.SemanticSegmentationDataset 41 | dataset_name: "scannet" 42 | data_dir: data/processed/scannet 43 | image_augmentations_path: null 44 | volume_augmentations_path: null 45 | label_db_filepath: data/processed/scannet/label_database.yaml 46 | color_mean_std: data/processed/scannet/color_mean_std.yaml 47 | data_percent: 1.0 48 | mode: ${data.validation_mode} 49 | ignore_label: ${data.ignore_label} 50 | num_labels: ${data.num_labels} 51 | add_raw_coordinates: ${data.add_raw_coordinates} 52 | add_colors: ${data.add_colors} 53 | add_normals: ${data.add_normals} 54 | add_instance: ${data.add_instance} 55 | cropping: false 56 | is_tta: false 57 | crop_min_size: ${data.crop_min_size} 58 | crop_length: ${data.crop_length} 59 | filter_out_classes: [0, 1] 60 | label_offset: 2 61 | layer0_path: ${data.layer0_path} 62 | part_path: ${data.part_path} 63 | objrct_path: ${data.objrct_path} 64 | df: ${data.df} 65 | 66 | test_dataset: 67 | _target_: datasets.semseg.SemanticSegmentationDataset 68 | dataset_name: "scannet" 69 | data_dir: data/processed/scannet 70 | image_augmentations_path: null 71 | volume_augmentations_path: null 72 | label_db_filepath: data/processed/scannet/label_database.yaml 73 | color_mean_std: data/processed/scannet/color_mean_std.yaml 74 | data_percent: 1.0 75 | mode: ${data.test_mode} 76 | ignore_label: ${data.ignore_label} 77 | num_labels: ${data.num_labels} 78 | add_raw_coordinates: ${data.add_raw_coordinates} 79 | add_colors: ${data.add_colors} 80 | add_normals: ${data.add_normals} 81 | add_instance: ${data.add_instance} 82 | cropping: false 83 | is_tta: false 84 | crop_min_size: ${data.crop_min_size} 85 | crop_length: ${data.crop_length} 86 | filter_out_classes: [0, 1] 87 | label_offset: 2 88 | layer0_path: ${data.layer0_path} 89 | part_path: ${data.part_path} 90 | objrct_path: ${data.objrct_path} 91 | df: ${data.df} 92 | -------------------------------------------------------------------------------- /conf/data/datasets/scannet200.yaml: -------------------------------------------------------------------------------- 1 | # @package data 2 | train_dataset: 3 | _target_: datasets.semseg.SemanticSegmentationDataset 4 | dataset_name: "scannet200" 5 | data_dir: data/processed/scannet200 6 | image_augmentations_path: conf/augmentation/albumentations_aug.yaml 7 | volume_augmentations_path: conf/augmentation/volumentations_aug.yaml 8 | label_db_filepath: data/processed/scannet200/label_database.yaml 9 | color_mean_std: data/processed/scannet200/color_mean_std.yaml 10 | data_percent: 1.0 11 | mode: ${data.train_mode} 12 | ignore_label: ${data.ignore_label} 13 | num_labels: ${data.num_labels} 14 | add_raw_coordinates: ${data.add_raw_coordinates} 15 | add_colors: ${data.add_colors} 16 | add_normals: ${data.add_normals} 17 | add_instance: ${data.add_instance} 18 | # different augs experiments 19 | instance_oversampling: 0.0 20 | place_around_existing: false 21 | point_per_cut: 0 22 | max_cut_region: 0 23 | flip_in_center: false 24 | noise_rate: 0 25 | resample_points: 0 26 | add_unlabeled_pc: false 27 | cropping: ${data.cropping} 28 | cropping_args: ${data.cropping_args} 29 | is_tta: false 30 | crop_min_size: ${data.crop_min_size} 31 | crop_length: ${data.crop_length} 32 | filter_out_classes: [0, 2] 33 | label_offset: 2 34 | 35 | validation_dataset: 36 | _target_: datasets.semseg.SemanticSegmentationDataset 37 | dataset_name: "scannet200" 38 | data_dir: data/processed/scannet200 39 | image_augmentations_path: null 40 | volume_augmentations_path: null 41 | label_db_filepath: data/processed/scannet200/label_database.yaml 42 | color_mean_std: data/processed/scannet200/color_mean_std.yaml 43 | data_percent: 1.0 44 | mode: ${data.validation_mode} 45 | ignore_label: ${data.ignore_label} 46 | num_labels: ${data.num_labels} 47 | add_raw_coordinates: ${data.add_raw_coordinates} 48 | add_colors: ${data.add_colors} 49 | add_normals: ${data.add_normals} 50 | add_instance: ${data.add_instance} 51 | cropping: false 52 | is_tta: false 53 | crop_min_size: ${data.crop_min_size} 54 | crop_length: ${data.crop_length} 55 | filter_out_classes: [0, 2] 56 | label_offset: 2 57 | 58 | test_dataset: 59 | _target_: datasets.semseg.SemanticSegmentationDataset 60 | dataset_name: "scannet200" 61 | data_dir: data/processed/scannet200 62 | image_augmentations_path: null 63 | volume_augmentations_path: null 64 | label_db_filepath: data/processed/scannet200/label_database.yaml 65 | color_mean_std: data/processed/scannet200/color_mean_std.yaml 66 | data_percent: 1.0 67 | mode: ${data.test_mode} 68 | ignore_label: ${data.ignore_label} 69 | num_labels: ${data.num_labels} 70 | add_raw_coordinates: ${data.add_raw_coordinates} 71 | add_colors: ${data.add_colors} 72 | add_normals: ${data.add_normals} 73 | add_instance: ${data.add_instance} 74 | cropping: false 75 | is_tta: false 76 | crop_min_size: ${data.crop_min_size} 77 | crop_length: ${data.crop_length} 78 | filter_out_classes: [0, 2] 79 | label_offset: 2 80 | -------------------------------------------------------------------------------- /conf/data/datasets/semantic_kitti.yaml: -------------------------------------------------------------------------------- 1 | # @package data 2 | train_dataset: 3 | _target_: mix3d.datasets.outdoor_semseg.LidarDataset 4 | data_dir: data/processed/semantic_kitti 5 | label_db_filepath: data/processed/semantic_kitti/label_database.yaml 6 | mode: ${data.train_mode} 7 | add_reflection: ${data.add_reflection} 8 | add_distance: ${data.add_distance} 9 | add_instance: ${data.add_instance} 10 | num_labels: ${data.num_labels} 11 | sweep: ${data.sweep} 12 | data_percent: 1.0 13 | ignore_label: ${data.ignore_label} 14 | volume_augmentations_path: mix3d/conf/augmentation/volumentations_aug.yaml 15 | 16 | validation_dataset: 17 | _target_: mix3d.datasets.outdoor_semseg.LidarDataset 18 | data_dir: data/processed/semantic_kitti 19 | label_db_filepath: data/processed/semantic_kitti/label_database.yaml 20 | mode: ${data.validation_mode} 21 | add_reflection: ${data.add_reflection} 22 | add_distance: ${data.add_distance} 23 | add_instance: ${data.add_instance} 24 | num_labels: ${data.num_labels} 25 | sweep: ${data.sweep} 26 | data_percent: 1.0 27 | ignore_label: ${data.ignore_label} 28 | volume_augmentations_path: null 29 | 30 | test_dataset: 31 | _target_: mix3d.datasets.outdoor_semseg.LidarDataset 32 | data_dir: data/processed/semantic_kitti 33 | label_db_filepath: data/processed/semantic_kitti/label_database.yaml 34 | mode: ${data.test_mode} 35 | add_reflection: ${data.add_reflection} 36 | add_distance: ${data.add_distance} 37 | add_instance: ${data.add_instance} 38 | num_labels: ${data.num_labels} 39 | sweep: ${data.sweep} 40 | data_percent: 1.0 41 | ignore_label: ${data.ignore_label} 42 | volume_augmentations_path: null 43 | -------------------------------------------------------------------------------- /conf/data/datasets/stpls3d.yaml: -------------------------------------------------------------------------------- 1 | # @package data 2 | train_dataset: 3 | _target_: datasets.semseg.SemanticSegmentationDataset 4 | dataset_name: "stpls3d" 5 | data_dir: data/processed/stpls3d 6 | image_augmentations_path: conf/augmentation/albumentations_aug.yaml 7 | volume_augmentations_path: conf/augmentation/volumentations_aug.yaml 8 | label_db_filepath: data/processed/stpls3d/label_database.yaml 9 | color_mean_std: data/processed/stpls3d/color_mean_std.yaml 10 | data_percent: 1.0 11 | mode: ${data.train_mode} 12 | ignore_label: ${data.ignore_label} 13 | num_labels: ${data.num_labels} 14 | add_raw_coordinates: ${data.add_raw_coordinates} 15 | add_colors: ${data.add_colors} 16 | add_normals: ${data.add_normals} 17 | add_instance: ${data.add_instance} 18 | cache_data: ${data.cache_data} 19 | # different augs experiments 20 | instance_oversampling: 0.0 21 | place_around_existing: False 22 | point_per_cut: 0 23 | max_cut_region: 0 24 | flip_in_center: false 25 | noise_rate: 0 26 | resample_points: 0 27 | cropping: ${data.cropping} 28 | cropping_args: ${data.cropping_args} 29 | is_tta: false 30 | crop_min_size: ${data.crop_min_size} 31 | crop_length: ${data.crop_length} 32 | cropping_v1: ${data.cropping_v1} 33 | area: ${general.area} 34 | reps_per_epoch: ${general.reps_per_epoch} 35 | eval_inner_core: ${general.eval_inner_core} 36 | filter_out_classes: [0] 37 | label_offset: 1 38 | is_elastic_distortion: true 39 | color_drop: 0.0 40 | 41 | validation_dataset: 42 | _target_: datasets.semseg.SemanticSegmentationDataset 43 | dataset_name: "stpls3d" 44 | data_dir: data/processed/stpls3d 45 | image_augmentations_path: null 46 | volume_augmentations_path: null 47 | label_db_filepath: data/processed/stpls3d/label_database.yaml 48 | color_mean_std: data/processed/stpls3d/color_mean_std.yaml 49 | data_percent: 1.0 50 | mode: ${data.validation_mode} 51 | ignore_label: ${data.ignore_label} 52 | num_labels: ${data.num_labels} 53 | add_raw_coordinates: ${data.add_raw_coordinates} 54 | add_colors: ${data.add_colors} 55 | add_normals: ${data.add_normals} 56 | add_instance: ${data.add_instance} 57 | cache_data: ${data.cache_data} 58 | cropping: false 59 | is_tta: false 60 | crop_min_size: ${data.crop_min_size} 61 | crop_length: ${data.crop_length} 62 | cropping_v1: ${data.cropping_v1} 63 | area: ${general.area} 64 | on_crops: ${general.on_crops} 65 | eval_inner_core: ${general.eval_inner_core} 66 | filter_out_classes: [0] 67 | label_offset: 1 68 | 69 | test_dataset: 70 | _target_: datasets.semseg.SemanticSegmentationDataset 71 | dataset_name: "stpls3d" 72 | data_dir: data/processed/stpls3d 73 | image_augmentations_path: null 74 | volume_augmentations_path: null 75 | label_db_filepath: data/processed/stpls3d/label_database.yaml 76 | color_mean_std: data/processed/stpls3d/color_mean_std.yaml 77 | data_percent: 1.0 78 | mode: ${data.test_mode} 79 | ignore_label: ${data.ignore_label} 80 | num_labels: ${data.num_labels} 81 | add_raw_coordinates: ${data.add_raw_coordinates} 82 | add_colors: ${data.add_colors} 83 | add_normals: ${data.add_normals} 84 | add_instance: ${data.add_instance} 85 | cache_data: ${data.cache_data} 86 | cropping: false 87 | is_tta: false 88 | crop_min_size: ${data.crop_min_size} 89 | crop_length: ${data.crop_length} 90 | cropping_v1: ${data.cropping_v1} 91 | area: ${general.area} 92 | on_crops: ${general.on_crops} 93 | eval_inner_core: ${general.eval_inner_core} 94 | filter_out_classes: [0] 95 | label_offset: 1 96 | -------------------------------------------------------------------------------- /conf/data/indoor.yaml: -------------------------------------------------------------------------------- 1 | # @package _group_ 2 | 3 | # these parameters are inherited by datasets, data_loaders and collators 4 | # but they might be overwritten 5 | 6 | # splits 7 | train_mode: train 8 | validation_mode: validation 9 | test_mode: validation # test # validation 10 | 11 | # dataset 12 | ignore_label: 255 13 | add_raw_coordinates: true # 3dim 14 | add_colors: true # 3dim 15 | add_normals: false # 3dim 16 | in_channels: 3 # in_channels = 3 * (add_normals + add_colors + add_raw_coordinates) 17 | num_labels: 20 18 | # num_labels: 41 19 | add_instance: ${general.add_instance} 20 | task: ${general.task} 21 | layer0_path: "" 22 | part_path: "" 23 | objrct_path: "" 24 | df: 0 25 | 26 | 27 | # data loader 28 | pin_memory: false 29 | num_workers: 3 30 | batch_size: 2 31 | test_batch_size: 1 32 | cache_data: True 33 | 34 | # collation 35 | voxel_size: 0.02 36 | 37 | 38 | reps_per_epoch: ${general.reps_per_epoch} 39 | 40 | cropping: false 41 | cropping_args: 42 | min_points: 30000 43 | aspect: 0.8 44 | min_crop: 0.5 45 | max_crop: 1.0 46 | 47 | crop_min_size: 20000 48 | crop_length: 6.0 49 | cropping_v1: true -------------------------------------------------------------------------------- /conf/data/outdoor.yaml: -------------------------------------------------------------------------------- 1 | # @package _group_ 2 | 3 | # these parameters are inherited by datasets, data_loaders and collators 4 | # but they might be overwritten 5 | 6 | # splits 7 | train_mode: train 8 | validation_mode: validation 9 | test_mode: validation 10 | 11 | # dataset 12 | ignore_label: 255 13 | add_distance: true # 1dim 14 | add_reflection: true # 1dim 15 | in_channels: 2 # in_channels = add_distance + add_reflection 16 | num_labels: 19 17 | add_instance: false 18 | 19 | # data loader 20 | pin_memory: true 21 | num_workers: 4 22 | batch_size: 18 23 | sweep: 1 24 | 25 | # collation 26 | voxel_size: 0.15 27 | -------------------------------------------------------------------------------- /conf/logging/base.yaml: -------------------------------------------------------------------------------- 1 | # @package _group_ 2 | - _target_: pytorch_lightning.loggers.NeptuneLogger 3 | project_name: ${general.workspace}/${general.project_name} 4 | experiment_name: ${general.experiment_name} 5 | offline_mode: false 6 | 7 | - _target_: pytorch_lightning.loggers.CSVLogger 8 | save_dir: ${general.save_dir} 9 | name: ${general.experiment_id} 10 | version: ${general.version} 11 | -------------------------------------------------------------------------------- /conf/logging/full.yaml: -------------------------------------------------------------------------------- 1 | # @package _group_ 2 | - _target_: pytorch_lightning.loggers.WandbLogger 3 | project: ${general.project_name} 4 | name: ${general.experiment_name} 5 | save_dir: ${general.save_dir} 6 | entity: "schult" 7 | resume: "allow" 8 | id: ${general.experiment_name} 9 | -------------------------------------------------------------------------------- /conf/logging/minimal.yaml: -------------------------------------------------------------------------------- 1 | # @package _group_ 2 | - _target_: pytorch_lightning.loggers.CSVLogger 3 | save_dir: ${general.save_dir} 4 | name: ${general.experiment_id} 5 | version: ${general.version} 6 | -------------------------------------------------------------------------------- /conf/logging/offline.yaml: -------------------------------------------------------------------------------- 1 | # @package _group_ 2 | - _target_: pytorch_lightning.loggers.TensorBoardLogger 3 | name: ${general.experiment_id} 4 | version: ${general.version} 5 | save_dir: ${general.save_dir} 6 | 7 | # - _target_: pytorch_lightning.loggers.CSVLogger 8 | # name: ${general.experiment_id} 9 | # version: ${general.version} 10 | # save_dir: ${general.save_dir} -------------------------------------------------------------------------------- /conf/loss/cross_entropy.yaml: -------------------------------------------------------------------------------- 1 | # @package _group_ 2 | _target_: torch.nn.CrossEntropyLoss 3 | ignore_index: ${data.ignore_label} 4 | -------------------------------------------------------------------------------- /conf/loss/set_criterion.yaml: -------------------------------------------------------------------------------- 1 | # @package _group_ 2 | _target_: models.criterion.SetCriterion 3 | num_classes: ${general.num_targets} 4 | eos_coef: 0.1 5 | losses: 6 | - "labels" 7 | - "masks" 8 | num_points: ${matcher.num_points} 9 | oversample_ratio: 3.0 10 | importance_sample_ratio: 0.75 11 | class_weights: -1 12 | -------------------------------------------------------------------------------- /conf/loss/set_criterion_custom_weights_1.yaml: -------------------------------------------------------------------------------- 1 | # @package _group_ 2 | _target_: models.criterion.SetCriterion 3 | num_classes: ${general.num_targets} 4 | eos_coef: 0.1 5 | losses: 6 | - "labels" 7 | - "masks" 8 | num_points: ${matcher.num_points} 9 | oversample_ratio: 3.0 10 | importance_sample_ratio: 0.75 11 | class_weights: [1.0,1.5,10.0,1.0,1.0,1.0,1.0,1.0,10.0,10.0,1.0,10.0,1.0,1.0] 12 | -------------------------------------------------------------------------------- /conf/matcher/hungarian_matcher.yaml: -------------------------------------------------------------------------------- 1 | # @package _group_ 2 | _target_: models.matcher.HungarianMatcher 3 | cost_class: 2. 4 | cost_mask: 5. 5 | cost_dice: 2. 6 | num_points: -1 7 | -------------------------------------------------------------------------------- /conf/metrics/miou.yaml: -------------------------------------------------------------------------------- 1 | # @package _group_ 2 | _target_: models.metrics.ConfusionMatrix 3 | num_classes: ${data.num_labels} 4 | ignore_label: ${data.ignore_label} 5 | -------------------------------------------------------------------------------- /conf/model/mask3d.yaml: -------------------------------------------------------------------------------- 1 | # @package _group_ 2 | _target_: models.Mask3D 3 | 4 | # transformer parameters 5 | hidden_dim: 128 6 | dim_feedforward: 1024 7 | num_queries: 150 8 | num_heads: 8 9 | num_decoders: 3 10 | dropout: 0.0 11 | pre_norm: false 12 | use_level_embed: false 13 | normalize_pos_enc: true 14 | positional_encoding_type: "fourier" 15 | gauss_scale: 1.0 16 | hlevels: [0,1,2,3] 17 | 18 | # queries 19 | non_parametric_queries: true 20 | random_query_both: false 21 | random_normal: false 22 | random_queries: false 23 | use_np_features: false 24 | 25 | # sampling 26 | sample_sizes: [200, 800, 3200, 12800, 51200] 27 | max_sample_size: false # change false means sampling activated 28 | 29 | shared_decoder: true 30 | num_classes: ${general.num_targets} 31 | train_on_segments: ${general.train_on_segments} 32 | scatter_type: "mean" 33 | 34 | voxel_size: ${data.voxel_size} 35 | 36 | config: 37 | backbone: 38 | _target_: models.Res16UNet34C 39 | config: 40 | dialations: [ 1, 1, 1, 1 ] 41 | conv1_kernel_size: 5 42 | bn_momentum: 0.02 43 | # depends on normals, color, raw_coordinates 44 | # varies from 3 to 9 45 | in_channels: ${data.in_channels} 46 | out_channels: ${data.num_labels} 47 | out_fpn: true 48 | -------------------------------------------------------------------------------- /conf/optimizer/adamw.yaml: -------------------------------------------------------------------------------- 1 | # @package _group_ 2 | _target_: torch.optim.AdamW 3 | lr: 0.00005 -------------------------------------------------------------------------------- /conf/optimizer/adamw_lower.yaml: -------------------------------------------------------------------------------- 1 | # @package _group_ 2 | _target_: torch.optim.AdamW 3 | lr: 0.005 4 | -------------------------------------------------------------------------------- /conf/scheduler/exponentiallr.yaml: -------------------------------------------------------------------------------- 1 | # @package _group_ 2 | 3 | scheduler: 4 | _target_: torch.optim.lr_scheduler.ExponentialLR 5 | gamma: 0.99999 6 | last_epoch: -1 # ${trainer.max_epochs} 7 | # need to set to number because of tensorboard logger 8 | # steps_per_epoch: -1 9 | 10 | pytorch_lightning_params: 11 | interval: step 12 | -------------------------------------------------------------------------------- /conf/scheduler/lambdalr.yaml: -------------------------------------------------------------------------------- 1 | # @package _group_ 2 | 3 | scheduler: 4 | _target_: torch.optim.lr_scheduler.StepLR 5 | step_size: 99999 6 | 7 | pytorch_lightning_params: 8 | interval: epoch 9 | -------------------------------------------------------------------------------- /conf/scheduler/onecyclelr.yaml: -------------------------------------------------------------------------------- 1 | # @package _group_ 2 | 3 | scheduler: 4 | _target_: torch.optim.lr_scheduler.OneCycleLR 5 | max_lr: ${optimizer.lr} 6 | epochs: ${trainer.max_epochs} 7 | # need to set to number because of tensorboard logger 8 | steps_per_epoch: -1 9 | 10 | pytorch_lightning_params: 11 | interval: step 12 | -------------------------------------------------------------------------------- /conf/trainer/trainer.yaml: -------------------------------------------------------------------------------- 1 | # @package _group_ 2 | deterministic: false 3 | max_epochs: 1000 4 | min_epochs: 1 5 | resume_from_checkpoint: null 6 | check_val_every_n_epoch: 10 7 | num_sanity_val_steps: -1 8 | -------------------------------------------------------------------------------- /conf/trainer/trainer600.yaml: -------------------------------------------------------------------------------- 1 | # @package _group_ 2 | deterministic: false 3 | max_epochs: 601 4 | min_epochs: 1 5 | resume_from_checkpoint: null 6 | check_val_every_n_epoch: 50 7 | num_sanity_val_steps: 2 8 | -------------------------------------------------------------------------------- /datasets/__pycache__/random_cuboid.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/datasets/__pycache__/random_cuboid.cpython-310.pyc -------------------------------------------------------------------------------- /datasets/__pycache__/random_cuboid.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/datasets/__pycache__/random_cuboid.cpython-38.pyc -------------------------------------------------------------------------------- /datasets/__pycache__/semseg.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/datasets/__pycache__/semseg.cpython-310.pyc -------------------------------------------------------------------------------- /datasets/__pycache__/semseg.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/datasets/__pycache__/semseg.cpython-38.pyc -------------------------------------------------------------------------------- /datasets/__pycache__/utils.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/datasets/__pycache__/utils.cpython-310.pyc -------------------------------------------------------------------------------- /datasets/__pycache__/utils.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/datasets/__pycache__/utils.cpython-38.pyc -------------------------------------------------------------------------------- /datasets/preprocessing/__pycache__/base_preprocessing.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/datasets/preprocessing/__pycache__/base_preprocessing.cpython-310.pyc -------------------------------------------------------------------------------- /datasets/preprocessing/__pycache__/base_preprocessing.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/datasets/preprocessing/__pycache__/base_preprocessing.cpython-38.pyc -------------------------------------------------------------------------------- /datasets/random_cuboid.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. 2 | import numpy as np 3 | import torch 4 | 5 | 6 | def check_aspect(crop_range, aspect_min): 7 | xy_aspect = np.min(crop_range[:2]) / np.max(crop_range[:2]) 8 | xz_aspect = np.min(crop_range[[0, 2]]) / np.max(crop_range[[0, 2]]) 9 | yz_aspect = np.min(crop_range[1:]) / np.max(crop_range[1:]) 10 | return ( 11 | (xy_aspect >= aspect_min) 12 | or (xz_aspect >= aspect_min) 13 | or (yz_aspect >= aspect_min) 14 | ) 15 | 16 | 17 | class RandomCuboid(object): 18 | """ 19 | RandomCuboid augmentation from DepthContrast [https://arxiv.org/abs/2101.02691] 20 | We slightly modify this operation to account for object detection. 21 | This augmentation randomly crops a cuboid from the input and 22 | ensures that the cropped cuboid contains at least one bounding box 23 | """ 24 | 25 | def __init__( 26 | self, 27 | min_points, 28 | #aspect=0.8, 29 | crop_length=6.0, 30 | version1=True 31 | ): 32 | #self.aspect = aspect 33 | self.crop_length = crop_length 34 | self.min_points = min_points 35 | self.version1 = version1 36 | 37 | def __call__(self, point_cloud): 38 | if point_cloud.shape[0] < self.min_points: 39 | print("too small pcd") 40 | return np.ones(point_cloud.shape[0], dtype=np.bool) 41 | 42 | range_xyz = np.max(point_cloud[:, :2], axis=0) - np.min( 43 | point_cloud[:, :2], axis=0 44 | ) 45 | 46 | for _ in range(100): 47 | #crop_range = self.min_crop + np.random.rand(3) * ( 48 | # self.max_crop - self.min_crop 49 | #) 50 | #crop_range[-1] = 999. 51 | # if not check_aspect(crop_range, self.aspect): 52 | # continue 53 | 54 | sample_center = point_cloud[:, :2].min(axis=0) + range_xyz/2 55 | 56 | if self.version1: 57 | offset_x = np.random.uniform(-range_xyz[0]/4,range_xyz[0]/4) 58 | offset_y = np.random.uniform(-range_xyz[1]/4,range_xyz[1]/4) 59 | else: 60 | offset_x = np.random.uniform(-(range_xyz[0]/2) + self.crop_length / 4, 61 | +(range_xyz[0]/2) - self.crop_length / 4) 62 | offset_y = np.random.uniform(-(range_xyz[1]/2) + self.crop_length / 4, 63 | +(range_xyz[1]/2) - self.crop_length / 4) 64 | 65 | sample_center[0] = sample_center[0] + offset_x 66 | sample_center[1] = sample_center[1] + offset_y 67 | 68 | min_xy = sample_center - self.crop_length / 2 69 | max_xy = sample_center + self.crop_length / 2 70 | 71 | upper_idx = ( 72 | np.sum((point_cloud[:, :2] <= max_xy).astype(np.int32), 1) == 2 73 | ) 74 | lower_idx = ( 75 | np.sum((point_cloud[:, :2] >= min_xy).astype(np.int32), 1) == 2 76 | ) 77 | 78 | new_pointidx = (upper_idx) & (lower_idx) 79 | 80 | if np.sum(new_pointidx) < self.min_points: 81 | print("TOO SMALL") 82 | continue 83 | 84 | return new_pointidx 85 | 86 | # fallback 87 | print("FALLBACK") 88 | return np.ones(point_cloud.shape[0], dtype=np.bool) 89 | -------------------------------------------------------------------------------- /datasets/scannet200/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/datasets/scannet200/__init__.py -------------------------------------------------------------------------------- /datasets/scannet200/__pycache__/__init__.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/datasets/scannet200/__pycache__/__init__.cpython-310.pyc -------------------------------------------------------------------------------- /datasets/scannet200/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/datasets/scannet200/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /datasets/scannet200/__pycache__/scannet200_constants.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/datasets/scannet200/__pycache__/scannet200_constants.cpython-310.pyc -------------------------------------------------------------------------------- /datasets/scannet200/__pycache__/scannet200_constants.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/datasets/scannet200/__pycache__/scannet200_constants.cpython-38.pyc -------------------------------------------------------------------------------- /datasets/scannet200/__pycache__/scannet200_splits.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/datasets/scannet200/__pycache__/scannet200_splits.cpython-310.pyc -------------------------------------------------------------------------------- /datasets/scannet200/__pycache__/scannet200_splits.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/datasets/scannet200/__pycache__/scannet200_splits.cpython-38.pyc -------------------------------------------------------------------------------- /models/__init__.py: -------------------------------------------------------------------------------- 1 | import models.resunet as resunet 2 | import models.res16unet as res16unet 3 | from models.res16unet import Res16UNet34C, Res16UNet34A, Res16UNet14A, Res16UNet34D, Res16UNet18D, Res16UNet18B, Custom30M 4 | from models.mask3d import Mask3D 5 | 6 | MODELS = [] 7 | 8 | 9 | def add_models(module): 10 | MODELS.extend([getattr(module, a) for a in dir(module) if "Net" in a]) 11 | 12 | 13 | add_models(resunet) 14 | add_models(res16unet) 15 | add_models(mask3d) 16 | 17 | 18 | def get_models(): 19 | """Returns a tuple of sample models.""" 20 | return MODELS 21 | 22 | 23 | def load_model(name): 24 | """Creates and returns an instance of the model given its class name.""" 25 | # Find the model class from its name 26 | all_models = get_models() 27 | mdict = {model.__name__: model for model in all_models} 28 | if name not in mdict: 29 | print("Invalid model index. Options are:") 30 | # Display a list of valid model names 31 | for model in all_models: 32 | print(f"\t* {model.__name__}") 33 | return None 34 | NetClass = mdict[name] 35 | 36 | return NetClass 37 | -------------------------------------------------------------------------------- /models/metrics/__init__.py: -------------------------------------------------------------------------------- 1 | from .confusionmatrix import ConfusionMatrix 2 | from .metrics import IoU 3 | 4 | __all__ = ["ConfusionMatrix", "IoU"] 5 | -------------------------------------------------------------------------------- /models/metrics/metrics.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | 4 | class IoU: 5 | """Computes the intersection over union (IoU) per class and corresponding 6 | mean (mIoU). 7 | 8 | Intersection over union (IoU) is a common evaluation metric for semantic 9 | segmentation. The predictions are first accumulated in a confusion matrix 10 | and the IoU is computed from it as follows: 11 | 12 | IoU = true_positive / (true_positive + false_positive + false_negative). 13 | 14 | Keyword arguments: 15 | - num_classes (int): number of classes in the classification problem 16 | - normalized (boolean, optional): Determines whether or not the confusion 17 | matrix is normalized or not. Default: False. 18 | - ignore_index (int or iterable, optional): Index of the classes to ignore 19 | when computing the IoU. Can be an int, or any iterable of ints. 20 | 21 | Modified from: https://github.com/pytorch/tnt/blob/master/torchnet/meter 22 | 23 | """ 24 | 25 | def __init__(self): 26 | super().__init__() 27 | 28 | def value(self, conf_matrix): 29 | """Computes the IoU and mean IoU. 30 | 31 | The mean computation ignores NaN elements of the IoU array. 32 | 33 | Returns: 34 | Tuple: (IoU, mIoU). The first output is the per class IoU, 35 | for K classes it's numpy.ndarray with K elements. The second output, 36 | is the mean IoU. 37 | """ 38 | true_positive = np.diag(conf_matrix) 39 | false_positive = np.sum(conf_matrix, 0) - true_positive 40 | false_negative = np.sum(conf_matrix, 1) - true_positive 41 | 42 | # Just in case we get a division by 0, ignore/hide the error 43 | with np.errstate(divide="ignore", invalid="ignore"): 44 | iou = true_positive / (true_positive + false_positive + false_negative) 45 | 46 | return iou 47 | -------------------------------------------------------------------------------- /models/model.py: -------------------------------------------------------------------------------- 1 | from MinkowskiEngine import MinkowskiNetwork 2 | 3 | 4 | class Model(MinkowskiNetwork): 5 | """ 6 | Base network for all sparse convnet 7 | 8 | By default, all networks are segmentation networks. 9 | """ 10 | 11 | OUT_PIXEL_DIST = -1 12 | 13 | def __init__(self, in_channels, out_channels, config, D, **kwargs): 14 | super().__init__(D) 15 | self.in_channels = in_channels 16 | self.out_channels = out_channels 17 | self.config = config 18 | 19 | 20 | class HighDimensionalModel(Model): 21 | """ 22 | Base network for all spatio (temporal) chromatic sparse convnet 23 | """ 24 | 25 | def __init__(self, in_channels, out_channels, config, D, **kwargs): 26 | assert D > 4, "Num dimension smaller than 5" 27 | super().__init__(in_channels, out_channels, config, D, **kwargs) 28 | -------------------------------------------------------------------------------- /models/modules/3detr_helpers.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. 2 | import torch.nn as nn 3 | from functools import partial 4 | import copy 5 | 6 | 7 | class BatchNormDim1Swap(nn.BatchNorm1d): 8 | """ 9 | Used for nn.Transformer that uses a HW x N x C rep 10 | """ 11 | 12 | def forward(self, x): 13 | """ 14 | x: HW x N x C 15 | permute to N x C x HW 16 | Apply BN on C 17 | permute back 18 | """ 19 | hw, n, c = x.shape 20 | x = x.permute(1, 2, 0) 21 | x = super(BatchNormDim1Swap, self).forward(x) 22 | # x: n x c x hw -> hw x n x c 23 | x = x.permute(2, 0, 1) 24 | return x 25 | 26 | 27 | NORM_DICT = { 28 | "bn": BatchNormDim1Swap, 29 | "bn1d": nn.BatchNorm1d, 30 | "id": nn.Identity, 31 | "ln": nn.LayerNorm, 32 | } 33 | 34 | ACTIVATION_DICT = { 35 | "relu": nn.ReLU, 36 | "gelu": nn.GELU, 37 | "leakyrelu": partial(nn.LeakyReLU, negative_slope=0.1), 38 | } 39 | 40 | WEIGHT_INIT_DICT = { 41 | "xavier_uniform": nn.init.xavier_uniform_, 42 | } 43 | 44 | 45 | class GenericMLP(nn.Module): 46 | def __init__( 47 | self, 48 | input_dim, 49 | hidden_dims, 50 | output_dim, 51 | norm_fn_name=None, 52 | activation="relu", 53 | use_conv=False, 54 | dropout=None, 55 | hidden_use_bias=False, 56 | output_use_bias=True, 57 | output_use_activation=False, 58 | output_use_norm=False, 59 | weight_init_name=None, 60 | ): 61 | super().__init__() 62 | activation = ACTIVATION_DICT[activation] 63 | norm = None 64 | if norm_fn_name is not None: 65 | norm = NORM_DICT[norm_fn_name] 66 | if norm_fn_name == "ln" and use_conv: 67 | norm = lambda x: nn.GroupNorm(1, x) # easier way to use LayerNorm 68 | 69 | if dropout is not None: 70 | if not isinstance(dropout, list): 71 | dropout = [dropout for _ in range(len(hidden_dims))] 72 | 73 | layers = [] 74 | prev_dim = input_dim 75 | for idx, x in enumerate(hidden_dims): 76 | if use_conv: 77 | layer = nn.Conv1d(prev_dim, x, 1, bias=hidden_use_bias) 78 | else: 79 | layer = nn.Linear(prev_dim, x, bias=hidden_use_bias) 80 | layers.append(layer) 81 | if norm: 82 | layers.append(norm(x)) 83 | layers.append(activation()) 84 | if dropout is not None: 85 | layers.append(nn.Dropout(p=dropout[idx])) 86 | prev_dim = x 87 | if use_conv: 88 | layer = nn.Conv1d(prev_dim, output_dim, 1, bias=output_use_bias) 89 | else: 90 | layer = nn.Linear(prev_dim, output_dim, bias=output_use_bias) 91 | layers.append(layer) 92 | 93 | if output_use_norm: 94 | layers.append(norm(output_dim)) 95 | 96 | if output_use_activation: 97 | layers.append(activation()) 98 | 99 | self.layers = nn.Sequential(*layers) 100 | 101 | if weight_init_name is not None: 102 | self.do_weight_init(weight_init_name) 103 | 104 | def do_weight_init(self, weight_init_name): 105 | func = WEIGHT_INIT_DICT[weight_init_name] 106 | for (_, param) in self.named_parameters(): 107 | if param.dim() > 1: # skips batchnorm/layernorm 108 | func(param) 109 | 110 | def forward(self, x): 111 | output = self.layers(x) 112 | return output 113 | 114 | 115 | def get_clones(module, N): 116 | return nn.ModuleList([copy.deepcopy(module) for i in range(N)]) -------------------------------------------------------------------------------- /models/modules/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/models/modules/__init__.py -------------------------------------------------------------------------------- /models/modules/helpers_3detr.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. 2 | import torch.nn as nn 3 | from functools import partial 4 | import copy 5 | 6 | 7 | class BatchNormDim1Swap(nn.BatchNorm1d): 8 | """ 9 | Used for nn.Transformer that uses a HW x N x C rep 10 | """ 11 | 12 | def forward(self, x): 13 | """ 14 | x: HW x N x C 15 | permute to N x C x HW 16 | Apply BN on C 17 | permute back 18 | """ 19 | hw, n, c = x.shape 20 | x = x.permute(1, 2, 0) 21 | x = super(BatchNormDim1Swap, self).forward(x) 22 | # x: n x c x hw -> hw x n x c 23 | x = x.permute(2, 0, 1) 24 | return x 25 | 26 | 27 | NORM_DICT = { 28 | "bn": BatchNormDim1Swap, 29 | "bn1d": nn.BatchNorm1d, 30 | "id": nn.Identity, 31 | "ln": nn.LayerNorm, 32 | } 33 | 34 | ACTIVATION_DICT = { 35 | "relu": nn.ReLU, 36 | "gelu": nn.GELU, 37 | "leakyrelu": partial(nn.LeakyReLU, negative_slope=0.1), 38 | } 39 | 40 | WEIGHT_INIT_DICT = { 41 | "xavier_uniform": nn.init.xavier_uniform_, 42 | } 43 | 44 | 45 | class GenericMLP(nn.Module): 46 | def __init__( 47 | self, 48 | input_dim, 49 | hidden_dims, 50 | output_dim, 51 | norm_fn_name=None, 52 | activation="relu", 53 | use_conv=False, 54 | dropout=None, 55 | hidden_use_bias=False, 56 | output_use_bias=True, 57 | output_use_activation=False, 58 | output_use_norm=False, 59 | weight_init_name=None, 60 | ): 61 | super().__init__() 62 | activation = ACTIVATION_DICT[activation] 63 | norm = None 64 | if norm_fn_name is not None: 65 | norm = NORM_DICT[norm_fn_name] 66 | if norm_fn_name == "ln" and use_conv: 67 | norm = lambda x: nn.GroupNorm(1, x) # easier way to use LayerNorm 68 | 69 | if dropout is not None: 70 | if not isinstance(dropout, list): 71 | dropout = [dropout for _ in range(len(hidden_dims))] 72 | 73 | layers = [] 74 | prev_dim = input_dim 75 | for idx, x in enumerate(hidden_dims): 76 | if use_conv: 77 | layer = nn.Conv1d(prev_dim, x, 1, bias=hidden_use_bias) 78 | else: 79 | layer = nn.Linear(prev_dim, x, bias=hidden_use_bias) 80 | layers.append(layer) 81 | if norm: 82 | layers.append(norm(x)) 83 | layers.append(activation()) 84 | if dropout is not None: 85 | layers.append(nn.Dropout(p=dropout[idx])) 86 | prev_dim = x 87 | if use_conv: 88 | layer = nn.Conv1d(prev_dim, output_dim, 1, bias=output_use_bias) 89 | else: 90 | layer = nn.Linear(prev_dim, output_dim, bias=output_use_bias) 91 | layers.append(layer) 92 | 93 | if output_use_norm: 94 | layers.append(norm(output_dim)) 95 | 96 | if output_use_activation: 97 | layers.append(activation()) 98 | 99 | self.layers = nn.Sequential(*layers) 100 | 101 | if weight_init_name is not None: 102 | self.do_weight_init(weight_init_name) 103 | 104 | def do_weight_init(self, weight_init_name): 105 | func = WEIGHT_INIT_DICT[weight_init_name] 106 | for (_, param) in self.named_parameters(): 107 | if param.dim() > 1: # skips batchnorm/layernorm 108 | func(param) 109 | 110 | def forward(self, x): 111 | output = self.layers(x) 112 | return output 113 | 114 | 115 | def get_clones(module, N): 116 | return nn.ModuleList([copy.deepcopy(module) for i in range(N)]) -------------------------------------------------------------------------------- /models/modules/resnet_block.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | from MinkowskiEngine import MinkowskiReLU 3 | 4 | from models.modules.common import ConvType, NormType, conv, get_norm 5 | 6 | 7 | class BasicBlockBase(nn.Module): 8 | expansion = 1 9 | NORM_TYPE = NormType.BATCH_NORM 10 | 11 | def __init__( 12 | self, 13 | inplanes, 14 | planes, 15 | stride=1, 16 | dilation=1, 17 | downsample=None, 18 | conv_type=ConvType.HYPERCUBE, 19 | bn_momentum=0.1, 20 | D=3, 21 | ): 22 | super().__init__() 23 | 24 | self.conv1 = conv( 25 | inplanes, 26 | planes, 27 | kernel_size=3, 28 | stride=stride, 29 | dilation=dilation, 30 | conv_type=conv_type, 31 | D=D, 32 | ) 33 | self.norm1 = get_norm(self.NORM_TYPE, planes, D, bn_momentum=bn_momentum) 34 | self.conv2 = conv( 35 | planes, 36 | planes, 37 | kernel_size=3, 38 | stride=1, 39 | dilation=dilation, 40 | bias=False, 41 | conv_type=conv_type, 42 | D=D, 43 | ) 44 | self.norm2 = get_norm(self.NORM_TYPE, planes, D, bn_momentum=bn_momentum) 45 | self.relu = MinkowskiReLU(inplace=True) 46 | self.downsample = downsample 47 | 48 | def forward(self, x): 49 | residual = x 50 | 51 | out = self.conv1(x) 52 | out = self.norm1(out) 53 | out = self.relu(out) 54 | 55 | out = self.conv2(out) 56 | out = self.norm2(out) 57 | 58 | if self.downsample is not None: 59 | residual = self.downsample(x) 60 | 61 | out += residual 62 | out = self.relu(out) 63 | 64 | return out 65 | 66 | 67 | class BasicBlock(BasicBlockBase): 68 | NORM_TYPE = NormType.BATCH_NORM 69 | 70 | 71 | class BasicBlockIN(BasicBlockBase): 72 | NORM_TYPE = NormType.INSTANCE_NORM 73 | 74 | 75 | class BasicBlockINBN(BasicBlockBase): 76 | NORM_TYPE = NormType.INSTANCE_BATCH_NORM 77 | 78 | 79 | class BottleneckBase(nn.Module): 80 | expansion = 4 81 | NORM_TYPE = NormType.BATCH_NORM 82 | 83 | def __init__( 84 | self, 85 | inplanes, 86 | planes, 87 | stride=1, 88 | dilation=1, 89 | downsample=None, 90 | conv_type=ConvType.HYPERCUBE, 91 | bn_momentum=0.1, 92 | D=3, 93 | ): 94 | super().__init__() 95 | self.conv1 = conv(inplanes, planes, kernel_size=1, D=D) 96 | self.norm1 = get_norm(self.NORM_TYPE, planes, D, bn_momentum=bn_momentum) 97 | 98 | self.conv2 = conv( 99 | planes, 100 | planes, 101 | kernel_size=3, 102 | stride=stride, 103 | dilation=dilation, 104 | conv_type=conv_type, 105 | D=D, 106 | ) 107 | self.norm2 = get_norm(self.NORM_TYPE, planes, D, bn_momentum=bn_momentum) 108 | 109 | self.conv3 = conv(planes, planes * self.expansion, kernel_size=1, D=D) 110 | self.norm3 = get_norm( 111 | self.NORM_TYPE, planes * self.expansion, D, bn_momentum=bn_momentum 112 | ) 113 | 114 | self.relu = MinkowskiReLU(inplace=True) 115 | self.downsample = downsample 116 | 117 | def forward(self, x): 118 | residual = x 119 | 120 | out = self.conv1(x) 121 | out = self.norm1(out) 122 | out = self.relu(out) 123 | 124 | out = self.conv2(out) 125 | out = self.norm2(out) 126 | out = self.relu(out) 127 | 128 | out = self.conv3(out) 129 | out = self.norm3(out) 130 | 131 | if self.downsample is not None: 132 | residual = self.downsample(x) 133 | 134 | out += residual 135 | out = self.relu(out) 136 | 137 | return out 138 | 139 | 140 | class Bottleneck(BottleneckBase): 141 | NORM_TYPE = NormType.BATCH_NORM 142 | 143 | 144 | class BottleneckIN(BottleneckBase): 145 | NORM_TYPE = NormType.INSTANCE_NORM 146 | 147 | 148 | class BottleneckINBN(BottleneckBase): 149 | NORM_TYPE = NormType.INSTANCE_BATCH_NORM 150 | -------------------------------------------------------------------------------- /models/modules/senet_block.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import MinkowskiEngine as ME 3 | 4 | from mix3d.models.modules.common import ConvType, NormType 5 | from mix3d.models.modules.resnet_block import BasicBlock, Bottleneck 6 | 7 | 8 | class SELayer(nn.Module): 9 | def __init__(self, channel, reduction=16, D=-1): 10 | # Global coords does not require coords_key 11 | super().__init__() 12 | self.fc = nn.Sequential( 13 | ME.MinkowskiLinear(channel, channel // reduction), 14 | ME.MinkowskiReLU(inplace=True), 15 | ME.MinkowskiLinear(channel // reduction, channel), 16 | ME.MinkowskiSigmoid(), 17 | ) 18 | self.pooling = ME.MinkowskiGlobalPooling(dimension=D) 19 | self.broadcast_mul = ME.MinkowskiBroadcastMultiplication(dimension=D) 20 | 21 | def forward(self, x): 22 | y = self.pooling(x) 23 | y = self.fc(y) 24 | return self.broadcast_mul(x, y) 25 | 26 | 27 | class SEBasicBlock(BasicBlock): 28 | def __init__( 29 | self, 30 | inplanes, 31 | planes, 32 | stride=1, 33 | dilation=1, 34 | downsample=None, 35 | conv_type=ConvType.HYPERCUBE, 36 | reduction=16, 37 | D=-1, 38 | ): 39 | super().__init__( 40 | inplanes, 41 | planes, 42 | stride=stride, 43 | dilation=dilation, 44 | downsample=downsample, 45 | conv_type=conv_type, 46 | D=D, 47 | ) 48 | self.se = SELayer(planes, reduction=reduction, D=D) 49 | 50 | def forward(self, x): 51 | residual = x 52 | 53 | out = self.conv1(x) 54 | out = self.norm1(out) 55 | out = self.relu(out) 56 | 57 | out = self.conv2(out) 58 | out = self.norm2(out) 59 | out = self.se(out) 60 | 61 | if self.downsample is not None: 62 | residual = self.downsample(x) 63 | 64 | out += residual 65 | out = self.relu(out) 66 | 67 | return out 68 | 69 | 70 | class SEBasicBlockSN(SEBasicBlock): 71 | NORM_TYPE = NormType.SPARSE_SWITCH_NORM 72 | 73 | 74 | class SEBasicBlockIN(SEBasicBlock): 75 | NORM_TYPE = NormType.SPARSE_INSTANCE_NORM 76 | 77 | 78 | class SEBasicBlockLN(SEBasicBlock): 79 | NORM_TYPE = NormType.SPARSE_LAYER_NORM 80 | 81 | 82 | class SEBottleneck(Bottleneck): 83 | def __init__( 84 | self, 85 | inplanes, 86 | planes, 87 | stride=1, 88 | dilation=1, 89 | downsample=None, 90 | conv_type=ConvType.HYPERCUBE, 91 | D=3, 92 | reduction=16, 93 | ): 94 | super().__init__( 95 | inplanes, 96 | planes, 97 | stride=stride, 98 | dilation=dilation, 99 | downsample=downsample, 100 | conv_type=conv_type, 101 | D=D, 102 | ) 103 | self.se = SELayer(planes * self.expansion, reduction=reduction, D=D) 104 | 105 | def forward(self, x): 106 | residual = x 107 | 108 | out = self.conv1(x) 109 | out = self.norm1(out) 110 | out = self.relu(out) 111 | 112 | out = self.conv2(out) 113 | out = self.norm2(out) 114 | out = self.relu(out) 115 | 116 | out = self.conv3(out) 117 | out = self.norm3(out) 118 | out = self.se(out) 119 | 120 | if self.downsample is not None: 121 | residual = self.downsample(x) 122 | 123 | out += residual 124 | out = self.relu(out) 125 | 126 | return out 127 | 128 | 129 | class SEBottleneckSN(SEBottleneck): 130 | NORM_TYPE = NormType.SPARSE_SWITCH_NORM 131 | 132 | 133 | class SEBottleneckIN(SEBottleneck): 134 | NORM_TYPE = NormType.SPARSE_INSTANCE_NORM 135 | 136 | 137 | class SEBottleneckLN(SEBottleneck): 138 | NORM_TYPE = NormType.SPARSE_LAYER_NORM 139 | -------------------------------------------------------------------------------- /models/wrapper.py: -------------------------------------------------------------------------------- 1 | import random 2 | 3 | from torch.nn import Module 4 | from MinkowskiEngine import SparseTensor 5 | 6 | 7 | class Wrapper(Module): 8 | """ 9 | Wrapper for the segmentation networks. 10 | """ 11 | 12 | OUT_PIXEL_DIST = -1 13 | 14 | def __init__(self, NetClass, in_nchannel, out_nchannel, config): 15 | super().__init__() 16 | self.initialize_filter(NetClass, in_nchannel, out_nchannel, config) 17 | 18 | def initialize_filter(self, NetClass, in_nchannel, out_nchannel, config): 19 | raise NotImplementedError("Must initialize a model and a filter") 20 | 21 | def forward(self, x, coords, colors=None): 22 | soutput = self.model(x) 23 | 24 | # During training, make the network invariant to the filter 25 | if not self.training or random.random() < 0.5: 26 | # Filter requires the model to finish the forward pass 27 | wrapper_coords = self.filter.initialize_coords(self.model, coords, colors) 28 | finput = SparseTensor(soutput.F, wrapper_coords) 29 | soutput = self.filter(finput) 30 | return soutput 31 | -------------------------------------------------------------------------------- /pseudo_mask_gen/README.md: -------------------------------------------------------------------------------- 1 | 2 | # Data Preprocessing 3 |
4 | 5 | Download the ScanNet dataset from [here](http://kaldir.vc.in.tum.de/scannet_benchmark/documentation) , and move it to ./data 6 | 7 | ### Apply VCCS Algorithm 8 | 9 | ```bash 10 | python vccs/data_prepare_ScanNet.py --data_path "PATH_TO_RAW_SCANNET_DATASET" 11 | ``` 12 | This code will preprcocess ScanNet and put it under ./data/ScanNet/processed 13 | 14 | ```bash 15 | python vccs/initialSP_prepare_ScanNet.py 16 | ``` 17 | This code will construct init segments on ScanNet and put it under ./data/ScanNet/initial_superpoints 18 | 19 | ### Get DINO features of Pointclouds 20 | Download DINO model and put it in ./DINO 21 | ```bash 22 | python project_feature.py 23 | ``` 24 | This code will get the DINO features for each point cloud and store them in ./DINO_point_feats 25 | 26 | ### Get Point Distances 27 | ```bash 28 | python get_dis_matrix.py 29 | ``` 30 | This code calculates the shortest distance between the initial segments and stores it in ./dis_matrixes_initseg 31 | 32 | ### Get Bbox Prior 33 | ```bash 34 | python get bbox prior/scene_frame_ncut.py --scannet_frames data/Scannetv2/frames_square --immedia_data_dir PATH_TO_STORE_INTERMEDIATE_RESULTS 35 | 36 | python get bbox prior/scene_frame_merge.py --immedia_data_dir PATH_TO_STORE_INTERMEDIATE_RESULTS --scannet_frames data/Scannetv2/frames_square --val_data_dir data/Scannetv2 --output_dir bbox_prior 37 | ``` 38 | This code will calculate Bbox Prior and store it in ./bbox_prior 39 | 40 | ### Data Prepare 41 | 42 | 1. Create a new directory called `processed_data`: 43 | 44 | ```bash 45 | mkdir processed_data 46 | ``` 47 | 48 | 2. Move or create a symbolic link to the folders completed in **Data Preprocessing** into the `processed_data` directory. 49 | 50 | 3. Your directory structure of `processed_data` should look like this: 51 | 52 | ```bash 53 | pseudo_mask_gen/ 54 | ├── processed_data/ 55 | │ └── processed/ 56 | │ └── DINO_point_feats/ 57 | │ └── dis_matrixes_initseg/ 58 | │ └── initial_superpoints/ 59 | │ └── bbox_prior/ 60 | ├── other_directories_and_files 61 | ``` 62 | 63 | 64 | 65 | # Pseudo Mask Generation 66 |
67 | 68 | ### Part2Object : hierarchical clustering 69 | 70 | You can obtain the hierarchical clustering results using `topk_merge.py`. To do this, you need to specify the file path to the results produced in the Data Preprocessing section and also specify an output directory. 71 | 72 | Here's an example: 73 | 74 | ```bash 75 | python topk_merge.py --input /path/to/processed_data --output /path/to/output/directory 76 | ``` 77 | 78 | Replace `/path/to/processed_data` with the path to your preprocessed data and `/path/to/output/directory` with the path to the directory where you want to save the results. 79 | 80 | ### Post Processing 81 | 82 | You can use `post_pro.py` for post-processing to eliminate noise points. In addition to specifying the `input` and `output` parameters as with `topk_merge.py`, you also need to specify `output-processed` to store the post-processing results. Here's how you can do it: 83 | 84 | Here's an example: 85 | 86 | ```bash 87 | python post_pro.py --input /path/to/processed_data --output /path/to/output/directory --output-processed /path/to/post_processed/results 88 | ``` 89 | Replace `/path/to/input/data` with the path to your input data, `/path/to/output/directory` with the path to the directory where you want to save the intermediate results, and `/path/to/post_processed/results` with the path to the directory where you want to save the post-processing results. 90 | -------------------------------------------------------------------------------- /pseudo_mask_gen/get_dis_matrix.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | 4 | from tqdm import tqdm 5 | import torch 6 | 7 | from multiprocessing import Pool 8 | 9 | def naive_distance(super_point_A, super_point_B): 10 | 11 | distances = np.sqrt(np.sum((super_point_A[:, np.newaxis] - super_point_B) ** 2, axis=2)) 12 | 13 | min_distances = np.min(distances, axis=1) 14 | 15 | min_distance = np.min(min_distances) 16 | return min_distance 17 | 18 | def calculate_distance(pair): 19 | i, j, super_points, super_points_ids, coords = pair 20 | mask_A = super_points == super_points_ids[i] 21 | mask_B = super_points == super_points_ids[j] 22 | distance = naive_distance(coords[mask_A], coords[mask_B]) 23 | 24 | return (i, j, distance) 25 | 26 | 27 | scene_list = sorted(os.listdir("data/frames_square/")) 28 | 29 | for scan in tqdm(scene_list): 30 | 31 | try: 32 | points_path = f"data/val/{scan}.pth" 33 | points = torch.load(points_path) 34 | except: 35 | points_path = f"data/train/{scan}.pth" 36 | points = torch.load(points_path) 37 | 38 | superpoints_path = f"init_segments/{scan}.pth" 39 | super_points = torch.load(superpoints_path) 40 | 41 | super_points_ids = np.unique(super_points) 42 | 43 | coords = points["coord"] 44 | colors = points["color"] 45 | 46 | 47 | super_point_coors = [] 48 | super_point_colors = [] 49 | super_point_feats = [] 50 | 51 | distance_matrix = np.zeros((super_points_ids.shape[0], super_points_ids.shape[0])) 52 | 53 | with Pool() as pool: 54 | pairs = [(i, j, super_points, super_points_ids, coords) for i in range(super_points_ids.shape[0]) for j in range(i+1, super_points_ids.shape[0])] 55 | results = pool.map(calculate_distance, pairs) 56 | 57 | for result in results: 58 | i, j, distance = result 59 | distance_matrix[i, j] = distance 60 | distance_matrix[j, i] = distance 61 | 62 | torch.save(distance_matrix, f"dis_matrixes_initseg/{scan}.pth") -------------------------------------------------------------------------------- /pseudo_mask_gen/vccs/data_prepare_ScanNet.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | from os.path import join, exists, dirname, abspath 3 | import os, sys, glob 4 | 5 | import numpy as np 6 | BASE_DIR = dirname(abspath(__file__)) 7 | ROOT_DIR = dirname(BASE_DIR) 8 | sys.path.append(BASE_DIR) 9 | sys.path.append(ROOT_DIR) 10 | from lib.helper_ply import read_ply, write_ply 11 | from concurrent.futures import ProcessPoolExecutor 12 | import argparse 13 | 14 | parser = argparse.ArgumentParser() 15 | parser.add_argument('--data_path', type=str, default='/remote-home/share/Datasets/Scannetv2/', help='raw data path') 16 | parser.add_argument('--processed_data_path', type=str, default='data/ScanNet/processed') 17 | args = parser.parse_args() 18 | 19 | SCANNET_RAW_PATH = Path(args.data_path) 20 | SCANNET_OUT_PATH = Path(args.processed_data_path) 21 | TRAIN_DEST = 'train' 22 | TEST_DEST = 'test' 23 | SUBSETS = {TRAIN_DEST: 'scans', TEST_DEST: 'scans_test'} 24 | POINTCLOUD_FILE = '_vh_clean_2.ply' 25 | BUGS = { 26 | 'scene0270_00': 50, 27 | 'scene0270_02': 50, 28 | 'scene0384_00': 149, 29 | } 30 | CLASS_LABELS = ('wall', 'floor', 'cabinet', 'bed', 'chair', 'sofa', 'table', 'door', 'window', 31 | 'bookshelf', 'picture', 'counter', 'desk', 'curtain', 'refrigerator', 32 | 'shower curtain', 'toilet', 'sink', 'bathtub', 'otherfurniture') 33 | VALID_CLASS_IDS = (1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 16, 24, 28, 33, 34, 36, 39) 34 | NUM_LABELS = 41 35 | IGNORE_LABELS = tuple(set(range(41)) - set(VALID_CLASS_IDS)) 36 | 37 | ''' Set Invalid Label to -1''' 38 | label_map = {} 39 | n_used = 0 40 | for l in range(NUM_LABELS): 41 | if l in IGNORE_LABELS: 42 | label_map[l] = -1#ignore label 43 | else: 44 | label_map[l] = n_used 45 | n_used += 1 46 | 47 | def handle_process(path): 48 | f = Path(path.split(',')[0]) 49 | phase_out_path = Path(path.split(',')[1]) 50 | # pointcloud = read_ply(f) 51 | data = read_ply(str(f), triangular_mesh=True)[0] 52 | coords = np.vstack((data['x'], data['y'], data['z'])).T 53 | colors = np.vstack((data['red'], data['green'], data['blue'])).T 54 | # Load label file. 55 | label_f = f.parent / (f.stem + '.labels' + f.suffix) 56 | if label_f.is_file(): 57 | label = read_ply(str(label_f), triangular_mesh=True)[0]['label'].T.squeeze() 58 | else: # Label may not exist in test case. 59 | label = -np.zeros(data.shape) 60 | out_f = phase_out_path / (f.name[:-len(POINTCLOUD_FILE)] + f.suffix) 61 | 62 | '''Alignment''' 63 | txtfile = str(f.parent) + '/'+ str(f.parts[-2]) +'.txt' 64 | print(txtfile) 65 | with open(txtfile) as txtfile: 66 | lines = txtfile.readlines() 67 | for line in lines: 68 | line = line.split() 69 | if line[0] == 'axisAlignment': 70 | align_mat = np.array([float(x) for x in line[2:]]).reshape([4, 4]).astype(np.float32) 71 | R = align_mat[:3, :3] 72 | T = align_mat[:3, 3] 73 | coords = coords.dot(R.T) + T 74 | 75 | '''Fix Data Bug''' 76 | for item, bug_index in BUGS.items(): 77 | if item in path: 78 | print('Fixing {} bugged label'.format(item)) 79 | bug_mask = label == bug_index 80 | label[bug_mask] = 0 81 | 82 | label = np.array([label_map[x] for x in label]) 83 | write_ply(str(out_f), [coords.astype(np.float64), colors, label[:, None].astype(np.float64)], ['x', 'y', 'z', 'red', 'green', 'blue', 'class']) 84 | 85 | 86 | print('start preprocess') 87 | 88 | path_list = [] 89 | for out_path, in_path in SUBSETS.items(): 90 | # phase_out_path = SCANNET_OUT_PATH / out_path 91 | phase_out_path = SCANNET_OUT_PATH 92 | phase_out_path.mkdir(parents=True, exist_ok=True) 93 | for f in (SCANNET_RAW_PATH / in_path).glob('*/*' + POINTCLOUD_FILE): 94 | path_list.append(str(f) + ',' + str(phase_out_path)) 95 | 96 | pool = ProcessPoolExecutor(max_workers=10) 97 | result = list(pool.map(handle_process, path_list)) 98 | -------------------------------------------------------------------------------- /scenes/1.txt: -------------------------------------------------------------------------------- 1 | scene0239_00 2 | scene0013_02 3 | scene0473_01 4 | scene0416_00 5 | scene0572_00 6 | scene0047_00 7 | scene0209_02 8 | scene0408_01 9 | scene0279_01 10 | scene0656_01 11 | scene0044_01 12 | scene0150_00 13 | -------------------------------------------------------------------------------- /scenes/10.txt: -------------------------------------------------------------------------------- 1 | scene0557_00 2 | scene0012_00 3 | scene0324_01 4 | scene0627_00 5 | scene0523_00 6 | scene0060_01 7 | scene0322_00 8 | scene0134_01 9 | scene0089_00 10 | scene0548_02 11 | scene0132_00 12 | scene0472_02 13 | scene0554_00 14 | scene0023_00 15 | scene0509_02 16 | scene0603_01 17 | scene0121_01 18 | scene0078_00 19 | scene0515_01 20 | scene0276_00 21 | scene0253_00 22 | scene0505_04 23 | scene0172_00 24 | scene0347_02 25 | scene0247_00 26 | scene0171_00 27 | scene0501_01 28 | scene0043_00 29 | scene0261_00 30 | scene0176_00 31 | scene0451_04 32 | scene0137_02 33 | scene0443_00 34 | scene0404_00 35 | scene0229_01 36 | scene0529_02 37 | scene0126_01 38 | scene0305_01 39 | scene0421_01 40 | scene0520_01 41 | scene0192_01 42 | scene0492_01 43 | scene0421_02 44 | scene0310_02 45 | scene0083_00 46 | scene0166_00 47 | scene0197_01 48 | scene0562_00 49 | scene0255_02 50 | scene0673_05 51 | scene0642_00 52 | scene0471_00 53 | scene0492_00 54 | scene0434_00 55 | scene0466_00 56 | scene0291_01 57 | scene0152_02 58 | scene0099_00 59 | scene0026_00 60 | scene0053_00 61 | scene0143_01 62 | scene0155_01 63 | scene0195_01 64 | scene0035_01 65 | scene0634_00 66 | scene0447_00 67 | scene0588_03 68 | scene0123_02 69 | scene0361_01 70 | scene0272_01 71 | scene0301_00 72 | scene0160_01 73 | scene0166_02 74 | scene0038_00 75 | scene0530_00 76 | scene0673_00 77 | scene0010_00 78 | scene0287_00 79 | scene0439_01 80 | scene0613_01 81 | scene0587_03 82 | scene0286_02 83 | scene0182_01 84 | scene0509_00 85 | scene0449_02 86 | scene0493_01 87 | scene0703_00 88 | scene0031_01 89 | scene0489_02 90 | scene0646_02 91 | scene0080_01 92 | scene0119_00 93 | scene0292_00 94 | scene0260_01 95 | scene0301_02 96 | scene0630_02 97 | scene0505_02 98 | scene0312_01 99 | scene0161_01 100 | scene0320_01 101 | scene0078_01 102 | scene0478_00 103 | scene0173_01 104 | scene0078_02 105 | scene0289_00 106 | scene0349_00 107 | scene0367_01 108 | scene0055_02 109 | scene0110_01 110 | scene0572_01 111 | scene0594_00 112 | scene0013_01 113 | scene0416_04 114 | scene0160_03 115 | scene0524_00 116 | scene0170_02 117 | scene0369_02 118 | scene0434_01 119 | scene0250_01 120 | scene0349_01 121 | -------------------------------------------------------------------------------- /scenes/5.txt: -------------------------------------------------------------------------------- 1 | scene0573_01 2 | scene0331_00 3 | scene0312_00 4 | scene0379_00 5 | scene0677_01 6 | scene0080_02 7 | scene0545_01 8 | scene0163_01 9 | scene0209_02 10 | scene0440_00 11 | scene0054_00 12 | scene0467_00 13 | scene0270_00 14 | scene0091_00 15 | scene0468_01 16 | scene0051_00 17 | scene0586_01 18 | scene0457_01 19 | scene0375_01 20 | scene0504_00 21 | scene0282_00 22 | scene0012_02 23 | scene0027_01 24 | scene0058_00 25 | scene0018_00 26 | scene0140_01 27 | scene0705_00 28 | scene0396_02 29 | scene0455_00 30 | scene0191_02 31 | scene0124_01 32 | scene0352_01 33 | scene0582_02 34 | scene0147_00 35 | scene0536_02 36 | scene0465_00 37 | scene0675_00 38 | scene0227_00 39 | scene0547_01 40 | scene0142_01 41 | scene0321_00 42 | scene0364_01 43 | scene0596_02 44 | scene0611_01 45 | scene0627_01 46 | scene0060_01 47 | scene0136_00 48 | scene0079_01 49 | scene0012_00 50 | scene0012_01 51 | scene0698_01 52 | scene0173_02 53 | scene0275_00 54 | scene0261_03 55 | scene0138_00 56 | scene0625_01 57 | scene0581_00 58 | scene0424_02 59 | scene0390_00 60 | scene0451_00 61 | -------------------------------------------------------------------------------- /scripts/s3dis/s3dis_from_scratch.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | export OMP_NUM_THREADS=3 # speeds up MinkowskiEngine 3 | 4 | CURR_AREA=1 # set the area number accordingly [1,6] 5 | CURR_DBSCAN=0.6 6 | CURR_TOPK=-1 7 | CURR_QUERY=100 8 | 9 | python main_instance_segmentation.py \ 10 | general.project_name="s3dis" \ 11 | general.experiment_name="area${CURR_AREA}_from_scratch" \ 12 | data.batch_size=4 \ 13 | data/datasets=s3dis \ 14 | general.num_targets=14 \ 15 | data.num_labels=13 \ 16 | trainer.max_epochs=1001 \ 17 | general.area=${CURR_AREA} \ 18 | trainer.check_val_every_n_epoch=10 19 | 20 | python main_instance_segmentation.py \ 21 | general.project_name="s3dis_eval" \ 22 | general.experiment_name="area${CURR_AREA}_from_scratch_eps_${CURR_DBSCAN}_topk_${CURR_TOPK}_q_${CURR_QUERY}" \ 23 | general.checkpoint="checkpoints/s3dis/from_scratch/area${CURR_AREA}.ckpt" \ 24 | general.train_mode=false \ 25 | data.batch_size=4 \ 26 | data/datasets=s3dis \ 27 | general.num_targets=14 \ 28 | data.num_labels=13 \ 29 | general.area=${CURR_AREA} \ 30 | model.num_queries=${CURR_QUERY} \ 31 | general.topk_per_image=${CURR_TOPK} \ 32 | general.use_dbscan=true \ 33 | general.dbscan_eps=${CURR_DBSCAN} 34 | -------------------------------------------------------------------------------- /scripts/s3dis/s3dis_pretrained.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | export OMP_NUM_THREADS=3 # speeds up MinkowskiEngine 3 | 4 | CURR_AREA=1 # set the area number accordingly [1,6] 5 | CURR_DBSCAN=0.6 6 | CURR_TOPK=-1 7 | CURR_QUERY=100 8 | 9 | python main_instance_segmentation.py \ 10 | general.project_name="s3dis" \ 11 | general.experiment_name="area${CURR_AREA}_pretrained" \ 12 | data.batch_size=4 \ 13 | data/datasets=s3dis \ 14 | general.num_targets=14 \ 15 | data.num_labels=13 \ 16 | general.area=${CURR_AREA} \ 17 | general.checkpoint="checkpoints/s3dis/scannet_pretrained/scannet_pretrained.ckpt" \ 18 | trainer.check_val_every_n_epoch=10 \ 19 | optimizer.lr=0.00001 20 | 21 | python main_instance_segmentation.py \ 22 | general.project_name="s3dis_eval" \ 23 | general.experiment_name="area${CURR_AREA}_pretrained_eps_${CURR_DBSCAN}_topk_${CURR_TOPK}_q_${CURR_QUERY}" \ 24 | general.checkpoint="checkpoints/s3dis/scannet_pretrained/area${CURR_AREA}.ckpt" \ 25 | general.train_mode=false \ 26 | data.batch_size=4 \ 27 | data/datasets=s3dis \ 28 | general.num_targets=14 \ 29 | data.num_labels=13 \ 30 | general.area=${CURR_AREA} \ 31 | model.num_queries=${CURR_QUERY} \ 32 | general.topk_per_image=${CURR_TOPK} \ 33 | general.use_dbscan=true \ 34 | general.dbscan_eps=${CURR_DBSCAN} 35 | -------------------------------------------------------------------------------- /scripts/scannet/scannet_benchmark.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | export OMP_NUM_THREADS=3 # speeds up MinkowskiEngine 3 | 4 | CURR_DBSCAN=0.95 5 | CURR_TOPK=300 6 | CURR_QUERY=150 7 | 8 | # TRAIN 9 | python main_instance_segmentation.py \ 10 | general.experiment_name="benchmark" \ 11 | general.eval_on_segments=true \ 12 | general.train_on_segments=true \ 13 | data.train_mode=train_validation 14 | 15 | # TEST 16 | python main_instance_segmentation.py \ 17 | general.experiment_name="benchmark_query_${CURR_QUERY}_topk_${CURR_TOPK}_dbscan_${CURR_DBSCAN}" \ 18 | general.project_name="scannet_eval" \ 19 | general.checkpoint='checkpoints/scannet/scannet_benchmark.ckpt' \ 20 | general.eval_on_segments=true \ 21 | general.train_on_segments=true \ 22 | general.train_mode=false \ 23 | general.export=true \ 24 | data.test_mode=test \ 25 | model.num_queries=${CURR_QUERY} \ 26 | general.topk_per_image=${CURR_TOPK} \ 27 | general.use_dbscan=true \ 28 | general.dbscan_eps=${CURR_DBSCAN} 29 | -------------------------------------------------------------------------------- /scripts/scannet/scannet_df.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # export OMP_NUM_THREADS=3 # speeds up MinkowskiEngine 3 | 4 | CURR_DBSCAN=0.95 5 | CURR_TOPK=-1 6 | CURR_QUERY=150 7 | # # CURR_AREA=5 8 | 9 | # TRAIN 10 | CUDA_VISIBLE_DEVICES=1 python main_instance_segmentation.py \ 11 | general.experiment_name="train" \ 12 | general.eval_on_segments=true \ 13 | general.train_on_segments=true \ 14 | general.checkpoint="base model path" \ 15 | data.df=1 16 | # data.df represents the percentage of annotated data to be used. 17 | # It can be set to 1, 5, 10, or 20 to indicate the desired percentage. -------------------------------------------------------------------------------- /scripts/scannet/scannet_pretrain_for_s3dis.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | export OMP_NUM_THREADS=3 # speeds up MinkowskiEngine 3 | 4 | # TRAIN 5 | python main_instance_segmentation.py \ 6 | general.experiment_name="pretrain_for_s3dis" \ 7 | data.train_mode=train_validation -------------------------------------------------------------------------------- /scripts/scannet/scannet_val.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # export OMP_NUM_THREADS=3 # speeds up MinkowskiEngine 3 | 4 | CURR_DBSCAN=0.95 5 | CURR_TOPK=-1 6 | CURR_QUERY=150 7 | # # CURR_AREA=5 8 | 9 | # TRAIN 10 | CUDA_VISIBLE_DEVICES=1 python main_instance_segmentation.py \ 11 | general.experiment_name="train" \ 12 | general.eval_on_segments=true \ 13 | general.train_on_segments=true \ 14 | data.layer0_path="path of init segments" \ 15 | data.part_path="path of part pseudo label" \ 16 | data.objrct_path="path of object pseudo label" 17 | 18 | 19 | 20 | 21 | # # TEST 22 | CUDA_VISIBLE_DEVICES=0 python main_instance_segmentation.py \ 23 | general.experiment_name="test" \ 24 | general.project_name="scannet_eval" \ 25 | general.checkpoint='path of ckpt' \ 26 | data/datasets=scannet \ 27 | data.layer0_path="path of init segments" \ 28 | general.train_mode=false \ 29 | general.eval_on_segments=true \ 30 | general.train_on_segments=true \ 31 | model.num_queries=${CURR_QUERY} \ 32 | general.topk_per_image=${CURR_TOPK} \ 33 | general.use_dbscan=true \ 34 | general.dbscan_eps=${CURR_DBSCAN} 35 | -------------------------------------------------------------------------------- /scripts/scannet200/scannet200_benchmark.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | export OMP_NUM_THREADS=3 # speeds up MinkowskiEngine 3 | 4 | CURR_DBSCAN=0.95 5 | CURR_TOPK=300 6 | CURR_QUERY=150 7 | CURR_T=0.001 8 | 9 | # TRAIN 10 | python main_instance_segmentation.py \ 11 | general.experiment_name="scannet200_benchmark_train_validation_test" \ 12 | general.project_name="scannet200" \ 13 | data/datasets=scannet200 \ 14 | general.num_targets=201 \ 15 | data.num_labels=200 \ 16 | general.eval_on_segments=true \ 17 | general.train_on_segments=true \ 18 | data.train_mode=train_validation 19 | 20 | # TEST 21 | python main_instance_segmentation.py \ 22 | general.experiment_name="scannet200_benchmark_query_${CURR_QUERY}_topk_${CURR_TOPK}_dbscan_${CURR_DBSCAN}_export_${CURR_T}" \ 23 | general.project_name="scannet200_eval" \ 24 | general.checkpoint="checkpoints/scannet200/scannet200_benchmark.ckpt" \ 25 | data/datasets=scannet200 \ 26 | general.num_targets=201 \ 27 | data.num_labels=200 \ 28 | general.eval_on_segments=true \ 29 | general.train_on_segments=true \ 30 | general.train_mode=false \ 31 | model.num_queries=${CURR_QUERY} \ 32 | general.topk_per_image=${CURR_TOPK} \ 33 | general.use_dbscan=true \ 34 | general.dbscan_eps=${CURR_DBSCAN} \ 35 | general.export=true \ 36 | data.test_mode=test \ 37 | general.export_threshold=${CURR_T} 38 | -------------------------------------------------------------------------------- /scripts/scannet200/scannet200_val.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | export OMP_NUM_THREADS=3 # speeds up MinkowskiEngine 3 | 4 | CURR_DBSCAN=0.95 5 | CURR_TOPK=750 6 | CURR_QUERY=150 7 | 8 | # TRAIN 9 | # python main_instance_segmentation.py \ 10 | # general.experiment_name="scannet200_val" \ 11 | # general.project_name="scannet200" \ 12 | # data/datasets=scannet200 \ 13 | # general.num_targets=201 \ 14 | # data.num_labels=200 \ 15 | # general.eval_on_segments=true \ 16 | # general.train_on_segments=true 17 | 18 | # TEST 19 | python main_instance_segmentation.py \ 20 | general.experiment_name="scannet200_val_query_${CURR_QUERY}_topk_${CURR_TOPK}_dbscan_${CURR_DBSCAN}" \ 21 | general.project_name="scannet200_eval" \ 22 | general.checkpoint="checkpoints/scannet200/scannet200_val.ckpt" \ 23 | data/datasets=scannet200 \ 24 | general.num_targets=201 \ 25 | data.num_labels=200 \ 26 | general.eval_on_segments=true \ 27 | general.train_on_segments=true \ 28 | general.train_mode=false \ 29 | model.num_queries=${CURR_QUERY} \ 30 | general.topk_per_image=${CURR_TOPK} \ 31 | general.use_dbscan=true \ 32 | general.dbscan_eps=${CURR_DBSCAN} 33 | -------------------------------------------------------------------------------- /scripts/stpls3d/merge_exports.py: -------------------------------------------------------------------------------- 1 | import os 2 | import shutil 3 | from glob import glob 4 | from tqdm import tqdm 5 | 6 | base_path = "INSERT_WORKING_DIRECTORY" 7 | vs03 = f"{base_path}/benchmark_03" 8 | vs02 = f"{base_path}/benchmark_02" 9 | 10 | target_path = "INSERT_TARGET_DIRECTORY" 11 | 12 | print("COPY MASKS FILES 1/2 ...") 13 | shutil.copytree(f"{vs02}/pred_mask", f"{target_path}/pred_mask_02") 14 | print("COPY MASKS FILES 2/2 ...") 15 | shutil.copytree(f"{vs03}/pred_mask", f"{target_path}/pred_mask_03") 16 | 17 | for scene03 in tqdm(glob(f"{vs03}/*.txt")): 18 | instances = [] 19 | with open(scene03, "r") as file03: 20 | while line := file03.readline().rstrip(): 21 | mask_path, class_id, score = line.split(" ") 22 | 23 | if int(class_id) in [1, 3, 4, 7, 8, 11, 12, 13]: 24 | instances.append(f'{mask_path.replace("pred_mask", "pred_mask_03")} {class_id} {score}') 25 | print(instances[-1]) 26 | else: 27 | print(f'DELETE {target_path}/{mask_path.replace("pred_mask", "pred_mask_03")}') 28 | os.remove(f'{target_path}/{mask_path.replace("pred_mask", "pred_mask_03")}') 29 | 30 | with open(f'{vs02}/{scene03.split("/")[-1]}', "r") as file02: 31 | while line := file02.readline().rstrip(): 32 | mask_path, class_id, score = line.split(" ") 33 | 34 | if int(class_id) not in [1, 3, 4, 7, 8, 11, 12, 13]: 35 | instances.append(f'{mask_path.replace("pred_mask", "pred_mask_02")} {class_id} {score}') 36 | print(instances[-1]) 37 | else: 38 | print(f'DELETE {target_path}/{mask_path.replace("pred_mask", "pred_mask_02")}') 39 | os.remove(f'{target_path}/{mask_path.replace("pred_mask", "pred_mask_02")}') 40 | 41 | with open(f'{target_path}/{scene03.split("/")[-1]}', 'w') as fout: 42 | for line in instances: 43 | fout.write(f"{line}\n") 44 | -------------------------------------------------------------------------------- /scripts/stpls3d/stpls3d_benchmark.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | export OMP_NUM_THREADS=3 3 | 4 | CURR_DBSCAN=12.5 5 | CURR_TOPK=200 6 | CURR_QUERY=160 7 | CURR_SIZE=54 8 | CURR_THRESHOLD=0.01 9 | 10 | # TRAIN network 1 with voxel size 0.333 11 | python main_instance_segmentation.py \ 12 | general.experiment_name="benchmark_03" \ 13 | general.project_name="stpls3d" \ 14 | data/datasets=stpls3d \ 15 | general.num_targets=15 \ 16 | data.num_labels=15 \ 17 | data.voxel_size=0.333 \ 18 | data.num_workers=10 \ 19 | data.cache_data=true \ 20 | data.cropping_v1=false \ 21 | general.reps_per_epoch=100 \ 22 | model.num_queries=${CURR_QUERY} \ 23 | general.on_crops=true \ 24 | model.config.backbone._target_=models.Res16UNet18B \ 25 | data.crop_length=${CURR_SIZE} \ 26 | general.eval_inner_core=50.0 \ 27 | data.train_mode=train_validation 28 | 29 | # TRAIN network 2 with voxel size 0.2 and larger backbone 30 | python main_instance_segmentation.py \ 31 | general.experiment_name="benchmark_02" \ 32 | general.project_name="stpls3d" \ 33 | data/datasets=stpls3d \ 34 | general.num_targets=15 \ 35 | data.num_labels=15 \ 36 | data.voxel_size=0.2 \ 37 | data.num_workers=10 \ 38 | data.cache_data=true \ 39 | data.cropping_v1=false \ 40 | general.reps_per_epoch=100 \ 41 | model.num_queries=${CURR_QUERY} \ 42 | general.on_crops=true \ 43 | data.crop_length=${CURR_SIZE} \ 44 | general.eval_inner_core=50.0 \ 45 | data.train_mode=train_validation 46 | 47 | # TEST network 1 48 | python main_instance_segmentation.py \ 49 | general.experiment_name="benchmark_03_query_${CURR_QUERY}_topk_${CURR_TOPK}_dbscan_${CURR_DBSCAN}_size_${CURR_SIZE}_T_${CURR_THRESHOLD}" \ 50 | general.project_name="stpls3d_eval" \ 51 | data/datasets=stpls3d \ 52 | general.num_targets=15 \ 53 | data.num_labels=15 \ 54 | data.voxel_size=0.333 \ 55 | data.num_workers=10 \ 56 | data.cache_data=true \ 57 | data.cropping_v1=false \ 58 | general.reps_per_epoch=100 \ 59 | model.num_queries=${CURR_QUERY} \ 60 | general.on_crops=true \ 61 | model.config.backbone._target_=models.Res16UNet18B \ 62 | general.train_mode=false \ 63 | general.checkpoint="checkpoints/stpls3d/stpls3d_benchmark_03.ckpt" \ 64 | data.crop_length=${CURR_SIZE} \ 65 | general.eval_inner_core=50.0 \ 66 | general.topk_per_image=${CURR_TOPK} \ 67 | general.use_dbscan=true \ 68 | general.dbscan_eps=${CURR_DBSCAN} \ 69 | data.test_mode=test \ 70 | general.export=true 71 | 72 | # TEST network 2 73 | python main_instance_segmentation.py \ 74 | general.experiment_name="benchmark_02_query_${CURR_QUERY}_topk_${CURR_TOPK}_dbscan_${CURR_DBSCAN}_size_${CURR_SIZE}_T_${CURR_THRESHOLD}" \ 75 | general.project_name="stpls3d_eval" \ 76 | data/datasets=stpls3d \ 77 | general.num_targets=15 \ 78 | data.num_labels=15 \ 79 | data.voxel_size=0.2 \ 80 | data.num_workers=10 \ 81 | data.cache_data=true \ 82 | data.cropping_v1=false \ 83 | general.reps_per_epoch=100 \ 84 | model.num_queries=${CURR_QUERY} \ 85 | general.on_crops=true \ 86 | general.train_mode=false \ 87 | general.checkpoint="checkpoints/stpls3d/stpls3d_benchmark_02.ckpt" \ 88 | data.crop_length=${CURR_SIZE} \ 89 | general.eval_inner_core=50.0 \ 90 | general.topk_per_image=${CURR_TOPK} \ 91 | general.use_dbscan=true \ 92 | general.dbscan_eps=${CURR_DBSCAN} \ 93 | data.test_mode=test \ 94 | general.export=true 95 | 96 | # COMBINE OUTPUTS OF ENSEMBLE 97 | # VOXEL SIZE 0.2 FOR OBJECTS OF SMALL CLASSES; VOXEL SIZE 0.333 FOR OBJECTS OF LARGE CLASS CATEGORIES 98 | # TODO FILL IN PATHS 99 | python merge_exports.py 100 | -------------------------------------------------------------------------------- /scripts/stpls3d/stpls3d_val.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | export OMP_NUM_THREADS=3 3 | 4 | CURR_DBSCAN=14.0 5 | CURR_TOPK=750 6 | CURR_QUERY=160 7 | CURR_SIZE=54 8 | 9 | # TRAIN 10 | python main_instance_segmentation.py \ 11 | general.experiment_name="validation" \ 12 | general.project_name="stpls3d" \ 13 | data/datasets=stpls3d \ 14 | general.num_targets=15 \ 15 | data.num_labels=15 \ 16 | data.voxel_size=0.333 \ 17 | data.num_workers=10 \ 18 | data.cache_data=true \ 19 | data.cropping_v1=false \ 20 | general.reps_per_epoch=100 \ 21 | model.num_queries=${CURR_QUERY} \ 22 | general.on_crops=true \ 23 | model.config.backbone._target_=models.Res16UNet18B \ 24 | data.crop_length=${CURR_SIZE} \ 25 | general.eval_inner_core=50.0 26 | 27 | # TEST 28 | python main_instance_segmentation.py \ 29 | general.experiment_name="validation_query_${CURR_QUERY}_topk_${CURR_TOPK}_dbscan_${CURR_DBSCAN}_size_${CURR_SIZE}" \ 30 | general.project_name="stpls3d_eval" \ 31 | data/datasets=stpls3d \ 32 | general.num_targets=15 \ 33 | data.num_labels=15 \ 34 | data.voxel_size=0.333 \ 35 | data.num_workers=10 \ 36 | data.cache_data=true \ 37 | data.cropping_v1=false \ 38 | general.reps_per_epoch=100 \ 39 | model.num_queries=${CURR_QUERY} \ 40 | general.on_crops=true \ 41 | model.config.backbone._target_=models.Res16UNet18B \ 42 | general.train_mode=false \ 43 | general.checkpoint="checkpoints/stpls3d/stpls3d_val.ckpt" \ 44 | data.crop_length=${CURR_SIZE} \ 45 | general.eval_inner_core=50.0 \ 46 | general.topk_per_image=${CURR_TOPK} \ 47 | general.use_dbscan=true \ 48 | general.dbscan_eps=${CURR_DBSCAN} 49 | -------------------------------------------------------------------------------- /third_party/pointnet2/__pycache__/pointnet2_utils.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/third_party/pointnet2/__pycache__/pointnet2_utils.cpython-310.pyc -------------------------------------------------------------------------------- /third_party/pointnet2/__pycache__/pointnet2_utils.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/third_party/pointnet2/__pycache__/pointnet2_utils.cpython-38.pyc -------------------------------------------------------------------------------- /third_party/pointnet2/__pycache__/pytorch_utils.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/third_party/pointnet2/__pycache__/pytorch_utils.cpython-310.pyc -------------------------------------------------------------------------------- /third_party/pointnet2/__pycache__/pytorch_utils.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/third_party/pointnet2/__pycache__/pytorch_utils.cpython-38.pyc -------------------------------------------------------------------------------- /third_party/pointnet2/_ext_src/include/ball_query.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | 3 | #pragma once 4 | #include 5 | 6 | at::Tensor ball_query(at::Tensor new_xyz, at::Tensor xyz, const float radius, 7 | const int nsample); 8 | -------------------------------------------------------------------------------- /third_party/pointnet2/_ext_src/include/cuda_utils.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | 3 | #ifndef _CUDA_UTILS_H 4 | #define _CUDA_UTILS_H 5 | 6 | #include 7 | #include 8 | #include 9 | 10 | #include 11 | #include 12 | 13 | #include 14 | 15 | #define TOTAL_THREADS 512 16 | 17 | inline int opt_n_threads(int work_size) { 18 | const int pow_2 = std::log(static_cast(work_size)) / std::log(2.0); 19 | 20 | return max(min(1 << pow_2, TOTAL_THREADS), 1); 21 | } 22 | 23 | inline dim3 opt_block_config(int x, int y) { 24 | const int x_threads = opt_n_threads(x); 25 | const int y_threads = 26 | max(min(opt_n_threads(y), TOTAL_THREADS / x_threads), 1); 27 | dim3 block_config(x_threads, y_threads, 1); 28 | 29 | return block_config; 30 | } 31 | 32 | #define CUDA_CHECK_ERRORS() \ 33 | do { \ 34 | cudaError_t err = cudaGetLastError(); \ 35 | if (cudaSuccess != err) { \ 36 | fprintf(stderr, "CUDA kernel failed : %s\n%s at L:%d in %s\n", \ 37 | cudaGetErrorString(err), __PRETTY_FUNCTION__, __LINE__, \ 38 | __FILE__); \ 39 | exit(-1); \ 40 | } \ 41 | } while (0) 42 | 43 | #endif 44 | -------------------------------------------------------------------------------- /third_party/pointnet2/_ext_src/include/group_points.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | 3 | 4 | #pragma once 5 | #include 6 | 7 | at::Tensor group_points(at::Tensor points, at::Tensor idx); 8 | at::Tensor group_points_grad(at::Tensor grad_out, at::Tensor idx, const int n); 9 | -------------------------------------------------------------------------------- /third_party/pointnet2/_ext_src/include/interpolate.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | 3 | #pragma once 4 | 5 | #include 6 | #include 7 | 8 | std::vector three_nn(at::Tensor unknowns, at::Tensor knows); 9 | at::Tensor three_interpolate(at::Tensor points, at::Tensor idx, 10 | at::Tensor weight); 11 | at::Tensor three_interpolate_grad(at::Tensor grad_out, at::Tensor idx, 12 | at::Tensor weight, const int m); 13 | -------------------------------------------------------------------------------- /third_party/pointnet2/_ext_src/include/sampling.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | 3 | 4 | #pragma once 5 | #include 6 | 7 | at::Tensor gather_points(at::Tensor points, at::Tensor idx); 8 | at::Tensor gather_points_grad(at::Tensor grad_out, at::Tensor idx, const int n); 9 | at::Tensor furthest_point_sampling(at::Tensor points, const int nsamples); 10 | -------------------------------------------------------------------------------- /third_party/pointnet2/_ext_src/include/utils.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | 3 | 4 | #pragma once 5 | #include 6 | #include 7 | 8 | #define CHECK_CUDA(x) \ 9 | do { \ 10 | AT_ASSERT(x.is_cuda(), #x " must be a CUDA tensor"); \ 11 | } while (0) 12 | 13 | #define CHECK_CONTIGUOUS(x) \ 14 | do { \ 15 | AT_ASSERT(x.is_contiguous(), #x " must be a contiguous tensor"); \ 16 | } while (0) 17 | 18 | #define CHECK_IS_INT(x) \ 19 | do { \ 20 | AT_ASSERT(x.scalar_type() == at::ScalarType::Int, \ 21 | #x " must be an int tensor"); \ 22 | } while (0) 23 | 24 | #define CHECK_IS_FLOAT(x) \ 25 | do { \ 26 | AT_ASSERT(x.scalar_type() == at::ScalarType::Float, \ 27 | #x " must be a float tensor"); \ 28 | } while (0) 29 | -------------------------------------------------------------------------------- /third_party/pointnet2/_ext_src/src/ball_query.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | 3 | 4 | #include "ball_query.h" 5 | #include "utils.h" 6 | 7 | void query_ball_point_kernel_wrapper(int b, int n, int m, float radius, 8 | int nsample, const float *new_xyz, 9 | const float *xyz, int *idx); 10 | 11 | at::Tensor ball_query(at::Tensor new_xyz, at::Tensor xyz, const float radius, 12 | const int nsample) { 13 | CHECK_CONTIGUOUS(new_xyz); 14 | CHECK_CONTIGUOUS(xyz); 15 | CHECK_IS_FLOAT(new_xyz); 16 | CHECK_IS_FLOAT(xyz); 17 | 18 | if (new_xyz.is_cuda()) { 19 | CHECK_CUDA(xyz); 20 | } 21 | 22 | at::Tensor idx = 23 | torch::zeros({new_xyz.size(0), new_xyz.size(1), nsample}, 24 | at::device(new_xyz.device()).dtype(at::ScalarType::Int)); 25 | 26 | if (new_xyz.is_cuda()) { 27 | query_ball_point_kernel_wrapper(xyz.size(0), xyz.size(1), new_xyz.size(1), 28 | radius, nsample, new_xyz.data(), 29 | xyz.data(), idx.data()); 30 | } else { 31 | AT_ASSERT(false, "CPU not supported"); 32 | } 33 | 34 | return idx; 35 | } 36 | -------------------------------------------------------------------------------- /third_party/pointnet2/_ext_src/src/ball_query_gpu.cu: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | 3 | 4 | #include 5 | #include 6 | #include 7 | 8 | #include "cuda_utils.h" 9 | 10 | // input: new_xyz(b, m, 3) xyz(b, n, 3) 11 | // output: idx(b, m, nsample) 12 | __global__ void query_ball_point_kernel(int b, int n, int m, float radius, 13 | int nsample, 14 | const float *__restrict__ new_xyz, 15 | const float *__restrict__ xyz, 16 | int *__restrict__ idx) { 17 | int batch_index = blockIdx.x; 18 | xyz += batch_index * n * 3; 19 | new_xyz += batch_index * m * 3; 20 | idx += m * nsample * batch_index; 21 | 22 | int index = threadIdx.x; 23 | int stride = blockDim.x; 24 | 25 | float radius2 = radius * radius; 26 | for (int j = index; j < m; j += stride) { 27 | float new_x = new_xyz[j * 3 + 0]; 28 | float new_y = new_xyz[j * 3 + 1]; 29 | float new_z = new_xyz[j * 3 + 2]; 30 | for (int k = 0, cnt = 0; k < n && cnt < nsample; ++k) { 31 | float x = xyz[k * 3 + 0]; 32 | float y = xyz[k * 3 + 1]; 33 | float z = xyz[k * 3 + 2]; 34 | float d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) + 35 | (new_z - z) * (new_z - z); 36 | if (d2 < radius2) { 37 | if (cnt == 0) { 38 | for (int l = 0; l < nsample; ++l) { 39 | idx[j * nsample + l] = k; 40 | } 41 | } 42 | idx[j * nsample + cnt] = k; 43 | ++cnt; 44 | } 45 | } 46 | } 47 | } 48 | 49 | void query_ball_point_kernel_wrapper(int b, int n, int m, float radius, 50 | int nsample, const float *new_xyz, 51 | const float *xyz, int *idx) { 52 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 53 | query_ball_point_kernel<<>>( 54 | b, n, m, radius, nsample, new_xyz, xyz, idx); 55 | 56 | CUDA_CHECK_ERRORS(); 57 | } 58 | -------------------------------------------------------------------------------- /third_party/pointnet2/_ext_src/src/bindings.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | 3 | 4 | #include "ball_query.h" 5 | #include "group_points.h" 6 | #include "interpolate.h" 7 | #include "sampling.h" 8 | 9 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 10 | m.def("gather_points", &gather_points); 11 | m.def("gather_points_grad", &gather_points_grad); 12 | m.def("furthest_point_sampling", &furthest_point_sampling); 13 | 14 | m.def("three_nn", &three_nn); 15 | m.def("three_interpolate", &three_interpolate); 16 | m.def("three_interpolate_grad", &three_interpolate_grad); 17 | 18 | m.def("ball_query", &ball_query); 19 | 20 | m.def("group_points", &group_points); 21 | m.def("group_points_grad", &group_points_grad); 22 | } 23 | -------------------------------------------------------------------------------- /third_party/pointnet2/_ext_src/src/group_points.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | 3 | 4 | #include "group_points.h" 5 | #include "utils.h" 6 | 7 | void group_points_kernel_wrapper(int b, int c, int n, int npoints, int nsample, 8 | const float *points, const int *idx, 9 | float *out); 10 | 11 | void group_points_grad_kernel_wrapper(int b, int c, int n, int npoints, 12 | int nsample, const float *grad_out, 13 | const int *idx, float *grad_points); 14 | 15 | at::Tensor group_points(at::Tensor points, at::Tensor idx) { 16 | CHECK_CONTIGUOUS(points); 17 | CHECK_CONTIGUOUS(idx); 18 | CHECK_IS_FLOAT(points); 19 | CHECK_IS_INT(idx); 20 | 21 | if (points.is_cuda()) { 22 | CHECK_CUDA(idx); 23 | } 24 | 25 | at::Tensor output = 26 | torch::zeros({points.size(0), points.size(1), idx.size(1), idx.size(2)}, 27 | at::device(points.device()).dtype(at::ScalarType::Float)); 28 | 29 | if (points.is_cuda()) { 30 | group_points_kernel_wrapper(points.size(0), points.size(1), points.size(2), 31 | idx.size(1), idx.size(2), points.data(), 32 | idx.data(), output.data()); 33 | } else { 34 | AT_ASSERT(false, "CPU not supported"); 35 | } 36 | 37 | return output; 38 | } 39 | 40 | at::Tensor group_points_grad(at::Tensor grad_out, at::Tensor idx, const int n) { 41 | CHECK_CONTIGUOUS(grad_out); 42 | CHECK_CONTIGUOUS(idx); 43 | CHECK_IS_FLOAT(grad_out); 44 | CHECK_IS_INT(idx); 45 | 46 | if (grad_out.is_cuda()) { 47 | CHECK_CUDA(idx); 48 | } 49 | 50 | at::Tensor output = 51 | torch::zeros({grad_out.size(0), grad_out.size(1), n}, 52 | at::device(grad_out.device()).dtype(at::ScalarType::Float)); 53 | 54 | if (grad_out.is_cuda()) { 55 | group_points_grad_kernel_wrapper( 56 | grad_out.size(0), grad_out.size(1), n, idx.size(1), idx.size(2), 57 | grad_out.data(), idx.data(), output.data()); 58 | } else { 59 | AT_ASSERT(false, "CPU not supported"); 60 | } 61 | 62 | return output; 63 | } 64 | -------------------------------------------------------------------------------- /third_party/pointnet2/_ext_src/src/group_points_gpu.cu: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | 3 | 4 | #include 5 | #include 6 | 7 | #include "cuda_utils.h" 8 | 9 | // input: points(b, c, n) idx(b, npoints, nsample) 10 | // output: out(b, c, npoints, nsample) 11 | __global__ void group_points_kernel(int b, int c, int n, int npoints, 12 | int nsample, 13 | const float *__restrict__ points, 14 | const int *__restrict__ idx, 15 | float *__restrict__ out) { 16 | int batch_index = blockIdx.x; 17 | points += batch_index * n * c; 18 | idx += batch_index * npoints * nsample; 19 | out += batch_index * npoints * nsample * c; 20 | 21 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 22 | const int stride = blockDim.y * blockDim.x; 23 | for (int i = index; i < c * npoints; i += stride) { 24 | const int l = i / npoints; 25 | const int j = i % npoints; 26 | for (int k = 0; k < nsample; ++k) { 27 | int ii = idx[j * nsample + k]; 28 | out[(l * npoints + j) * nsample + k] = points[l * n + ii]; 29 | } 30 | } 31 | } 32 | 33 | void group_points_kernel_wrapper(int b, int c, int n, int npoints, int nsample, 34 | const float *points, const int *idx, 35 | float *out) { 36 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 37 | 38 | group_points_kernel<<>>( 39 | b, c, n, npoints, nsample, points, idx, out); 40 | 41 | CUDA_CHECK_ERRORS(); 42 | } 43 | 44 | // input: grad_out(b, c, npoints, nsample), idx(b, npoints, nsample) 45 | // output: grad_points(b, c, n) 46 | __global__ void group_points_grad_kernel(int b, int c, int n, int npoints, 47 | int nsample, 48 | const float *__restrict__ grad_out, 49 | const int *__restrict__ idx, 50 | float *__restrict__ grad_points) { 51 | int batch_index = blockIdx.x; 52 | grad_out += batch_index * npoints * nsample * c; 53 | idx += batch_index * npoints * nsample; 54 | grad_points += batch_index * n * c; 55 | 56 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 57 | const int stride = blockDim.y * blockDim.x; 58 | for (int i = index; i < c * npoints; i += stride) { 59 | const int l = i / npoints; 60 | const int j = i % npoints; 61 | for (int k = 0; k < nsample; ++k) { 62 | int ii = idx[j * nsample + k]; 63 | atomicAdd(grad_points + l * n + ii, 64 | grad_out[(l * npoints + j) * nsample + k]); 65 | } 66 | } 67 | } 68 | 69 | void group_points_grad_kernel_wrapper(int b, int c, int n, int npoints, 70 | int nsample, const float *grad_out, 71 | const int *idx, float *grad_points) { 72 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 73 | 74 | group_points_grad_kernel<<>>( 75 | b, c, n, npoints, nsample, grad_out, idx, grad_points); 76 | 77 | CUDA_CHECK_ERRORS(); 78 | } 79 | -------------------------------------------------------------------------------- /third_party/pointnet2/_ext_src/src/interpolate.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | 3 | #include "interpolate.h" 4 | #include "utils.h" 5 | 6 | void three_nn_kernel_wrapper(int b, int n, int m, const float *unknown, 7 | const float *known, float *dist2, int *idx); 8 | void three_interpolate_kernel_wrapper(int b, int c, int m, int n, 9 | const float *points, const int *idx, 10 | const float *weight, float *out); 11 | void three_interpolate_grad_kernel_wrapper(int b, int c, int n, int m, 12 | const float *grad_out, 13 | const int *idx, const float *weight, 14 | float *grad_points); 15 | 16 | std::vector three_nn(at::Tensor unknowns, at::Tensor knows) { 17 | CHECK_CONTIGUOUS(unknowns); 18 | CHECK_CONTIGUOUS(knows); 19 | CHECK_IS_FLOAT(unknowns); 20 | CHECK_IS_FLOAT(knows); 21 | 22 | if (unknowns.is_cuda()) { 23 | CHECK_CUDA(knows); 24 | } 25 | 26 | at::Tensor idx = 27 | torch::zeros({unknowns.size(0), unknowns.size(1), 3}, 28 | at::device(unknowns.device()).dtype(at::ScalarType::Int)); 29 | at::Tensor dist2 = 30 | torch::zeros({unknowns.size(0), unknowns.size(1), 3}, 31 | at::device(unknowns.device()).dtype(at::ScalarType::Float)); 32 | 33 | if (unknowns.is_cuda()) { 34 | three_nn_kernel_wrapper(unknowns.size(0), unknowns.size(1), knows.size(1), 35 | unknowns.data(), knows.data(), 36 | dist2.data(), idx.data()); 37 | } else { 38 | AT_ASSERT(false, "CPU not supported"); 39 | } 40 | 41 | return {dist2, idx}; 42 | } 43 | 44 | at::Tensor three_interpolate(at::Tensor points, at::Tensor idx, 45 | at::Tensor weight) { 46 | CHECK_CONTIGUOUS(points); 47 | CHECK_CONTIGUOUS(idx); 48 | CHECK_CONTIGUOUS(weight); 49 | CHECK_IS_FLOAT(points); 50 | CHECK_IS_INT(idx); 51 | CHECK_IS_FLOAT(weight); 52 | 53 | if (points.is_cuda()) { 54 | CHECK_CUDA(idx); 55 | CHECK_CUDA(weight); 56 | } 57 | 58 | at::Tensor output = 59 | torch::zeros({points.size(0), points.size(1), idx.size(1)}, 60 | at::device(points.device()).dtype(at::ScalarType::Float)); 61 | 62 | if (points.is_cuda()) { 63 | three_interpolate_kernel_wrapper( 64 | points.size(0), points.size(1), points.size(2), idx.size(1), 65 | points.data(), idx.data(), weight.data(), 66 | output.data()); 67 | } else { 68 | AT_ASSERT(false, "CPU not supported"); 69 | } 70 | 71 | return output; 72 | } 73 | at::Tensor three_interpolate_grad(at::Tensor grad_out, at::Tensor idx, 74 | at::Tensor weight, const int m) { 75 | CHECK_CONTIGUOUS(grad_out); 76 | CHECK_CONTIGUOUS(idx); 77 | CHECK_CONTIGUOUS(weight); 78 | CHECK_IS_FLOAT(grad_out); 79 | CHECK_IS_INT(idx); 80 | CHECK_IS_FLOAT(weight); 81 | 82 | if (grad_out.is_cuda()) { 83 | CHECK_CUDA(idx); 84 | CHECK_CUDA(weight); 85 | } 86 | 87 | at::Tensor output = 88 | torch::zeros({grad_out.size(0), grad_out.size(1), m}, 89 | at::device(grad_out.device()).dtype(at::ScalarType::Float)); 90 | 91 | if (grad_out.is_cuda()) { 92 | three_interpolate_grad_kernel_wrapper( 93 | grad_out.size(0), grad_out.size(1), grad_out.size(2), m, 94 | grad_out.data(), idx.data(), weight.data(), 95 | output.data()); 96 | } else { 97 | AT_ASSERT(false, "CPU not supported"); 98 | } 99 | 100 | return output; 101 | } 102 | -------------------------------------------------------------------------------- /third_party/pointnet2/_ext_src/src/sampling.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | 3 | #include "sampling.h" 4 | #include "utils.h" 5 | 6 | void gather_points_kernel_wrapper(int b, int c, int n, int npoints, 7 | const float *points, const int *idx, 8 | float *out); 9 | void gather_points_grad_kernel_wrapper(int b, int c, int n, int npoints, 10 | const float *grad_out, const int *idx, 11 | float *grad_points); 12 | 13 | void furthest_point_sampling_kernel_wrapper(int b, int n, int m, 14 | const float *dataset, float *temp, 15 | int *idxs); 16 | 17 | at::Tensor gather_points(at::Tensor points, at::Tensor idx) { 18 | CHECK_CONTIGUOUS(points); 19 | CHECK_CONTIGUOUS(idx); 20 | CHECK_IS_FLOAT(points); 21 | CHECK_IS_INT(idx); 22 | 23 | if (points.is_cuda()) { 24 | CHECK_CUDA(idx); 25 | } 26 | 27 | at::Tensor output = 28 | torch::zeros({points.size(0), points.size(1), idx.size(1)}, 29 | at::device(points.device()).dtype(at::ScalarType::Float)); 30 | 31 | if (points.is_cuda()) { 32 | gather_points_kernel_wrapper(points.size(0), points.size(1), points.size(2), 33 | idx.size(1), points.data(), 34 | idx.data(), output.data()); 35 | } else { 36 | AT_ASSERT(false, "CPU not supported"); 37 | } 38 | 39 | return output; 40 | } 41 | 42 | at::Tensor gather_points_grad(at::Tensor grad_out, at::Tensor idx, 43 | const int n) { 44 | CHECK_CONTIGUOUS(grad_out); 45 | CHECK_CONTIGUOUS(idx); 46 | CHECK_IS_FLOAT(grad_out); 47 | CHECK_IS_INT(idx); 48 | 49 | if (grad_out.is_cuda()) { 50 | CHECK_CUDA(idx); 51 | } 52 | 53 | at::Tensor output = 54 | torch::zeros({grad_out.size(0), grad_out.size(1), n}, 55 | at::device(grad_out.device()).dtype(at::ScalarType::Float)); 56 | 57 | if (grad_out.is_cuda()) { 58 | gather_points_grad_kernel_wrapper(grad_out.size(0), grad_out.size(1), n, 59 | idx.size(1), grad_out.data(), 60 | idx.data(), output.data()); 61 | } else { 62 | AT_ASSERT(false, "CPU not supported"); 63 | } 64 | 65 | return output; 66 | } 67 | at::Tensor furthest_point_sampling(at::Tensor points, const int nsamples) { 68 | CHECK_CONTIGUOUS(points); 69 | CHECK_IS_FLOAT(points); 70 | 71 | at::Tensor output = 72 | torch::zeros({points.size(0), nsamples}, 73 | at::device(points.device()).dtype(at::ScalarType::Int)); 74 | 75 | at::Tensor tmp = 76 | torch::full({points.size(0), points.size(1)}, 1e10, 77 | at::device(points.device()).dtype(at::ScalarType::Float)); 78 | 79 | if (points.is_cuda()) { 80 | furthest_point_sampling_kernel_wrapper( 81 | points.size(0), points.size(1), nsamples, points.data(), 82 | tmp.data(), output.data()); 83 | } else { 84 | AT_ASSERT(false, "CPU not supported"); 85 | } 86 | 87 | return output; 88 | } 89 | -------------------------------------------------------------------------------- /third_party/pointnet2/pointnet2_test.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. 2 | 3 | ''' Testing customized ops. ''' 4 | 5 | import torch 6 | from torch.autograd import gradcheck 7 | import numpy as np 8 | 9 | import os 10 | import sys 11 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 12 | sys.path.append(BASE_DIR) 13 | import pointnet2_utils 14 | 15 | def test_interpolation_grad(): 16 | batch_size = 1 17 | feat_dim = 2 18 | m = 4 19 | feats = torch.randn(batch_size, feat_dim, m, requires_grad=True).float().cuda() 20 | 21 | def interpolate_func(inputs): 22 | idx = torch.from_numpy(np.array([[[0,1,2],[1,2,3]]])).int().cuda() 23 | weight = torch.from_numpy(np.array([[[1,1,1],[2,2,2]]])).float().cuda() 24 | interpolated_feats = pointnet2_utils.three_interpolate(inputs, idx, weight) 25 | return interpolated_feats 26 | 27 | assert (gradcheck(interpolate_func, feats, atol=1e-1, rtol=1e-1)) 28 | 29 | if __name__=='__main__': 30 | test_interpolation_grad() 31 | -------------------------------------------------------------------------------- /third_party/pointnet2/setup.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | from setuptools import setup 7 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 8 | import glob 9 | import os.path as osp 10 | 11 | this_dir = osp.dirname(osp.abspath(__file__)) 12 | 13 | _ext_src_root = "_ext_src" 14 | _ext_sources = glob.glob("{}/src/*.cpp".format(_ext_src_root)) + glob.glob( 15 | "{}/src/*.cu".format(_ext_src_root) 16 | ) 17 | _ext_headers = glob.glob("{}/include/*".format(_ext_src_root)) 18 | 19 | setup( 20 | name='pointnet2', 21 | ext_modules=[ 22 | CUDAExtension( 23 | name='pointnet2._ext', 24 | sources=_ext_sources, 25 | extra_compile_args={ 26 | "cxx": ["-O2", "-I{}".format("{}/include".format(_ext_src_root))], 27 | "nvcc": ["-O2", "-I{}".format("{}/include".format(_ext_src_root))], 28 | }, 29 | include_dirs=[osp.join(this_dir, _ext_src_root, "include")], 30 | ) 31 | ], 32 | cmdclass={ 33 | 'build_ext': BuildExtension 34 | } 35 | ) 36 | -------------------------------------------------------------------------------- /trainer/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/trainer/__init__.py -------------------------------------------------------------------------------- /utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/utils/__init__.py -------------------------------------------------------------------------------- /utils/gradflow_check.py: -------------------------------------------------------------------------------- 1 | """ https://github.com/alwynmathew/gradflow-check """ 2 | import matplotlib.pyplot as plt 3 | import numpy as np 4 | from matplotlib.lines import Line2D 5 | 6 | 7 | def plot_grad_flow(named_parameters): 8 | ave_grads = [] 9 | layers = [] 10 | for n, p in named_parameters: 11 | if (p.requires_grad) and ("bias" not in n): 12 | if p.grad: 13 | layers.append(n) 14 | ave_grads.append(p.grad.abs().mean()) 15 | else: 16 | print(f"{n} - doesn't have gradient computed") 17 | 18 | plt.plot(ave_grads, alpha=0.3, color="b") 19 | plt.hlines(0, 0, len(ave_grads) + 1, linewidth=1, color="k") 20 | plt.xticks(range(0, len(ave_grads), 1), layers, rotation="vertical") 21 | plt.xlim(xmin=0, xmax=len(ave_grads)) 22 | plt.xlabel("Layers") 23 | plt.ylabel("average gradient") 24 | plt.title("Gradient flow") 25 | plt.grid(True) 26 | 27 | 28 | def plot_grad_flow_v2(named_parameters): 29 | """Plots the gradients flowing through different layers in the net during training. 30 | Can be used for checking for possible gradient vanishing / exploding problems. 31 | 32 | Usage: Plug this function in Trainer class after loss.backwards() as 33 | "plot_grad_flow(self.model.named_parameters())" to visualize the gradient flow""" 34 | ave_grads = [] 35 | max_grads = [] 36 | layers = [] 37 | for n, p in named_parameters: 38 | if (p.requires_grad) and ("bias" not in n): 39 | layers.append(n) 40 | if p.grad: 41 | ave_grads.append(p.grad.abs().mean()) 42 | max_grads.append(p.grad.abs().max()) 43 | else: 44 | print(f"{n} - doesn't have gradient computed") 45 | plt.bar(np.arange(len(max_grads)), max_grads, alpha=0.1, lw=1, color="c") 46 | plt.bar(np.arange(len(max_grads)), ave_grads, alpha=0.1, lw=1, color="b") 47 | plt.hlines(0, 0, len(ave_grads) + 1, lw=2, color="k") 48 | plt.xticks(range(0, len(ave_grads), 1), layers, rotation="vertical") 49 | plt.xlim(left=0, right=len(ave_grads)) 50 | plt.ylim(bottom=-0.001, top=0.02) # zoom in on the lower gradient regions 51 | plt.xlabel("Layers") 52 | plt.ylabel("average gradient") 53 | plt.title("Gradient flow") 54 | plt.grid(True) 55 | plt.legend( 56 | [ 57 | Line2D([0], [0], color="c", lw=4), 58 | Line2D([0], [0], color="b", lw=4), 59 | Line2D([0], [0], color="k", lw=4), 60 | ], 61 | ["max-gradient", "mean-gradient", "zero-gradient"], 62 | ) 63 | -------------------------------------------------------------------------------- /utils/kfold.py: -------------------------------------------------------------------------------- 1 | """ Author: https://github.com/yk-szk/stratified_group_kfold """ 2 | import random 3 | import numpy as np 4 | 5 | 6 | class StratifiedGroupKFold: 7 | """ 8 | Stratified Group K-fold with sklearn.model_selection.KFold compabitility. 9 | 10 | Split dataset into k folds with balanced label distribution (stratified) and non-overlapping group. 11 | 12 | Args: 13 | n_splits (int): # of splits 14 | shuffle (bool): Shuffle 15 | seed (int): Seed value for random number generator 16 | """ 17 | 18 | def __init__(self, n_splits, shuffle=True, random_state=None): 19 | self.n_splits = n_splits 20 | self.shuffle = shuffle 21 | self.seed = random_state 22 | 23 | def split(self, X, labels, groups): 24 | assert len(X) == len(labels) == len(groups), "Invalid input length" 25 | assert ( 26 | len(set(groups)) >= self.n_splits 27 | ), "The number of groups needs to be larger than n_splits" 28 | 29 | def encode(v): 30 | s = set(v) 31 | d = {l: i for i, l in enumerate(s)} 32 | return [d[e] for e in v] 33 | 34 | labels, groups = encode(labels), encode(groups) 35 | num_labels, num_groups = max(labels) + 1, max(groups) + 1 36 | label_counts_per_group = np.zeros((num_groups, num_labels), dtype=int) 37 | global_label_dist = np.bincount(labels) 38 | for label, g in zip(labels, groups): 39 | label_counts_per_group[g][label] += 1 40 | 41 | label_counts_per_fold = np.zeros((self.n_splits, num_labels), dtype=int) 42 | groups_per_fold = [set() for _ in range(self.n_splits)] 43 | 44 | def eval_label_counts_per_fold(y_counts, fold): 45 | fold += y_counts 46 | std_per_label = np.std(label_counts_per_fold, axis=0) / global_label_dist 47 | fold -= y_counts 48 | return np.mean(std_per_label) 49 | 50 | groups_and_label_counts = list(enumerate(label_counts_per_group)) 51 | if self.shuffle: 52 | rng = random.Random(self.seed) 53 | mean_std = np.mean(np.std(label_counts_per_group, axis=1)) 54 | groups_and_label_counts.sort( 55 | key=lambda g_counts: -np.std(g_counts[1]) + rng.gauss(0, mean_std) 56 | ) # add rng.gauss to increase the randomness 57 | else: 58 | groups_and_label_counts.sort(key=lambda g_counts: -np.std(g_counts[1])) 59 | 60 | for g, label_counts in groups_and_label_counts: 61 | evals = [ 62 | eval_label_counts_per_fold(label_counts, label_counts_per_fold[i]) 63 | for i in range(self.n_splits) 64 | ] 65 | best_fold = np.argmin(evals) 66 | label_counts_per_fold[best_fold] += label_counts 67 | groups_per_fold[best_fold].add(g) 68 | 69 | all_groups = set(groups) 70 | for test_groups in groups_per_fold: 71 | train_groups = all_groups - test_groups 72 | 73 | train_indices = [i for i, g in enumerate(groups) if g in train_groups] 74 | test_indices = [i for i, g in enumerate(groups) if g in test_groups] 75 | 76 | yield train_indices, test_indices 77 | -------------------------------------------------------------------------------- /utils/point_cloud_utils.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | from typing import List, Optional, Tuple 3 | 4 | import numpy as np 5 | import open3d 6 | from plyfile import PlyData, PlyElement 7 | 8 | 9 | def load_ply(filepath): 10 | with open(filepath, "rb") as f: 11 | plydata = PlyData.read(f) 12 | data = plydata.elements[0].data 13 | coords = np.array([data["x"], data["y"], data["z"]], dtype=np.float32).T 14 | feats = None 15 | labels = None 16 | if ({"red", "green", "blue"} - set(data.dtype.names)) == set(): 17 | feats = np.array([data["red"], data["green"], data["blue"]], dtype=np.uint8).T 18 | if "label" in data.dtype.names: 19 | labels = np.array(data["label"], dtype=np.uint32) 20 | return coords, feats, labels 21 | 22 | 23 | def load_ply_with_normals(filepath): 24 | mesh = open3d.io.read_triangle_mesh(str(filepath)) 25 | if not mesh.has_vertex_normals(): 26 | mesh.compute_vertex_normals() 27 | vertices = np.asarray(mesh.vertices) 28 | normals = np.asarray(mesh.vertex_normals) 29 | 30 | coords, feats, labels = load_ply(filepath) 31 | assert np.allclose(coords, vertices), "different coordinates" 32 | feats = np.hstack((feats, normals)) 33 | 34 | return coords, feats, labels 35 | 36 | 37 | def load_obj_with_normals(filepath): 38 | mesh = open3d.io.read_triangle_mesh(str(filepath)) 39 | if not mesh.has_vertex_normals(): 40 | mesh.compute_vertex_normals() 41 | coords = np.asarray(mesh.vertices) 42 | normals = np.asarray(mesh.vertex_normals) 43 | colors = np.asarray(mesh.vertex_colors) 44 | feats = np.hstack((colors, normals)) 45 | 46 | return coords, feats 47 | 48 | 49 | def write_point_cloud_in_ply( 50 | filepath: Path, 51 | coords: np.ndarray, 52 | feats: Optional[np.ndarray] = None, 53 | labels: Optional[np.ndarray] = None, 54 | dtypes: Optional[List[Tuple[str, str]]] = [ 55 | ("x", " 2 | #include 3 | #include 4 | #include 5 | #include "aggregation_cuda_kernel.h" 6 | 7 | 8 | void aggregation_forward_cuda(int n, int nsample, int c, int w_c, at::Tensor input_tensor, at::Tensor position_tensor, at::Tensor weight_tensor, at::Tensor idx_tensor, at::Tensor output_tensor) 9 | { 10 | const float *input = input_tensor.data_ptr(); 11 | const float *position = position_tensor.data_ptr(); 12 | const float *weight = weight_tensor.data_ptr(); 13 | const int *idx = idx_tensor.data_ptr(); 14 | float *output = output_tensor.data_ptr(); 15 | aggregation_forward_cuda_launcher(n, nsample, c, w_c, input, position, weight, idx, output); 16 | } 17 | 18 | void aggregation_backward_cuda(int n, int nsample, int c, int w_c, at::Tensor input_tensor, at::Tensor position_tensor, at::Tensor weight_tensor, at::Tensor idx_tensor, at::Tensor grad_output_tensor, at::Tensor grad_input_tensor, at::Tensor grad_position_tensor, at::Tensor grad_weight_tensor) 19 | { 20 | const float *input = input_tensor.data_ptr(); 21 | const float *position = position_tensor.data_ptr(); 22 | const float *weight = weight_tensor.data_ptr(); 23 | const int *idx = idx_tensor.data_ptr(); 24 | const float *grad_output = grad_output_tensor.data_ptr(); 25 | float *grad_input = grad_input_tensor.data_ptr(); 26 | float *grad_position = grad_position_tensor.data_ptr(); 27 | float *grad_weight = grad_weight_tensor.data_ptr(); 28 | aggregation_backward_cuda_launcher(n, nsample, c, w_c, input, position, weight, idx, grad_output, grad_input, grad_position, grad_weight); 29 | } 30 | -------------------------------------------------------------------------------- /utils/pointops2/src/aggregation/aggregation_cuda_kernel.cu: -------------------------------------------------------------------------------- 1 | #include "../cuda_utils.h" 2 | #include "aggregation_cuda_kernel.h" 3 | 4 | 5 | __global__ void aggregation_forward_cuda_kernel(int n, int nsample, int c, int w_c, const float *input, const float *position, const float *weight, const int *idx, float *output) { 6 | // input: input: (n, c), position: (n, nsample, c), weight: (n, nsample, w_c), idx: (n, nsample), output: (n, c) 7 | int index = blockIdx.x * blockDim.x + threadIdx.x; 8 | if (index >= n * c) return; 9 | const int c_idx = index % c; 10 | const int n_idx = index / c; 11 | const int w_c_idx = c_idx % w_c; 12 | for (int nsample_idx = 0; nsample_idx < nsample; nsample_idx++) 13 | { 14 | int idx_idx = n_idx * nsample + nsample_idx; 15 | int input_idx = idx[idx_idx] * c + c_idx; 16 | int position_idx = n_idx * nsample * c + nsample_idx * c + c_idx; 17 | int weight_idx = n_idx * nsample * w_c + nsample_idx * w_c + w_c_idx; 18 | output[index] += (input[input_idx] + position[position_idx]) * weight[weight_idx]; 19 | } 20 | } 21 | 22 | __global__ void aggregation_backward_cuda_kernel(int n, int nsample, int c, int w_c, const float *input, const float *position, const float *weight, const int *idx, const float *grad_output, float *grad_input, float *grad_position, float *grad_weight) { 23 | // input: grad_output: (n, c), output: grad_input: (n, c), grad_position: (n, nsample, c), grad_weight: (n, nsample, w_c) 24 | int index = blockIdx.x * blockDim.x + threadIdx.x; 25 | if (index >= n * c) return; 26 | const int c_idx = index % c; 27 | const int n_idx = index / c; 28 | const int w_c_idx = c_idx % w_c; 29 | for (int nsample_idx = 0; nsample_idx < nsample; nsample_idx++) 30 | { 31 | int idx_idx = n_idx * nsample + nsample_idx; 32 | int input_idx = idx[idx_idx] * c + c_idx; 33 | int position_idx = n_idx * nsample * c + nsample_idx * c + c_idx; 34 | int weight_idx = n_idx * nsample * w_c + nsample_idx * w_c + w_c_idx; 35 | atomicAdd(grad_input + input_idx, grad_output[index] * weight[weight_idx]); 36 | grad_position[position_idx] = grad_output[index] * weight[weight_idx]; 37 | atomicAdd(grad_weight + weight_idx, grad_output[index] * (input[input_idx] + position[position_idx])); 38 | } 39 | } 40 | 41 | void aggregation_forward_cuda_launcher(int n, int nsample, int c, int w_c, const float *input, const float *position, const float *weight, const int *idx, float *output) { 42 | // input: input: (n, c), position: (n, nsample, c), weight: (n, nsample, w_c), idx: (n, nsample), output: (n, c) 43 | dim3 blocks(DIVUP(n * c, THREADS_PER_BLOCK)); 44 | dim3 threads(THREADS_PER_BLOCK); 45 | aggregation_forward_cuda_kernel<<>>(n, nsample, c, w_c, input, position, weight, idx, output); 46 | } 47 | 48 | void aggregation_backward_cuda_launcher(int n, int nsample, int c, int w_c, const float *input, const float *position, const float *weight, const int *idx, const float *grad_output, float *grad_input, float *grad_position, float *grad_weight) { 49 | // input: grad_output: (n, c), output: grad_input: (n, c), grad_position: (n, nsample, c), grad_weight: (n, nsample, w_c) 50 | dim3 blocks(DIVUP(n * c, THREADS_PER_BLOCK)); 51 | dim3 threads(THREADS_PER_BLOCK); 52 | aggregation_backward_cuda_kernel<<>>(n, nsample, c, w_c, input, position, weight, idx, grad_output, grad_input, grad_position, grad_weight); 53 | } 54 | -------------------------------------------------------------------------------- /utils/pointops2/src/aggregation/aggregation_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _AGGREGATION_CUDA_KERNEL 2 | #define _AGGREGATION_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void aggregation_forward_cuda(int n, int nsample, int c, int w_c, at::Tensor input_tensor, at::Tensor position_tensor, at::Tensor weight_tensor, at::Tensor idx_tensor, at::Tensor output_tensor); 8 | void aggregation_backward_cuda(int n, int nsample, int c, int w_c, at::Tensor input_tensor, at::Tensor position_tensor, at::Tensor weight_tensor, at::Tensor idx_tensor, at::Tensor grad_output_tensor, at::Tensor grad_input_tensor, at::Tensor grad_position_tensor, at::Tensor grad_weight_tensor); 9 | 10 | #ifdef __cplusplus 11 | extern "C" { 12 | #endif 13 | 14 | void aggregation_forward_cuda_launcher(int n, int nsample, int c, int w_c, const float *input, const float *position, const float *weight, const int *idx, float *output); 15 | void aggregation_backward_cuda_launcher(int n, int nsample, int c, int w_c, const float *input, const float *position, const float *weight, const int *idx, const float *grad_output, float *grad_input, float *grad_position, float *grad_weight); 16 | 17 | #ifdef __cplusplus 18 | } 19 | #endif 20 | #endif 21 | -------------------------------------------------------------------------------- /utils/pointops2/src/attention/attention_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | #include "attention_cuda_kernel.h" 6 | 7 | void attention_step1_forward_cuda(int N, int M, int h, int C, at::Tensor q_tensor, at::Tensor k_tensor, 8 | at::Tensor index0_tensor, at::Tensor index1_tensor, at::Tensor attn_tensor) 9 | { 10 | const float *q = q_tensor.data_ptr(); 11 | const float *k = k_tensor.data_ptr(); 12 | const int *index0 = index0_tensor.data_ptr(); 13 | const int *index1 = index1_tensor.data_ptr(); 14 | float *attn = attn_tensor.data_ptr(); 15 | attention_step1_forward_cuda_launcher(N, M, h, C, q, k, index0, index1, attn); 16 | } 17 | 18 | void attention_step1_backward_cuda(int N, int M, int h, int C, at::Tensor grad_out_tensor, 19 | at::Tensor index0_tensor, at::Tensor index1_tensor, at::Tensor q_tensor, at::Tensor k_tensor, 20 | at::Tensor grad_q_tensor, at::Tensor grad_k_tensor) 21 | { 22 | const float *grad_out = grad_out_tensor.data_ptr(); 23 | const int *index0 = index0_tensor.data_ptr(); 24 | const int *index1 = index1_tensor.data_ptr(); 25 | const float *q = q_tensor.data_ptr(); 26 | const float *k = k_tensor.data_ptr(); 27 | float *grad_q = grad_q_tensor.data_ptr(); 28 | float *grad_k = grad_k_tensor.data_ptr(); 29 | attention_step1_backward_cuda_launcher(N, M, h, C, grad_out, index0, index1, q, k, grad_q, grad_k); 30 | } 31 | 32 | void attention_step2_forward_cuda(int N, int M, int h, int C, at::Tensor attn_tensor, at::Tensor v_tensor, 33 | at::Tensor index0_tensor, at::Tensor index1_tensor, at::Tensor output_tensor) 34 | { 35 | const float *attn = attn_tensor.data_ptr(); 36 | const float *v = v_tensor.data_ptr(); 37 | const int *index0 = index0_tensor.data_ptr(); 38 | const int *index1 = index1_tensor.data_ptr(); 39 | float *output = output_tensor.data_ptr(); 40 | attention_step2_forward_cuda_launcher(N, M, h, C, attn, v, index0, index1, output); 41 | } 42 | 43 | 44 | void attention_step2_backward_cuda(int N, int M, int h, int C, at::Tensor grad_out_tensor, 45 | at::Tensor index0_tensor, at::Tensor index1_tensor, at::Tensor attn_tensor, at::Tensor v_tensor, 46 | at::Tensor grad_attn_tensor, at::Tensor grad_v_tensor) 47 | { 48 | const float *grad_out = grad_out_tensor.data_ptr(); 49 | const int *index0 = index0_tensor.data_ptr(); 50 | const int *index1 = index1_tensor.data_ptr(); 51 | const float *attn = attn_tensor.data_ptr(); 52 | const float *v = v_tensor.data_ptr(); 53 | float *grad_attn = grad_attn_tensor.data_ptr(); 54 | float *grad_v = grad_v_tensor.data_ptr(); 55 | attention_step2_backward_cuda_launcher(N, M, h, C, grad_out, index0, index1, attn, v, grad_attn, grad_v); 56 | } 57 | -------------------------------------------------------------------------------- /utils/pointops2/src/attention/attention_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _ATTENTION_CUDA_KERNEL 2 | #define _ATTENTION_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void attention_step1_forward_cuda(int N, int M, int h, int C, at::Tensor q_tensor, at::Tensor k_tensor, at::Tensor index0_tensor, at::Tensor index1_tensor, at::Tensor attn_tensor); 8 | void attention_step1_backward_cuda(int N, int M, int h, int C, at::Tensor grad_out_tensor, at::Tensor index0_tensor, at::Tensor index1_tensor, at::Tensor q_tensor, at::Tensor k_tensor, at::Tensor grad_q_tensor, at::Tensor grad_k_tensor); 9 | 10 | void attention_step2_forward_cuda(int N, int M, int h, int C, at::Tensor attn_tensor, at::Tensor v_tensor, at::Tensor index0_tensor, at::Tensor index1_tensor, at::Tensor output_tensor); 11 | void attention_step2_backward_cuda(int N, int M, int h, int C, at::Tensor grad_out_tensor, at::Tensor index0_tensor, at::Tensor index1_tensor, at::Tensor attn_tensor, at::Tensor v_tensor, at::Tensor grad_attn_tensor, at::Tensor grad_v_tensor); 12 | 13 | #ifdef __cplusplus 14 | extern "C" { 15 | #endif 16 | 17 | void attention_step1_forward_cuda_launcher(int N, int M, int h, int C, const float *q, const float *k, const int *index0, const int *index1, float *attn); 18 | void attention_step1_backward_cuda_launcher(int N, int M, int h, int C, const float *grad_out, const int *index0, const int *index1, const float *q, const float *k, float *grad_q, float *grad_k); 19 | 20 | void attention_step2_forward_cuda_launcher(int N, int M, int h, int C, const float *attn, const float *v, const int *index0, const int *index1, float *output); 21 | void attention_step2_backward_cuda_launcher(int N, int M, int h, int C, const float *grad_out, const int *index0, const int *index1, const float *attn, const float *v, float *grad_attn, float *grad_v); 22 | 23 | #ifdef __cplusplus 24 | } 25 | #endif 26 | #endif 27 | -------------------------------------------------------------------------------- /utils/pointops2/src/attention_v2/attention_cuda_kernel_v2.h: -------------------------------------------------------------------------------- 1 | #ifndef _ATTENTION_V2_CUDA_KERNEL 2 | #define _ATTENTION_V2_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void attention_step1_forward_cuda_v2(int N, int M, int h, int C, const unsigned int n_max, at::Tensor q_tensor, at::Tensor k_tensor, at::Tensor index0_tensor_offsets, at::Tensor index1_tensor, at::Tensor attn_tensor); 8 | void attention_step1_backward_cuda_v2(int N, int M, int h, int C, const unsigned int n_max, at::Tensor grad_out_tensor, at::Tensor index0_tensor_offsets, at::Tensor index1_tensor, at::Tensor q_tensor, at::Tensor k_tensor, at::Tensor grad_q_tensor, at::Tensor grad_k_tensor); 9 | 10 | void attention_step2_forward_cuda_v2(int N, int M, int h, int C, at::Tensor attn_tensor, at::Tensor v_tensor, at::Tensor index0_tensor, at::Tensor index1_tensor, at::Tensor output_tensor); 11 | void attention_step2_backward_cuda_v2(int N, int M, int h, int C, at::Tensor grad_out_tensor, at::Tensor index0_tensor, at::Tensor index1_tensor, at::Tensor attn_tensor, at::Tensor v_tensor, at::Tensor grad_attn_tensor, at::Tensor grad_v_tensor); 12 | 13 | #ifdef __cplusplus 14 | extern "C" { 15 | #endif 16 | 17 | void attention_step1_forward_cuda_launcher_v2(int N, int M, int h, int C, const unsigned int n_max, const float *q, const float *k, const int *index0_offsets, const int *index1, float *attn); 18 | void attention_step1_backward_cuda_launcher_v2(int N, int M, int h, int C, const unsigned int n_max, const float *grad_out, const int *index0_offsets, const int *index1, const float *q, const float *k, float *grad_q, float *grad_k); 19 | 20 | void attention_step2_forward_cuda_launcher_v2(int N, int M, int h, int C, const float *attn, const float *v, const int *index0, const int *index1, float *output); 21 | void attention_step2_backward_cuda_launcher_v2(int N, int M, int h, int C, const float *grad_out, const int *index0, const int *index1, const float *attn, const float *v, float *grad_attn, float *grad_v); 22 | 23 | #ifdef __cplusplus 24 | } 25 | #endif 26 | #endif 27 | -------------------------------------------------------------------------------- /utils/pointops2/src/attention_v2/attention_cuda_v2.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | #include "attention_cuda_kernel_v2.h" 6 | 7 | void attention_step1_forward_cuda_v2(int N, int M, int h, int C, const unsigned int n_max, at::Tensor q_tensor, at::Tensor k_tensor, 8 | at::Tensor index0_tensor_offsets, at::Tensor index1_tensor, at::Tensor attn_tensor) 9 | { 10 | const float *q = q_tensor.data_ptr(); 11 | const float *k = k_tensor.data_ptr(); 12 | const int *index0_offsets = index0_tensor_offsets.data_ptr(); 13 | const int *index1 = index1_tensor.data_ptr(); 14 | float *attn = attn_tensor.data_ptr(); 15 | attention_step1_forward_cuda_launcher_v2(N, M, h, C, n_max, q, k, index0_offsets, index1, attn); 16 | } 17 | 18 | void attention_step1_backward_cuda_v2(int N, int M, int h, int C, const unsigned int n_max, at::Tensor grad_out_tensor, 19 | at::Tensor index0_tensor_offsets, at::Tensor index1_tensor, at::Tensor q_tensor, at::Tensor k_tensor, 20 | at::Tensor grad_q_tensor, at::Tensor grad_k_tensor) 21 | { 22 | const float *grad_out = grad_out_tensor.data_ptr(); 23 | const int *index0_offsets = index0_tensor_offsets.data_ptr(); 24 | const int *index1 = index1_tensor.data_ptr(); 25 | const float *q = q_tensor.data_ptr(); 26 | const float *k = k_tensor.data_ptr(); 27 | float *grad_q = grad_q_tensor.data_ptr(); 28 | float *grad_k = grad_k_tensor.data_ptr(); 29 | attention_step1_backward_cuda_launcher_v2(N, M, h, C, n_max, grad_out, index0_offsets, index1, q, k, grad_q, grad_k); 30 | } 31 | 32 | void attention_step2_forward_cuda_v2(int N, int M, int h, int C, at::Tensor attn_tensor, at::Tensor v_tensor, 33 | at::Tensor index0_tensor, at::Tensor index1_tensor, at::Tensor output_tensor) 34 | { 35 | const float *attn = attn_tensor.data_ptr(); 36 | const float *v = v_tensor.data_ptr(); 37 | const int *index0 = index0_tensor.data_ptr(); 38 | const int *index1 = index1_tensor.data_ptr(); 39 | float *output = output_tensor.data_ptr(); 40 | attention_step2_forward_cuda_launcher_v2(N, M, h, C, attn, v, index0, index1, output); 41 | } 42 | 43 | 44 | void attention_step2_backward_cuda_v2(int N, int M, int h, int C, at::Tensor grad_out_tensor, 45 | at::Tensor index0_tensor, at::Tensor index1_tensor, at::Tensor attn_tensor, at::Tensor v_tensor, 46 | at::Tensor grad_attn_tensor, at::Tensor grad_v_tensor) 47 | { 48 | const float *grad_out = grad_out_tensor.data_ptr(); 49 | const int *index0 = index0_tensor.data_ptr(); 50 | const int *index1 = index1_tensor.data_ptr(); 51 | const float *attn = attn_tensor.data_ptr(); 52 | const float *v = v_tensor.data_ptr(); 53 | float *grad_attn = grad_attn_tensor.data_ptr(); 54 | float *grad_v = grad_v_tensor.data_ptr(); 55 | attention_step2_backward_cuda_launcher_v2(N, M, h, C, grad_out, index0, index1, attn, v, grad_attn, grad_v); 56 | } 57 | -------------------------------------------------------------------------------- /utils/pointops2/src/cuda_utils.h: -------------------------------------------------------------------------------- 1 | #ifndef _CUDA_UTILS_H 2 | #define _CUDA_UTILS_H 3 | 4 | #include 5 | #include 6 | 7 | #define TOTAL_THREADS 1024 8 | #define THREADS_PER_BLOCK 256 9 | #define DIVUP(m, n) ((m) / (n) + ((m) % (n) > 0)) 10 | 11 | inline int opt_n_threads(int work_size) { 12 | const int pow_2 = std::log(static_cast(work_size)) / std::log(2.0); 13 | return std::max(std::min(1 << pow_2, TOTAL_THREADS), 1); 14 | } 15 | 16 | inline dim3 opt_block_config(int x, int y) { 17 | const int x_threads = opt_n_threads(x); 18 | const int y_threads = std::max(std::min(opt_n_threads(y), TOTAL_THREADS / x_threads), 1); 19 | dim3 block_config(x_threads, y_threads, 1); 20 | return block_config; 21 | } 22 | 23 | #endif 24 | -------------------------------------------------------------------------------- /utils/pointops2/src/grouping/grouping_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | #include "grouping_cuda_kernel.h" 6 | 7 | 8 | void grouping_forward_cuda(int m, int nsample, int c, at::Tensor input_tensor, at::Tensor idx_tensor, at::Tensor output_tensor) 9 | { 10 | const float *input = input_tensor.data_ptr(); 11 | const int *idx = idx_tensor.data_ptr(); 12 | float *output = output_tensor.data_ptr(); 13 | grouping_forward_cuda_launcher(m, nsample, c, input, idx, output); 14 | } 15 | 16 | void grouping_backward_cuda(int m, int nsample, int c, at::Tensor grad_output_tensor, at::Tensor idx_tensor, at::Tensor grad_input_tensor) 17 | { 18 | const float *grad_output = grad_output_tensor.data_ptr(); 19 | const int *idx = idx_tensor.data_ptr(); 20 | float *grad_input = grad_input_tensor.data_ptr(); 21 | grouping_backward_cuda_launcher(m, nsample, c, grad_output, idx, grad_input); 22 | } 23 | -------------------------------------------------------------------------------- /utils/pointops2/src/grouping/grouping_cuda_kernel.cu: -------------------------------------------------------------------------------- 1 | #include "../cuda_utils.h" 2 | #include "grouping_cuda_kernel.h" 3 | 4 | 5 | __global__ void grouping_forward_cuda_kernel(int m, int nsample, int c, const float *__restrict__ input, const int *__restrict__ idx, float *__restrict__ output) { 6 | // input: input: (n, c), idx: (m, nsample), output: (m, nsample, c) 7 | int index = blockIdx.x * blockDim.x + threadIdx.x; 8 | if (index >= m * nsample * c) return; 9 | const int c_idx = index % c; 10 | const int nsample_idx = (index / c) % nsample; 11 | const int m_idx = index / nsample / c; 12 | const int input_idx = idx[m_idx * nsample + nsample_idx] * c + c_idx; 13 | output[index] = input[input_idx]; 14 | } 15 | 16 | __global__ void grouping_backward_cuda_kernel(int m, int nsample, int c, const float *__restrict__ grad_output, const int *__restrict__ idx, float *__restrict__ grad_input) { 17 | // input: grad_output: (m, nsample, c), idx: (m, nsample), output: grad_input: (n, c) 18 | int index = blockIdx.x * blockDim.x + threadIdx.x; 19 | if (index >= m * nsample * c) return; 20 | const int c_idx = index % c; 21 | const int nsample_idx = (index / c) % nsample; 22 | const int m_idx = index / nsample / c; 23 | const int input_idx = idx[m_idx * nsample + nsample_idx] * c + c_idx; 24 | atomicAdd(grad_input + input_idx, grad_output[index]); 25 | } 26 | 27 | void grouping_forward_cuda_launcher(int m, int nsample, int c, const float *input, const int *idx, float *output) { 28 | // input: input: (n, c), idx: (m, nsample), output: (m, nsample, c) 29 | dim3 blocks(DIVUP(m * nsample * c, THREADS_PER_BLOCK)); 30 | dim3 threads(THREADS_PER_BLOCK); 31 | grouping_forward_cuda_kernel<<>>(m, nsample, c, input, idx, output); 32 | } 33 | 34 | void grouping_backward_cuda_launcher(int m, int nsample, int c, const float *grad_output, const int *idx, float *grad_input) 35 | { 36 | // input: grad_output: (m, nsample, c), idx: (m, nsample), output: grad_input: (n, c) 37 | dim3 blocks(DIVUP(m * nsample * c, THREADS_PER_BLOCK)); 38 | dim3 threads(THREADS_PER_BLOCK); 39 | grouping_backward_cuda_kernel<<>>(m, nsample, c, grad_output, idx, grad_input); 40 | } 41 | -------------------------------------------------------------------------------- /utils/pointops2/src/grouping/grouping_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _GROUPING_CUDA_KERNEL 2 | #define _GROUPING_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void grouping_forward_cuda(int m, int nsample, int c, at::Tensor input_tensor, at::Tensor idx_tensor, at::Tensor output_tensor); 8 | void grouping_backward_cuda(int m, int nsample, int c, at::Tensor grad_output_tensor, at::Tensor idx_tensor, at::Tensor grad_input_tensor); 9 | 10 | #ifdef __cplusplus 11 | extern "C" { 12 | #endif 13 | 14 | void grouping_forward_cuda_launcher(int m, int nsample, int c, const float *input, const int *idx, float *output); 15 | void grouping_backward_cuda_launcher(int m, int nsample, int c, const float *grad_output, const int *idx, float *grad_input); 16 | 17 | #ifdef __cplusplus 18 | } 19 | #endif 20 | #endif 21 | -------------------------------------------------------------------------------- /utils/pointops2/src/interpolation/interpolation_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | #include "interpolation_cuda_kernel.h" 6 | 7 | 8 | void interpolation_forward_cuda(int n, int c, int k, at::Tensor input_tensor, at::Tensor idx_tensor, at::Tensor weight_tensor, at::Tensor output_tensor) 9 | { 10 | const float *input = input_tensor.data_ptr(); 11 | const int *idx = idx_tensor.data_ptr(); 12 | const float *weight = weight_tensor.data_ptr(); 13 | float *output = output_tensor.data_ptr(); 14 | interpolation_forward_cuda_launcher(n, c, k, input, idx, weight, output); 15 | } 16 | 17 | void interpolation_backward_cuda(int n, int c, int k, at::Tensor grad_output_tensor, at::Tensor idx_tensor, at::Tensor weight_tensor, at::Tensor grad_input_tensor) 18 | { 19 | const float *grad_output = grad_output_tensor.data_ptr(); 20 | const int *idx = idx_tensor.data_ptr(); 21 | const float *weight = weight_tensor.data_ptr(); 22 | float *grad_input = grad_input_tensor.data_ptr(); 23 | interpolation_backward_cuda_launcher(n, c, k, grad_output, idx, weight, grad_input); 24 | } 25 | -------------------------------------------------------------------------------- /utils/pointops2/src/interpolation/interpolation_cuda_kernel.cu: -------------------------------------------------------------------------------- 1 | #include "../cuda_utils.h" 2 | #include "interpolation_cuda_kernel.h" 3 | 4 | 5 | __global__ void interpolation_forward_cuda_kernel(int n, int c, int k, const float *input, const int *idx, const float *weight, float *output) 6 | { 7 | // input: input: (m, c), idx: (n, k), weight: (n, k), output: output (n, c) 8 | int index = blockIdx.x * blockDim.x + threadIdx.x; 9 | if (index >= n * c) return; 10 | int c_idx = index % c; 11 | int n_idx = index / c; 12 | for (int i = 0; i < k; i++) 13 | { 14 | int idx_idx = n_idx * k + i; 15 | int input_idx = idx[idx_idx] * c + c_idx; 16 | output[index] += input[input_idx] * weight[idx_idx]; 17 | } 18 | } 19 | 20 | __global__ void interpolation_backward_cuda_kernel(int n, int c, int k, const float *grad_output, const int *idx, const float *weight, float *grad_input) 21 | { 22 | // input: grad_output: (n, c), idx: (n, k), weight: (n, k), output: grad_input (m, c) 23 | int index = blockIdx.x * blockDim.x + threadIdx.x; 24 | if (index >= n * c) return; 25 | int c_idx = index % c; 26 | int n_idx = index / c; 27 | for (int i = 0; i < k; i++) 28 | { 29 | int idx_idx = n_idx * k + i; 30 | int input_idx = idx[idx_idx] * c + c_idx; 31 | atomicAdd(grad_input + input_idx, grad_output[index] * weight[idx_idx]); 32 | } 33 | } 34 | 35 | void interpolation_forward_cuda_launcher(int n, int c, int k, const float *input, const int *idx, const float *weight, float *output) { 36 | // input: input: (m, c), idx: (n, k), weight: (n, k), output: output (n, c) 37 | dim3 blocks(DIVUP(n * c, THREADS_PER_BLOCK)); 38 | dim3 threads(THREADS_PER_BLOCK); 39 | interpolation_forward_cuda_kernel<<>>(n, c, k, input, idx, weight, output); 40 | } 41 | 42 | void interpolation_backward_cuda_launcher(int n, int c, int k, const float *grad_output, const int *idx, const float *weight, float *grad_input) { 43 | // input: grad_output: (n, c), idx: (n, k), weight: (n, k), output: grad_input (m, c) 44 | dim3 blocks(DIVUP(n * c, THREADS_PER_BLOCK)); 45 | dim3 threads(THREADS_PER_BLOCK); 46 | interpolation_backward_cuda_kernel<<>>(n, c, k, grad_output, idx, weight, grad_input); 47 | } 48 | -------------------------------------------------------------------------------- /utils/pointops2/src/interpolation/interpolation_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _INTERPOLATION_CUDA_KERNEL 2 | #define _INTERPOLATION_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void interpolation_forward_cuda(int n, int c, int k, at::Tensor input_tensor, at::Tensor idx_tensor, at::Tensor weight_tensor, at::Tensor output_tensor); 8 | void interpolation_backward_cuda(int n, int c, int k, at::Tensor grad_output_tensor, at::Tensor idx_tensor, at::Tensor weight_tensor, at::Tensor grad_input_tensor); 9 | 10 | #ifdef __cplusplus 11 | extern "C" { 12 | #endif 13 | 14 | void interpolation_forward_cuda_launcher(int n, int c, int k, const float *input, const int *idx, const float *weight, float *output); 15 | void interpolation_backward_cuda_launcher(int n, int c, int k, const float *grad_output, const int *idx, const float *weight, float *grad_input); 16 | 17 | #ifdef __cplusplus 18 | } 19 | #endif 20 | #endif 21 | -------------------------------------------------------------------------------- /utils/pointops2/src/knnquery/knnquery_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | #include "knnquery_cuda_kernel.h" 6 | 7 | 8 | void knnquery_cuda(int m, int nsample, at::Tensor xyz_tensor, at::Tensor new_xyz_tensor, at::Tensor offset_tensor, at::Tensor new_offset_tensor, at::Tensor idx_tensor, at::Tensor dist2_tensor) 9 | { 10 | const float *xyz = xyz_tensor.data_ptr(); 11 | const float *new_xyz = new_xyz_tensor.data_ptr(); 12 | const int *offset = offset_tensor.data_ptr(); 13 | const int *new_offset = new_offset_tensor.data_ptr(); 14 | int *idx = idx_tensor.data_ptr(); 15 | float *dist2 = dist2_tensor.data_ptr(); 16 | knnquery_cuda_launcher(m, nsample, xyz, new_xyz, offset, new_offset, idx, dist2); 17 | } 18 | -------------------------------------------------------------------------------- /utils/pointops2/src/knnquery/knnquery_cuda_kernel.cu: -------------------------------------------------------------------------------- 1 | #include "../cuda_utils.h" 2 | #include "knnquery_cuda_kernel.h" 3 | 4 | 5 | __device__ void swap_float(float *x, float *y) 6 | { 7 | float tmp = *x; 8 | *x = *y; 9 | *y = tmp; 10 | } 11 | 12 | 13 | __device__ void swap_int(int *x, int *y) 14 | { 15 | int tmp = *x; 16 | *x = *y; 17 | *y = tmp; 18 | } 19 | 20 | 21 | __device__ void reheap(float *dist, int *idx, int k) 22 | { 23 | int root = 0; 24 | int child = root * 2 + 1; 25 | while (child < k) 26 | { 27 | if(child + 1 < k && dist[child+1] > dist[child]) 28 | child++; 29 | if(dist[root] > dist[child]) 30 | return; 31 | swap_float(&dist[root], &dist[child]); 32 | swap_int(&idx[root], &idx[child]); 33 | root = child; 34 | child = root * 2 + 1; 35 | } 36 | } 37 | 38 | 39 | __device__ void heap_sort(float *dist, int *idx, int k) 40 | { 41 | int i; 42 | for (i = k - 1; i > 0; i--) 43 | { 44 | swap_float(&dist[0], &dist[i]); 45 | swap_int(&idx[0], &idx[i]); 46 | reheap(dist, idx, i); 47 | } 48 | } 49 | 50 | 51 | __device__ int get_bt_idx(int idx, const int *offset) 52 | { 53 | int i = 0; 54 | while (1) 55 | { 56 | if (idx < offset[i]) 57 | break; 58 | else 59 | i++; 60 | } 61 | return i; 62 | } 63 | 64 | 65 | __global__ void knnquery_cuda_kernel(int m, int nsample, const float *__restrict__ xyz, const float *__restrict__ new_xyz, const int *__restrict__ offset, const int *__restrict__ new_offset, int *__restrict__ idx, float *__restrict__ dist2) { 66 | // input: xyz (n, 3) new_xyz (m, 3) 67 | // output: idx (m, nsample) dist2 (m, nsample) 68 | int pt_idx = blockIdx.x * blockDim.x + threadIdx.x; 69 | if (pt_idx >= m) return; 70 | 71 | new_xyz += pt_idx * 3; 72 | idx += pt_idx * nsample; 73 | dist2 += pt_idx * nsample; 74 | int bt_idx = get_bt_idx(pt_idx, new_offset); 75 | int start; 76 | if (bt_idx == 0) 77 | start = 0; 78 | else 79 | start = offset[bt_idx - 1]; 80 | int end = offset[bt_idx]; 81 | 82 | float new_x = new_xyz[0]; 83 | float new_y = new_xyz[1]; 84 | float new_z = new_xyz[2]; 85 | 86 | float best_dist[100]; 87 | int best_idx[100]; 88 | for(int i = 0; i < nsample; i++){ 89 | best_dist[i] = 1e10; 90 | best_idx[i] = start; 91 | } 92 | for(int i = start; i < end; i++){ 93 | float x = xyz[i * 3 + 0]; 94 | float y = xyz[i * 3 + 1]; 95 | float z = xyz[i * 3 + 2]; 96 | float d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) + (new_z - z) * (new_z - z); 97 | if (d2 < best_dist[0]){ 98 | best_dist[0] = d2; 99 | best_idx[0] = i; 100 | reheap(best_dist, best_idx, nsample); 101 | } 102 | } 103 | heap_sort(best_dist, best_idx, nsample); 104 | for(int i = 0; i < nsample; i++){ 105 | idx[i] = best_idx[i]; 106 | dist2[i] = best_dist[i]; 107 | } 108 | } 109 | 110 | 111 | void knnquery_cuda_launcher(int m, int nsample, const float *xyz, const float *new_xyz, const int *offset, const int *new_offset, int *idx, float *dist2) { 112 | // input: new_xyz: (m, 3), xyz: (n, 3), idx: (m, nsample) 113 | dim3 blocks(DIVUP(m, THREADS_PER_BLOCK)); 114 | dim3 threads(THREADS_PER_BLOCK); 115 | knnquery_cuda_kernel<<>>(m, nsample, xyz, new_xyz, offset, new_offset, idx, dist2); 116 | } 117 | -------------------------------------------------------------------------------- /utils/pointops2/src/knnquery/knnquery_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _KNNQUERY_CUDA_KERNEL 2 | #define _KNNQUERY_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void knnquery_cuda(int m, int nsample, at::Tensor xyz_tensor, at::Tensor new_xyz_tensor, at::Tensor offset_tensor, at::Tensor new_offset_tensor, at::Tensor idx_tensor, at::Tensor dist2_tensor); 8 | 9 | #ifdef __cplusplus 10 | extern "C" { 11 | #endif 12 | 13 | void knnquery_cuda_launcher(int m, int nsample, const float *xyz, const float *new_xyz, const int *offset, const int *new_offset, int *idx, float *dist2); 14 | 15 | #ifdef __cplusplus 16 | } 17 | #endif 18 | #endif 19 | -------------------------------------------------------------------------------- /utils/pointops2/src/pointops_api.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #include "knnquery/knnquery_cuda_kernel.h" 5 | #include "sampling/sampling_cuda_kernel.h" 6 | #include "grouping/grouping_cuda_kernel.h" 7 | #include "interpolation/interpolation_cuda_kernel.h" 8 | #include "aggregation/aggregation_cuda_kernel.h" 9 | #include "subtraction/subtraction_cuda_kernel.h" 10 | #include "attention/attention_cuda_kernel.h" 11 | #include "rpe/relative_pos_encoding_cuda_kernel.h" 12 | #include "attention_v2/attention_cuda_kernel_v2.h" 13 | #include "rpe_v2/relative_pos_encoding_cuda_kernel_v2.h" 14 | 15 | 16 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 17 | m.def("knnquery_cuda", &knnquery_cuda, "knnquery_cuda"); 18 | m.def("furthestsampling_cuda", &furthestsampling_cuda, "furthestsampling_cuda"); 19 | m.def("grouping_forward_cuda", &grouping_forward_cuda, "grouping_forward_cuda"); 20 | m.def("grouping_backward_cuda", &grouping_backward_cuda, "grouping_backward_cuda"); 21 | m.def("interpolation_forward_cuda", &interpolation_forward_cuda, "interpolation_forward_cuda"); 22 | m.def("interpolation_backward_cuda", &interpolation_backward_cuda, "interpolation_backward_cuda"); 23 | m.def("subtraction_forward_cuda", &subtraction_forward_cuda, "subtraction_forward_cuda"); 24 | m.def("subtraction_backward_cuda", &subtraction_backward_cuda, "subtraction_backward_cuda"); 25 | m.def("aggregation_forward_cuda", &aggregation_forward_cuda, "aggregation_forward_cuda"); 26 | m.def("aggregation_backward_cuda", &aggregation_backward_cuda, "aggregation_backward_cuda"); 27 | m.def("attention_step1_forward_cuda", &attention_step1_forward_cuda, "attention_step1_forward_cuda"); 28 | m.def("attention_step1_backward_cuda", &attention_step1_backward_cuda, "attention_step1_backward_cuda"); 29 | m.def("attention_step2_forward_cuda", &attention_step2_forward_cuda, "attention_step2_forward_cuda"); 30 | m.def("attention_step2_backward_cuda", &attention_step2_backward_cuda, "attention_step2_backward_cuda"); 31 | m.def("dot_prod_with_idx_forward_cuda", &dot_prod_with_idx_forward_cuda, "dot_prod_with_idx_forward_cuda"); 32 | m.def("dot_prod_with_idx_backward_cuda", &dot_prod_with_idx_backward_cuda, "dot_prod_with_idx_backward_cuda"); 33 | m.def("attention_step2_with_rel_pos_value_forward_cuda", &attention_step2_with_rel_pos_value_forward_cuda, "attention_step2_with_rel_pos_value_forward_cuda"); 34 | m.def("attention_step2_with_rel_pos_value_backward_cuda", &attention_step2_with_rel_pos_value_backward_cuda, "attention_step2_with_rel_pos_value_backward_cuda"); 35 | m.def("attention_step1_forward_cuda_v2", &attention_step1_forward_cuda_v2, "attention_step1_forward_cuda_v2"); 36 | m.def("attention_step1_backward_cuda_v2", &attention_step1_backward_cuda_v2, "attention_step1_backward_cuda_v2"); 37 | m.def("attention_step2_forward_cuda_v2", &attention_step2_forward_cuda_v2, "attention_step2_forward_cuda_v2"); 38 | m.def("attention_step2_backward_cuda_v2", &attention_step2_backward_cuda_v2, "attention_step2_backward_cuda_v2"); 39 | m.def("dot_prod_with_idx_forward_cuda_v2", &dot_prod_with_idx_forward_cuda_v2, "dot_prod_with_idx_forward_cuda_v2"); 40 | m.def("dot_prod_with_idx_backward_cuda_v2", &dot_prod_with_idx_backward_cuda_v2, "dot_prod_with_idx_backward_cuda_v2"); 41 | m.def("attention_step2_with_rel_pos_value_forward_cuda_v2", &attention_step2_with_rel_pos_value_forward_cuda_v2, "attention_step2_with_rel_pos_value_forward_cuda_v2"); 42 | m.def("attention_step2_with_rel_pos_value_backward_cuda_v2", &attention_step2_with_rel_pos_value_backward_cuda_v2, "attention_step2_with_rel_pos_value_backward_cuda_v2"); 43 | m.def("dot_prod_with_idx_forward_cuda_v3", &dot_prod_with_idx_forward_cuda_v3, "dot_prod_with_idx_forward_cuda_v3"); 44 | m.def("dot_prod_with_idx_backward_cuda_v3", &dot_prod_with_idx_backward_cuda_v3, "dot_prod_with_idx_backward_cuda_v3"); 45 | } 46 | -------------------------------------------------------------------------------- /utils/pointops2/src/rpe/relative_pos_encoding_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | #include "relative_pos_encoding_cuda_kernel.h" 6 | 7 | void dot_prod_with_idx_forward_cuda(int N, int M, int h, int hdim, at::Tensor q_tensor, at::Tensor index_tensor, 8 | at::Tensor table_tensor, at::Tensor rel_idx_tensor, at::Tensor output_tensor) 9 | { 10 | const float *q = q_tensor.data_ptr(); 11 | const float *table = table_tensor.data_ptr(); 12 | const int *index = index_tensor.data_ptr(); 13 | const int *rel_idx = rel_idx_tensor.data_ptr(); 14 | float *output = output_tensor.data_ptr(); 15 | dot_prod_with_idx_forward_cuda_launcher(N, M, h, hdim, q, index, table, rel_idx, output); 16 | } 17 | 18 | void dot_prod_with_idx_backward_cuda(int N, int M, int h, int hdim, at::Tensor grad_out_tensor, 19 | at::Tensor q_tensor, at::Tensor index_tensor, at::Tensor table_tensor, at::Tensor rel_idx_tensor, 20 | at::Tensor grad_q_tensor, at::Tensor grad_table_tensor) 21 | { 22 | const float *grad_out = grad_out_tensor.data_ptr(); 23 | const float *q = q_tensor.data_ptr(); 24 | const int *index = index_tensor.data_ptr(); 25 | const float *table = table_tensor.data_ptr(); 26 | const int *rel_idx = rel_idx_tensor.data_ptr(); 27 | float *grad_q = grad_q_tensor.data_ptr(); 28 | float *grad_table = grad_table_tensor.data_ptr(); 29 | dot_prod_with_idx_backward_cuda_launcher(N, M, h, hdim, grad_out, q, index, table, rel_idx, grad_q, grad_table); 30 | } 31 | 32 | void attention_step2_with_rel_pos_value_forward_cuda(int N, int M, int h, int hdim, at::Tensor attn_tensor, at::Tensor v_tensor, 33 | at::Tensor index0_tensor, at::Tensor index1_tensor, at::Tensor table_tensor, at::Tensor rel_idx_tensor, at::Tensor output_tensor) 34 | { 35 | const float *attn = attn_tensor.data_ptr(); 36 | const float *v = v_tensor.data_ptr(); 37 | const int *index0 = index0_tensor.data_ptr(); 38 | const int *index1 = index1_tensor.data_ptr(); 39 | const float *table = table_tensor.data_ptr(); 40 | const int *rel_idx = rel_idx_tensor.data_ptr(); 41 | float *output = output_tensor.data_ptr(); 42 | attention_step2_with_rel_pos_value_forward_cuda_launcher(N, M, h, hdim, attn, v, index0, index1, table, rel_idx, output); 43 | } 44 | 45 | void attention_step2_with_rel_pos_value_backward_cuda(int N, int M, int h, int hdim, at::Tensor grad_out_tensor, 46 | at::Tensor index0_tensor, at::Tensor index1_tensor, at::Tensor attn_tensor, at::Tensor v_tensor, at::Tensor table_tensor, 47 | at::Tensor rel_idx_tensor, at::Tensor grad_attn_tensor, at::Tensor grad_v_tensor, at::Tensor grad_table_tensor) 48 | { 49 | const float *grad_out = grad_out_tensor.data_ptr(); 50 | const int *index0 = index0_tensor.data_ptr(); 51 | const int *index1 = index1_tensor.data_ptr(); 52 | const float *attn = attn_tensor.data_ptr(); 53 | const float *v = v_tensor.data_ptr(); 54 | const float *table = table_tensor.data_ptr(); 55 | const int *rel_idx = rel_idx_tensor.data_ptr(); 56 | float *grad_attn = grad_attn_tensor.data_ptr(); 57 | float *grad_v = grad_v_tensor.data_ptr(); 58 | float *grad_table = grad_table_tensor.data_ptr(); 59 | attention_step2_with_rel_pos_value_backward_cuda_launcher(N, M, h, hdim, grad_out, index0, index1, attn, v, table, rel_idx, grad_attn, grad_v, grad_table); 60 | } 61 | -------------------------------------------------------------------------------- /utils/pointops2/src/rpe/relative_pos_encoding_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _RPE_CUDA_KERNEL 2 | #define _RPE_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void dot_prod_with_idx_forward_cuda(int N, int M, int h, int hdim, at::Tensor q_tensor, at::Tensor index_tensor, at::Tensor table_tensor, at::Tensor rel_idx_tensor, at::Tensor output_tensor); 8 | void dot_prod_with_idx_backward_cuda(int N, int M, int h, int hdim, at::Tensor grad_out_tensor, at::Tensor q_tensor, at::Tensor index_tensor, at::Tensor table_tensor, at::Tensor rel_idx_tensor, at::Tensor grad_q_tensor, at::Tensor grad_table_tensor); 9 | 10 | void attention_step2_with_rel_pos_value_forward_cuda(int N, int M, int h, int hdim, at::Tensor attn_tensor, at::Tensor v_tensor, at::Tensor index0_tensor, at::Tensor index1_tensor, at::Tensor table_tensor, at::Tensor rel_idx_tensor, at::Tensor output_tensor); 11 | void attention_step2_with_rel_pos_value_backward_cuda(int N, int M, int h, int hdim, at::Tensor grad_out_tensor, at::Tensor index0_tensor, at::Tensor index1_tensor, at::Tensor attn_tensor, at::Tensor v_tensor, at::Tensor table_tensor, at::Tensor rel_idx_tensor, at::Tensor grad_attn_tensor, at::Tensor grad_v_tensor, at::Tensor grad_table_tensor); 12 | 13 | #ifdef __cplusplus 14 | extern "C" { 15 | #endif 16 | 17 | void dot_prod_with_idx_forward_cuda_launcher(int N, int M, int h, int hdim, const float *q, const int *index, const float *table, const int *rel_idx, float *output); 18 | void dot_prod_with_idx_backward_cuda_launcher(int N, int M, int h, int hdim, const float *grad_out, const float *q, const int *index, const float *table, const int *rel_idx, float *grad_q, float *grad_table); 19 | 20 | void attention_step2_with_rel_pos_value_forward_cuda_launcher(int N, int M, int h, int hdim, const float *attn, const float *v, const int *index0, const int *index1, const float *table, const int *rel_idx, float *output); 21 | void attention_step2_with_rel_pos_value_backward_cuda_launcher(int N, int M, int h, int hdim, const float *grad_out, const int *index0, const int *index1, const float *attn, const float *v, const float *table, const int *rel_idx, float *grad_attn, float *grad_v, float *grad_table); 22 | 23 | #ifdef __cplusplus 24 | } 25 | #endif 26 | #endif 27 | -------------------------------------------------------------------------------- /utils/pointops2/src/sampling/sampling_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | #include "sampling_cuda_kernel.h" 6 | 7 | 8 | void furthestsampling_cuda(int b, int n, at::Tensor xyz_tensor, at::Tensor offset_tensor, at::Tensor new_offset_tensor, at::Tensor tmp_tensor, at::Tensor idx_tensor) 9 | { 10 | const float *xyz = xyz_tensor.data_ptr(); 11 | const int *offset = offset_tensor.data_ptr(); 12 | const int *new_offset = new_offset_tensor.data_ptr(); 13 | float *tmp = tmp_tensor.data_ptr(); 14 | int *idx = idx_tensor.data_ptr(); 15 | furthestsampling_cuda_launcher(b, n, xyz, offset, new_offset, tmp, idx); 16 | } 17 | -------------------------------------------------------------------------------- /utils/pointops2/src/sampling/sampling_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _SAMPLING_CUDA_KERNEL 2 | #define _SAMPLING_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void furthestsampling_cuda(int b, int n, at::Tensor xyz_tensor, at::Tensor offset_tensor, at::Tensor new_offset_tensor, at::Tensor tmp_tensor, at::Tensor idx_tensor); 8 | 9 | #ifdef __cplusplus 10 | extern "C" { 11 | #endif 12 | 13 | void furthestsampling_cuda_launcher(int b, int n, const float *xyz, const int *offset, const int *new_offset, float *tmp, int *idx); 14 | 15 | #ifdef __cplusplus 16 | } 17 | #endif 18 | #endif 19 | -------------------------------------------------------------------------------- /utils/pointops2/src/subtraction/subtraction_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | #include "subtraction_cuda_kernel.h" 6 | 7 | 8 | void subtraction_forward_cuda(int n, int nsample, int c, at::Tensor input1_tensor, at::Tensor input2_tensor, at::Tensor idx_tensor, at::Tensor output_tensor) 9 | { 10 | const float *input1 = input1_tensor.data_ptr(); 11 | const float *input2 = input2_tensor.data_ptr(); 12 | const int *idx = idx_tensor.data_ptr(); 13 | float *output = output_tensor.data_ptr(); 14 | subtraction_forward_cuda_launcher(n, nsample, c, input1, input2, idx, output); 15 | } 16 | 17 | void subtraction_backward_cuda(int n, int nsample, int c, at::Tensor idx_tensor, at::Tensor grad_output_tensor, at::Tensor grad_input1_tensor, at::Tensor grad_input2_tensor) 18 | { 19 | const int *idx = idx_tensor.data_ptr(); 20 | const float *grad_output = grad_output_tensor.data_ptr(); 21 | float *grad_input1 = grad_input1_tensor.data_ptr(); 22 | float *grad_input2 = grad_input2_tensor.data_ptr(); 23 | subtraction_backward_cuda_launcher(n, nsample, c, idx, grad_output, grad_input1, grad_input2); 24 | } 25 | -------------------------------------------------------------------------------- /utils/pointops2/src/subtraction/subtraction_cuda_kernel.cu: -------------------------------------------------------------------------------- 1 | #include "../cuda_utils.h" 2 | #include "subtraction_cuda_kernel.h" 3 | 4 | 5 | __global__ void subtraction_forward_cuda_kernel(int n, int nsample, int c, const float *input1, const float *input2, const int *idx, float *output) { 6 | // input: input1: (n, c), input2: (n, c), idx: (n, nsample), output: (n, nsample, c) 7 | int index = blockIdx.x * blockDim.x + threadIdx.x; 8 | if (index >= n * nsample * c) return; 9 | const int c_idx = index % c; 10 | const int nsample_idx = (index / c) % nsample; 11 | const int n_idx = index / nsample / c; 12 | const int idx_idx = n_idx * nsample + nsample_idx; 13 | const int input1_idx = n_idx * c + c_idx; 14 | const int input2_idx = idx[idx_idx] * c + c_idx; 15 | output[index] = input1[input1_idx] - input2[input2_idx]; 16 | } 17 | 18 | __global__ void subtraction_backward_cuda_kernel(int n, int nsample, int c, const int *idx, const float *grad_output, float *grad_input1, float *grad_input2) { 19 | // input: grad_output: (n, nsample, c), output: grad_input1: (n, c), grad_input2: (n, c) 20 | int index = blockIdx.x * blockDim.x + threadIdx.x; 21 | if (index >= n * nsample * c) return; 22 | const int c_idx = index % c; 23 | const int nsample_idx = (index / c) % nsample; 24 | const int n_idx = index / nsample / c; 25 | const int idx_idx = n_idx * nsample + nsample_idx; 26 | const int input1_idx = n_idx * c + c_idx; 27 | const int input2_idx = idx[idx_idx] * c + c_idx; 28 | atomicAdd(grad_input1 + input1_idx, grad_output[index]); 29 | atomicAdd(grad_input2 + input2_idx, -grad_output[index]); 30 | } 31 | 32 | void subtraction_forward_cuda_launcher(int n, int nsample, int c, const float *input1, const float *input2, const int *idx, float *output) { 33 | // input: input1: (n, c), input2: (n, c), idx: (n, nsample), output: (n, nsample, c) 34 | dim3 blocks(DIVUP(n * nsample * c, THREADS_PER_BLOCK)); 35 | dim3 threads(THREADS_PER_BLOCK); 36 | subtraction_forward_cuda_kernel<<>>(n, nsample, c, input1, input2, idx, output); 37 | } 38 | 39 | void subtraction_backward_cuda_launcher(int n, int nsample, int c, const int *idx, const float *grad_output, float *grad_input1, float *grad_input2) { 40 | // input: grad_output: (n, nsample, c), output: grad_input1: (n, c), grad_input2: (n, c) 41 | dim3 blocks(DIVUP(n * nsample * c, THREADS_PER_BLOCK)); 42 | dim3 threads(THREADS_PER_BLOCK); 43 | subtraction_backward_cuda_kernel<<>>(n, nsample, c, idx, grad_output, grad_input1, grad_input2); 44 | } 45 | -------------------------------------------------------------------------------- /utils/pointops2/src/subtraction/subtraction_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _SUBTRACTION_CUDA_KERNEL 2 | #define _SUBTRACTION_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void subtraction_forward_cuda(int n, int nsample, int c, at::Tensor input1_tensor, at::Tensor input2_tensor, at::Tensor idx_tensor, at::Tensor output_tensor); 8 | void subtraction_backward_cuda(int n, int nsample, int c, at::Tensor idx_tensor, at::Tensor grad_output_tensor, at::Tensor grad_input1_tensor, at::Tensor grad_input2_tensor); 9 | 10 | #ifdef __cplusplus 11 | extern "C" { 12 | #endif 13 | 14 | void subtraction_forward_cuda_launcher(int n, int nsample, int c, const float *input1, const float *input2, const int *idx, float *output); 15 | void subtraction_backward_cuda_launcher(int n, int nsample, int c, const int *idx, const float *grad_output, float *grad_input1, float *grad_input2); 16 | 17 | #ifdef __cplusplus 18 | } 19 | #endif 20 | #endif 21 | -------------------------------------------------------------------------------- /utils/votenet_utils/__pycache__/box_util.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/utils/votenet_utils/__pycache__/box_util.cpython-310.pyc -------------------------------------------------------------------------------- /utils/votenet_utils/__pycache__/box_util.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/utils/votenet_utils/__pycache__/box_util.cpython-38.pyc -------------------------------------------------------------------------------- /utils/votenet_utils/__pycache__/eval_det.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/utils/votenet_utils/__pycache__/eval_det.cpython-310.pyc -------------------------------------------------------------------------------- /utils/votenet_utils/__pycache__/eval_det.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/utils/votenet_utils/__pycache__/eval_det.cpython-38.pyc -------------------------------------------------------------------------------- /utils/votenet_utils/__pycache__/metric_util.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/utils/votenet_utils/__pycache__/metric_util.cpython-310.pyc -------------------------------------------------------------------------------- /utils/votenet_utils/__pycache__/metric_util.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SooLab/Part2Object/65d728edecdb461bf6fd88c173996874ae862c45/utils/votenet_utils/__pycache__/metric_util.cpython-38.pyc -------------------------------------------------------------------------------- /utils/votenet_utils/nn_distance.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | """ Chamfer distance in Pytorch. 7 | Author: Charles R. Qi 8 | """ 9 | 10 | import torch 11 | import torch.nn as nn 12 | import numpy as np 13 | 14 | 15 | def huber_loss(error, delta=1.0): 16 | """ 17 | Args: 18 | error: Torch tensor (d1,d2,...,dk) 19 | Returns: 20 | loss: Torch tensor (d1,d2,...,dk) 21 | 22 | x = error = pred - gt or dist(pred,gt) 23 | 0.5 * |x|^2 if |x|<=d 24 | 0.5 * d^2 + d * (|x|-d) if |x|>d 25 | Ref: https://github.com/charlesq34/frustum-pointnets/blob/master/models/model_util.py 26 | """ 27 | abs_error = torch.abs(error) 28 | #quadratic = torch.min(abs_error, torch.FloatTensor([delta])) 29 | quadratic = torch.clamp(abs_error, max=delta) 30 | linear = (abs_error - quadratic) 31 | loss = 0.5 * quadratic**2 + delta * linear 32 | return loss 33 | 34 | def nn_distance(pc1, pc2, l1smooth=False, delta=1.0, l1=False): 35 | """ 36 | Input: 37 | pc1: (B,N,C) torch tensor 38 | pc2: (B,M,C) torch tensor 39 | l1smooth: bool, whether to use l1smooth loss 40 | delta: scalar, the delta used in l1smooth loss 41 | Output: 42 | dist1: (B,N) torch float32 tensor 43 | idx1: (B,N) torch int64 tensor 44 | dist2: (B,M) torch float32 tensor 45 | idx2: (B,M) torch int64 tensor 46 | """ 47 | N = pc1.shape[1] 48 | M = pc2.shape[1] 49 | pc1_expand_tile = pc1.unsqueeze(2).repeat(1,1,M,1) 50 | pc2_expand_tile = pc2.unsqueeze(1).repeat(1,N,1,1) 51 | pc_diff = pc1_expand_tile - pc2_expand_tile 52 | 53 | if l1smooth: 54 | pc_dist = torch.sum(huber_loss(pc_diff, delta), dim=-1) # (B,N,M) 55 | elif l1: 56 | pc_dist = torch.sum(torch.abs(pc_diff), dim=-1) # (B,N,M) 57 | else: 58 | pc_dist = torch.sum(pc_diff**2, dim=-1) # (B,N,M) 59 | dist1, idx1 = torch.min(pc_dist, dim=2) # (B,N) 60 | dist2, idx2 = torch.min(pc_dist, dim=1) # (B,M) 61 | return dist1, idx1, dist2, idx2 62 | 63 | def demo_nn_distance(): 64 | np.random.seed(0) 65 | pc1arr = np.random.random((1,5,3)) 66 | pc2arr = np.random.random((1,6,3)) 67 | pc1 = torch.from_numpy(pc1arr.astype(np.float32)) 68 | pc2 = torch.from_numpy(pc2arr.astype(np.float32)) 69 | dist1, idx1, dist2, idx2 = nn_distance(pc1, pc2) 70 | print(dist1) 71 | print(idx1) 72 | dist = np.zeros((5,6)) 73 | for i in range(5): 74 | for j in range(6): 75 | dist[i,j] = np.sum((pc1arr[0,i,:] - pc2arr[0,j,:]) ** 2) 76 | print(dist) 77 | print('-'*30) 78 | print('L1smooth dists:') 79 | dist1, idx1, dist2, idx2 = nn_distance(pc1, pc2, True) 80 | print(dist1) 81 | print(idx1) 82 | dist = np.zeros((5,6)) 83 | for i in range(5): 84 | for j in range(6): 85 | error = np.abs(pc1arr[0,i,:] - pc2arr[0,j,:]) 86 | quad = np.minimum(error, 1.0) 87 | linear = error - quad 88 | loss = 0.5*quad**2 + 1.0*linear 89 | dist[i,j] = np.sum(loss) 90 | print(dist) 91 | 92 | 93 | if __name__ == '__main__': 94 | demo_nn_distance() 95 | -------------------------------------------------------------------------------- /utils/votenet_utils/tf_logger.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | import tensorflow as tf 7 | import numpy as np 8 | import scipy.misc 9 | try: 10 | from StringIO import StringIO # Python 2.7 11 | except ImportError: 12 | from io import BytesIO # Python 3.x 13 | 14 | 15 | class Logger(object): 16 | 17 | def __init__(self, log_dir): 18 | """Create a summary writer logging to log_dir.""" 19 | self.writer = tf.summary.FileWriter(log_dir) 20 | 21 | def scalar_summary(self, tag, value, step): 22 | """Log a scalar variable.""" 23 | summary = tf.Summary(value=[tf.Summary.Value(tag=tag, simple_value=value)]) 24 | self.writer.add_summary(summary, step) 25 | 26 | def image_summary(self, tag, images, step): 27 | """Log a list of images.""" 28 | 29 | img_summaries = [] 30 | for i, img in enumerate(images): 31 | # Write the image to a string 32 | try: 33 | s = StringIO() 34 | except: 35 | s = BytesIO() 36 | scipy.misc.toimage(img).save(s, format="png") 37 | 38 | # Create an Image object 39 | img_sum = tf.Summary.Image(encoded_image_string=s.getvalue(), 40 | height=img.shape[0], 41 | width=img.shape[1]) 42 | # Create a Summary value 43 | img_summaries.append(tf.Summary.Value(tag='%s/%d' % (tag, i), image=img_sum)) 44 | 45 | # Create and write Summary 46 | summary = tf.Summary(value=img_summaries) 47 | self.writer.add_summary(summary, step) 48 | 49 | def histo_summary(self, tag, values, step, bins=1000): 50 | """Log a histogram of the tensor of values.""" 51 | 52 | # Create a histogram using numpy 53 | counts, bin_edges = np.histogram(values, bins=bins) 54 | 55 | # Fill the fields of the histogram proto 56 | hist = tf.HistogramProto() 57 | hist.min = float(np.min(values)) 58 | hist.max = float(np.max(values)) 59 | hist.num = int(np.prod(values.shape)) 60 | hist.sum = float(np.sum(values)) 61 | hist.sum_squares = float(np.sum(values**2)) 62 | 63 | # Drop the start of the first bin 64 | bin_edges = bin_edges[1:] 65 | 66 | # Add bin edges and counts 67 | for edge in bin_edges: 68 | hist.bucket_limit.append(edge) 69 | for c in counts: 70 | hist.bucket.append(c) 71 | 72 | # Create and write Summary 73 | summary = tf.Summary(value=[tf.Summary.Value(tag=tag, histo=hist)]) 74 | self.writer.add_summary(summary, step) 75 | self.writer.flush() 76 | -------------------------------------------------------------------------------- /utils/votenet_utils/tf_visualizer.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | '''Code adapted from https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix''' 7 | import os 8 | import time 9 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 10 | import sys 11 | sys.path.append(BASE_DIR) 12 | import tf_logger 13 | 14 | 15 | class Visualizer(): 16 | def __init__(self, opt, name='train'): 17 | # self.opt = opt 18 | #self.logger = tf_logger.Logger(os.path.join(opt.logging_dir, opt.name)) 19 | #self.log_name = os.path.join(opt.checkpoint_dir, opt.name, 'loss_log.txt') 20 | self.logger = tf_logger.Logger(os.path.join(opt.log_dir, name)) 21 | self.log_name = os.path.join(opt.log_dir, 'tf_visualizer_log.txt') 22 | with open(self.log_name, "a") as log_file: 23 | now = time.strftime("%c") 24 | log_file.write('================ Training Loss (%s) ================\n' % now) 25 | 26 | # |visuals|: dictionary of images to save 27 | def log_images(self, visuals, step): 28 | for label, image_numpy in visuals.items(): 29 | self.logger.image_summary( 30 | label, [image_numpy], step) 31 | 32 | # scalars: dictionary of scalar labels and values 33 | def log_scalars(self, scalars, step): 34 | for label, val in scalars.items(): 35 | self.logger.scalar_summary(label, val, step) 36 | 37 | # scatter plots 38 | def plot_current_points(self, points, disp_offset=10): 39 | pass 40 | 41 | # scalars: same format as |scalars| of plot_current_scalars 42 | def print_current_scalars(self, epoch, i, scalars): 43 | message = '(epoch: %d, iters: %d) ' % (epoch, i) 44 | for k, v in scalars.items(): 45 | message += '%s: %.3f ' % (k, v) 46 | 47 | print(message) 48 | with open(self.log_name, "a") as log_file: 49 | log_file.write('%s\n' % message) 50 | --------------------------------------------------------------------------------