├── DATASET.md ├── KITTI_metric.py ├── README.md ├── cfgs ├── KITTI_models │ ├── FoldingNet.yaml │ ├── GRNet.yaml │ ├── PCN.yaml │ ├── PoinTr.yaml │ └── TopNet.yaml ├── PCN_models │ ├── FoldingNet.yaml │ ├── GRNet.yaml │ ├── PCN.yaml │ ├── PT_model.yaml │ ├── PoinTr.yaml │ ├── PoinTr_Encoder.yaml │ ├── TopNet.yaml │ └── our_model.yaml ├── ShapeNet34_models │ ├── FoldingNet.yaml │ ├── GRNet.yaml │ ├── PCN.yaml │ ├── PoinTr.yaml │ ├── TopNet.yaml │ └── our_model.yaml ├── ShapeNet55_models │ ├── FoldingNet.yaml │ ├── GRNet.yaml │ ├── PCN.yaml │ ├── PT_model.yaml │ ├── PoinTr.yaml │ ├── TopNet.yaml │ └── our_model.yaml ├── ShapeNetUnseen21_models │ ├── FoldingNet.yaml │ ├── GRNet.yaml │ ├── PCN.yaml │ ├── PoinTr.yaml │ ├── TopNet.yaml │ └── our_model.yaml └── dataset_configs │ ├── KITTI.yaml │ ├── PCN.yaml │ ├── PCNCars.yaml │ ├── ShapeNet-34.yaml │ ├── ShapeNet-55.yaml │ └── ShapeNet-Unseen21.yaml ├── datasets ├── KITTIDataset.py ├── PCN_Dataloader.py ├── ShapeNet55Dataset.py ├── build.py ├── data_transforms.py └── io.py ├── extensions ├── chamfer_dist │ ├── __init__.py │ ├── __pycache__ │ │ └── __init__.cpython-37.pyc │ ├── build │ │ ├── lib.linux-x86_64-3.7 │ │ │ └── chamfer.cpython-37m-x86_64-linux-gnu.so │ │ └── temp.linux-x86_64-3.7 │ │ │ ├── chamfer.o │ │ │ └── chamfer_cuda.o │ ├── chamfer.cu │ ├── chamfer.egg-info │ │ ├── PKG-INFO │ │ ├── SOURCES.txt │ │ ├── dependency_links.txt │ │ └── top_level.txt │ ├── chamfer_cuda.cpp │ ├── dist │ │ └── chamfer-2.0.0-py3.7-linux-x86_64.egg │ ├── setup.py │ └── test.py ├── cubic_feature_sampling │ ├── __init__.py │ ├── __pycache__ │ │ └── __init__.cpython-37.pyc │ ├── cubic_feature_sampling.cu │ ├── cubic_feature_sampling_cuda.cpp │ ├── setup.py │ └── test.py ├── gridding │ ├── __init__.py │ ├── __pycache__ │ │ └── __init__.cpython-37.pyc │ ├── gridding.cu │ ├── gridding_cuda.cpp │ ├── gridding_reverse.cu │ ├── setup.py │ └── test.py └── gridding_loss │ ├── __init__.py │ ├── __pycache__ │ └── __init__.cpython-37.pyc │ ├── gridding_distance.cu │ ├── gridding_distance_cuda.cpp │ └── setup.py ├── figs ├── More_Results_On_ShapNet55_1.pdf ├── PCN_Dataset_Result_Picture.pdf ├── PCN_results.png ├── Pipeline.png ├── Pipeline2.pdf ├── ShapeNet55_Results.pdf ├── ShapeNet55_results.png └── more_shapenet55_results.png ├── install.sh ├── main.py ├── other_models ├── FoldingNet.py ├── GRNet.py ├── PCN.py ├── PoinTr.py ├── TopNet.py ├── Transformer.py ├── build.py └── dgcnn_group.py ├── our_models ├── PT_model.py ├── PoinTr_Encoder.py ├── PointTransformer.py ├── model_Transformer.py ├── model_Transformer2.py ├── model_Transformer3.py └── pointnet_util.py ├── requirements.txt ├── scripts ├── dist_train.sh ├── test.sh └── train.sh ├── tools ├── __init__.py ├── __pycache__ │ ├── __init__.cpython-37.pyc │ ├── __init__.cpython-39.pyc │ ├── builder.cpython-37.pyc │ ├── runner.cpython-37.pyc │ ├── runner2.cpython-37.pyc │ ├── runner3.cpython-37.pyc │ ├── runner_model.cpython-37.pyc │ ├── runner_reconstruct.cpython-37.pyc │ └── runner_reconstruct.cpython-39.pyc ├── builder.py ├── runner.py ├── runner2.py ├── runner3.py ├── runner_model.py └── runner_reconstruct.py └── utils ├── AverageMeter.py ├── __pycache__ ├── AverageMeter.cpython-37.pyc ├── config.cpython-37.pyc ├── dist_utils.cpython-37.pyc ├── eulerangles.cpython-37.pyc ├── logger.cpython-37.pyc ├── metrics.cpython-37.pyc ├── misc.cpython-37.pyc ├── other_utils.cpython-37.pyc ├── parser.cpython-37.pyc ├── pc_util.cpython-37.pyc ├── plyfile.cpython-37.pyc └── registry.cpython-37.pyc ├── config.py ├── dist_utils.py ├── eulerangles.py ├── logger.py ├── metrics.py ├── misc.py ├── other_utils.py ├── parser.py ├── pc_util.py ├── plyfile.py ├── pointnet2_ops_lib ├── build │ ├── lib.linux-x86_64-3.7 │ │ └── pointnet2_ops │ │ │ ├── __init__.py │ │ │ ├── _ext-src │ │ │ └── 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 │ │ │ ├── _ext.cpython-37m-x86_64-linux-gnu.so │ │ │ ├── _version.py │ │ │ ├── pointnet2_modules.py │ │ │ └── pointnet2_utils.py │ └── temp.linux-x86_64-3.7 │ │ ├── .ninja_deps │ │ ├── .ninja_log │ │ ├── build.ninja │ │ └── pointnet2_ops │ │ └── _ext-src │ │ └── src │ │ ├── ball_query.o │ │ ├── ball_query_gpu.o │ │ ├── bindings.o │ │ ├── group_points.o │ │ ├── group_points_gpu.o │ │ ├── interpolate.o │ │ ├── interpolate_gpu.o │ │ ├── sampling.o │ │ └── sampling_gpu.o ├── dist │ └── pointnet2_ops-3.0.0-py3.7-linux-x86_64.egg ├── pointnet2_ops.egg-info │ ├── PKG-INFO │ ├── SOURCES.txt │ ├── dependency_links.txt │ ├── requires.txt │ └── top_level.txt ├── pointnet2_ops │ ├── __init__.py │ ├── _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 │ ├── _version.py │ ├── pointnet2_modules.py │ └── pointnet2_utils.py └── setup.py └── registry.py /DATASET.md: -------------------------------------------------------------------------------- 1 | ## Dataset 2 | 3 | ### ShapNet55/34 4 | 5 | Thanks to the authors of PoinTr for these two extremely challenging datasets. 6 | 7 | 8 | ### Data Preparation 9 | The overall directory structure should be: 10 | 11 | ``` 12 | │PoinTr/ 13 | ├──cfgs/ 14 | ├──datasets/ 15 | ├──data/ 16 | │ ├──ShapeNet55-34/ 17 | │ ├──PCN/ 18 | │ ├──KITTI/ 19 | ├──....... 20 | ``` 21 | **ShapeNet55/34 Dataset**: You can download the processed ShapeNet55/34 dataset at [[BaiduCloud](https://pan.baidu.com/s/16Q-GsEXEHkXRhmcSZTY86A)] (code:le04) or [[Google Drive](https://drive.google.com/file/d/1jUB5yD7DP97-EqqU2A9mmr61JpNwZBVK/view?usp=sharing)]. Unzip the file under `ShapeNet55-34/`. The directory structure should be 22 | 23 | ``` 24 | │ShapeNet55-34/ 25 | ├──shapenet_pc/ 26 | │ ├── 02691156-1a04e3eab45ca15dd86060f189eb133.npy 27 | │ ├── 02691156-1a6ad7a24bb89733f412783097373bdc.npy 28 | │ ├── ....... 29 | ├──ShapeNet-34/ 30 | │ ├── train.txt 31 | │ └── test.txt 32 | ├──ShapeNet-34/ 33 | │ ├── train.txt 34 | │ └── test.txt 35 | ├──ShapeNet-Unseen21/ 36 | └── test.txt 37 | ``` 38 | 39 | **PCN Dataset**: You can download the processed PCN dataset from this [url](https://gateway.infinitescript.com/?fileName=ShapeNetCompletion). The directory structure should be 40 | 41 | ``` 42 | │PCN/ 43 | ├──train/ 44 | │ ├── complete 45 | │ │ ├── 02691156 46 | │ │ │ ├── 1a04e3eab45ca15dd86060f189eb133.pcd 47 | │ │ │ ├── ....... 48 | │ │ ├── ....... 49 | │ ├── partial 50 | │ │ ├── 02691156 51 | │ │ │ ├── 1a04e3eab45ca15dd86060f189eb133 52 | │ │ │ │ ├── 00.pcd 53 | │ │ │ │ ├── 01.pcd 54 | │ │ │ │ ├── ....... 55 | │ │ │ │ └── 07.pcd 56 | │ │ │ ├── ....... 57 | │ │ ├── ....... 58 | ├──test/ 59 | │ ├── complete 60 | │ │ ├── 02691156 61 | │ │ │ ├── 1d63eb2b1f78aa88acf77e718d93f3e1.pcd 62 | │ │ │ ├── ....... 63 | │ │ ├── ....... 64 | │ ├── partial 65 | │ │ ├── 02691156 66 | │ │ │ ├── 1d63eb2b1f78aa88acf77e718d93f3e1 67 | │ │ │ │ └── 00.pcd 68 | │ │ │ ├── ....... 69 | │ │ ├── ....... 70 | ├──val/ 71 | │ ├── complete 72 | │ │ ├── 02691156 73 | │ │ │ ├── 4bae467a3dad502b90b1d6deb98feec6.pcd 74 | │ │ │ ├── ....... 75 | │ │ ├── ....... 76 | │ ├── partial 77 | │ │ ├── 02691156 78 | │ │ │ ├── 4bae467a3dad502b90b1d6deb98feec6 79 | │ │ │ │ └── 00.pcd 80 | │ │ │ ├── ....... 81 | │ │ ├── ....... 82 | ├──PCN.json 83 | └──category.txt 84 | ``` 85 | 86 | **KITTI**: You can download the KITTI dataset from this [url](https://drive.google.com/drive/folders/1fSu0_huWhticAlzLh3Ejpg8zxzqO1z-F). The directory structure should be 87 | 88 | ``` 89 | │KITTI/ 90 | ├──bboxes/ 91 | │ ├── frame_0_car_0.txt 92 | │ ├── ....... 93 | ├──cars/ 94 | │ ├── frame_0_car_0.pcd 95 | │ ├── ....... 96 | ├──tracklets/ 97 | │ ├── tracklet_0.txt 98 | │ ├── ....... 99 | ├──KITTI.json 100 | ``` 101 | -------------------------------------------------------------------------------- /KITTI_metric.py: -------------------------------------------------------------------------------- 1 | import os 2 | import yaml 3 | from datasets import build_dataset_from_cfg 4 | 5 | from easydict import EasyDict 6 | import argparse 7 | from extensions.chamfer_dist import ChamferDistanceL2_split, ChamferDistanceL2 8 | import torch 9 | import numpy as np 10 | 11 | from tqdm import tqdm 12 | 13 | def build_ShapeNetCars(): 14 | ShapeNetCars_config = val = yaml.load(open('cfgs/dataset_configs/PCNCars.yaml', 'r'), Loader=yaml.FullLoader) 15 | train_dataset = build_dataset_from_cfg(EasyDict(ShapeNetCars_config), EasyDict(subset = 'train')) 16 | test_dataset = build_dataset_from_cfg(EasyDict(ShapeNetCars_config), EasyDict(subset = 'test')) 17 | val_dataset = build_dataset_from_cfg(EasyDict(ShapeNetCars_config), EasyDict(subset = 'val')) 18 | CarsDataset = train_dataset + test_dataset + val_dataset 19 | return CarsDataset 20 | def get_args(): 21 | parser = argparse.ArgumentParser() 22 | parser.add_argument( 23 | '--vis_path', 24 | type = str, 25 | help = 'KITTI visualize path') 26 | args = parser.parse_args() 27 | return args 28 | 29 | def get_Fidelity(): 30 | # Fidelity Error 31 | criterion = ChamferDistanceL2_split(ignore_zeros=True) 32 | 33 | metric = [] 34 | for sample in Samples: 35 | input_data = torch.from_numpy(np.load(os.path.join(Data_path, sample, 'input.npy'))).unsqueeze(0).cuda() 36 | pred_data = torch.from_numpy(np.load(os.path.join(Data_path, sample, 'pred.npy'))).unsqueeze(0).cuda() 37 | metric.append(criterion(input_data, pred_data)[0]) 38 | print('Fidelity is %f' % (sum(metric)/len(metric))) 39 | 40 | def get_Consistency(): 41 | #Consistency 42 | criterion = ChamferDistanceL2(ignore_zeros=True) 43 | Cars_dict = {} 44 | for sample in Samples: 45 | all_elements = sample.split('_') # example sample = 'frame_1_car_3_647' 46 | frame_id = int(all_elements[1]) 47 | car_id = int(all_elements[-2]) 48 | sample_id = int(all_elements[-1]) 49 | 50 | if Cars_dict.get(car_id) is None: 51 | Cars_dict[car_id] = [f'frame_{frame_id:03d}_car_{car_id:02d}_{sample_id:03d}'] 52 | else: 53 | Cars_dict[car_id].append(f'frame_{frame_id:03d}_car_{car_id:02d}_{sample_id:03d}') # example sample = 'frame_001_car_003_647' 54 | 55 | Consistency = [] 56 | for key, car_list in Cars_dict.items(): 57 | car_list = sorted(car_list) 58 | Each_Car_Consistency = [] 59 | for i, this_car in enumerate(car_list): 60 | if i == len(car_list) - 1: 61 | break 62 | this_elements = this_car.split('_') 63 | this_frame =int(this_elements[1]) 64 | 65 | next_car = car_list[i + 1] 66 | next_elements = next_car.split('_') 67 | next_frame = int(next_elements[1]) 68 | 69 | if next_frame - 1 != this_frame: 70 | continue 71 | 72 | this_car = torch.from_numpy(np.load(os.path.join(Data_path, f'frame_{this_frame}_car_{int(this_elements[3])}_{int(this_elements[4]):03d}', 'pred.npy'))).unsqueeze(0).cuda() 73 | next_car = torch.from_numpy(np.load(os.path.join(Data_path, f'frame_{next_frame}_car_{int(next_elements[3])}_{int(next_elements[4]):03d}', 'pred.npy'))).unsqueeze(0).cuda() 74 | cd = criterion(this_car, next_car) 75 | Each_Car_Consistency.append(cd) 76 | 77 | MeanCD = sum(Each_Car_Consistency) / len(Each_Car_Consistency) 78 | Consistency.append(MeanCD) 79 | MeanCD = sum(Consistency) / len(Consistency) 80 | print(f'Consistency is {MeanCD:.6f}') 81 | 82 | def get_MMD(): 83 | criterion = ChamferDistanceL2(ignore_zeros=True) 84 | #MMD 85 | metric = [] 86 | for item in tqdm(sorted(Samples)): 87 | pred_data = torch.from_numpy(np.load(os.path.join(Data_path, item, 'pred.npy'))).unsqueeze(0).cuda() 88 | batch_cd = [] 89 | for index in range(len(ShapeNetCars_dataset)): 90 | gt = ShapeNetCars_dataset[index][-1][1].cuda().unsqueeze(0) 91 | # for index, (taxonomy_ids, model_ids, data) in enumerate(CarsDataloader): 92 | # gt = data[1].cuda() 93 | # batch_pred_data = pred_data.expand(gt.size(0), -1, -1).contiguous() 94 | min_cd = criterion(gt, pred_data) 95 | batch_cd.append(min_cd) 96 | min_cd = min(batch_cd).item() 97 | metric.append(min_cd) 98 | print('This item %s CD %f, MMD %f' % (item, min_cd, sum(metric)*1.0 / len(metric) )) 99 | print('MMD is %f' % (sum(metric)/len(metric))) 100 | 101 | if __name__ == '__main__': 102 | args = get_args() 103 | ShapeNetCars_dataset = build_ShapeNetCars() 104 | CarsDataloader = torch.utils.data.DataLoader(ShapeNetCars_dataset, batch_size=1, shuffle = False, num_workers = 8) 105 | 106 | # Your data 107 | Data_path = args.vis_path 108 | Samples = [item for item in os.listdir(Data_path) if os.path.isdir(Data_path + '/' + item)] 109 | criterion = ChamferDistanceL2_split(ignore_zeros=True) 110 | 111 | get_Fidelity() 112 | get_MMD() 113 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # ProxyFormer: Proxy Alignment Assisted Point Cloud Completion with Missing Part Sensitive Transformer 2 | 3 | Created by [Shanshan Li](https://github.com/MarkLiSS), [Pan Gao]()\*, [Xiaoyang Tan](), [Mingqiang Wei]() 4 | 5 | [[arXiv]](https://arxiv.org/abs/2302.14435) [[Dataset]](./DATASET.md) [[supp]](https://github.com/MarkLiSS/MyPapers/blob/main/Appendix_for_ProxyFormer_camera_ready.pdf) 6 | 7 | This repository contains PyTorch implementation for __ProxyFormer: Proxy Alignment Assisted Point Cloud Completion with Missing Part Sensitive Transformer__ (Accepted by CVPR 2023). 8 | 9 | Problems such as equipment defects or limited viewpoints will lead the captured point clouds to be incomplete. Therefore, recovering the complete point clouds from the partial ones plays an vital role in many practical tasks, and one of the keys lies in the prediction of the missing part. In this paper, we propose a novel point cloud completion approach namely ProxyFormer that divides point clouds into existing (input) and missing (to be predicted) parts and each part communicates information through its proxies. Specifically, we fuse information into point proxy via feature and position extractor, and generate features for missing point proxies from the features of existing point proxies. Then, in order to better perceive the position of missing points, we design a missing part sensitive transformer, which converts random normal distribution into reasonable position information, and uses proxy alignment to refine the missing proxies. It makes the predicted point proxies more sensitive to the features and positions of the missing part, and thus make these proxies more suitable for subsequent coarse-to-fine processes. Experimental results show that our method outperforms state-of-the-art completion networks on several benchmark datasets and has the fastest inference speed. 10 | 11 | ![intro](figs/Pipeline.png) 12 | 13 | ## 🔥News 14 | - **2023-02-28** ProxyFormer is accepted by CVPR 2023. 15 | - **2022-11-05** ProxyFormer outperforms state-of-the-art completion networks on several benchmark datasets and has the fastest inference speed (PCN: CD 6.77 & DCD 0.577 etc.). 16 | 17 | 18 | ## Usage 19 | 20 | ### Requirements 21 | 22 | - PyTorch >= 1.9.0 23 | - python >= 3.7 24 | - CUDA >= 11.0 25 | - GCC >= 4.9 26 | - torchvision 27 | - open3d 28 | - plyfile 29 | - tensorboardX 30 | 31 | ``` 32 | pip install -r requirements.txt 33 | ``` 34 | 35 | #### Building Pytorch Extensions for Chamfer Distance, PointNet++ and kNN 36 | 37 | *NOTE:* PyTorch >= 1.9 and GCC >= 4.9 are required. 38 | 39 | ``` 40 | # Using a bash to install them 41 | bash install.sh 42 | ``` 43 | The solution for a common bug in chamfer distance installation can be found in Issue [#6] of PoinTr (https://github.com/yuxumin/PoinTr/issues/6) 44 | 45 | 46 | ### Dataset 47 | 48 | The details of ***ShapeNet-55/34*** datasets and other existing datasets can be found in [DATASET.md](./DATASET.md). 49 | 50 | ### Evaluation 51 | 52 | To evaluate a pre-trained ProxyFormer model on the Three Dataset with single GPU, run: 53 | 54 | ``` 55 | bash ./scripts/test.sh \ 56 | --ckpts \ 57 | --config \ 58 | --exp_name \ 59 | [--mode ] 60 | ``` 61 | 62 | #### Some examples: 63 | Test the ProxyFormer pretrained model on the PCN benchmark: 64 | ``` 65 | bash ./scripts/test.sh 0 \ 66 | --ckpts ./pretrained/ckpt-best.pth \ 67 | --config ./cfgs/PCN_models/our_model.yaml \ 68 | --exp_name example 69 | ``` 70 | Test the ProxyFormer pretrained model on ShapeNet55 benchmark (*easy* mode): 71 | ``` 72 | bash ./scripts/test.sh 0 \ 73 | --ckpts ./pretrained/ckpt-best.pth \ 74 | --config ./cfgs/ShapeNet55_models/our_model.yaml \ 75 | --mode easy \ 76 | --exp_name example 77 | ``` 78 | 79 | ### Training 80 | 81 | To train ProxyFormer from scratch, run: 82 | 83 | ``` 84 | # Use DistributedDataParallel (DDP) 85 | bash ./scripts/dist_train.sh \ 86 | --config \ 87 | --exp_name \ 88 | [--resume] \ 89 | [--start_ckpts ] \ 90 | [--val_freq ] 91 | # or just use DataParallel (DP) 92 | bash ./scripts/train.sh \ 93 | --config \ 94 | --exp_name \ 95 | [--resume] \ 96 | [--start_ckpts ] \ 97 | [--val_freq ] 98 | ``` 99 | #### Some examples: 100 | Train a ProxyFormer model with a single GPU: 101 | ``` 102 | bash ./scripts/train.sh 0 \ 103 | --config ./cfgs/PCN_models/our_model.yaml \ 104 | --exp_name example 105 | ``` 106 | 107 | ### Completion Results on PCN and ShapeNet-55 108 | 109 | #### PCN 110 | ![results](figs/PCN_results.png) 111 | #### ShapeNet-55 112 | ![results](figs/ShapeNet55_results.png) 113 | ![results](figs/more_shapenet55_results.png) 114 | 115 | ## Acknowledgements 116 | 117 | Our code is inspired by [PoinTr](https://github.com/yuxumin/PoinTr) and [SeedFormer](https://github.com/hrzhou2/seedformer). Thanks for their excellent works. 118 | 119 | ## Citation 120 | If you find our work useful in your research, please consider citing: 121 | ``` 122 | 123 | ``` 124 | -------------------------------------------------------------------------------- /cfgs/KITTI_models/FoldingNet.yaml: -------------------------------------------------------------------------------- 1 | optimizer : { 2 | type: Adam, 3 | kwargs: { 4 | lr : 0.0001, 5 | weight_decay : 0 6 | }} 7 | scheduler: { 8 | type: StepLR, 9 | kwargs: { 10 | step_size: 50, 11 | gamma : 0.5 12 | }} 13 | 14 | dataset : { 15 | train : { _base_: cfgs/dataset_configs/PCNCars.yaml, 16 | others: {subset: 'train'}}, 17 | val : { _base_: cfgs/dataset_configs/PCNCars.yaml, 18 | others: {subset: 'test'}}, 19 | test : { _base_: cfgs/dataset_configs/KITTI.yaml, 20 | others: {subset: 'test'}}} 21 | 22 | model : { 23 | NAME: FoldingNet, num_pred: 16384, encoder_channel: 1024} 24 | total_bs : 32 25 | step_per_update : 1 26 | max_epoch : 150 27 | 28 | consider_metric: CDL1 -------------------------------------------------------------------------------- /cfgs/KITTI_models/GRNet.yaml: -------------------------------------------------------------------------------- 1 | optimizer : { 2 | type: Adam, 3 | kwargs: { 4 | lr : 0.0001, 5 | weight_decay : 0 6 | }} 7 | scheduler: { 8 | type: StepLR, 9 | kwargs: { 10 | step_size: 50, 11 | gamma : 0.5 12 | }} 13 | 14 | dataset : { 15 | train : { _base_: cfgs/dataset_configs/PCNCars.yaml, 16 | others: {subset: 'train'}}, 17 | val : { _base_: cfgs/dataset_configs/PCNCars.yaml, 18 | others: {subset: 'test'}}, 19 | test : { _base_: cfgs/dataset_configs/KITTI.yaml, 20 | others: {subset: 'test'}}} 21 | 22 | model : { 23 | NAME: GRNet, num_pred: 16384, gridding_loss_scales: 128, gridding_loss_alphas: 0.1} 24 | total_bs : 32 25 | step_per_update : 1 26 | max_epoch : 150 27 | 28 | consider_metric: CDL1 -------------------------------------------------------------------------------- /cfgs/KITTI_models/PCN.yaml: -------------------------------------------------------------------------------- 1 | optimizer : { 2 | type: Adam, 3 | kwargs: { 4 | lr : 0.0001, 5 | weight_decay : 0 6 | }} 7 | scheduler: { 8 | type: StepLR, 9 | kwargs: { 10 | step_size: 50, 11 | gamma : 0.5 12 | }} 13 | 14 | dataset : { 15 | train : { _base_: cfgs/dataset_configs/PCNCars.yaml, 16 | others: {subset: 'train'}}, 17 | val : { _base_: cfgs/dataset_configs/PCNCars.yaml, 18 | others: {subset: 'test'}}, 19 | test : { _base_: cfgs/dataset_configs/KITTI.yaml, 20 | others: {subset: 'test'}}} 21 | 22 | model : { 23 | NAME: PCN, num_pred: 16384, encoder_channel: 1024} 24 | total_bs : 32 25 | step_per_update : 1 26 | max_epoch : 150 27 | 28 | consider_metric: CDL1 -------------------------------------------------------------------------------- /cfgs/KITTI_models/PoinTr.yaml: -------------------------------------------------------------------------------- 1 | optimizer : { 2 | type: AdamW, 3 | kwargs: { 4 | lr : 0.0001, 5 | weight_decay : 0.0005 6 | }} 7 | scheduler: { 8 | type: LambdaLR, 9 | kwargs: { 10 | decay_step: 21, 11 | lr_decay: 0.9, 12 | lowest_decay: 0.02 # min lr = lowest_decay * lr 13 | }} 14 | bnmscheduler: { 15 | type: Lambda, 16 | kwargs: { 17 | decay_step: 21, 18 | bn_decay: 0.5, 19 | bn_momentum: 0.9, 20 | lowest_decay: 0.01 21 | }} 22 | 23 | dataset : { 24 | train : { _base_: cfgs/dataset_configs/PCNCars.yaml, 25 | others: {subset: 'train'}}, 26 | val : { _base_: cfgs/dataset_configs/PCNCars.yaml, 27 | others: {subset: 'test'}}, 28 | test : { _base_: cfgs/dataset_configs/KITTI.yaml, 29 | others: {subset: 'test'}}} 30 | 31 | model : { 32 | NAME: PoinTr, num_pred: 14336, num_query: 224, knn_layer: 1, trans_dim: 384} 33 | total_bs : 64 34 | step_per_update : 1 35 | max_epoch : 600 36 | 37 | consider_metric: CDL1 38 | -------------------------------------------------------------------------------- /cfgs/KITTI_models/TopNet.yaml: -------------------------------------------------------------------------------- 1 | optimizer : { 2 | type: Adam, 3 | kwargs: { 4 | lr : 0.001, 5 | weight_decay : 0 6 | }} 7 | scheduler: { 8 | type: StepLR, 9 | kwargs: { 10 | step_size: 50, 11 | gamma : 0.5 12 | }} 13 | 14 | dataset : { 15 | train : { _base_: cfgs/dataset_configs/PCNCars.yaml, 16 | others: {subset: 'train'}}, 17 | val : { _base_: cfgs/dataset_configs/PCNCars.yaml, 18 | others: {subset: 'test'}}, 19 | test : { _base_: cfgs/dataset_configs/KITTI.yaml, 20 | others: {subset: 'test'}}} 21 | 22 | model : { 23 | NAME: TopNet, node_feature: 8, num_pred: 16384, encoder_channel: 1024, nlevels: 8} 24 | total_bs : 32 25 | step_per_update : 1 26 | max_epoch : 150 27 | 28 | consider_metric: CDL1 -------------------------------------------------------------------------------- /cfgs/PCN_models/FoldingNet.yaml: -------------------------------------------------------------------------------- 1 | optimizer : { 2 | type: Adam, 3 | kwargs: { 4 | lr : 0.0001, 5 | weight_decay : 0 6 | }} 7 | scheduler: { 8 | type: StepLR, 9 | kwargs: { 10 | step_size: 50, 11 | gamma : 0.5 12 | }} 13 | 14 | dataset : { 15 | train : { _base_: cfgs/dataset_configs/PCN.yaml, 16 | others: {subset: 'train'}}, 17 | val : { _base_: cfgs/dataset_configs/PCN.yaml, 18 | others: {subset: 'test'}}, 19 | test : { _base_: cfgs/dataset_configs/PCN.yaml, 20 | others: {subset: 'test'}}} 21 | 22 | model : { 23 | NAME: FoldingNet, num_pred: 16384, encoder_channel: 1024} 24 | total_bs : 32 25 | step_per_update : 1 26 | max_epoch : 150 27 | 28 | consider_metric: CDL1 -------------------------------------------------------------------------------- /cfgs/PCN_models/GRNet.yaml: -------------------------------------------------------------------------------- 1 | optimizer : { 2 | type: Adam, 3 | kwargs: { 4 | lr : 0.0001, 5 | weight_decay : 0 6 | }} 7 | scheduler: { 8 | type: StepLR, 9 | kwargs: { 10 | step_size: 50, 11 | gamma : 0.5 12 | }} 13 | 14 | dataset : { 15 | train : { _base_: cfgs/dataset_configs/PCN.yaml, 16 | others: {subset: 'train'}}, 17 | val : { _base_: cfgs/dataset_configs/PCN.yaml, 18 | others: {subset: 'test'}}, 19 | test : { _base_: cfgs/dataset_configs/PCN.yaml, 20 | others: {subset: 'test'}}} 21 | model : { 22 | NAME: GRNet, num_pred: 16384, gridding_loss_scales: 128, gridding_loss_alphas: 0.1} 23 | total_bs : 32 24 | step_per_update : 1 25 | max_epoch : 150 26 | 27 | consider_metric: CDL1 -------------------------------------------------------------------------------- /cfgs/PCN_models/PCN.yaml: -------------------------------------------------------------------------------- 1 | optimizer : { 2 | type: Adam, 3 | kwargs: { 4 | lr : 0.0001, 5 | weight_decay : 0 6 | }} 7 | scheduler: { 8 | type: StepLR, 9 | kwargs: { 10 | step_size: 50, 11 | gamma : 0.5 12 | }} 13 | 14 | dataset : { 15 | train : { _base_: cfgs/dataset_configs/PCN.yaml, 16 | others: {subset: 'train'}}, 17 | val : { _base_: cfgs/dataset_configs/PCN.yaml, 18 | others: {subset: 'test'}}, 19 | test : { _base_: cfgs/dataset_configs/PCN.yaml, 20 | others: {subset: 'test'}}} 21 | 22 | model : { 23 | NAME: PCN, num_pred: 16384, encoder_channel: 1024} 24 | total_bs : 32 25 | step_per_update : 1 26 | max_epoch : 150 27 | 28 | consider_metric: CDL1 -------------------------------------------------------------------------------- /cfgs/PCN_models/PT_model.yaml: -------------------------------------------------------------------------------- 1 | optimizer : { 2 | type: AdamW, 3 | kwargs: { 4 | lr : 0.0005, 5 | weight_decay : 0.0005 6 | }} 7 | 8 | scheduler: { 9 | type: LambdaLR, 10 | kwargs: { 11 | decay_step: 21, 12 | lr_decay: 0.9, 13 | lowest_decay: 0.02 # min lr = lowest_decay * lr 14 | }} 15 | 16 | bnmscheduler: { 17 | type: Lambda, 18 | kwargs: { 19 | decay_step: 21, 20 | bn_decay: 0.5, 21 | bn_momentum: 0.9, 22 | lowest_decay: 0.01 23 | }} 24 | 25 | dataset : { 26 | train : { _base_: cfgs/dataset_configs/PCN.yaml, 27 | others: {subset: 'train'}}, 28 | val : { _base_: cfgs/dataset_configs/PCN.yaml, 29 | others: {subset: 'test'}}, 30 | test : { _base_: cfgs/dataset_configs/PCN.yaml, 31 | others: {subset: 'test'}}} 32 | model : { 33 | NAME: PT_model, num_pred: 14336, num_query: 224, knn_layer: 1, trans_dim: 256} 34 | 35 | total_bs : 64 36 | step_per_update : 1 37 | max_epoch : 300 38 | 39 | consider_metric: CDL1 -------------------------------------------------------------------------------- /cfgs/PCN_models/PoinTr.yaml: -------------------------------------------------------------------------------- 1 | optimizer : { 2 | type: AdamW, 3 | kwargs: { 4 | lr : 0.0005, 5 | weight_decay : 0.0005 6 | }} 7 | 8 | scheduler: { 9 | type: LambdaLR, 10 | kwargs: { 11 | decay_step: 21, 12 | lr_decay: 0.9, 13 | lowest_decay: 0.02 # min lr = lowest_decay * lr 14 | }} 15 | 16 | bnmscheduler: { 17 | type: Lambda, 18 | kwargs: { 19 | decay_step: 21, 20 | bn_decay: 0.5, 21 | bn_momentum: 0.9, 22 | lowest_decay: 0.01 23 | }} 24 | 25 | dataset : { 26 | train : { _base_: cfgs/dataset_configs/PCN.yaml, 27 | others: {subset: 'train'}}, 28 | val : { _base_: cfgs/dataset_configs/PCN.yaml, 29 | others: {subset: 'test'}}, 30 | test : { _base_: cfgs/dataset_configs/PCN.yaml, 31 | others: {subset: 'test'}}} 32 | model : { 33 | NAME: PoinTr, num_pred: 14336, num_query: 224, knn_layer: 1, trans_dim: 384} 34 | 35 | total_bs : 48 36 | step_per_update : 1 37 | max_epoch : 300 38 | 39 | consider_metric: CDL1 -------------------------------------------------------------------------------- /cfgs/PCN_models/PoinTr_Encoder.yaml: -------------------------------------------------------------------------------- 1 | optimizer : { 2 | type: AdamW, 3 | kwargs: { 4 | lr : 0.0005, 5 | weight_decay : 0.0005 6 | }} 7 | 8 | scheduler: { 9 | type: LambdaLR, 10 | kwargs: { 11 | decay_step: 21, 12 | lr_decay: 0.9, 13 | lowest_decay: 0.02 # min lr = lowest_decay * lr 14 | }} 15 | 16 | bnmscheduler: { 17 | type: Lambda, 18 | kwargs: { 19 | decay_step: 21, 20 | bn_decay: 0.5, 21 | bn_momentum: 0.9, 22 | lowest_decay: 0.01 23 | }} 24 | 25 | dataset : { 26 | train : { _base_: cfgs/dataset_configs/PCN.yaml, 27 | others: {subset: 'train'}}, 28 | val : { _base_: cfgs/dataset_configs/PCN.yaml, 29 | others: {subset: 'test'}}, 30 | test : { _base_: cfgs/dataset_configs/PCN.yaml, 31 | others: {subset: 'test'}}} 32 | model : { 33 | NAME: PoinTr_Encoder, num_pred: 14336, num_query: 224, knn_layer: 1, trans_dim: 384} 34 | 35 | total_bs : 48 36 | step_per_update : 1 37 | max_epoch : 300 38 | 39 | consider_metric: CDL1 -------------------------------------------------------------------------------- /cfgs/PCN_models/TopNet.yaml: -------------------------------------------------------------------------------- 1 | optimizer : { 2 | type: Adam, 3 | kwargs: { 4 | lr : 0.001, 5 | weight_decay : 0 6 | }} 7 | scheduler: { 8 | type: StepLR, 9 | kwargs: { 10 | step_size: 50, 11 | gamma : 0.5 12 | }} 13 | 14 | dataset : { 15 | train : { _base_: cfgs/dataset_configs/PCN.yaml, 16 | others: {subset: 'train'}}, 17 | val : { _base_: cfgs/dataset_configs/PCN.yaml, 18 | others: {subset: 'test'}}, 19 | test : { _base_: cfgs/dataset_configs/PCN.yaml, 20 | others: {subset: 'test'}}} 21 | 22 | model : { 23 | NAME: TopNet, node_feature: 8, num_pred: 16384, encoder_channel: 1024, nlevels: 8} 24 | total_bs : 32 25 | step_per_update : 1 26 | max_epoch : 150 27 | 28 | consider_metric: CDL1 -------------------------------------------------------------------------------- /cfgs/PCN_models/our_model.yaml: -------------------------------------------------------------------------------- 1 | optimizer : { 2 | type: AdamW, 3 | kwargs: { 4 | lr : 0.0005, 5 | weight_decay : 0.0005 6 | }} 7 | 8 | scheduler: { 9 | type: LambdaLR, 10 | kwargs: { 11 | decay_step: 21, 12 | lr_decay: 0.9, 13 | lowest_decay: 0.02 # min lr = lowest_decay * lr 14 | }} 15 | 16 | bnmscheduler: { 17 | type: Lambda, 18 | kwargs: { 19 | decay_step: 21, 20 | bn_decay: 0.5, 21 | bn_momentum: 0.9, 22 | lowest_decay: 0.01 23 | }} 24 | 25 | dataset : { 26 | train : { _base_: cfgs/dataset_configs/PCN.yaml, 27 | others: {subset: 'train'}}, 28 | val : { _base_: cfgs/dataset_configs/PCN.yaml, 29 | others: {subset: 'test'}}, 30 | test : { _base_: cfgs/dataset_configs/PCN.yaml, 31 | others: {subset: 'test'}}} 32 | model : { 33 | NAME: our_model, num_pred: 14336, num_query: 224, knn_layer: 1, trans_dim: 384} 34 | 35 | total_bs : 1 36 | step_per_update : 1 37 | max_epoch : 300 38 | 39 | consider_metric: CDL1 -------------------------------------------------------------------------------- /cfgs/ShapeNet34_models/FoldingNet.yaml: -------------------------------------------------------------------------------- 1 | optimizer : { 2 | type: Adam, 3 | kwargs: { 4 | lr : 0.0001, 5 | weight_decay : 0 6 | }} 7 | scheduler: { 8 | type: StepLR, 9 | kwargs: { 10 | step_size: 50, 11 | gamma : 0.5 12 | }} 13 | 14 | dataset : { 15 | train : { _base_: cfgs/dataset_configs/ShapeNet-34.yaml, 16 | others: {subset: 'train'}}, 17 | val : { _base_: cfgs/dataset_configs/ShapeNet-34.yaml, 18 | others: {subset: 'test'}}, 19 | test : { _base_: cfgs/dataset_configs/ShapeNet-34.yaml, 20 | others: {subset: 'test'}}} 21 | 22 | model : { 23 | NAME: FoldingNet, num_pred: 8192, encoder_channel: 1024} 24 | total_bs : 32 25 | step_per_update : 1 26 | max_epoch : 150 27 | 28 | consider_metric: CDL1 -------------------------------------------------------------------------------- /cfgs/ShapeNet34_models/GRNet.yaml: -------------------------------------------------------------------------------- 1 | optimizer : { 2 | type: Adam, 3 | kwargs: { 4 | lr : 0.0001, 5 | weight_decay : 0 6 | }} 7 | scheduler: { 8 | type: StepLR, 9 | kwargs: { 10 | step_size: 50, 11 | gamma : 0.5 12 | }} 13 | 14 | dataset : { 15 | train : { _base_: cfgs/dataset_configs/ShapeNet-34.yaml, 16 | others: {subset: 'train'}}, 17 | val : { _base_: cfgs/dataset_configs/ShapeNet-34.yaml, 18 | others: {subset: 'test'}}, 19 | test : { _base_: cfgs/dataset_configs/ShapeNet-34.yaml, 20 | others: {subset: 'test'}}} 21 | model : { 22 | NAME: GRNet, num_pred: 8192, gridding_loss_scales: 128, gridding_loss_alphas: 0.1} 23 | total_bs : 32 24 | step_per_update : 1 25 | max_epoch : 150 26 | 27 | consider_metric: CDL2 -------------------------------------------------------------------------------- /cfgs/ShapeNet34_models/PCN.yaml: -------------------------------------------------------------------------------- 1 | optimizer : { 2 | type: Adam, 3 | kwargs: { 4 | lr : 0.0001, 5 | weight_decay : 0 6 | }} 7 | scheduler: { 8 | type: StepLR, 9 | kwargs: { 10 | step_size: 50, 11 | gamma : 0.5 12 | }} 13 | 14 | dataset : { 15 | train : { _base_: cfgs/dataset_configs/ShapeNet-34.yaml, 16 | others: {subset: 'train'}}, 17 | val : { _base_: cfgs/dataset_configs/ShapeNet-34.yaml, 18 | others: {subset: 'test'}}, 19 | test : { _base_: cfgs/dataset_configs/ShapeNet-34.yaml, 20 | others: {subset: 'test'}}} 21 | model : { 22 | NAME: PCN, num_pred: 8192, encoder_channel: 1024} 23 | total_bs : 32 24 | step_per_update : 1 25 | max_epoch : 150 26 | 27 | consider_metric: CDL2 -------------------------------------------------------------------------------- /cfgs/ShapeNet34_models/PoinTr.yaml: -------------------------------------------------------------------------------- 1 | optimizer : { 2 | type: AdamW, 3 | kwargs: { 4 | lr : 0.0005, 5 | weight_decay : 0.0005 6 | }} 7 | scheduler: { 8 | type: LambdaLR, 9 | kwargs: { 10 | decay_step: 21, 11 | lr_decay: 0.9, 12 | lowest_decay: 0.02 # min lr = lowest_decay * lr 13 | }} 14 | bnmscheduler: { 15 | type: Lambda, 16 | kwargs: { 17 | decay_step: 21, 18 | bn_decay: 0.5, 19 | bn_momentum: 0.9, 20 | lowest_decay: 0.01 21 | }} 22 | 23 | dataset : { 24 | train : { _base_: cfgs/dataset_configs/ShapeNet-34.yaml, 25 | others: {subset: 'train'}}, 26 | val : { _base_: cfgs/dataset_configs/ShapeNet-34.yaml, 27 | others: {subset: 'test'}}, 28 | test : { _base_: cfgs/dataset_configs/ShapeNet-34.yaml, 29 | others: {subset: 'test'}}} 30 | model : { 31 | NAME: PoinTr, num_pred: 6144, num_query: 96, knn_layer: 1, trans_dim: 384} 32 | total_bs : 128 33 | step_per_update : 1 34 | max_epoch : 200 35 | 36 | consider_metric: CDL2 -------------------------------------------------------------------------------- /cfgs/ShapeNet34_models/TopNet.yaml: -------------------------------------------------------------------------------- 1 | optimizer : { 2 | type: Adam, 3 | kwargs: { 4 | lr : 0.001, 5 | weight_decay : 0 6 | }} 7 | scheduler: { 8 | type: StepLR, 9 | kwargs: { 10 | step_size: 50, 11 | gamma : 0.5 12 | }} 13 | 14 | dataset : { 15 | train : { _base_: cfgs/dataset_configs/ShapeNet-34.yaml, 16 | others: {subset: 'train'}}, 17 | val : { _base_: cfgs/dataset_configs/ShapeNet-34.yaml, 18 | others: {subset: 'test'}}, 19 | test : { _base_: cfgs/dataset_configs/ShapeNet-34.yaml, 20 | others: {subset: 'test'}}} 21 | 22 | model : { 23 | NAME: TopNet, node_feature: 8, num_pred: 8192, encoder_channel: 1024, nlevels: 8} 24 | total_bs : 32 25 | step_per_update : 1 26 | max_epoch : 150 27 | 28 | consider_metric: CDL1 -------------------------------------------------------------------------------- /cfgs/ShapeNet34_models/our_model.yaml: -------------------------------------------------------------------------------- 1 | optimizer : { 2 | type: AdamW, 3 | kwargs: { 4 | lr : 0.0005, 5 | weight_decay : 0.0005 6 | }} 7 | scheduler: { 8 | type: LambdaLR, 9 | kwargs: { 10 | decay_step: 21, 11 | lr_decay: 0.9, 12 | lowest_decay: 0.02 # min lr = lowest_decay * lr 13 | }} 14 | bnmscheduler: { 15 | type: Lambda, 16 | kwargs: { 17 | decay_step: 21, 18 | bn_decay: 0.5, 19 | bn_momentum: 0.9, 20 | lowest_decay: 0.01 21 | }} 22 | 23 | dataset : { 24 | train : { _base_: cfgs/dataset_configs/ShapeNet-34.yaml, 25 | others: {subset: 'train'}}, 26 | val : { _base_: cfgs/dataset_configs/ShapeNet-34.yaml, 27 | others: {subset: 'test'}}, 28 | test : { _base_: cfgs/dataset_configs/ShapeNet-34.yaml, 29 | others: {subset: 'test'}}} 30 | model : { 31 | NAME: our_model, num_pred: 6144, num_query: 96, knn_layer: 1, trans_dim: 384} 32 | total_bs : 128 33 | step_per_update : 1 34 | max_epoch : 200 35 | 36 | consider_metric: CDL2 -------------------------------------------------------------------------------- /cfgs/ShapeNet55_models/FoldingNet.yaml: -------------------------------------------------------------------------------- 1 | optimizer : { 2 | type: Adam, 3 | kwargs: { 4 | lr : 0.0001, 5 | weight_decay : 0 6 | }} 7 | scheduler: { 8 | type: StepLR, 9 | kwargs: { 10 | step_size: 50, 11 | gamma : 0.5 12 | }} 13 | 14 | dataset : { 15 | train : { _base_: cfgs/dataset_configs/ShapeNet-55.yaml, 16 | others: {subset: 'train'}}, 17 | val : { _base_: cfgs/dataset_configs/ShapeNet-55.yaml, 18 | others: {subset: 'test'}}, 19 | test : { _base_: cfgs/dataset_configs/ShapeNet-55.yaml, 20 | others: {subset: 'test'}}} 21 | 22 | model : { 23 | NAME: FoldingNet, num_pred: 8192, encoder_channel: 1024} 24 | total_bs : 32 25 | step_per_update : 1 26 | max_epoch : 150 27 | 28 | consider_metric: CDL1 -------------------------------------------------------------------------------- /cfgs/ShapeNet55_models/GRNet.yaml: -------------------------------------------------------------------------------- 1 | optimizer : { 2 | type: Adam, 3 | kwargs: { 4 | lr : 0.0001, 5 | weight_decay : 0 6 | }} 7 | scheduler: { 8 | type: StepLR, 9 | kwargs: { 10 | step_size: 50, 11 | gamma : 0.5 12 | }} 13 | 14 | dataset : { 15 | train : { _base_: cfgs/dataset_configs/ShapeNet-55.yaml, 16 | others: {subset: 'train'}}, 17 | val : { _base_: cfgs/dataset_configs/ShapeNet-55.yaml, 18 | others: {subset: 'test'}}, 19 | test : { _base_: cfgs/dataset_configs/ShapeNet-55.yaml, 20 | others: {subset: 'test'}}} 21 | model : { 22 | NAME: GRNet, num_pred: 8192, gridding_loss_scales: 128, gridding_loss_alphas: 0.1} 23 | total_bs : 32 24 | step_per_update : 1 25 | max_epoch : 150 26 | 27 | consider_metric: CDL2 -------------------------------------------------------------------------------- /cfgs/ShapeNet55_models/PCN.yaml: -------------------------------------------------------------------------------- 1 | optimizer : { 2 | type: Adam, 3 | kwargs: { 4 | lr : 0.0001, 5 | weight_decay : 0 6 | }} 7 | scheduler: { 8 | type: StepLR, 9 | kwargs: { 10 | step_size: 50, 11 | gamma : 0.5 12 | }} 13 | 14 | dataset : { 15 | train : { _base_: cfgs/dataset_configs/ShapeNet-55.yaml, 16 | others: {subset: 'train'}}, 17 | val : { _base_: cfgs/dataset_configs/ShapeNet-55.yaml, 18 | others: {subset: 'test'}}, 19 | test : { _base_: cfgs/dataset_configs/ShapeNet-55.yaml, 20 | others: {subset: 'test'}}} 21 | model : { 22 | NAME: PCN, num_pred: 8192, encoder_channel: 1024} 23 | total_bs : 32 24 | step_per_update : 1 25 | max_epoch : 150 26 | 27 | consider_metric: CDL2 -------------------------------------------------------------------------------- /cfgs/ShapeNet55_models/PT_model.yaml: -------------------------------------------------------------------------------- 1 | optimizer : { 2 | type: AdamW, 3 | kwargs: { 4 | lr : 0.0005, 5 | weight_decay : 0.0005 6 | }} 7 | scheduler: { 8 | type: LambdaLR, 9 | kwargs: { 10 | decay_step: 21, 11 | lr_decay: 0.76, 12 | lowest_decay: 0.02 # min lr = lowest_decay * lr 13 | }} 14 | bnmscheduler: { 15 | type: Lambda, 16 | kwargs: { 17 | decay_step: 21, 18 | bn_decay: 0.5, 19 | bn_momentum: 0.9, 20 | lowest_decay: 0.01 21 | }} 22 | 23 | dataset : { 24 | train : { _base_: cfgs/dataset_configs/ShapeNet-55.yaml, 25 | others: {subset: 'train'}}, 26 | val : { _base_: cfgs/dataset_configs/ShapeNet-55.yaml, 27 | others: {subset: 'test'}}, 28 | test : { _base_: cfgs/dataset_configs/ShapeNet-55.yaml, 29 | others: {subset: 'test'}}} 30 | model : { 31 | NAME: PT_model, num_pred: 6144, num_query: 96, knn_layer: 1, trans_dim: 512} 32 | total_bs : 64 33 | step_per_update : 1 34 | max_epoch : 200 35 | 36 | consider_metric: CDL2 -------------------------------------------------------------------------------- /cfgs/ShapeNet55_models/PoinTr.yaml: -------------------------------------------------------------------------------- 1 | optimizer : { 2 | type: AdamW, 3 | kwargs: { 4 | lr : 0.0005, 5 | weight_decay : 0.0005 6 | }} 7 | scheduler: { 8 | type: LambdaLR, 9 | kwargs: { 10 | decay_step: 21, 11 | lr_decay: 0.76, 12 | lowest_decay: 0.02 # min lr = lowest_decay * lr 13 | }} 14 | bnmscheduler: { 15 | type: Lambda, 16 | kwargs: { 17 | decay_step: 21, 18 | bn_decay: 0.5, 19 | bn_momentum: 0.9, 20 | lowest_decay: 0.01 21 | }} 22 | 23 | dataset : { 24 | train : { _base_: cfgs/dataset_configs/ShapeNet-55.yaml, 25 | others: {subset: 'train'}}, 26 | val : { _base_: cfgs/dataset_configs/ShapeNet-55.yaml, 27 | others: {subset: 'test'}}, 28 | test : { _base_: cfgs/dataset_configs/ShapeNet-55.yaml, 29 | others: {subset: 'test'}}} 30 | model : { 31 | NAME: PoinTr, num_pred: 6144, num_query: 96, knn_layer: 1, trans_dim: 384} 32 | total_bs : 96 33 | step_per_update : 1 34 | max_epoch : 200 35 | 36 | consider_metric: CDL2 -------------------------------------------------------------------------------- /cfgs/ShapeNet55_models/TopNet.yaml: -------------------------------------------------------------------------------- 1 | optimizer : { 2 | type: Adam, 3 | kwargs: { 4 | lr : 0.001, 5 | weight_decay : 0 6 | }} 7 | scheduler: { 8 | type: StepLR, 9 | kwargs: { 10 | step_size: 50, 11 | gamma : 0.5 12 | }} 13 | 14 | dataset : { 15 | train : { _base_: cfgs/dataset_configs/ShapeNet-55.yaml, 16 | others: {subset: 'train'}}, 17 | val : { _base_: cfgs/dataset_configs/ShapeNet-55.yaml, 18 | others: {subset: 'test'}}, 19 | test : { _base_: cfgs/dataset_configs/ShapeNet-55.yaml, 20 | others: {subset: 'test'}}} 21 | 22 | model : { 23 | NAME: TopNet, node_feature: 8, num_pred: 8192, encoder_channel: 1024, nlevels: 8} 24 | total_bs : 32 25 | step_per_update : 1 26 | max_epoch : 150 27 | 28 | consider_metric: CDL1 -------------------------------------------------------------------------------- /cfgs/ShapeNet55_models/our_model.yaml: -------------------------------------------------------------------------------- 1 | optimizer : { 2 | type: AdamW, 3 | kwargs: { 4 | lr : 0.0005, 5 | weight_decay : 0.0005 6 | }} 7 | scheduler: { 8 | type: LambdaLR, 9 | kwargs: { 10 | decay_step: 21, 11 | lr_decay: 0.76, 12 | lowest_decay: 0.02 # min lr = lowest_decay * lr 13 | }} 14 | bnmscheduler: { 15 | type: Lambda, 16 | kwargs: { 17 | decay_step: 21, 18 | bn_decay: 0.5, 19 | bn_momentum: 0.9, 20 | lowest_decay: 0.01 21 | }} 22 | 23 | dataset : { 24 | train : { _base_: cfgs/dataset_configs/ShapeNet-55.yaml, 25 | others: {subset: 'train'}}, 26 | val : { _base_: cfgs/dataset_configs/ShapeNet-55.yaml, 27 | others: {subset: 'test'}}, 28 | test : { _base_: cfgs/dataset_configs/ShapeNet-55.yaml, 29 | others: {subset: 'test'}}} 30 | model : { 31 | NAME: our_model, num_pred: 6144, num_query: 96, knn_layer: 1, trans_dim: 384} 32 | total_bs : 96 33 | step_per_update : 1 34 | max_epoch : 200 35 | 36 | consider_metric: CDL2 -------------------------------------------------------------------------------- /cfgs/ShapeNetUnseen21_models/FoldingNet.yaml: -------------------------------------------------------------------------------- 1 | optimizer : { 2 | type: Adam, 3 | kwargs: { 4 | lr : 0.0001, 5 | weight_decay : 0 6 | }} 7 | scheduler: { 8 | type: StepLR, 9 | kwargs: { 10 | step_size: 50, 11 | gamma : 0.5 12 | }} 13 | 14 | dataset : { 15 | train : { _base_: cfgs/dataset_configs/ShapeNet-34.yaml, 16 | others: {subset: 'train'}}, 17 | val : { _base_: cfgs/dataset_configs/ShapeNet-34.yaml, 18 | others: {subset: 'test'}}, 19 | test : { _base_: cfgs/dataset_configs/ShapeNet-Unseen21.yaml, 20 | others: {subset: 'test'}}} 21 | 22 | model : { 23 | NAME: FoldingNet, num_pred: 8192, encoder_channel: 1024} 24 | total_bs : 32 25 | step_per_update : 1 26 | max_epoch : 150 27 | 28 | consider_metric: CDL1 -------------------------------------------------------------------------------- /cfgs/ShapeNetUnseen21_models/GRNet.yaml: -------------------------------------------------------------------------------- 1 | optimizer : { 2 | type: Adam, 3 | kwargs: { 4 | lr : 0.0001, 5 | weight_decay : 0 6 | }} 7 | scheduler: { 8 | type: StepLR, 9 | kwargs: { 10 | step_size: 50, 11 | gamma : 0.5 12 | }} 13 | 14 | dataset : { 15 | train : { _base_: cfgs/dataset_configs/ShapeNet-34.yaml, 16 | others: {subset: 'train'}}, 17 | val : { _base_: cfgs/dataset_configs/ShapeNet-34.yaml, 18 | others: {subset: 'test'}}, 19 | test : { _base_: cfgs/dataset_configs/ShapeNet-Unseen21.yaml, 20 | others: {subset: 'test'}}} 21 | model : { 22 | NAME: GRNet, num_pred: 8192, gridding_loss_scales: 128, gridding_loss_alphas: 0.1} 23 | total_bs : 32 24 | step_per_update : 1 25 | max_epoch : 150 26 | 27 | consider_metric: CDL2 -------------------------------------------------------------------------------- /cfgs/ShapeNetUnseen21_models/PCN.yaml: -------------------------------------------------------------------------------- 1 | optimizer : { 2 | type: Adam, 3 | kwargs: { 4 | lr : 0.0001, 5 | weight_decay : 0 6 | }} 7 | scheduler: { 8 | type: StepLR, 9 | kwargs: { 10 | step_size: 50, 11 | gamma : 0.5 12 | }} 13 | 14 | dataset : { 15 | train : { _base_: cfgs/dataset_configs/ShapeNet-34.yaml, 16 | others: {subset: 'train'}}, 17 | val : { _base_: cfgs/dataset_configs/ShapeNet-34.yaml, 18 | others: {subset: 'test'}}, 19 | test : { _base_: cfgs/dataset_configs/ShapeNet-Unseen21.yaml, 20 | others: {subset: 'test'}}} 21 | model : { 22 | NAME: PCN, num_pred: 8192, encoder_channel: 1024} 23 | total_bs : 32 24 | step_per_update : 1 25 | max_epoch : 150 26 | 27 | consider_metric: CDL2 -------------------------------------------------------------------------------- /cfgs/ShapeNetUnseen21_models/PoinTr.yaml: -------------------------------------------------------------------------------- 1 | optimizer : { 2 | type: AdamW, 3 | kwargs: { 4 | lr : 0.0005, 5 | weight_decay : 0.0005 6 | }} 7 | scheduler: { 8 | type: LambdaLR, 9 | kwargs: { 10 | decay_step: 21, 11 | lr_decay: 0.9, 12 | lowest_decay: 0.02 # min lr = lowest_decay * lr 13 | }} 14 | bnmscheduler: { 15 | type: Lambda, 16 | kwargs: { 17 | decay_step: 21, 18 | bn_decay: 0.5, 19 | bn_momentum: 0.9, 20 | lowest_decay: 0.01 21 | }} 22 | 23 | dataset : { 24 | train : { _base_: cfgs/dataset_configs/ShapeNet-34.yaml, 25 | others: {subset: 'train'}}, 26 | val : { _base_: cfgs/dataset_configs/ShapeNet-34.yaml, 27 | others: {subset: 'test'}}, 28 | test : { _base_: cfgs/dataset_configs/ShapeNet-Unseen21.yaml, 29 | others: {subset: 'test'}}} 30 | model : { 31 | NAME: PoinTr, num_pred: 6144, num_query: 96, knn_layer: 1, trans_dim: 384} 32 | total_bs : 128 33 | step_per_update : 1 34 | max_epoch : 200 35 | 36 | consider_metric: CDL2 -------------------------------------------------------------------------------- /cfgs/ShapeNetUnseen21_models/TopNet.yaml: -------------------------------------------------------------------------------- 1 | optimizer : { 2 | type: Adam, 3 | kwargs: { 4 | lr : 0.001, 5 | weight_decay : 0 6 | }} 7 | scheduler: { 8 | type: StepLR, 9 | kwargs: { 10 | step_size: 50, 11 | gamma : 0.5 12 | }} 13 | 14 | dataset : { 15 | train : { _base_: cfgs/dataset_configs/ShapeNet-34.yaml, 16 | others: {subset: 'train'}}, 17 | val : { _base_: cfgs/dataset_configs/ShapeNet-34.yaml, 18 | others: {subset: 'test'}}, 19 | test : { _base_: cfgs/dataset_configs/ShapeNet-Unseen21.yaml, 20 | others: {subset: 'test'}}} 21 | 22 | model : { 23 | NAME: TopNet, node_feature: 8, num_pred: 8192, encoder_channel: 1024, nlevels: 8} 24 | total_bs : 32 25 | step_per_update : 1 26 | max_epoch : 150 27 | 28 | consider_metric: CDL1 -------------------------------------------------------------------------------- /cfgs/ShapeNetUnseen21_models/our_model.yaml: -------------------------------------------------------------------------------- 1 | optimizer : { 2 | type: AdamW, 3 | kwargs: { 4 | lr : 0.0005, 5 | weight_decay : 0.0005 6 | }} 7 | scheduler: { 8 | type: LambdaLR, 9 | kwargs: { 10 | decay_step: 21, 11 | lr_decay: 0.9, 12 | lowest_decay: 0.02 # min lr = lowest_decay * lr 13 | }} 14 | bnmscheduler: { 15 | type: Lambda, 16 | kwargs: { 17 | decay_step: 21, 18 | bn_decay: 0.5, 19 | bn_momentum: 0.9, 20 | lowest_decay: 0.01 21 | }} 22 | 23 | dataset : { 24 | train : { _base_: cfgs/dataset_configs/ShapeNet-34.yaml, 25 | others: {subset: 'train'}}, 26 | val : { _base_: cfgs/dataset_configs/ShapeNet-34.yaml, 27 | others: {subset: 'test'}}, 28 | test : { _base_: cfgs/dataset_configs/ShapeNet-Unseen21.yaml, 29 | others: {subset: 'test'}}} 30 | model : { 31 | NAME: our_model, num_pred: 6144, num_query: 96, knn_layer: 1, trans_dim: 384} 32 | total_bs : 128 33 | step_per_update : 1 34 | max_epoch : 200 35 | 36 | consider_metric: CDL2 -------------------------------------------------------------------------------- /cfgs/dataset_configs/KITTI.yaml: -------------------------------------------------------------------------------- 1 | NAME: KITTI 2 | CATEGORY_FILE_PATH: data/KITTI/KITTI.json 3 | N_POINTS: 16384 4 | CLOUD_PATH: /home/lss/Project/RL-GAN-Net/datasets/kitti/cars/%s.pcd 5 | BBOX_PATH: /home/lss/Project/RL-GAN-Net/datasets/kitti/bboxes/%s.txt -------------------------------------------------------------------------------- /cfgs/dataset_configs/PCN.yaml: -------------------------------------------------------------------------------- 1 | NAME: PCN 2 | CATEGORY_FILE_PATH: data/PCN/PCN.json 3 | N_POINTS: 16384 4 | N_RENDERINGS: 8 5 | PARTIAL_POINTS_PATH: data/PCN/%s/partial/%s/%s/%02d.pcd 6 | COMPLETE_POINTS_PATH: data/PCN/%s/complete/%s/%s.pcd 7 | CARS: FALSE -------------------------------------------------------------------------------- /cfgs/dataset_configs/PCNCars.yaml: -------------------------------------------------------------------------------- 1 | NAME: PCN 2 | CATEGORY_FILE_PATH: data/PCN/PCN.json 3 | N_POINTS: 16384 4 | N_RENDERINGS: 8 5 | PARTIAL_POINTS_PATH: data/PCN/%s/partial/%s/%s/%02d.pcd 6 | COMPLETE_POINTS_PATH: data/PCN/%s/complete/%s/%s.pcd 7 | CARS: TRUE -------------------------------------------------------------------------------- /cfgs/dataset_configs/ShapeNet-34.yaml: -------------------------------------------------------------------------------- 1 | NAME: ShapeNet 2 | DATA_PATH: data/ShapeNet55-34/ShapeNet-34 3 | N_POINTS: 8192 4 | PC_PATH: data/ShapeNet55-34/shapenet_pc -------------------------------------------------------------------------------- /cfgs/dataset_configs/ShapeNet-55.yaml: -------------------------------------------------------------------------------- 1 | NAME: ShapeNet 2 | DATA_PATH: data/ShapeNet55-34/ShapeNet-55 3 | N_POINTS: 8192 4 | PC_PATH: data/ShapeNet55-34/shapenet_pc 5 | -------------------------------------------------------------------------------- /cfgs/dataset_configs/ShapeNet-Unseen21.yaml: -------------------------------------------------------------------------------- 1 | NAME: ShapeNet 2 | DATA_PATH: data/ShapeNet55-34/ShapeNet-Unseen21 3 | N_POINTS: 8192 4 | PC_PATH: data/ShapeNet55-34/shapenet_pc -------------------------------------------------------------------------------- /datasets/KITTIDataset.py: -------------------------------------------------------------------------------- 1 | import torch.utils.data as data 2 | import numpy as np 3 | import os, sys 4 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 5 | sys.path.append(BASE_DIR) 6 | import data_transforms 7 | from .io import IO 8 | import json 9 | from .build import DATASETS 10 | 11 | # References: 12 | # - https://github.com/hzxie/GRNet/blob/master/utils/data_loaders.py 13 | @DATASETS.register_module() 14 | class KITTI(data.Dataset): 15 | def __init__(self, config): 16 | self.cloud_path = config.CLOUD_PATH 17 | self.bbox_path = config.BBOX_PATH 18 | self.category_file = config.CATEGORY_FILE_PATH 19 | self.npoints = config.N_POINTS 20 | self.subset = config.subset 21 | assert self.subset == 'test' 22 | 23 | self.dataset_categories = [] 24 | with open(self.category_file) as f: 25 | self.dataset_categories = json.loads(f.read()) 26 | self.transforms = data_transforms.Compose([{ 27 | 'callback': 'NormalizeObjectPose', 28 | 'parameters': { 29 | 'input_keys': { 30 | 'ptcloud': 'partial_cloud', 31 | 'bbox': 'bounding_box' 32 | } 33 | }, 34 | 'objects': ['partial_cloud', 'bounding_box'] 35 | }, { 36 | 'callback': 'RandomSamplePoints', 37 | 'parameters': { 38 | 'n_points': 2048 39 | }, 40 | 'objects': ['partial_cloud'] 41 | }, { 42 | 'callback': 'ToTensor', 43 | 'objects': ['partial_cloud', 'bounding_box'] 44 | }]) 45 | self.file_list = self._get_file_list(self.subset) 46 | 47 | def _get_file_list(self, subset): 48 | """Prepare file list for the dataset""" 49 | file_list = [] 50 | for dc in self.dataset_categories: 51 | samples = dc[subset] 52 | for s in samples: 53 | file_list.append({ 54 | 'taxonomy_id': dc['taxonomy_id'], 55 | 'model_id': s, 56 | 'partial_cloud_path': self.cloud_path % s, 57 | 'bounding_box_path': self.bbox_path % s, 58 | }) 59 | return file_list 60 | 61 | def __len__(self): 62 | return len(self.file_list) 63 | 64 | def __getitem__(self, idx): 65 | sample = self.file_list[idx] 66 | data = {} 67 | 68 | for ri in ['partial_cloud', 'bounding_box']: 69 | file_path = sample['%s_path' % ri] 70 | data[ri] = IO.get(file_path).astype(np.float32) 71 | 72 | if self.transforms is not None: 73 | data = self.transforms(data) 74 | 75 | return sample['taxonomy_id'], sample['model_id'], data['partial_cloud'] 76 | -------------------------------------------------------------------------------- /datasets/ShapeNet55Dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | import numpy as np 4 | import torch.utils.data as data 5 | from .io import IO 6 | from .build import DATASETS 7 | import logging 8 | 9 | @DATASETS.register_module() 10 | class ShapeNet(data.Dataset): 11 | def __init__(self, config): 12 | self.data_root = config.DATA_PATH 13 | self.pc_path = config.PC_PATH 14 | self.subset = config.subset 15 | self.npoints = config.N_POINTS 16 | self.data_list_file = os.path.join(self.data_root, f'{self.subset}.txt') 17 | 18 | print(f'[DATASET] Open file {self.data_list_file}') 19 | with open(self.data_list_file, 'r') as f: 20 | lines = f.readlines() 21 | 22 | self.file_list = [] 23 | for line in lines: 24 | line = line.strip() 25 | taxonomy_id = line.split('-')[0] 26 | model_id = line.split('-')[1].split('.')[0] 27 | self.file_list.append({ 28 | 'taxonomy_id': taxonomy_id, 29 | 'model_id': model_id, 30 | 'file_path': line 31 | }) 32 | print(f'[DATASET] {len(self.file_list)} instances were loaded') 33 | def pc_norm(self, pc): 34 | """ pc: NxC, return NxC """ 35 | centroid = np.mean(pc, axis=0) 36 | pc = pc - centroid 37 | m = np.max(np.sqrt(np.sum(pc**2, axis=1))) 38 | pc = pc / m 39 | return pc 40 | 41 | def __getitem__(self, idx): 42 | sample = self.file_list[idx] 43 | 44 | data = IO.get(os.path.join(self.pc_path, sample['file_path'])).astype(np.float32) 45 | data = self.pc_norm(data) 46 | data = torch.from_numpy(data).float() 47 | 48 | return sample['taxonomy_id'], sample['model_id'], data 49 | 50 | def __len__(self): 51 | return len(self.file_list) -------------------------------------------------------------------------------- /datasets/build.py: -------------------------------------------------------------------------------- 1 | from utils import registry 2 | 3 | 4 | DATASETS = registry.Registry('dataset') 5 | 6 | 7 | def build_dataset_from_cfg(cfg, default_args = None): 8 | """ 9 | Build a dataset, defined by `dataset_name`. 10 | Args: 11 | cfg (eDICT): 12 | Returns: 13 | Dataset: a constructed dataset specified by dataset_name. 14 | """ 15 | return DATASETS.build(cfg, default_args = default_args) 16 | 17 | 18 | -------------------------------------------------------------------------------- /datasets/data_transforms.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Haozhe Xie 3 | # @Date: 2019-08-02 14:38:36 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2020-07-03 09:23:07 6 | # @Email: cshzxie@gmail.com 7 | 8 | import numpy as np 9 | import torch 10 | import transforms3d 11 | 12 | class Compose(object): 13 | def __init__(self, transforms): 14 | self.transformers = [] 15 | for tr in transforms: 16 | transformer = eval(tr['callback']) 17 | parameters = tr['parameters'] if 'parameters' in tr else None 18 | self.transformers.append({ 19 | 'callback': transformer(parameters), 20 | 'objects': tr['objects'] 21 | }) # yapf: disable 22 | 23 | def __call__(self, data): 24 | for tr in self.transformers: 25 | transform = tr['callback'] 26 | objects = tr['objects'] 27 | rnd_value = np.random.uniform(0, 1) 28 | if transform.__class__ in [NormalizeObjectPose]: 29 | data = transform(data) 30 | else: 31 | for k, v in data.items(): 32 | if k in objects and k in data: 33 | if transform.__class__ in [ 34 | RandomMirrorPoints 35 | ]: 36 | data[k] = transform(v, rnd_value) 37 | else: 38 | data[k] = transform(v) 39 | 40 | return data 41 | 42 | class ToTensor(object): 43 | def __init__(self, parameters): 44 | pass 45 | 46 | def __call__(self, arr): 47 | shape = arr.shape 48 | if len(shape) == 3: # RGB/Depth Images 49 | arr = arr.transpose(2, 0, 1) 50 | 51 | # Ref: https://discuss.pytorch.org/t/torch-from-numpy-not-support-negative-strides/3663/2 52 | return torch.from_numpy(arr.copy()).float() 53 | 54 | 55 | class RandomSamplePoints(object): 56 | def __init__(self, parameters): 57 | self.n_points = parameters['n_points'] 58 | 59 | def __call__(self, ptcloud): 60 | choice = np.random.permutation(ptcloud.shape[0]) 61 | ptcloud = ptcloud[choice[:self.n_points]] 62 | 63 | if ptcloud.shape[0] < self.n_points: 64 | zeros = np.zeros((self.n_points - ptcloud.shape[0], 3)) 65 | ptcloud = np.concatenate([ptcloud, zeros]) 66 | 67 | return ptcloud 68 | 69 | 70 | class RandomMirrorPoints(object): 71 | def __init__(self, parameters): 72 | pass 73 | 74 | def __call__(self, ptcloud, rnd_value): 75 | trfm_mat = transforms3d.zooms.zfdir2mat(1) 76 | trfm_mat_x = np.dot(transforms3d.zooms.zfdir2mat(-1, [1, 0, 0]), trfm_mat) 77 | trfm_mat_z = np.dot(transforms3d.zooms.zfdir2mat(-1, [0, 0, 1]), trfm_mat) 78 | if rnd_value <= 0.25: 79 | trfm_mat = np.dot(trfm_mat_x, trfm_mat) 80 | trfm_mat = np.dot(trfm_mat_z, trfm_mat) 81 | elif rnd_value > 0.25 and rnd_value <= 0.5: # lgtm [py/redundant-comparison] 82 | trfm_mat = np.dot(trfm_mat_x, trfm_mat) 83 | elif rnd_value > 0.5 and rnd_value <= 0.75: 84 | trfm_mat = np.dot(trfm_mat_z, trfm_mat) 85 | 86 | ptcloud[:, :3] = np.dot(ptcloud[:, :3], trfm_mat.T) 87 | return ptcloud 88 | 89 | 90 | class NormalizeObjectPose(object): 91 | def __init__(self, parameters): 92 | input_keys = parameters['input_keys'] 93 | self.ptcloud_key = input_keys['ptcloud'] 94 | self.bbox_key = input_keys['bbox'] 95 | 96 | def __call__(self, data): 97 | ptcloud = data[self.ptcloud_key] 98 | bbox = data[self.bbox_key] 99 | 100 | # Calculate center, rotation and scale 101 | # References: 102 | # - https://github.com/wentaoyuan/pcn/blob/master/test_kitti.py#L40-L52 103 | center = (bbox.min(0) + bbox.max(0)) / 2 104 | bbox -= center 105 | yaw = np.arctan2(bbox[3, 1] - bbox[0, 1], bbox[3, 0] - bbox[0, 0]) 106 | rotation = np.array([[np.cos(yaw), -np.sin(yaw), 0], [np.sin(yaw), np.cos(yaw), 0], [0, 0, 1]]) 107 | bbox = np.dot(bbox, rotation) 108 | scale = bbox[3, 0] - bbox[0, 0] 109 | bbox /= scale 110 | ptcloud = np.dot(ptcloud - center, rotation) / scale 111 | ptcloud = np.dot(ptcloud, [[1, 0, 0], [0, 0, 1], [0, 1, 0]]) 112 | 113 | data[self.ptcloud_key] = ptcloud 114 | return data 115 | -------------------------------------------------------------------------------- /datasets/io.py: -------------------------------------------------------------------------------- 1 | import h5py 2 | import numpy as np 3 | import open3d 4 | import os 5 | 6 | class IO: 7 | @classmethod 8 | def get(cls, file_path): 9 | _, file_extension = os.path.splitext(file_path) 10 | 11 | if file_extension in ['.npy']: 12 | return cls._read_npy(file_path) 13 | elif file_extension in ['.pcd']: 14 | return cls._read_pcd(file_path) 15 | elif file_extension in ['.h5']: 16 | return cls._read_h5(file_path) 17 | elif file_extension in ['.txt']: 18 | return cls._read_txt(file_path) 19 | else: 20 | raise Exception('Unsupported file extension: %s' % file_extension) 21 | 22 | # References: https://github.com/numpy/numpy/blob/master/numpy/lib/format.py 23 | @classmethod 24 | def _read_npy(cls, file_path): 25 | return np.load(file_path) 26 | 27 | # References: https://github.com/dimatura/pypcd/blob/master/pypcd/pypcd.py#L275 28 | # Support PCD files without compression ONLY! 29 | @classmethod 30 | def _read_pcd(cls, file_path): 31 | pc = open3d.io.read_point_cloud(file_path) 32 | ptcloud = np.array(pc.points) 33 | return ptcloud 34 | 35 | @classmethod 36 | def _read_txt(cls, file_path): 37 | return np.loadtxt(file_path) 38 | 39 | @classmethod 40 | def _read_h5(cls, file_path): 41 | f = h5py.File(file_path, 'r') 42 | return f['data'][()] -------------------------------------------------------------------------------- /extensions/chamfer_dist/__init__.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Thibault GROUEIX 3 | # @Date: 2019-08-07 20:54:24 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2019-12-18 15:06:25 6 | # @Email: cshzxie@gmail.com 7 | 8 | import torch 9 | 10 | import chamfer 11 | 12 | 13 | class ChamferFunction(torch.autograd.Function): 14 | @staticmethod 15 | def forward(ctx, xyz1, xyz2): 16 | dist1, dist2, idx1, idx2 = chamfer.forward(xyz1, xyz2) 17 | ctx.save_for_backward(xyz1, xyz2, idx1, idx2) 18 | 19 | return dist1, dist2 20 | 21 | @staticmethod 22 | def backward(ctx, grad_dist1, grad_dist2): 23 | xyz1, xyz2, idx1, idx2 = ctx.saved_tensors 24 | grad_xyz1, grad_xyz2 = chamfer.backward(xyz1, xyz2, idx1, idx2, grad_dist1, grad_dist2) 25 | return grad_xyz1, grad_xyz2 26 | 27 | 28 | class ChamferDistanceL2(torch.nn.Module): 29 | f''' Chamder Distance L2 30 | ''' 31 | def __init__(self, ignore_zeros=False): 32 | super().__init__() 33 | self.ignore_zeros = ignore_zeros 34 | 35 | def forward(self, xyz1, xyz2): 36 | batch_size = xyz1.size(0) 37 | if batch_size == 1 and self.ignore_zeros: 38 | non_zeros1 = torch.sum(xyz1, dim=2).ne(0) 39 | non_zeros2 = torch.sum(xyz2, dim=2).ne(0) 40 | xyz1 = xyz1[non_zeros1].unsqueeze(dim=0) 41 | xyz2 = xyz2[non_zeros2].unsqueeze(dim=0) 42 | 43 | dist1, dist2 = ChamferFunction.apply(xyz1, xyz2) 44 | return torch.mean(dist1) + torch.mean(dist2) 45 | 46 | class ChamferDistanceL2_split(torch.nn.Module): 47 | f''' Chamder Distance L2 48 | ''' 49 | def __init__(self, ignore_zeros=False): 50 | super().__init__() 51 | self.ignore_zeros = ignore_zeros 52 | 53 | def forward(self, xyz1, xyz2): 54 | batch_size = xyz1.size(0) 55 | if batch_size == 1 and self.ignore_zeros: 56 | non_zeros1 = torch.sum(xyz1, dim=2).ne(0) 57 | non_zeros2 = torch.sum(xyz2, dim=2).ne(0) 58 | xyz1 = xyz1[non_zeros1].unsqueeze(dim=0) 59 | xyz2 = xyz2[non_zeros2].unsqueeze(dim=0) 60 | 61 | dist1, dist2 = ChamferFunction.apply(xyz1, xyz2) 62 | return torch.mean(dist1), torch.mean(dist2) 63 | 64 | class ChamferDistanceL1(torch.nn.Module): 65 | f''' Chamder Distance L1 66 | ''' 67 | def __init__(self, ignore_zeros=False): 68 | super().__init__() 69 | self.ignore_zeros = ignore_zeros 70 | 71 | def forward(self, xyz1, xyz2): 72 | batch_size = xyz1.size(0) 73 | if batch_size == 1 and self.ignore_zeros: 74 | non_zeros1 = torch.sum(xyz1, dim=2).ne(0) 75 | non_zeros2 = torch.sum(xyz2, dim=2).ne(0) 76 | xyz1 = xyz1[non_zeros1].unsqueeze(dim=0) 77 | xyz2 = xyz2[non_zeros2].unsqueeze(dim=0) 78 | 79 | dist1, dist2 = ChamferFunction.apply(xyz1, xyz2) 80 | # import pdb 81 | # pdb.set_trace() 82 | dist1 = torch.sqrt(dist1) 83 | dist2 = torch.sqrt(dist2) 84 | return (torch.mean(dist1) + torch.mean(dist2))/2 85 | 86 | -------------------------------------------------------------------------------- /extensions/chamfer_dist/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/extensions/chamfer_dist/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /extensions/chamfer_dist/build/lib.linux-x86_64-3.7/chamfer.cpython-37m-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/extensions/chamfer_dist/build/lib.linux-x86_64-3.7/chamfer.cpython-37m-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /extensions/chamfer_dist/build/temp.linux-x86_64-3.7/chamfer.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/extensions/chamfer_dist/build/temp.linux-x86_64-3.7/chamfer.o -------------------------------------------------------------------------------- /extensions/chamfer_dist/build/temp.linux-x86_64-3.7/chamfer_cuda.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/extensions/chamfer_dist/build/temp.linux-x86_64-3.7/chamfer_cuda.o -------------------------------------------------------------------------------- /extensions/chamfer_dist/chamfer.egg-info/PKG-INFO: -------------------------------------------------------------------------------- 1 | Metadata-Version: 2.1 2 | Name: chamfer 3 | Version: 2.0.0 4 | Summary: UNKNOWN 5 | License: UNKNOWN 6 | Platform: UNKNOWN 7 | 8 | UNKNOWN 9 | 10 | -------------------------------------------------------------------------------- /extensions/chamfer_dist/chamfer.egg-info/SOURCES.txt: -------------------------------------------------------------------------------- 1 | chamfer.cu 2 | chamfer_cuda.cpp 3 | setup.py 4 | chamfer.egg-info/PKG-INFO 5 | chamfer.egg-info/SOURCES.txt 6 | chamfer.egg-info/dependency_links.txt 7 | chamfer.egg-info/top_level.txt -------------------------------------------------------------------------------- /extensions/chamfer_dist/chamfer.egg-info/dependency_links.txt: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /extensions/chamfer_dist/chamfer.egg-info/top_level.txt: -------------------------------------------------------------------------------- 1 | chamfer 2 | -------------------------------------------------------------------------------- /extensions/chamfer_dist/chamfer_cuda.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * @Author: Haozhe Xie 3 | * @Date: 2019-08-07 20:54:24 4 | * @Last Modified by: Haozhe Xie 5 | * @Last Modified time: 2019-12-10 10:33:50 6 | * @Email: cshzxie@gmail.com 7 | */ 8 | 9 | #include 10 | #include 11 | 12 | std::vector chamfer_cuda_forward(torch::Tensor xyz1, 13 | torch::Tensor xyz2); 14 | 15 | std::vector chamfer_cuda_backward(torch::Tensor xyz1, 16 | torch::Tensor xyz2, 17 | torch::Tensor idx1, 18 | torch::Tensor idx2, 19 | torch::Tensor grad_dist1, 20 | torch::Tensor grad_dist2); 21 | 22 | std::vector chamfer_forward(torch::Tensor xyz1, 23 | torch::Tensor xyz2) { 24 | return chamfer_cuda_forward(xyz1, xyz2); 25 | } 26 | 27 | std::vector chamfer_backward(torch::Tensor xyz1, 28 | torch::Tensor xyz2, 29 | torch::Tensor idx1, 30 | torch::Tensor idx2, 31 | torch::Tensor grad_dist1, 32 | torch::Tensor grad_dist2) { 33 | return chamfer_cuda_backward(xyz1, xyz2, idx1, idx2, grad_dist1, grad_dist2); 34 | } 35 | 36 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 37 | m.def("forward", &chamfer_forward, "Chamfer forward (CUDA)"); 38 | m.def("backward", &chamfer_backward, "Chamfer backward (CUDA)"); 39 | } 40 | -------------------------------------------------------------------------------- /extensions/chamfer_dist/dist/chamfer-2.0.0-py3.7-linux-x86_64.egg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/extensions/chamfer_dist/dist/chamfer-2.0.0-py3.7-linux-x86_64.egg -------------------------------------------------------------------------------- /extensions/chamfer_dist/setup.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Haozhe Xie 3 | # @Date: 2019-08-07 20:54:24 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2019-12-10 10:04:25 6 | # @Email: cshzxie@gmail.com 7 | 8 | from setuptools import setup 9 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 10 | 11 | setup(name='chamfer', 12 | version='2.0.0', 13 | ext_modules=[ 14 | CUDAExtension('chamfer', [ 15 | 'chamfer_cuda.cpp', 16 | 'chamfer.cu', 17 | ]), 18 | ], 19 | cmdclass={'build_ext': BuildExtension}) 20 | -------------------------------------------------------------------------------- /extensions/chamfer_dist/test.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Haozhe Xie 3 | # @Date: 2019-12-10 10:38:01 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2019-12-26 14:21:36 6 | # @Email: cshzxie@gmail.com 7 | # 8 | # Note: 9 | # - Replace float -> double, kFloat -> kDouble in chamfer.cu 10 | 11 | import os 12 | import sys 13 | import torch 14 | import unittest 15 | 16 | 17 | from torch.autograd import gradcheck 18 | 19 | sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))) 20 | from extensions.chamfer_dist import ChamferFunction 21 | 22 | 23 | class ChamferDistanceTestCase(unittest.TestCase): 24 | def test_chamfer_dist(self): 25 | x = torch.rand(4, 64, 3).double() 26 | y = torch.rand(4, 128, 3).double() 27 | x.requires_grad = True 28 | y.requires_grad = True 29 | print(gradcheck(ChamferFunction.apply, [x.cuda(), y.cuda()])) 30 | 31 | 32 | 33 | if __name__ == '__main__': 34 | # unittest.main() 35 | import pdb 36 | x = torch.rand(32,128,3) 37 | y = torch.rand(32,128,3) 38 | pdb.set_trace() 39 | -------------------------------------------------------------------------------- /extensions/cubic_feature_sampling/__init__.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Haozhe Xie 3 | # @Date: 2019-12-19 16:55:15 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2019-12-26 13:15:14 6 | # @Email: cshzxie@gmail.com 7 | 8 | import torch 9 | 10 | import cubic_feature_sampling 11 | 12 | 13 | class CubicFeatureSamplingFunction(torch.autograd.Function): 14 | @staticmethod 15 | def forward(ctx, ptcloud, cubic_features, neighborhood_size=1): 16 | scale = cubic_features.size(2) 17 | point_features, grid_pt_indexes = cubic_feature_sampling.forward(scale, neighborhood_size, ptcloud, 18 | cubic_features) 19 | ctx.save_for_backward(torch.Tensor([scale]), torch.Tensor([neighborhood_size]), grid_pt_indexes) 20 | return point_features 21 | 22 | @staticmethod 23 | def backward(ctx, grad_point_features): 24 | scale, neighborhood_size, grid_pt_indexes = ctx.saved_tensors 25 | scale = int(scale.item()) 26 | neighborhood_size = int(neighborhood_size.item()) 27 | grad_point_features = grad_point_features.contiguous() 28 | grad_ptcloud, grad_cubic_features = cubic_feature_sampling.backward(scale, neighborhood_size, 29 | grad_point_features, grid_pt_indexes) 30 | return grad_ptcloud, grad_cubic_features, None 31 | 32 | 33 | class CubicFeatureSampling(torch.nn.Module): 34 | def __init__(self): 35 | super(CubicFeatureSampling, self).__init__() 36 | 37 | def forward(self, ptcloud, cubic_features, neighborhood_size=1): 38 | h_scale = cubic_features.size(2) / 2 39 | ptcloud = ptcloud * h_scale + h_scale 40 | return CubicFeatureSamplingFunction.apply(ptcloud, cubic_features, neighborhood_size) 41 | -------------------------------------------------------------------------------- /extensions/cubic_feature_sampling/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/extensions/cubic_feature_sampling/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /extensions/cubic_feature_sampling/cubic_feature_sampling_cuda.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * @Author: Haozhe Xie 3 | * @Date: 2019-12-19 17:04:38 4 | * @Last Modified by: Haozhe Xie 5 | * @Last Modified time: 2020-06-17 14:50:22 6 | * @Email: cshzxie@gmail.com 7 | */ 8 | 9 | #include 10 | #include 11 | 12 | // NOTE: AT_ASSERT has become AT_CHECK on master after 0.4. 13 | #define CHECK_CUDA(x) AT_ASSERTM(x.is_cuda(), #x " must be a CUDA tensor") 14 | #define CHECK_CONTIGUOUS(x) \ 15 | AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") 16 | #define CHECK_INPUT(x) \ 17 | CHECK_CUDA(x); \ 18 | CHECK_CONTIGUOUS(x) 19 | 20 | std::vector cubic_feature_sampling_cuda_forward( 21 | int scale, 22 | int neighborhood_size, 23 | torch::Tensor ptcloud, 24 | torch::Tensor cubic_features, 25 | cudaStream_t stream); 26 | 27 | std::vector cubic_feature_sampling_cuda_backward( 28 | int scale, 29 | int neighborhood_size, 30 | torch::Tensor grad_point_features, 31 | torch::Tensor grid_pt_indexes, 32 | cudaStream_t stream); 33 | 34 | std::vector cubic_feature_sampling_forward( 35 | int scale, 36 | int neighborhood_size, 37 | torch::Tensor ptcloud, 38 | torch::Tensor cubic_features) { 39 | CHECK_INPUT(ptcloud); 40 | CHECK_INPUT(cubic_features); 41 | 42 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 43 | return cubic_feature_sampling_cuda_forward(scale, neighborhood_size, ptcloud, 44 | cubic_features, stream); 45 | } 46 | 47 | std::vector cubic_feature_sampling_backward( 48 | int scale, 49 | int neighborhood_size, 50 | torch::Tensor grad_point_features, 51 | torch::Tensor grid_pt_indexes) { 52 | CHECK_INPUT(grad_point_features); 53 | CHECK_INPUT(grid_pt_indexes); 54 | 55 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 56 | return cubic_feature_sampling_cuda_backward( 57 | scale, neighborhood_size, grad_point_features, grid_pt_indexes, stream); 58 | } 59 | 60 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 61 | m.def("forward", &cubic_feature_sampling_forward, 62 | "Cubic Feature Sampling forward (CUDA)"); 63 | m.def("backward", &cubic_feature_sampling_backward, 64 | "Cubic Feature Sampling backward (CUDA)"); 65 | } -------------------------------------------------------------------------------- /extensions/cubic_feature_sampling/setup.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Haozhe Xie 3 | # @Date: 2019-12-19 17:03:06 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2019-12-26 14:02:06 6 | # @Email: cshzxie@gmail.com 7 | 8 | from setuptools import setup 9 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 10 | 11 | setup(name='cubic_feature_sampling', 12 | version='1.1.0', 13 | ext_modules=[ 14 | CUDAExtension('cubic_feature_sampling', ['cubic_feature_sampling_cuda.cpp', 'cubic_feature_sampling.cu']), 15 | ], 16 | cmdclass={'build_ext': BuildExtension}) 17 | -------------------------------------------------------------------------------- /extensions/cubic_feature_sampling/test.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Haozhe Xie 3 | # @Date: 2019-12-20 11:50:50 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2019-12-26 13:52:33 6 | # @Email: cshzxie@gmail.com 7 | # 8 | # Note: 9 | # - Replace float -> double, kFloat -> kDouble in cubic_feature_sampling.cu 10 | 11 | import os 12 | import sys 13 | import torch 14 | import unittest 15 | 16 | from torch.autograd import gradcheck 17 | 18 | sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))) 19 | from extensions.cubic_feature_sampling import CubicFeatureSamplingFunction 20 | 21 | 22 | class CubicFeatureSamplingTestCase(unittest.TestCase): 23 | def test_neighborhood_size_1(self): 24 | ptcloud = torch.rand(2, 64, 3) * 2 - 1 25 | cubic_features = torch.rand(2, 4, 8, 8, 8) 26 | ptcloud.requires_grad = True 27 | cubic_features.requires_grad = True 28 | self.assertTrue( 29 | gradcheck(CubicFeatureSamplingFunction.apply, 30 | [ptcloud.double().cuda(), cubic_features.double().cuda()])) 31 | 32 | def test_neighborhood_size_2(self): 33 | ptcloud = torch.rand(2, 32, 3) * 2 - 1 34 | cubic_features = torch.rand(2, 2, 8, 8, 8) 35 | ptcloud.requires_grad = True 36 | cubic_features.requires_grad = True 37 | self.assertTrue( 38 | gradcheck(CubicFeatureSamplingFunction.apply, 39 | [ptcloud.double().cuda(), cubic_features.double().cuda(), 2])) 40 | 41 | def test_neighborhood_size_3(self): 42 | ptcloud = torch.rand(1, 32, 3) * 2 - 1 43 | cubic_features = torch.rand(1, 2, 16, 16, 16) 44 | ptcloud.requires_grad = True 45 | cubic_features.requires_grad = True 46 | self.assertTrue( 47 | gradcheck(CubicFeatureSamplingFunction.apply, 48 | [ptcloud.double().cuda(), cubic_features.double().cuda(), 3])) 49 | 50 | 51 | if __name__ == '__main__': 52 | unittest.main() 53 | -------------------------------------------------------------------------------- /extensions/gridding/__init__.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Haozhe Xie 3 | # @Date: 2019-11-15 20:33:52 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2019-12-30 09:55:53 6 | # @Email: cshzxie@gmail.com 7 | 8 | import torch 9 | 10 | import gridding 11 | 12 | 13 | class GriddingFunction(torch.autograd.Function): 14 | @staticmethod 15 | def forward(ctx, scale, ptcloud): 16 | grid, grid_pt_weights, grid_pt_indexes = gridding.forward(-scale, scale - 1, -scale, scale - 1, -scale, 17 | scale - 1, ptcloud) 18 | # print(grid.size()) # torch.Size(batch_size, n_grid_vertices) 19 | # print(grid_pt_weights.size()) # torch.Size(batch_size, n_pts, 8, 3) 20 | # print(grid_pt_indexes.size()) # torch.Size(batch_size, n_pts, 8) 21 | ctx.save_for_backward(grid_pt_weights, grid_pt_indexes) 22 | 23 | return grid 24 | 25 | @staticmethod 26 | def backward(ctx, grad_grid): 27 | grid_pt_weights, grid_pt_indexes = ctx.saved_tensors 28 | grad_ptcloud = gridding.backward(grid_pt_weights, grid_pt_indexes, grad_grid) 29 | # print(grad_ptcloud.size()) # torch.Size(batch_size, n_pts, 3) 30 | 31 | return None, grad_ptcloud 32 | 33 | 34 | class Gridding(torch.nn.Module): 35 | def __init__(self, scale=1): 36 | super(Gridding, self).__init__() 37 | self.scale = scale // 2 38 | 39 | def forward(self, ptcloud): 40 | ptcloud = ptcloud * self.scale 41 | _ptcloud = torch.split(ptcloud, 1, dim=0) 42 | grids = [] 43 | for p in _ptcloud: 44 | non_zeros = torch.sum(p, dim=2).ne(0) 45 | p = p[non_zeros].unsqueeze(dim=0) 46 | grids.append(GriddingFunction.apply(self.scale, p.contiguous())) 47 | 48 | return torch.cat(grids, dim=0).contiguous() 49 | 50 | 51 | class GriddingReverseFunction(torch.autograd.Function): 52 | @staticmethod 53 | def forward(ctx, scale, grid): 54 | ptcloud = gridding.rev_forward(scale, grid) 55 | ctx.save_for_backward(torch.Tensor([scale]), grid, ptcloud) 56 | return ptcloud 57 | 58 | @staticmethod 59 | def backward(ctx, grad_ptcloud): 60 | scale, grid, ptcloud = ctx.saved_tensors 61 | scale = int(scale.item()) 62 | grad_grid = gridding.rev_backward(ptcloud, grid, grad_ptcloud) 63 | grad_grid = grad_grid.view(-1, scale, scale, scale) 64 | return None, grad_grid 65 | 66 | 67 | class GriddingReverse(torch.nn.Module): 68 | def __init__(self, scale=1): 69 | super(GriddingReverse, self).__init__() 70 | self.scale = scale 71 | 72 | def forward(self, grid): 73 | ptcloud = GriddingReverseFunction.apply(self.scale, grid) 74 | return ptcloud / self.scale * 2 75 | -------------------------------------------------------------------------------- /extensions/gridding/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/extensions/gridding/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /extensions/gridding/gridding_cuda.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * @Author: Haozhe Xie 3 | * @Date: 2019-11-13 10:52:53 4 | * @Last Modified by: Haozhe Xie 5 | * @Last Modified time: 2020-06-17 14:52:32 6 | * @Email: cshzxie@gmail.com 7 | */ 8 | 9 | #include 10 | #include 11 | 12 | // NOTE: AT_ASSERT has become AT_CHECK on master after 0.4. 13 | #define CHECK_CUDA(x) AT_ASSERTM(x.is_cuda(), #x " must be a CUDA tensor") 14 | #define CHECK_CONTIGUOUS(x) \ 15 | AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") 16 | #define CHECK_INPUT(x) \ 17 | CHECK_CUDA(x); \ 18 | CHECK_CONTIGUOUS(x) 19 | 20 | std::vector gridding_cuda_forward(float min_x, 21 | float max_x, 22 | float min_y, 23 | float max_y, 24 | float min_z, 25 | float max_z, 26 | torch::Tensor ptcloud, 27 | cudaStream_t stream); 28 | 29 | torch::Tensor gridding_cuda_backward(torch::Tensor grid_pt_weights, 30 | torch::Tensor grid_pt_indexes, 31 | torch::Tensor grad_grid, 32 | cudaStream_t stream); 33 | 34 | torch::Tensor gridding_reverse_cuda_forward(int scale, 35 | torch::Tensor grid, 36 | cudaStream_t stream); 37 | 38 | torch::Tensor gridding_reverse_cuda_backward(torch::Tensor ptcloud, 39 | torch::Tensor grid, 40 | torch::Tensor grad_ptcloud, 41 | cudaStream_t stream); 42 | 43 | std::vector gridding_forward(float min_x, 44 | float max_x, 45 | float min_y, 46 | float max_y, 47 | float min_z, 48 | float max_z, 49 | torch::Tensor ptcloud) { 50 | CHECK_INPUT(ptcloud); 51 | 52 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 53 | return gridding_cuda_forward(min_x, max_x, min_y, max_y, min_z, max_z, 54 | ptcloud, stream); 55 | } 56 | 57 | torch::Tensor gridding_backward(torch::Tensor grid_pt_weights, 58 | torch::Tensor grid_pt_indexes, 59 | torch::Tensor grad_grid) { 60 | CHECK_INPUT(grid_pt_weights); 61 | CHECK_INPUT(grid_pt_indexes); 62 | CHECK_INPUT(grad_grid); 63 | 64 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 65 | return gridding_cuda_backward(grid_pt_weights, grid_pt_indexes, grad_grid, 66 | stream); 67 | } 68 | 69 | torch::Tensor gridding_reverse_forward(int scale, torch::Tensor grid) { 70 | CHECK_INPUT(grid); 71 | 72 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 73 | return gridding_reverse_cuda_forward(scale, grid, stream); 74 | } 75 | 76 | torch::Tensor gridding_reverse_backward(torch::Tensor ptcloud, 77 | torch::Tensor grid, 78 | torch::Tensor grad_ptcloud) { 79 | CHECK_INPUT(ptcloud); 80 | CHECK_INPUT(grid); 81 | CHECK_INPUT(grad_ptcloud); 82 | 83 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 84 | return gridding_reverse_cuda_backward(ptcloud, grid, grad_ptcloud, stream); 85 | } 86 | 87 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 88 | m.def("forward", &gridding_forward, "Gridding forward (CUDA)"); 89 | m.def("backward", &gridding_backward, "Gridding backward (CUDA)"); 90 | m.def("rev_forward", &gridding_reverse_forward, 91 | "Gridding Reverse forward (CUDA)"); 92 | m.def("rev_backward", &gridding_reverse_backward, 93 | "Gridding Reverse backward (CUDA)"); 94 | } 95 | -------------------------------------------------------------------------------- /extensions/gridding/setup.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Haozhe Xie 3 | # @Date: 2019-11-13 10:51:33 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2019-12-02 17:02:16 6 | # @Email: cshzxie@gmail.com 7 | 8 | from setuptools import setup 9 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 10 | 11 | setup(name='gridding', 12 | version='2.1.0', 13 | ext_modules=[ 14 | CUDAExtension('gridding', ['gridding_cuda.cpp', 'gridding.cu', 'gridding_reverse.cu']), 15 | ], 16 | cmdclass={'build_ext': BuildExtension}) 17 | -------------------------------------------------------------------------------- /extensions/gridding/test.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Haozhe Xie 3 | # @Date: 2019-12-10 10:48:55 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2019-12-26 14:20:42 6 | # @Email: cshzxie@gmail.com 7 | # 8 | # Note: 9 | # - Replace float -> double, kFloat -> kDouble in gridding.cu and gridding_reverse.cu 10 | 11 | import os 12 | import sys 13 | import torch 14 | import unittest 15 | 16 | from torch.autograd import gradcheck 17 | 18 | sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))) 19 | from extensions.gridding import GriddingFunction, GriddingReverseFunction 20 | 21 | 22 | class GriddingTestCase(unittest.TestCase): 23 | def test_gridding_reverse_function_4(self): 24 | x = torch.rand(2, 4, 4, 4) 25 | x.requires_grad = True 26 | self.assertTrue(gradcheck(GriddingReverseFunction.apply, [4, x.double().cuda()])) 27 | 28 | def test_gridding_reverse_function_8(self): 29 | x = torch.rand(4, 8, 8, 8) 30 | x.requires_grad = True 31 | self.assertTrue(gradcheck(GriddingReverseFunction.apply, [8, x.double().cuda()])) 32 | 33 | def test_gridding_reverse_function_16(self): 34 | x = torch.rand(1, 16, 16, 16) 35 | x.requires_grad = True 36 | self.assertTrue(gradcheck(GriddingReverseFunction.apply, [16, x.double().cuda()])) 37 | 38 | def test_gridding_function_32pts(self): 39 | x = torch.rand(1, 32, 3) 40 | x.requires_grad = True 41 | self.assertTrue(gradcheck(GriddingFunction.apply, [x.double().cuda()])) 42 | 43 | 44 | if __name__ == '__main__': 45 | unittest.main() 46 | -------------------------------------------------------------------------------- /extensions/gridding_loss/__init__.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Haozhe Xie 3 | # @Date: 2019-12-30 09:56:06 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2020-02-22 19:19:43 6 | # @Email: cshzxie@gmail.com 7 | 8 | import torch 9 | 10 | import gridding_distance 11 | 12 | 13 | class GriddingDistanceFunction(torch.autograd.Function): 14 | @staticmethod 15 | def forward(ctx, min_x, max_x, min_y, max_y, min_z, max_z, pred_cloud, gt_cloud): 16 | pred_grid, pred_grid_pt_weights, pred_grid_pt_indexes = gridding_distance.forward( 17 | min_x, max_x, min_y, max_y, min_z, max_z, pred_cloud) 18 | # print(pred_grid.size()) # torch.Size(batch_size, n_grid_vertices, 8) 19 | # print(pred_grid_pt_weights.size()) # torch.Size(batch_size, n_pts, 8, 3) 20 | # print(pred_grid_pt_indexes.size()) # torch.Size(batch_size, n_pts, 8) 21 | gt_grid, gt_grid_pt_weights, gt_grid_pt_indexes = gridding_distance.forward( 22 | min_x, max_x, min_y, max_y, min_z, max_z, gt_cloud) 23 | # print(gt_grid.size()) # torch.Size(batch_size, n_grid_vertices, 8) 24 | # print(gt_grid_pt_weights.size()) # torch.Size(batch_size, n_pts, 8, 3) 25 | # print(gt_grid_pt_indexes.size()) # torch.Size(batch_size, n_pts, 8) 26 | 27 | ctx.save_for_backward(pred_grid_pt_weights, pred_grid_pt_indexes, gt_grid_pt_weights, gt_grid_pt_indexes) 28 | return pred_grid, gt_grid 29 | 30 | @staticmethod 31 | def backward(ctx, grad_pred_grid, grad_gt_grid): 32 | pred_grid_pt_weights, pred_grid_pt_indexes, gt_grid_pt_weights, gt_grid_pt_indexes = ctx.saved_tensors 33 | 34 | grad_pred_cloud = gridding_distance.backward(pred_grid_pt_weights, pred_grid_pt_indexes, grad_pred_grid) 35 | # print(grad_pred_cloud.size()) # torch.Size(batch_size, n_pts, 3) 36 | grad_gt_cloud = gridding_distance.backward(gt_grid_pt_weights, gt_grid_pt_indexes, grad_gt_grid) 37 | # print(grad_gt_cloud.size()) # torch.Size(batch_size, n_pts, 3) 38 | 39 | return None, None, None, None, None, None, grad_pred_cloud, grad_gt_cloud 40 | 41 | 42 | class GriddingDistance(torch.nn.Module): 43 | def __init__(self, scale=1): 44 | super(GriddingDistance, self).__init__() 45 | self.scale = scale 46 | 47 | def forward(self, pred_cloud, gt_cloud): 48 | ''' 49 | pred_cloud(b, n_pts1, 3) 50 | gt_cloud(b, n_pts2, 3) 51 | ''' 52 | pred_cloud = pred_cloud * self.scale / 2 53 | gt_cloud = gt_cloud * self.scale / 2 54 | 55 | min_pred_x = torch.min(pred_cloud[:, :, 0]) 56 | max_pred_x = torch.max(pred_cloud[:, :, 0]) 57 | min_pred_y = torch.min(pred_cloud[:, :, 1]) 58 | max_pred_y = torch.max(pred_cloud[:, :, 1]) 59 | min_pred_z = torch.min(pred_cloud[:, :, 2]) 60 | max_pred_z = torch.max(pred_cloud[:, :, 2]) 61 | 62 | min_gt_x = torch.min(gt_cloud[:, :, 0]) 63 | max_gt_x = torch.max(gt_cloud[:, :, 0]) 64 | min_gt_y = torch.min(gt_cloud[:, :, 1]) 65 | max_gt_y = torch.max(gt_cloud[:, :, 1]) 66 | min_gt_z = torch.min(gt_cloud[:, :, 2]) 67 | max_gt_z = torch.max(gt_cloud[:, :, 2]) 68 | 69 | min_x = torch.floor(torch.min(min_pred_x, min_gt_x)) - 1 70 | max_x = torch.ceil(torch.max(max_pred_x, max_gt_x)) + 1 71 | min_y = torch.floor(torch.min(min_pred_y, min_gt_y)) - 1 72 | max_y = torch.ceil(torch.max(max_pred_y, max_gt_y)) + 1 73 | min_z = torch.floor(torch.min(min_pred_z, min_gt_z)) - 1 74 | max_z = torch.ceil(torch.max(max_pred_z, max_gt_z)) + 1 75 | 76 | _pred_clouds = torch.split(pred_cloud, 1, dim=0) 77 | _gt_clouds = torch.split(gt_cloud, 1, dim=0) 78 | pred_grids = [] 79 | gt_grids = [] 80 | for pc, gc in zip(_pred_clouds, _gt_clouds): 81 | non_zeros = torch.sum(pc, dim=2).ne(0) 82 | pc = pc[non_zeros].unsqueeze(dim=0) 83 | non_zeros = torch.sum(gc, dim=2).ne(0) 84 | gc = gc[non_zeros].unsqueeze(dim=0) 85 | pred_grid, gt_grid = GriddingDistanceFunction.apply(min_x, max_x, min_y, max_y, min_z, max_z, pc, gc) 86 | pred_grids.append(pred_grid) 87 | gt_grids.append(gt_grid) 88 | 89 | return torch.cat(pred_grids, dim=0).contiguous(), torch.cat(gt_grids, dim=0).contiguous() 90 | 91 | 92 | class GriddingLoss(torch.nn.Module): 93 | def __init__(self, scales=[], alphas=[]): 94 | super(GriddingLoss, self).__init__() 95 | self.scales = scales 96 | self.alphas = alphas 97 | self.gridding_dists = [GriddingDistance(scale=s) for s in scales] 98 | self.l1_loss = torch.nn.L1Loss() 99 | 100 | def forward(self, pred_cloud, gt_cloud): 101 | gridding_loss = None 102 | n_dists = len(self.scales) 103 | 104 | for i in range(n_dists): 105 | alpha = self.alphas[i] 106 | gdist = self.gridding_dists[i] 107 | pred_grid, gt_grid = gdist(pred_cloud, gt_cloud) 108 | 109 | if gridding_loss is None: 110 | gridding_loss = alpha * self.l1_loss(pred_grid, gt_grid) 111 | else: 112 | gridding_loss += alpha * self.l1_loss(pred_grid, gt_grid) 113 | 114 | return gridding_loss 115 | -------------------------------------------------------------------------------- /extensions/gridding_loss/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/extensions/gridding_loss/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /extensions/gridding_loss/gridding_distance_cuda.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * @Author: Haozhe Xie 3 | * @Date: 2019-12-30 10:59:31 4 | * @Last Modified by: Haozhe Xie 5 | * @Last Modified time: 2020-06-17 14:52:52 6 | * @Email: cshzxie@gmail.com 7 | */ 8 | 9 | #include 10 | #include 11 | 12 | // NOTE: AT_ASSERT has become AT_CHECK on master after 0.4. 13 | #define CHECK_CUDA(x) AT_ASSERTM(x.is_cuda(), #x " must be a CUDA tensor") 14 | #define CHECK_CONTIGUOUS(x) \ 15 | AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") 16 | #define CHECK_INPUT(x) \ 17 | CHECK_CUDA(x); \ 18 | CHECK_CONTIGUOUS(x) 19 | 20 | std::vector gridding_distance_cuda_forward(float min_x, 21 | float max_x, 22 | float min_y, 23 | float max_y, 24 | float min_z, 25 | float max_z, 26 | torch::Tensor ptcloud, 27 | cudaStream_t stream); 28 | 29 | torch::Tensor gridding_distance_cuda_backward(torch::Tensor grid_pt_weights, 30 | torch::Tensor grid_pt_indexes, 31 | torch::Tensor grad_grid, 32 | cudaStream_t stream); 33 | 34 | std::vector gridding_distance_forward(float min_x, 35 | float max_x, 36 | float min_y, 37 | float max_y, 38 | float min_z, 39 | float max_z, 40 | torch::Tensor ptcloud) { 41 | CHECK_INPUT(ptcloud); 42 | 43 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 44 | return gridding_distance_cuda_forward(min_x, max_x, min_y, max_y, min_z, 45 | max_z, ptcloud, stream); 46 | } 47 | 48 | torch::Tensor gridding_distance_backward(torch::Tensor grid_pt_weights, 49 | torch::Tensor grid_pt_indexes, 50 | torch::Tensor grad_grid) { 51 | CHECK_INPUT(grid_pt_weights); 52 | CHECK_INPUT(grid_pt_indexes); 53 | CHECK_INPUT(grad_grid); 54 | 55 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 56 | return gridding_distance_cuda_backward(grid_pt_weights, grid_pt_indexes, 57 | grad_grid, stream); 58 | } 59 | 60 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 61 | m.def("forward", &gridding_distance_forward, 62 | "Gridding Distance Forward (CUDA)"); 63 | m.def("backward", &gridding_distance_backward, 64 | "Gridding Distance Backward (CUDA)"); 65 | } 66 | -------------------------------------------------------------------------------- /extensions/gridding_loss/setup.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Haozhe Xie 3 | # @Date: 2019-12-30 11:03:55 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2019-12-30 11:13:39 6 | # @Email: cshzxie@gmail.com 7 | 8 | from setuptools import setup 9 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 10 | 11 | setup(name='gridding_distance', 12 | version='1.0.0', 13 | ext_modules=[ 14 | CUDAExtension('gridding_distance', ['gridding_distance_cuda.cpp', 'gridding_distance.cu']), 15 | ], 16 | cmdclass={'build_ext': BuildExtension}) 17 | -------------------------------------------------------------------------------- /figs/More_Results_On_ShapNet55_1.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/figs/More_Results_On_ShapNet55_1.pdf -------------------------------------------------------------------------------- /figs/PCN_Dataset_Result_Picture.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/figs/PCN_Dataset_Result_Picture.pdf -------------------------------------------------------------------------------- /figs/PCN_results.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/figs/PCN_results.png -------------------------------------------------------------------------------- /figs/Pipeline.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/figs/Pipeline.png -------------------------------------------------------------------------------- /figs/Pipeline2.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/figs/Pipeline2.pdf -------------------------------------------------------------------------------- /figs/ShapeNet55_Results.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/figs/ShapeNet55_Results.pdf -------------------------------------------------------------------------------- /figs/ShapeNet55_results.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/figs/ShapeNet55_results.png -------------------------------------------------------------------------------- /figs/more_shapenet55_results.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/figs/more_shapenet55_results.png -------------------------------------------------------------------------------- /install.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env sh 2 | HOME='/home/lss/Project/PoinTr' 3 | 4 | # Chamfer Distance 5 | cd $HOME/extensions/chamfer_dist 6 | python setup.py install --user 7 | 8 | # NOTE: For GRNet 9 | 10 | # Cubic Feature Sampling 11 | cd $HOME/extensions/cubic_feature_sampling 12 | python setup.py install --user 13 | 14 | # Gridding & Gridding Reverse 15 | cd $HOME/extensions/gridding 16 | python setup.py install --user 17 | 18 | # Gridding Loss 19 | cd $HOME/extensions/gridding_loss 20 | python setup.py install --user 21 | 22 | -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | from tools import run_net 2 | from tools import test_net 3 | from utils import parser, dist_utils, misc 4 | from utils.logger import * 5 | from utils.config import * 6 | import time 7 | import os 8 | import torch 9 | from tensorboardX import SummaryWriter 10 | 11 | def main(): 12 | # args 13 | args = parser.get_args() 14 | # CUDA 15 | args.use_gpu = torch.cuda.is_available() 16 | if args.use_gpu: 17 | torch.backends.cudnn.benchmark = True 18 | # init distributed env first, since logger depends on the dist info. 19 | if args.launcher == 'none': 20 | args.distributed = False 21 | else: 22 | args.distributed = True 23 | dist_utils.init_dist(args.launcher) 24 | # re-set gpu_ids with distributed training mode 25 | _, world_size = dist_utils.get_dist_info() 26 | args.world_size = world_size 27 | # logger 28 | timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime()) 29 | log_file = os.path.join(args.experiment_path, f'{timestamp}.log') 30 | logger = get_root_logger(log_file=log_file, name=args.log_name) 31 | # define the tensorboard writer 32 | if not args.test: 33 | if args.local_rank == 0: 34 | train_writer = SummaryWriter(os.path.join(args.tfboard_path, 'train')) 35 | val_writer = SummaryWriter(os.path.join(args.tfboard_path, 'test')) 36 | else: 37 | train_writer = None 38 | val_writer = None 39 | # config 40 | config = get_config(args, logger = logger) 41 | # batch size 42 | if args.distributed: 43 | assert config.total_bs % world_size == 0 44 | config.dataset.train.others.bs = config.total_bs // world_size 45 | else: 46 | config.dataset.train.others.bs = config.total_bs 47 | # log 48 | log_args_to_file(args, 'args', logger = logger) 49 | log_config_to_file(config, 'config', logger = logger) 50 | # exit() 51 | logger.info(f'Distributed training: {args.distributed}') 52 | # set random seeds 53 | if args.seed is not None: 54 | logger.info(f'Set random seed to {args.seed}, ' 55 | f'deterministic: {args.deterministic}') 56 | misc.set_random_seed(args.seed + args.local_rank, deterministic=args.deterministic) # seed + rank, for augmentation 57 | if args.distributed: 58 | assert args.local_rank == torch.distributed.get_rank() 59 | 60 | # run 61 | if args.test: 62 | test_net(args, config) 63 | else: 64 | run_net(args, config, train_writer, val_writer) 65 | 66 | 67 | if __name__ == '__main__': 68 | main() 69 | -------------------------------------------------------------------------------- /other_models/FoldingNet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from .build import MODELS 4 | from extensions.chamfer_dist import ChamferDistanceL2 5 | 6 | @MODELS.register_module() 7 | class FoldingNet(nn.Module): 8 | def __init__(self, config): 9 | super().__init__() 10 | self.num_pred = config.num_pred 11 | self.encoder_channel = config.encoder_channel 12 | self.grid_size = int(pow(self.num_pred,0.5) + 0.5) 13 | 14 | self.first_conv = nn.Sequential( 15 | nn.Conv1d(3,128,1), 16 | nn.BatchNorm1d(128), 17 | nn.ReLU(inplace=True), 18 | nn.Conv1d(128,256,1) 19 | ) 20 | self.second_conv = nn.Sequential( 21 | nn.Conv1d(512,512,1), 22 | nn.BatchNorm1d(512), 23 | nn.ReLU(inplace=True), 24 | nn.Conv1d(512,self.encoder_channel,1) 25 | ) 26 | 27 | self.folding1 = nn.Sequential( 28 | nn.Conv1d(self.encoder_channel + 2, 512, 1), 29 | nn.BatchNorm1d(512), 30 | nn.ReLU(inplace=True), 31 | nn.Conv1d(512, 512, 1), 32 | nn.BatchNorm1d(512), 33 | nn.ReLU(inplace=True), 34 | nn.Conv1d(512, 3, 1), 35 | ) 36 | 37 | self.folding2 = nn.Sequential( 38 | nn.Conv1d(self.encoder_channel + 3, 512, 1), 39 | nn.BatchNorm1d(512), 40 | nn.ReLU(inplace=True), 41 | nn.Conv1d(512, 512, 1), 42 | nn.BatchNorm1d(512), 43 | nn.ReLU(inplace=True), 44 | nn.Conv1d(512, 3, 1), 45 | ) 46 | 47 | a = torch.linspace(-0.5, 0.5, steps=self.grid_size, dtype=torch.float).view(1, self.grid_size).expand(self.grid_size, self.grid_size).reshape(1, -1) 48 | b = torch.linspace(-0.5, 0.5, steps=self.grid_size, dtype=torch.float).view(self.grid_size, 1).expand(self.grid_size, self.grid_size).reshape(1, -1) 49 | self.folding_seed = torch.cat([a, b], dim=0).view(1, 2, self.grid_size ** 2).cuda() # 1 2 N 50 | self.build_loss_func() 51 | 52 | def build_loss_func(self): 53 | self.loss_func = ChamferDistanceL2() 54 | 55 | def get_loss(self, ret, gt): 56 | loss_coarse = self.loss_func(ret[0], gt) 57 | loss_fine = self.loss_func(ret[1], gt) 58 | return loss_coarse, loss_fine 59 | 60 | def forward(self, xyz): 61 | bs , n , _ = xyz.shape 62 | # encoder 63 | feature = self.first_conv(xyz.transpose(2,1)) # B 256 n 64 | feature_global = torch.max(feature,dim=2,keepdim=True)[0] # B 256 1 65 | feature = torch.cat([feature_global.expand(-1,-1,n), feature], dim=1)# B 512 n 66 | feature = self.second_conv(feature) # B 1024 n 67 | feature_global = torch.max(feature,dim=2,keepdim=False)[0] # B 1024 68 | # folding decoder 69 | fd1, fd2 = self.decoder(feature_global) # B N 3 70 | return (fd2, fd2) # FoldingNet producing final result directly 71 | 72 | def decoder(self,x): 73 | num_sample = self.grid_size * self.grid_size 74 | bs = x.size(0) 75 | features = x.view(bs, self.encoder_channel, 1).expand(bs, self.encoder_channel, num_sample) 76 | seed = self.folding_seed.view(1, 2, num_sample).expand(bs, 2, num_sample).to(x.device) 77 | 78 | x = torch.cat([seed, features], dim=1) 79 | fd1 = self.folding1(x) 80 | x = torch.cat([fd1, features], dim=1) 81 | fd2 = self.folding2(x) 82 | 83 | return fd1.transpose(2,1).contiguous() , fd2.transpose(2,1).contiguous() 84 | # return fd2.transpose(2,1).contiguous() -------------------------------------------------------------------------------- /other_models/PCN.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from .build import MODELS 4 | from extensions.chamfer_dist import ChamferDistanceL2 5 | 6 | @MODELS.register_module() 7 | class PCN(nn.Module): 8 | def __init__(self, config): 9 | super().__init__() 10 | self.number_fine = config.num_pred 11 | self.encoder_channel = config.encoder_channel 12 | grid_size = 4 # set default 13 | self.grid_size = grid_size 14 | assert self.number_fine % grid_size**2 == 0 15 | self.number_coarse = self.number_fine // (grid_size ** 2 ) 16 | self.first_conv = nn.Sequential( 17 | nn.Conv1d(3,128,1), 18 | nn.BatchNorm1d(128), 19 | nn.ReLU(inplace=True), 20 | nn.Conv1d(128,256,1) 21 | ) 22 | self.second_conv = nn.Sequential( 23 | nn.Conv1d(512,512,1), 24 | nn.BatchNorm1d(512), 25 | nn.ReLU(inplace=True), 26 | nn.Conv1d(512,self.encoder_channel,1) 27 | ) 28 | self.mlp = nn.Sequential( 29 | nn.Linear(self.encoder_channel,1024), 30 | nn.ReLU(inplace=True), 31 | nn.Linear(1024,1024), 32 | nn.ReLU(inplace=True), 33 | nn.Linear(1024,3*self.number_coarse) 34 | ) 35 | self.final_conv = nn.Sequential( 36 | nn.Conv1d(1024+3+2,512,1), 37 | nn.BatchNorm1d(512), 38 | nn.ReLU(inplace=True), 39 | nn.Conv1d(512,512,1), 40 | nn.BatchNorm1d(512), 41 | nn.ReLU(inplace=True), 42 | nn.Conv1d(512,3,1) 43 | ) 44 | a = torch.linspace(-0.05, 0.05, steps=grid_size, dtype=torch.float).view(1, grid_size).expand(grid_size, grid_size).reshape(1, -1) 45 | b = torch.linspace(-0.05, 0.05, steps=grid_size, dtype=torch.float).view(grid_size, 1).expand(grid_size, grid_size).reshape(1, -1) 46 | self.folding_seed = torch.cat([a, b], dim=0).view(1, 2, grid_size ** 2).cuda() # 1 2 S 47 | self.build_loss_func() 48 | 49 | def build_loss_func(self): 50 | self.loss_func = ChamferDistanceL2() 51 | 52 | def get_loss(self, ret, gt): 53 | loss_coarse = self.loss_func(ret[0], gt) 54 | loss_fine = self.loss_func(ret[1], gt) 55 | return loss_coarse, loss_fine 56 | 57 | def forward(self, xyz): 58 | bs , n , _ = xyz.shape 59 | # encoder 60 | feature = self.first_conv(xyz.transpose(2,1)) # B 256 n 61 | feature_global = torch.max(feature,dim=2,keepdim=True)[0] # B 256 1 62 | feature = torch.cat([feature_global.expand(-1,-1,n), feature], dim=1)# B 512 n 63 | feature = self.second_conv(feature) # B 1024 n 64 | feature_global = torch.max(feature,dim=2,keepdim=False)[0] # B 1024 65 | # decoder 66 | coarse = self.mlp(feature_global).reshape(-1,self.number_coarse,3) # B M 3 67 | point_feat = coarse.unsqueeze(2).expand(-1,-1,self.grid_size**2,-1) # B M S 3 68 | point_feat = point_feat.reshape(-1,self.number_fine,3).transpose(2,1) # B 3 N 69 | 70 | seed = self.folding_seed.unsqueeze(2).expand(bs,-1,self.number_coarse, -1) # B 2 M S 71 | seed = seed.reshape(bs,-1,self.number_fine) # B 2 N 72 | 73 | feature_global = feature_global.unsqueeze(2).expand(-1,-1,self.number_fine) # B 1024 N 74 | feat = torch.cat([feature_global, seed, point_feat], dim=1) # B C N 75 | 76 | fine = self.final_conv(feat) + point_feat # B 3 N 77 | 78 | return (coarse.contiguous(), fine.transpose(1,2).contiguous()) 79 | -------------------------------------------------------------------------------- /other_models/PoinTr.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | 4 | from pointnet2_ops import pointnet2_utils 5 | from extensions.chamfer_dist import ChamferDistanceL1 6 | from .Transformer import PCTransformer 7 | from .build import MODELS 8 | 9 | 10 | def fps(pc, num): 11 | fps_idx = pointnet2_utils.furthest_point_sample(pc, num) 12 | sub_pc = pointnet2_utils.gather_operation(pc.transpose(1, 2).contiguous(), fps_idx).transpose(1,2).contiguous() 13 | return sub_pc 14 | 15 | 16 | class Fold(nn.Module): 17 | def __init__(self, in_channel , step , hidden_dim = 512): 18 | super().__init__() 19 | 20 | self.in_channel = in_channel 21 | self.step = step 22 | 23 | a = torch.linspace(-1., 1., steps=step, dtype=torch.float).view(1, step).expand(step, step).reshape(1, -1) 24 | b = torch.linspace(-1., 1., steps=step, dtype=torch.float).view(step, 1).expand(step, step).reshape(1, -1) 25 | self.folding_seed = torch.cat([a, b], dim=0).cuda() 26 | 27 | self.folding1 = nn.Sequential( 28 | nn.Conv1d(in_channel + 2, hidden_dim, 1), 29 | nn.BatchNorm1d(hidden_dim), 30 | nn.ReLU(inplace=True), 31 | nn.Conv1d(hidden_dim, hidden_dim//2, 1), 32 | nn.BatchNorm1d(hidden_dim//2), 33 | nn.ReLU(inplace=True), 34 | nn.Conv1d(hidden_dim//2, 3, 1), 35 | ) 36 | 37 | self.folding2 = nn.Sequential( 38 | nn.Conv1d(in_channel + 3, hidden_dim, 1), 39 | nn.BatchNorm1d(hidden_dim), 40 | nn.ReLU(inplace=True), 41 | nn.Conv1d(hidden_dim, hidden_dim//2, 1), 42 | nn.BatchNorm1d(hidden_dim//2), 43 | nn.ReLU(inplace=True), 44 | nn.Conv1d(hidden_dim//2, 3, 1), 45 | ) 46 | 47 | def forward(self, x): 48 | num_sample = self.step * self.step 49 | bs = x.size(0) 50 | features = x.view(bs, self.in_channel, 1).expand(bs, self.in_channel, num_sample) 51 | seed = self.folding_seed.view(1, 2, num_sample).expand(bs, 2, num_sample).to(x.device) 52 | 53 | x = torch.cat([seed, features], dim=1) 54 | fd1 = self.folding1(x) 55 | x = torch.cat([fd1, features], dim=1) 56 | fd2 = self.folding2(x) 57 | 58 | return fd2 59 | 60 | @MODELS.register_module() 61 | class PoinTr(nn.Module): 62 | def __init__(self, config, **kwargs): 63 | super().__init__() 64 | self.trans_dim = config.trans_dim 65 | self.knn_layer = config.knn_layer 66 | self.num_pred = config.num_pred 67 | self.num_query = config.num_query 68 | 69 | self.fold_step = int(pow(self.num_pred//self.num_query, 0.5) + 0.5) 70 | self.base_model = PCTransformer(in_chans = 3, embed_dim = self.trans_dim, depth = [6, 8], drop_rate = 0., num_query = self.num_query, knn_layer = self.knn_layer) 71 | 72 | self.foldingnet = Fold(self.trans_dim, step = self.fold_step, hidden_dim = 256) # rebuild a cluster point 73 | 74 | self.increase_dim = nn.Sequential( 75 | nn.Conv1d(self.trans_dim, 1024, 1), 76 | nn.BatchNorm1d(1024), 77 | nn.LeakyReLU(negative_slope=0.2), 78 | nn.Conv1d(1024, 1024, 1) 79 | ) 80 | self.reduce_map = nn.Linear(self.trans_dim + 1027, self.trans_dim) 81 | self.build_loss_func() 82 | 83 | def build_loss_func(self): 84 | self.loss_func = ChamferDistanceL1() 85 | 86 | def get_loss(self, ret, gt): 87 | loss_coarse = self.loss_func(ret[0], gt) 88 | loss_fine = self.loss_func(ret[1], gt) 89 | return loss_coarse, loss_fine 90 | 91 | def forward(self, xyz): 92 | q, coarse_point_cloud = self.base_model(xyz) # B M C and B M 3 93 | 94 | B, M ,C = q.shape 95 | 96 | global_feature = self.increase_dim(q.transpose(1,2)).transpose(1,2) # B M 1024 97 | global_feature = torch.max(global_feature, dim=1)[0] # B 1024 98 | 99 | rebuild_feature = torch.cat([ 100 | global_feature.unsqueeze(-2).expand(-1, M, -1), 101 | q, 102 | coarse_point_cloud], dim=-1) # B M 1027 + C 103 | 104 | rebuild_feature = self.reduce_map(rebuild_feature.reshape(B*M, -1)) # BM C 105 | # # NOTE: try to rebuild pc 106 | # coarse_point_cloud = self.refine_coarse(rebuild_feature).reshape(B, M, 3) 107 | 108 | # NOTE: foldingNet 109 | relative_xyz = self.foldingnet(rebuild_feature).reshape(B, M, 3, -1) # B M 3 S 110 | rebuild_points = (relative_xyz + coarse_point_cloud.unsqueeze(-1)).transpose(2,3).reshape(B, -1, 3) # B N 3 111 | 112 | # NOTE: fc 113 | # relative_xyz = self.refine(rebuild_feature) # BM 3S 114 | # rebuild_points = (relative_xyz.reshape(B,M,3,-1) + coarse_point_cloud.unsqueeze(-1)).transpose(2,3).reshape(B, -1, 3) 115 | 116 | # cat the input 117 | inp_sparse = fps(xyz, self.num_query) 118 | coarse_point_cloud = torch.cat([coarse_point_cloud, inp_sparse], dim=1).contiguous() 119 | # print(coarse_point_cloud.size()) 120 | 121 | rebuild_points = rebuild_points.contiguous() 122 | # rebuild_points = torch.cat([rebuild_points, xyz],dim=1).contiguous() 123 | 124 | ret = (coarse_point_cloud, rebuild_points) 125 | return ret 126 | 127 | -------------------------------------------------------------------------------- /other_models/TopNet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import math 4 | import numpy as np 5 | from .build import MODELS 6 | from extensions.chamfer_dist import ChamferDistanceL2 7 | 8 | 9 | 10 | tree_arch = {} 11 | tree_arch[2] = [32, 64] 12 | tree_arch[4] = [4, 8, 8, 8] 13 | tree_arch[6] = [2, 4, 4, 4, 4, 4] 14 | tree_arch[8] = [2, 2, 2, 2, 2, 4, 4, 4] 15 | 16 | def get_arch(nlevels, npts): 17 | logmult = int(math.log2(npts/2048)) 18 | assert 2048*(2**(logmult)) == npts, "Number of points is %d, expected 2048x(2^n)" % (npts) 19 | arch = tree_arch[nlevels] 20 | while logmult > 0: 21 | last_min_pos = np.where(arch==np.min(arch))[0][-1] 22 | arch[last_min_pos]*=2 23 | logmult -= 1 24 | return arch 25 | 26 | @MODELS.register_module() 27 | class TopNet(nn.Module): 28 | def __init__(self, config): # node_feature = 8, encoder_feature = 1024, nlevels = 8, num_pred = 2048 29 | super().__init__() 30 | self.node_feature = config.node_feature 31 | self.encoder_feature = config.encoder_feature 32 | self.nlevels = config.nlevels 33 | self.num_pred = config.num_pred 34 | 35 | self.tarch = get_arch(self.nlevels, self.num_pred) 36 | self.Top_in_channel = self.encoder_feature + self.node_feature 37 | self.Top_out_channel = self.node_feature 38 | self.first_conv = nn.Sequential( 39 | nn.Conv1d(3,128,1), 40 | nn.BatchNorm1d(128), 41 | nn.ReLU(inplace=True), 42 | nn.Conv1d(128,256,1) 43 | ) 44 | self.second_conv = nn.Sequential( 45 | nn.Conv1d(512,512,1), 46 | nn.BatchNorm1d(512), 47 | nn.ReLU(inplace=True), 48 | nn.Conv1d(512,self.encoder_feature,1) 49 | ) 50 | self.root_layer = nn.Sequential( 51 | nn.Linear(self.encoder_feature,256), 52 | nn.BatchNorm1d(256), 53 | nn.ReLU(inplace=True), 54 | nn.Linear(256,64), 55 | nn.BatchNorm1d(64), 56 | nn.ReLU(inplace=True), 57 | nn.Linear(64 , self.node_feature * int(self.tarch[0])), 58 | nn.Tanh() 59 | ) 60 | self.leaf_layer = self.get_tree_layer(self.Top_in_channel, 3, int(self.tarch[-1])) 61 | self.feature_layers = nn.ModuleList([self.get_tree_layer(self.Top_in_channel, self.Top_out_channel, int(self.tarch[d]) ) for d in range(1, self.nlevels - 1)]) 62 | self.build_loss_func() 63 | 64 | def build_loss_func(self): 65 | self.loss_func = ChamferDistanceL2() 66 | 67 | def get_loss(self, ret, gt): 68 | loss_coarse = self.loss_func(ret[0], gt) 69 | loss_fine = self.loss_func(ret[1], gt) 70 | return loss_coarse, loss_fine 71 | 72 | @staticmethod 73 | def get_tree_layer(in_channel, out_channel, node): 74 | return nn.Sequential( 75 | nn.Conv1d(in_channel, in_channel//2, 1), 76 | nn.BatchNorm1d(in_channel//2), 77 | nn.ReLU(inplace=True), 78 | nn.Conv1d(in_channel//2, in_channel//4, 1), 79 | nn.BatchNorm1d(in_channel//4), 80 | nn.ReLU(inplace=True), 81 | nn.Conv1d(in_channel//4, in_channel//8, 1), 82 | nn.BatchNorm1d(in_channel//8), 83 | nn.ReLU(inplace=True), 84 | nn.Conv1d(in_channel//8, out_channel * node, 1), 85 | ) 86 | 87 | def forward(self, xyz): 88 | bs , n , _ = xyz.shape 89 | # encoder 90 | feature = self.first_conv(xyz.transpose(2,1)) # B 256 n 91 | feature_global = torch.max(feature,dim=2,keepdim=True)[0] # B 256 1 92 | feature = torch.cat([feature_global.expand(-1,-1,n), feature], dim=1)# B 512 n 93 | feature = self.second_conv(feature) # B 1024 n 94 | feature_global = torch.max(feature,dim=2,keepdim=False)[0] # B 1024 95 | # decoder 96 | level10 = self.root_layer(feature_global).reshape(-1, self.node_feature ,int(self.tarch[0])) # B 8 node 97 | outs = [level10,] 98 | for i in range(1, self.nlevels): 99 | last_level = outs[-1] 100 | expand_feature = feature_global.unsqueeze(2).expand(-1,-1,last_level.shape[2]) 101 | if i == self.nlevels - 1: 102 | layer_feature = self.leaf_layer(torch.cat([expand_feature,last_level],dim=1)).reshape(bs, 3 ,-1) 103 | else: 104 | layer_feature = self.feature_layers[i-1](torch.cat([expand_feature,last_level],dim=1)).reshape(bs, self.node_feature, -1) 105 | outs.append(nn.Tanh()(layer_feature)) 106 | return (outs[-1].transpose(1,2).contiguous(), outs[-1].transpose(1,2).contiguous()) 107 | -------------------------------------------------------------------------------- /other_models/build.py: -------------------------------------------------------------------------------- 1 | from utils import registry 2 | 3 | 4 | MODELS = registry.Registry('models') 5 | 6 | 7 | def build_model_from_cfg(cfg, **kwargs): 8 | """ 9 | Build a dataset, defined by `dataset_name`. 10 | Args: 11 | cfg (eDICT): 12 | Returns: 13 | Dataset: a constructed dataset specified by dataset_name. 14 | """ 15 | return MODELS.build(cfg, **kwargs) 16 | 17 | 18 | -------------------------------------------------------------------------------- /other_models/dgcnn_group.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | from pointnet2_ops import pointnet2_utils 4 | from knn_cuda import KNN 5 | knn = KNN(k=16, transpose_mode=False) 6 | 7 | 8 | class DGCNN_Grouper(nn.Module): 9 | def __init__(self): 10 | super().__init__() 11 | ''' 12 | K has to be 16 13 | ''' 14 | self.input_trans = nn.Conv1d(3, 8, 1) 15 | 16 | self.layer1 = nn.Sequential(nn.Conv2d(16, 32, kernel_size=1, bias=False), 17 | nn.GroupNorm(4, 32), 18 | nn.LeakyReLU(negative_slope=0.2) 19 | ) 20 | 21 | self.layer2 = nn.Sequential(nn.Conv2d(64, 64, kernel_size=1, bias=False), 22 | nn.GroupNorm(4, 64), 23 | nn.LeakyReLU(negative_slope=0.2) 24 | ) 25 | 26 | self.layer3 = nn.Sequential(nn.Conv2d(128, 64, kernel_size=1, bias=False), 27 | nn.GroupNorm(4, 64), 28 | nn.LeakyReLU(negative_slope=0.2) 29 | ) 30 | 31 | self.layer4 = nn.Sequential(nn.Conv2d(128, 128, kernel_size=1, bias=False), 32 | nn.GroupNorm(4, 128), 33 | nn.LeakyReLU(negative_slope=0.2) 34 | ) 35 | 36 | 37 | @staticmethod 38 | def fps_downsample(coor, x, num_group): 39 | xyz = coor.transpose(1, 2).contiguous() # b, n, 3 40 | fps_idx = pointnet2_utils.furthest_point_sample(xyz, num_group) 41 | 42 | combined_x = torch.cat([coor, x], dim=1) 43 | 44 | new_combined_x = ( 45 | pointnet2_utils.gather_operation( 46 | combined_x, fps_idx 47 | ) 48 | ) 49 | 50 | new_coor = new_combined_x[:, :3] 51 | new_x = new_combined_x[:, 3:] 52 | 53 | return new_coor, new_x 54 | 55 | @staticmethod 56 | def get_graph_feature(coor_q, x_q, coor_k, x_k): 57 | 58 | # coor: bs, 3, np, x: bs, c, np 59 | 60 | k = 16 61 | batch_size = x_k.size(0) 62 | num_points_k = x_k.size(2) 63 | num_points_q = x_q.size(2) 64 | 65 | with torch.no_grad(): 66 | _, idx = knn(coor_k, coor_q) # bs k np 67 | assert idx.shape[1] == k 68 | idx_base = torch.arange(0, batch_size, device=x_q.device).view(-1, 1, 1) * num_points_k 69 | idx = idx + idx_base 70 | idx = idx.view(-1) 71 | num_dims = x_k.size(1) 72 | x_k = x_k.transpose(2, 1).contiguous() 73 | feature = x_k.view(batch_size * num_points_k, -1)[idx, :] 74 | feature = feature.view(batch_size, k, num_points_q, num_dims).permute(0, 3, 2, 1).contiguous() 75 | x_q = x_q.view(batch_size, num_dims, num_points_q, 1).expand(-1, -1, -1, k) 76 | feature = torch.cat((feature - x_q, x_q), dim=1) 77 | return feature 78 | 79 | def forward(self, x): 80 | 81 | # x: bs, 3, np 82 | 83 | # bs 3 N(128) bs C(224)128 N(128) 84 | _, _, N = x.shape 85 | coor = x 86 | f = self.input_trans(x) 87 | 88 | f = self.get_graph_feature(coor, f, coor, f) 89 | f = self.layer1(f) 90 | f = f.max(dim=-1, keepdim=False)[0] 91 | 92 | coor_q, f_q = self.fps_downsample(coor, f, N // 4) 93 | f = self.get_graph_feature(coor_q, f_q, coor, f) 94 | f = self.layer2(f) 95 | f = f.max(dim=-1, keepdim=False)[0] 96 | coor = coor_q 97 | 98 | f = self.get_graph_feature(coor, f, coor, f) 99 | f = self.layer3(f) 100 | f = f.max(dim=-1, keepdim=False)[0] 101 | 102 | coor_q, f_q = self.fps_downsample(coor, f, N // 16) 103 | f = self.get_graph_feature(coor_q, f_q, coor, f) 104 | f = self.layer4(f) 105 | f = f.max(dim=-1, keepdim=False)[0] 106 | coor = coor_q 107 | 108 | return coor, f -------------------------------------------------------------------------------- /our_models/PT_model.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | 4 | from pointnet2_ops import pointnet2_utils 5 | from extensions.chamfer_dist import ChamferDistanceL1 6 | from our_models.PT_Transformer_Coarse import PCTransformer 7 | from models.build import MODELS 8 | 9 | 10 | def fps(pc, num): 11 | fps_idx = pointnet2_utils.furthest_point_sample(pc, num) 12 | sub_pc = pointnet2_utils.gather_operation(pc.transpose(1, 2).contiguous(), fps_idx).transpose(1,2).contiguous() 13 | return sub_pc 14 | 15 | 16 | class Fold(nn.Module): 17 | def __init__(self, in_channel , step , hidden_dim = 512): 18 | super().__init__() 19 | 20 | self.in_channel = in_channel 21 | self.step = step 22 | 23 | a = torch.linspace(-1., 1., steps=step, dtype=torch.float).view(1, step).expand(step, step).reshape(1, -1) 24 | b = torch.linspace(-1., 1., steps=step, dtype=torch.float).view(step, 1).expand(step, step).reshape(1, -1) 25 | self.folding_seed = torch.cat([a, b], dim=0).cuda() 26 | 27 | self.folding1 = nn.Sequential( 28 | nn.Conv1d(in_channel + 2, hidden_dim, 1), 29 | nn.BatchNorm1d(hidden_dim), 30 | nn.ReLU(inplace=True), 31 | nn.Conv1d(hidden_dim, hidden_dim//2, 1), 32 | nn.BatchNorm1d(hidden_dim//2), 33 | nn.ReLU(inplace=True), 34 | nn.Conv1d(hidden_dim//2, 3, 1), 35 | ) 36 | 37 | self.folding2 = nn.Sequential( 38 | nn.Conv1d(in_channel + 3, hidden_dim, 1), 39 | nn.BatchNorm1d(hidden_dim), 40 | nn.ReLU(inplace=True), 41 | nn.Conv1d(hidden_dim, hidden_dim//2, 1), 42 | nn.BatchNorm1d(hidden_dim//2), 43 | nn.ReLU(inplace=True), 44 | nn.Conv1d(hidden_dim//2, 3, 1), 45 | ) 46 | 47 | def forward(self, x): 48 | num_sample = self.step * self.step 49 | bs = x.size(0) 50 | features = x.view(bs, self.in_channel, 1).expand(bs, self.in_channel, num_sample) 51 | seed = self.folding_seed.view(1, 2, num_sample).expand(bs, 2, num_sample).to(x.device) 52 | 53 | x = torch.cat([seed, features], dim=1) 54 | fd1 = self.folding1(x) 55 | x = torch.cat([fd1, features], dim=1) 56 | fd2 = self.folding2(x) 57 | 58 | return fd2 59 | 60 | @MODELS.register_module() 61 | class PT_model(nn.Module): 62 | def __init__(self, config, **kwargs): 63 | super().__init__() 64 | self.trans_dim = config.trans_dim 65 | self.knn_layer = config.knn_layer 66 | self.num_pred = config.num_pred 67 | self.num_query = config.num_query 68 | 69 | self.fold_step = int(pow(self.num_pred//self.num_query, 0.5) + 0.5) 70 | self.base_model = PCTransformer(in_chans = 3, embed_dim = self.trans_dim, depth = [6, 8], drop_rate = 0., num_query = self.num_query, knn_layer = self.knn_layer) 71 | 72 | self.foldingnet = Fold(self.trans_dim, step = self.fold_step, hidden_dim = 256) # rebuild a cluster point 73 | 74 | self.increase_dim = nn.Sequential( 75 | nn.Conv1d(self.trans_dim, 1024, 1), 76 | nn.BatchNorm1d(1024), 77 | nn.LeakyReLU(negative_slope=0.2), 78 | nn.Conv1d(1024, 1024, 1) 79 | ) 80 | self.reduce_map = nn.Linear(1027, self.trans_dim) 81 | self.build_loss_func() 82 | 83 | def build_loss_func(self): 84 | self.loss_func = ChamferDistanceL1() 85 | 86 | def get_loss(self, ret, gt): 87 | loss_coarse = self.loss_func(ret[0], gt) 88 | loss_fine = self.loss_func(ret[1], gt) 89 | return loss_coarse, loss_fine 90 | 91 | def forward(self, xyz): 92 | coarse, point_feature = self.base_model(xyz) # B 128 3, B 128 512 93 | # print(coarse.size()) 94 | # print(fps_xyz.size()) 95 | 96 | B, M, _ = coarse.shape 97 | 98 | # global_feature = self.increase_dim(point_feature.transpose(1,2)).transpose(1,2) # B M 1024 99 | # global_feature = torch.max(global_feature, dim=1)[0] # B 1024 100 | 101 | # print(global_feature.size()) 102 | # print(point_feature.size()) 103 | # print(coarse.size()) 104 | 105 | rebuild_feature = torch.cat([ 106 | point_feature.unsqueeze(-2).expand(-1, M, -1), 107 | # point_feature, 108 | coarse], dim=-1) # B M 1027 + C 109 | 110 | rebuild_feature = self.reduce_map(rebuild_feature.reshape(B*M, -1)) # B*M C 111 | # # NOTE: try to rebuild pc 112 | 113 | # NOTE: foldingNet 114 | relative_xyz = self.foldingnet(rebuild_feature).reshape(B, M, 3, -1) # B M 3 S 115 | rebuild_points = (relative_xyz + coarse.unsqueeze(-1)).transpose(2,3).reshape(B, -1, 3) # B N 3 116 | # print(rebuild_points.size()) 117 | # cat the input 118 | # inp_sparse = fps(xyz, self.num_query) 119 | # coarse_point_cloud = torch.cat([coarse_point_cloud, inp_sparse], dim=1).contiguous() 120 | # rebuild_points = torch.cat([rebuild_points, xyz],dim=1).contiguous() 121 | # print(coarse.size()) 122 | # print(rebuild_points.size()) 123 | ret = (coarse, rebuild_points) 124 | return ret 125 | 126 | 127 | 128 | -------------------------------------------------------------------------------- /our_models/PoinTr_Encoder.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | 4 | from pointnet2_ops import pointnet2_utils 5 | from extensions.chamfer_dist import ChamferDistanceL1 6 | from our_models.Transformer_Coarse import PCTransformer 7 | from models.build import MODELS 8 | 9 | 10 | def fps(pc, num): 11 | fps_idx = pointnet2_utils.furthest_point_sample(pc, num) 12 | sub_pc = pointnet2_utils.gather_operation(pc.transpose(1, 2).contiguous(), fps_idx).transpose(1,2).contiguous() 13 | return sub_pc 14 | 15 | 16 | class Fold(nn.Module): 17 | def __init__(self, in_channel , step , hidden_dim = 512): 18 | super().__init__() 19 | 20 | self.in_channel = in_channel 21 | self.step = step 22 | 23 | a = torch.linspace(-1., 1., steps=step, dtype=torch.float).view(1, step).expand(step, step).reshape(1, -1) 24 | b = torch.linspace(-1., 1., steps=step, dtype=torch.float).view(step, 1).expand(step, step).reshape(1, -1) 25 | self.folding_seed = torch.cat([a, b], dim=0).cuda() 26 | 27 | self.folding1 = nn.Sequential( 28 | nn.Conv1d(in_channel + 2, hidden_dim, 1), 29 | nn.BatchNorm1d(hidden_dim), 30 | nn.ReLU(inplace=True), 31 | nn.Conv1d(hidden_dim, hidden_dim//2, 1), 32 | nn.BatchNorm1d(hidden_dim//2), 33 | nn.ReLU(inplace=True), 34 | nn.Conv1d(hidden_dim//2, 3, 1), 35 | ) 36 | 37 | self.folding2 = nn.Sequential( 38 | nn.Conv1d(in_channel + 3, hidden_dim, 1), 39 | nn.BatchNorm1d(hidden_dim), 40 | nn.ReLU(inplace=True), 41 | nn.Conv1d(hidden_dim, hidden_dim//2, 1), 42 | nn.BatchNorm1d(hidden_dim//2), 43 | nn.ReLU(inplace=True), 44 | nn.Conv1d(hidden_dim//2, 3, 1), 45 | ) 46 | 47 | def forward(self, x): 48 | num_sample = self.step * self.step 49 | bs = x.size(0) 50 | features = x.view(bs, self.in_channel, 1).expand(bs, self.in_channel, num_sample) 51 | seed = self.folding_seed.view(1, 2, num_sample).expand(bs, 2, num_sample).to(x.device) 52 | 53 | x = torch.cat([seed, features], dim=1) 54 | fd1 = self.folding1(x) 55 | x = torch.cat([fd1, features], dim=1) 56 | fd2 = self.folding2(x) 57 | 58 | return fd2 59 | 60 | @MODELS.register_module() 61 | class PoinTr_Encoder(nn.Module): 62 | def __init__(self, config, **kwargs): 63 | super().__init__() 64 | self.trans_dim = config.trans_dim 65 | self.knn_layer = config.knn_layer 66 | self.num_pred = config.num_pred 67 | self.num_query = config.num_query 68 | 69 | self.fold_step = int(pow(self.num_pred//self.num_query, 0.5) + 0.5) 70 | self.base_model = PCTransformer(in_chans = 3, embed_dim = self.trans_dim, depth = [6, 8], drop_rate = 0., num_query = self.num_query, knn_layer = self.knn_layer) 71 | 72 | self.foldingnet = Fold(self.trans_dim, step = self.fold_step, hidden_dim = 256) # rebuild a cluster point 73 | 74 | self.increase_dim = nn.Sequential( 75 | nn.Conv1d(self.trans_dim, 1024, 1), 76 | nn.BatchNorm1d(1024), 77 | nn.LeakyReLU(negative_slope=0.2), 78 | nn.Conv1d(1024, 1024, 1) 79 | ) 80 | self.reduce_map = nn.Linear(self.trans_dim + 1027, self.trans_dim) 81 | self.build_loss_func() 82 | 83 | def build_loss_func(self): 84 | self.loss_func = ChamferDistanceL1() 85 | self.loss_f = nn.MSELoss() 86 | 87 | def get_loss(self, ret, gt): 88 | loss_coarse = self.loss_func(ret[0], gt) 89 | loss_fine = self.loss_func(ret[1], gt) 90 | return loss_coarse, loss_fine 91 | 92 | def get_feature_loss(self, miss_f, pre_miss_f): 93 | loss_f = self.loss_f(pre_miss_f, miss_f) 94 | return loss_f 95 | 96 | def forward(self, xyz, missing_part=None): 97 | if missing_part is None: 98 | y, coarse_point_cloud = self.base_model(xyz) # B M C and B M 3 99 | 100 | B, M, _ = coarse_point_cloud.shape 101 | 102 | global_feature = self.increase_dim(y.transpose(1,2)).transpose(1,2) # B M 1024 103 | global_feature = torch.max(global_feature, dim=1)[0] # B 1024 104 | 105 | rebuild_feature = torch.cat([ 106 | global_feature.unsqueeze(-2).expand(-1, M, -1), 107 | y, 108 | coarse_point_cloud], dim=-1) # B M 1027 + C 109 | 110 | rebuild_feature = self.reduce_map(rebuild_feature.reshape(B*M, -1)) # BM C 111 | # # NOTE: try to rebuild pc 112 | # coarse_point_cloud = self.refine_coarse(rebuild_feature).reshape(B, M, 3) 113 | 114 | # NOTE: foldingNet 115 | relative_xyz = self.foldingnet(rebuild_feature).reshape(B, M, 3, -1) # B M 3 S 116 | rebuild_points = (relative_xyz + coarse_point_cloud.unsqueeze(-1)).transpose(2,3).reshape(B, -1, 3) # B N 3 117 | 118 | # NOTE: fc 119 | # relative_xyz = self.refine(rebuild_feature) # BM 3S 120 | # rebuild_points = (relative_xyz.reshape(B,M,3,-1) + coarse_point_cloud.unsqueeze(-1)).transpose(2,3).reshape(B, -1, 3) 121 | 122 | # cat the input 123 | inp_sparse = fps(xyz, self.num_query) 124 | coarse_point_cloud = torch.cat([coarse_point_cloud, inp_sparse], dim=1).contiguous() 125 | rebuild_points = torch.cat([rebuild_points, xyz],dim=1).contiguous() 126 | 127 | ret = (coarse_point_cloud, rebuild_points) 128 | return ret, y 129 | else: 130 | x = self.base_model(xyz, missing_part) 131 | return x 132 | 133 | -------------------------------------------------------------------------------- /our_models/PointTransformer.py: -------------------------------------------------------------------------------- 1 | from our_models.pointnet_util import index_points, square_distance 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | import numpy as np 6 | 7 | class TransformerBlock(nn.Module): 8 | def __init__(self, d_points, d_model, k) -> None: 9 | super().__init__() 10 | self.fc1 = nn.Linear(d_points, d_model) 11 | self.fc2 = nn.Linear(d_model, d_points) 12 | self.fc_delta = nn.Sequential( 13 | nn.Linear(3, d_model), 14 | nn.ReLU(), 15 | nn.Linear(d_model, d_model) 16 | ) 17 | self.fc_gamma = nn.Sequential( 18 | nn.Linear(d_model, d_model), 19 | nn.ReLU(), 20 | nn.Linear(d_model, d_model) 21 | ) 22 | self.w_qs = nn.Linear(d_model, d_model, bias=False) 23 | self.w_ks = nn.Linear(d_model, d_model, bias=False) 24 | self.w_vs = nn.Linear(d_model, d_model, bias=False) 25 | self.k = k 26 | 27 | # xyz: b x n x 3, features: b x n x f 28 | def forward(self, xyz, features): 29 | dists = square_distance(xyz, xyz) 30 | knn_idx = dists.argsort()[:, :, :self.k] # b x n x k 31 | knn_xyz = index_points(xyz, knn_idx) 32 | 33 | pre = features 34 | x = self.fc1(features) 35 | q, k, v = self.w_qs(x), index_points(self.w_ks(x), knn_idx), index_points(self.w_vs(x), knn_idx) 36 | 37 | pos_enc = self.fc_delta(xyz[:, :, None] - knn_xyz) # b x n x k x f 38 | 39 | attn = self.fc_gamma(q[:, :, None] - k + pos_enc) 40 | attn = F.softmax(attn / np.sqrt(k.size(-1)), dim=-2) # b x n x k x f 41 | 42 | res = torch.einsum('bmnf,bmnf->bmf', attn, v + pos_enc) 43 | res = self.fc2(res) + pre 44 | return res, attn 45 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | # This file may be used to create an environment using: 2 | # $ conda create --name --file 3 | # platform: linux-64 4 | _libgcc_mutex=0.1=main 5 | _openmp_mutex=4.5=1_gnu 6 | absl-py=1.0.0=pypi_0 7 | argon2-cffi=21.3.0=pypi_0 8 | argon2-cffi-bindings=21.2.0=pypi_0 9 | argparse=1.4.0=pypi_0 10 | attrs=21.4.0=pypi_0 11 | backcall=0.2.0=pypi_0 12 | beautifulsoup4=4.11.1=pypi_0 13 | blas=1.0=mkl 14 | bleach=5.0.0=pypi_0 15 | brotlipy=0.7.0=py37h27cfd23_1003 16 | bzip2=1.0.8=h7b6447c_0 17 | ca-certificates=2022.4.26=h06a4308_0 18 | cached-property=1.5.2=pypi_0 19 | cachetools=5.1.0=pypi_0 20 | certifi=2021.10.8=py37h06a4308_2 21 | cffi=1.15.0=py37hd667e15_1 22 | charset-normalizer=2.0.4=pyhd3eb1b0_0 23 | cryptography=37.0.1=py37h9ce1e76_0 24 | cudatoolkit=11.3.1=h2bc3f7f_2 25 | cycler=0.11.0=pypi_0 26 | debugpy=1.6.0=pypi_0 27 | decorator=5.1.1=pypi_0 28 | defusedxml=0.7.1=pypi_0 29 | easydict=1.9=pypi_0 30 | einops=0.4.1=pypi_0 31 | entrypoints=0.4=pypi_0 32 | fastjsonschema=2.15.3=pypi_0 33 | ffmpeg=4.3=hf484d3e_0 34 | fonttools=4.33.3=pypi_0 35 | freetype=2.11.0=h70c0345_0 36 | future=0.18.2=pypi_0 37 | fvcore=0.1.5.post20220512=pypi_0 38 | giflib=5.2.1=h7b6447c_0 39 | gmp=6.2.1=h2531618_2 40 | gnutls=3.6.15=he1e5248_0 41 | google-auth=2.6.6=pypi_0 42 | google-auth-oauthlib=0.4.6=pypi_0 43 | grpcio=1.46.1=pypi_0 44 | h5py=3.6.0=pypi_0 45 | idna=3.3=pyhd3eb1b0_0 46 | importlib-metadata=4.11.3=pypi_0 47 | importlib-resources=5.7.1=pypi_0 48 | intel-openmp=2021.4.0=h06a4308_3561 49 | iopath=0.1.10=pypi_0 50 | ipykernel=6.13.0=pypi_0 51 | ipython=7.33.0=pypi_0 52 | ipython-genutils=0.2.0=pypi_0 53 | ipywidgets=7.7.0=pypi_0 54 | jedi=0.18.1=pypi_0 55 | jinja2=3.1.2=pypi_0 56 | joblib=1.1.0=pypi_0 57 | jpeg=9e=h7f8727e_0 58 | jsonschema=4.5.1=pypi_0 59 | jupyter-client=7.3.1=pypi_0 60 | jupyter-core=4.10.0=pypi_0 61 | jupyterlab-pygments=0.2.2=pypi_0 62 | jupyterlab-widgets=1.1.0=pypi_0 63 | kiwisolver=1.4.2=pypi_0 64 | knn-cuda=0.2=pypi_0 65 | lame=3.100=h7b6447c_0 66 | lcms2=2.12=h3be6417_0 67 | ld_impl_linux-64=2.38=h1181459_0 68 | libffi=3.3=he6710b0_2 69 | libgcc-ng=9.3.0=h5101ec6_17 70 | libgomp=9.3.0=h5101ec6_17 71 | libiconv=1.16=h7f8727e_2 72 | libidn2=2.3.2=h7f8727e_0 73 | libpng=1.6.37=hbc83047_0 74 | libstdcxx-ng=9.3.0=hd4cf53a_17 75 | libtasn1=4.16.0=h27cfd23_0 76 | libtiff=4.2.0=h85742a9_0 77 | libunistring=0.9.10=h27cfd23_0 78 | libuv=1.40.0=h7b6447c_0 79 | libwebp=1.2.2=h55f646e_0 80 | libwebp-base=1.2.2=h7f8727e_0 81 | lz4-c=1.9.3=h295c915_1 82 | markdown=3.3.7=pypi_0 83 | markupsafe=2.1.1=pypi_0 84 | matplotlib=3.5.2=pypi_0 85 | matplotlib-inline=0.1.3=pypi_0 86 | mistune=0.8.4=pypi_0 87 | mkl=2021.4.0=h06a4308_640 88 | mkl-service=2.4.0=py37h7f8727e_0 89 | mkl_fft=1.3.1=py37hd3c417c_0 90 | mkl_random=1.2.2=py37h51133e4_0 91 | nbclient=0.6.3=pypi_0 92 | nbconvert=6.5.0=pypi_0 93 | nbformat=5.4.0=pypi_0 94 | ncurses=6.3=h7f8727e_2 95 | nest-asyncio=1.5.5=pypi_0 96 | nettle=3.7.3=hbbd107a_1 97 | ninja=1.10.2.3=pypi_0 98 | notebook=6.4.11=pypi_0 99 | numpy=1.21.5=py37he7a7128_2 100 | numpy-base=1.21.5=py37hf524024_2 101 | oauthlib=3.2.0=pypi_0 102 | open3d=0.9.0.0=pypi_0 103 | opencv-python=4.5.5.64=pypi_0 104 | openexr=1.3.9=pypi_0 105 | openh264=2.1.1=h4ff587b_0 106 | openssl=1.1.1o=h7f8727e_0 107 | packaging=21.3=pypi_0 108 | pandas=1.3.5=pypi_0 109 | pandocfilters=1.5.0=pypi_0 110 | parso=0.8.3=pypi_0 111 | pexpect=4.8.0=pypi_0 112 | pickleshare=0.7.5=pypi_0 113 | pillow=9.0.1=py37h22f2fdc_0 114 | pip=21.2.2=py37h06a4308_0 115 | plyfile=0.7.4=pypi_0 116 | pointnet2-ops=3.0.0=pypi_0 117 | portalocker=2.5.1=pypi_0 118 | prometheus-client=0.14.1=pypi_0 119 | prompt-toolkit=3.0.29=pypi_0 120 | protobuf=3.20.1=pypi_0 121 | psutil=5.9.0=pypi_0 122 | ptyprocess=0.7.0=pypi_0 123 | pyasn1=0.4.8=pypi_0 124 | pyasn1-modules=0.2.8=pypi_0 125 | pycparser=2.21=pyhd3eb1b0_0 126 | pyexr=0.3.10=pypi_0 127 | pygments=2.12.0=pypi_0 128 | pyntcloud=0.3.1=pypi_0 129 | pyopenssl=22.0.0=pyhd3eb1b0_0 130 | pyparsing=3.0.9=pypi_0 131 | pyrsistent=0.18.1=pypi_0 132 | pysocks=1.7.1=py37_1 133 | python=3.7.13=h12debd9_0 134 | python-dateutil=2.8.2=pypi_0 135 | pytorch=1.11.0=py3.7_cuda11.3_cudnn8.2.0_0 136 | pytorch-mutex=1.0=cuda 137 | pytorch3d=0.6.2=pypi_0 138 | pytz=2022.1=pypi_0 139 | pyyaml=6.0=pypi_0 140 | pyzmq=22.3.0=pypi_0 141 | readline=8.1.2=h7f8727e_1 142 | requests=2.27.1=pyhd3eb1b0_0 143 | requests-oauthlib=1.3.1=pypi_0 144 | rsa=4.8=pypi_0 145 | scikit-learn=1.0.2=pypi_0 146 | scipy=1.7.3=pypi_0 147 | send2trash=1.8.0=pypi_0 148 | setuptools=61.2.0=py37h06a4308_0 149 | six=1.16.0=pyhd3eb1b0_1 150 | soupsieve=2.3.2.post1=pypi_0 151 | sqlite=3.38.3=hc218d9a_0 152 | tabulate=0.8.10=pypi_0 153 | tensorboard=2.9.0=pypi_0 154 | tensorboard-data-server=0.6.1=pypi_0 155 | tensorboard-plugin-wit=1.8.1=pypi_0 156 | tensorboardx=2.5=pypi_0 157 | termcolor=1.1.0=pypi_0 158 | terminado=0.15.0=pypi_0 159 | threadpoolctl=3.1.0=pypi_0 160 | timm=0.4.5=pypi_0 161 | tinycss2=1.1.1=pypi_0 162 | tk=8.6.11=h1ccaba5_1 163 | torch-tb-profiler=0.4.0=pypi_0 164 | torchaudio=0.11.0=py37_cu113 165 | torchstat=0.0.7=pypi_0 166 | torchvision=0.12.0=py37_cu113 167 | tornado=6.1=pypi_0 168 | tqdm=4.64.0=pypi_0 169 | traitlets=5.2.1.post0=pypi_0 170 | transforms3d=0.3.1=pypi_0 171 | typing_extensions=4.1.1=pyh06a4308_0 172 | urllib3=1.26.9=py37h06a4308_0 173 | wcwidth=0.2.5=pypi_0 174 | webencodings=0.5.1=pypi_0 175 | werkzeug=2.1.2=pypi_0 176 | wheel=0.37.1=pyhd3eb1b0_0 177 | widgetsnbextension=3.6.0=pypi_0 178 | xz=5.2.5=h7f8727e_1 179 | yacs=0.1.8=pypi_0 180 | zipp=3.8.0=pypi_0 181 | zlib=1.2.12=h7f8727e_2 182 | zstd=1.4.9=haebb681_0 183 | -------------------------------------------------------------------------------- /scripts/dist_train.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | set -x 4 | NGPUS=$1 5 | PORT=$2 6 | PY_ARGS=${@:3} 7 | 8 | python -m torch.distributed.launch --master_port=${PORT} --nproc_per_node=${NGPUS} main.py --launcher pytorch --sync_bn ${PY_ARGS} 9 | -------------------------------------------------------------------------------- /scripts/test.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | set -x 4 | GPUS=$1 5 | PY_ARGS=${@:2} 6 | 7 | CUDA_VISIBLE_DEVICES=${GPUS} python main.py --test ${PY_ARGS} -------------------------------------------------------------------------------- /scripts/train.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | set -x 4 | GPUS=$1 5 | 6 | PY_ARGS=${@:2} 7 | 8 | CUDA_VISIBLE_DEVICES=${GPUS} python main.py ${PY_ARGS} 9 | -------------------------------------------------------------------------------- /tools/__init__.py: -------------------------------------------------------------------------------- 1 | from .runner_model import run_net 2 | from .runner_model import test_net -------------------------------------------------------------------------------- /tools/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/tools/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /tools/__pycache__/__init__.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/tools/__pycache__/__init__.cpython-39.pyc -------------------------------------------------------------------------------- /tools/__pycache__/builder.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/tools/__pycache__/builder.cpython-37.pyc -------------------------------------------------------------------------------- /tools/__pycache__/runner.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/tools/__pycache__/runner.cpython-37.pyc -------------------------------------------------------------------------------- /tools/__pycache__/runner2.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/tools/__pycache__/runner2.cpython-37.pyc -------------------------------------------------------------------------------- /tools/__pycache__/runner3.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/tools/__pycache__/runner3.cpython-37.pyc -------------------------------------------------------------------------------- /tools/__pycache__/runner_model.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/tools/__pycache__/runner_model.cpython-37.pyc -------------------------------------------------------------------------------- /tools/__pycache__/runner_reconstruct.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/tools/__pycache__/runner_reconstruct.cpython-37.pyc -------------------------------------------------------------------------------- /tools/__pycache__/runner_reconstruct.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/tools/__pycache__/runner_reconstruct.cpython-39.pyc -------------------------------------------------------------------------------- /utils/AverageMeter.py: -------------------------------------------------------------------------------- 1 | 2 | class AverageMeter(object): 3 | def __init__(self, items=None): 4 | self.items = items 5 | self.n_items = 1 if items is None else len(items) 6 | self.reset() 7 | 8 | def reset(self): 9 | self._val = [0] * self.n_items 10 | self._sum = [0] * self.n_items 11 | self._count = [0] * self.n_items 12 | 13 | def update(self, values): 14 | if type(values).__name__ == 'list': 15 | for idx, v in enumerate(values): 16 | self._val[idx] = v 17 | self._sum[idx] += v 18 | self._count[idx] += 1 19 | else: 20 | self._val[0] = values 21 | self._sum[0] += values 22 | self._count[0] += 1 23 | 24 | def val(self, idx=None): 25 | if idx is None: 26 | return self._val[0] if self.items is None else [self._val[i] for i in range(self.n_items)] 27 | else: 28 | return self._val[idx] 29 | 30 | def count(self, idx=None): 31 | if idx is None: 32 | return self._count[0] if self.items is None else [self._count[i] for i in range(self.n_items)] 33 | else: 34 | return self._count[idx] 35 | 36 | def avg(self, idx=None): 37 | if idx is None: 38 | return self._sum[0] / self._count[0] if self.items is None else [ 39 | self._sum[i] / self._count[i] for i in range(self.n_items) 40 | ] 41 | else: 42 | return self._sum[idx] / self._count[idx] -------------------------------------------------------------------------------- /utils/__pycache__/AverageMeter.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/utils/__pycache__/AverageMeter.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/config.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/utils/__pycache__/config.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/dist_utils.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/utils/__pycache__/dist_utils.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/eulerangles.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/utils/__pycache__/eulerangles.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/logger.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/utils/__pycache__/logger.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/metrics.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/utils/__pycache__/metrics.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/misc.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/utils/__pycache__/misc.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/other_utils.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/utils/__pycache__/other_utils.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/parser.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/utils/__pycache__/parser.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/pc_util.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/utils/__pycache__/pc_util.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/plyfile.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/utils/__pycache__/plyfile.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/registry.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/utils/__pycache__/registry.cpython-37.pyc -------------------------------------------------------------------------------- /utils/config.py: -------------------------------------------------------------------------------- 1 | import yaml 2 | from easydict import EasyDict 3 | import os 4 | from .logger import print_log 5 | 6 | def log_args_to_file(args, pre='args', logger=None): 7 | for key, val in args.__dict__.items(): 8 | print_log(f'{pre}.{key} : {val}', logger = logger) 9 | 10 | def log_config_to_file(cfg, pre='cfg', logger=None): 11 | for key, val in cfg.items(): 12 | if isinstance(cfg[key], EasyDict): 13 | print_log(f'{pre}.{key} = edict()', logger = logger) 14 | log_config_to_file(cfg[key], pre=pre + '.' + key, logger=logger) 15 | continue 16 | print_log(f'{pre}.{key} : {val}', logger = logger) 17 | 18 | def merge_new_config(config, new_config): 19 | for key, val in new_config.items(): 20 | if not isinstance(val, dict): 21 | if key == '_base_': 22 | with open(new_config['_base_'], 'r') as f: 23 | try: 24 | val = yaml.load(f, Loader=yaml.FullLoader) 25 | except: 26 | val = yaml.load(f) 27 | config[key] = EasyDict() 28 | merge_new_config(config[key], val) 29 | else: 30 | config[key] = val 31 | continue 32 | if key not in config: 33 | config[key] = EasyDict() 34 | merge_new_config(config[key], val) 35 | return config 36 | 37 | def cfg_from_yaml_file(cfg_file): 38 | config = EasyDict() 39 | with open(cfg_file, 'r') as f: 40 | try: 41 | new_config = yaml.load(f, Loader=yaml.FullLoader) 42 | except: 43 | new_config = yaml.load(f) 44 | merge_new_config(config=config, new_config=new_config) 45 | return config 46 | 47 | def get_config(args, logger=None): 48 | if args.resume: 49 | cfg_path = os.path.join(args.experiment_path, 'config.yaml') 50 | if not os.path.exists(cfg_path): 51 | print_log("Failed to resume", logger = logger) 52 | raise FileNotFoundError() 53 | print_log(f'Resume yaml from {cfg_path}', logger = logger) 54 | args.config = cfg_path 55 | config = cfg_from_yaml_file(args.config) 56 | if not args.resume and args.local_rank == 0: 57 | save_experiment_config(args, config, logger) 58 | return config 59 | 60 | def save_experiment_config(args, config, logger = None): 61 | config_path = os.path.join(args.experiment_path, 'config.yaml') 62 | os.system('cp %s %s' % (args.config, config_path)) 63 | print_log(f'Copy the Config file from {args.config} to {config_path}',logger = logger ) -------------------------------------------------------------------------------- /utils/dist_utils.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import torch 4 | import torch.multiprocessing as mp 5 | from torch import distributed as dist 6 | 7 | 8 | 9 | def init_dist(launcher, backend='nccl', **kwargs): 10 | if mp.get_start_method(allow_none=True) is None: 11 | mp.set_start_method('spawn') 12 | if launcher == 'pytorch': 13 | _init_dist_pytorch(backend, **kwargs) 14 | else: 15 | raise ValueError(f'Invalid launcher type: {launcher}') 16 | 17 | 18 | def _init_dist_pytorch(backend, **kwargs): 19 | # TODO: use local_rank instead of rank % num_gpus 20 | rank = int(os.environ['RANK']) 21 | num_gpus = torch.cuda.device_count() 22 | torch.cuda.set_device(rank % num_gpus) 23 | dist.init_process_group(backend=backend, **kwargs) 24 | print(f'init distributed in rank {torch.distributed.get_rank()}') 25 | 26 | 27 | def get_dist_info(): 28 | if dist.is_available(): 29 | initialized = dist.is_initialized() 30 | else: 31 | initialized = False 32 | if initialized: 33 | rank = dist.get_rank() 34 | world_size = dist.get_world_size() 35 | else: 36 | rank = 0 37 | world_size = 1 38 | return rank, world_size 39 | 40 | 41 | def reduce_tensor(tensor, args): 42 | ''' 43 | for acc kind, get the mean in each gpu 44 | ''' 45 | rt = tensor.clone() 46 | torch.distributed.all_reduce(rt, op=torch.distributed.ReduceOp.SUM) 47 | rt /= args.world_size 48 | return rt 49 | 50 | def gather_tensor(tensor, args): 51 | output_tensors = [tensor.clone() for _ in range(args.world_size)] 52 | torch.distributed.all_gather(output_tensors, tensor) 53 | concat = torch.cat(output_tensors, dim=0) 54 | return concat 55 | -------------------------------------------------------------------------------- /utils/logger.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import torch.distributed as dist 3 | 4 | logger_initialized = {} 5 | 6 | def get_root_logger(log_file=None, log_level=logging.INFO, name='main'): 7 | """Get root logger and add a keyword filter to it. 8 | The logger will be initialized if it has not been initialized. By default a 9 | StreamHandler will be added. If `log_file` is specified, a FileHandler will 10 | also be added. The name of the root logger is the top-level package name, 11 | e.g., "mmdet3d". 12 | Args: 13 | log_file (str, optional): File path of log. Defaults to None. 14 | log_level (int, optional): The level of logger. 15 | Defaults to logging.INFO. 16 | name (str, optional): The name of the root logger, also used as a 17 | filter keyword. Defaults to 'mmdet3d'. 18 | Returns: 19 | :obj:`logging.Logger`: The obtained logger 20 | """ 21 | logger = get_logger(name=name, log_file=log_file, log_level=log_level) 22 | # add a logging filter 23 | logging_filter = logging.Filter(name) 24 | logging_filter.filter = lambda record: record.find(name) != -1 25 | 26 | return logger 27 | 28 | 29 | def get_logger(name, log_file=None, log_level=logging.INFO, file_mode='w'): 30 | """Initialize and get a logger by name. 31 | If the logger has not been initialized, this method will initialize the 32 | logger by adding one or two handlers, otherwise the initialized logger will 33 | be directly returned. During initialization, a StreamHandler will always be 34 | added. If `log_file` is specified and the process rank is 0, a FileHandler 35 | will also be added. 36 | Args: 37 | name (str): Logger name. 38 | log_file (str | None): The log filename. If specified, a FileHandler 39 | will be added to the logger. 40 | log_level (int): The logger level. Note that only the process of 41 | rank 0 is affected, and other processes will set the level to 42 | "Error" thus be silent most of the time. 43 | file_mode (str): The file mode used in opening log file. 44 | Defaults to 'w'. 45 | Returns: 46 | logging.Logger: The expected logger. 47 | """ 48 | logger = logging.getLogger(name) 49 | if name in logger_initialized: 50 | return logger 51 | # handle hierarchical names 52 | # e.g., logger "a" is initialized, then logger "a.b" will skip the 53 | # initialization since it is a child of "a". 54 | for logger_name in logger_initialized: 55 | if name.startswith(logger_name): 56 | return logger 57 | 58 | # handle duplicate logs to the console 59 | # Starting in 1.8.0, PyTorch DDP attaches a StreamHandler (NOTSET) 60 | # to the root logger. As logger.propagate is True by default, this root 61 | # level handler causes logging messages from rank>0 processes to 62 | # unexpectedly show up on the console, creating much unwanted clutter. 63 | # To fix this issue, we set the root logger's StreamHandler, if any, to log 64 | # at the ERROR level. 65 | for handler in logger.root.handlers: 66 | if type(handler) is logging.StreamHandler: 67 | handler.setLevel(logging.ERROR) 68 | 69 | stream_handler = logging.StreamHandler() 70 | handlers = [stream_handler] 71 | 72 | if dist.is_available() and dist.is_initialized(): 73 | rank = dist.get_rank() 74 | else: 75 | rank = 0 76 | 77 | # only rank 0 will add a FileHandler 78 | if rank == 0 and log_file is not None: 79 | # Here, the default behaviour of the official logger is 'a'. Thus, we 80 | # provide an interface to change the file mode to the default 81 | # behaviour. 82 | file_handler = logging.FileHandler(log_file, file_mode) 83 | handlers.append(file_handler) 84 | 85 | formatter = logging.Formatter( 86 | '%(asctime)s - %(name)s - %(levelname)s - %(message)s') 87 | for handler in handlers: 88 | handler.setFormatter(formatter) 89 | handler.setLevel(log_level) 90 | logger.addHandler(handler) 91 | 92 | if rank == 0: 93 | logger.setLevel(log_level) 94 | else: 95 | logger.setLevel(logging.ERROR) 96 | 97 | logger_initialized[name] = True 98 | 99 | 100 | return logger 101 | 102 | 103 | def print_log(msg, logger=None, level=logging.INFO): 104 | """Print a log message. 105 | Args: 106 | msg (str): The message to be logged. 107 | logger (logging.Logger | str | None): The logger to be used. 108 | Some special loggers are: 109 | - "silent": no message will be printed. 110 | - other str: the logger obtained with `get_root_logger(logger)`. 111 | - None: The `print()` method will be used to print log messages. 112 | level (int): Logging level. Only available when `logger` is a Logger 113 | object or "root". 114 | """ 115 | if logger is None: 116 | print(msg) 117 | elif isinstance(logger, logging.Logger): 118 | logger.log(level, msg) 119 | elif logger == 'silent': 120 | pass 121 | elif isinstance(logger, str): 122 | _logger = get_logger(logger) 123 | _logger.log(level, msg) 124 | else: 125 | raise TypeError( 126 | 'logger should be either a logging.Logger object, str, ' 127 | f'"silent" or None, but got {type(logger)}') -------------------------------------------------------------------------------- /utils/metrics.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: Haozhe Xie 3 | # @Date: 2019-08-08 14:31:30 4 | # @Last Modified by: Haozhe Xie 5 | # @Last Modified time: 2020-05-25 09:13:32 6 | # @Email: cshzxie@gmail.com 7 | 8 | import logging 9 | import open3d 10 | 11 | from extensions.chamfer_dist import ChamferDistanceL1, ChamferDistanceL2 12 | 13 | 14 | class Metrics(object): 15 | ITEMS = [{ 16 | 'name': 'F-Score', 17 | 'enabled': True, 18 | 'eval_func': 'cls._get_f_score', 19 | 'is_greater_better': True, 20 | 'init_value': 0 21 | }, { 22 | 'name': 'CDL1', 23 | 'enabled': True, 24 | 'eval_func': 'cls._get_chamfer_distancel1', 25 | 'eval_object': ChamferDistanceL1(ignore_zeros=True), 26 | 'is_greater_better': False, 27 | 'init_value': 32767 28 | }, { 29 | 'name': 'CDL2', 30 | 'enabled': True, 31 | 'eval_func': 'cls._get_chamfer_distancel2', 32 | 'eval_object': ChamferDistanceL2(ignore_zeros=True), 33 | 'is_greater_better': False, 34 | 'init_value': 32767 35 | }] 36 | 37 | @classmethod 38 | def get(cls, pred, gt): 39 | _items = cls.items() 40 | _values = [0] * len(_items) 41 | for i, item in enumerate(_items): 42 | eval_func = eval(item['eval_func']) 43 | _values[i] = eval_func(pred, gt) 44 | 45 | return _values 46 | 47 | @classmethod 48 | def items(cls): 49 | return [i for i in cls.ITEMS if i['enabled']] 50 | 51 | @classmethod 52 | def names(cls): 53 | _items = cls.items() 54 | return [i['name'] for i in _items] 55 | 56 | @classmethod 57 | def _get_f_score(cls, pred, gt, th=0.01): 58 | 59 | """References: https://github.com/lmb-freiburg/what3d/blob/master/util.py""" 60 | b = pred.size(0) 61 | assert pred.size(0) == gt.size(0) 62 | if b != 1: 63 | f_score_list = [] 64 | for idx in range(b): 65 | f_score_list.append(cls._get_f_score(pred[idx:idx+1], gt[idx:idx+1])) 66 | return sum(f_score_list)/len(f_score_list) 67 | else: 68 | pred = cls._get_open3d_ptcloud(pred) 69 | gt = cls._get_open3d_ptcloud(gt) 70 | 71 | dist1 = pred.compute_point_cloud_distance(gt) 72 | dist2 = gt.compute_point_cloud_distance(pred) 73 | 74 | recall = float(sum(d < th for d in dist2)) / float(len(dist2)) 75 | precision = float(sum(d < th for d in dist1)) / float(len(dist1)) 76 | return 2 * recall * precision / (recall + precision) if recall + precision else 0 77 | 78 | @classmethod 79 | def _get_open3d_ptcloud(cls, tensor): 80 | """pred and gt bs is 1""" 81 | tensor = tensor.squeeze().cpu().numpy() 82 | ptcloud = open3d.geometry.PointCloud() 83 | ptcloud.points = open3d.utility.Vector3dVector(tensor) 84 | 85 | return ptcloud 86 | 87 | @classmethod 88 | def _get_chamfer_distancel1(cls, pred, gt): 89 | chamfer_distance = cls.ITEMS[1]['eval_object'] 90 | return chamfer_distance(pred, gt).item() * 1000 91 | 92 | @classmethod 93 | def _get_chamfer_distancel2(cls, pred, gt): 94 | chamfer_distance = cls.ITEMS[2]['eval_object'] 95 | return chamfer_distance(pred, gt).item() * 1000 96 | 97 | def __init__(self, metric_name, values): 98 | self._items = Metrics.items() 99 | self._values = [item['init_value'] for item in self._items] 100 | self.metric_name = metric_name 101 | 102 | if type(values).__name__ == 'list': 103 | self._values = values 104 | elif type(values).__name__ == 'dict': 105 | metric_indexes = {} 106 | for idx, item in enumerate(self._items): 107 | item_name = item['name'] 108 | metric_indexes[item_name] = idx 109 | for k, v in values.items(): 110 | if k not in metric_indexes: 111 | logging.warn('Ignore Metric[Name=%s] due to disability.' % k) 112 | continue 113 | self._values[metric_indexes[k]] = v 114 | else: 115 | raise Exception('Unsupported value type: %s' % type(values)) 116 | 117 | def state_dict(self): 118 | _dict = dict() 119 | for i in range(len(self._items)): 120 | item = self._items[i]['name'] 121 | value = self._values[i] 122 | _dict[item] = value 123 | 124 | return _dict 125 | 126 | def __repr__(self): 127 | return str(self.state_dict()) 128 | 129 | def better_than(self, other): 130 | if other is None: 131 | return True 132 | 133 | _index = -1 134 | for i, _item in enumerate(self._items): 135 | if _item['name'] == self.metric_name: 136 | _index = i 137 | break 138 | if _index == -1: 139 | raise Exception('Invalid metric name to compare.') 140 | 141 | _metric = self._items[i] 142 | _value = self._values[_index] 143 | other_value = other._values[_index] 144 | return _value > other_value if _metric['is_greater_better'] else _value < other_value 145 | -------------------------------------------------------------------------------- /utils/parser.py: -------------------------------------------------------------------------------- 1 | import os 2 | import argparse 3 | from pathlib import Path 4 | 5 | def get_args(): 6 | parser = argparse.ArgumentParser() 7 | parser.add_argument( 8 | '--config', 9 | type = str, 10 | help = 'yaml config file') 11 | parser.add_argument( 12 | '--launcher', 13 | choices=['none', 'pytorch'], 14 | default='none', 15 | help='job launcher') 16 | parser.add_argument('--local_rank', type=int, default=0) 17 | parser.add_argument('--num_workers', type=int, default=4) 18 | # seed 19 | parser.add_argument('--seed', type=int, default=0, help='random seed') 20 | parser.add_argument( 21 | '--deterministic', 22 | action='store_true', 23 | help='whether to set deterministic options for CUDNN backend.') 24 | # bn 25 | parser.add_argument( 26 | '--sync_bn', 27 | action='store_true', 28 | default=False, 29 | help='whether to use sync bn') 30 | # some args 31 | parser.add_argument('--exp_name', type = str, default='default', help = 'experiment name') 32 | parser.add_argument('--start_ckpts', type = str, default=None, help = 'reload used ckpt path') 33 | parser.add_argument('--ckpts', type = str, default=None, help = 'test used ckpt path') 34 | parser.add_argument('--val_freq', type = int, default=1, help = 'test freq') 35 | parser.add_argument( 36 | '--resume', 37 | action='store_true', 38 | default=False, 39 | help = 'autoresume training (interrupted by accident)') 40 | parser.add_argument( 41 | '--test', 42 | action='store_true', 43 | default=False, 44 | help = 'test mode for certain ckpt') 45 | parser.add_argument( 46 | '--mode', 47 | choices=['easy', 'median', 'hard', None], 48 | default=None, 49 | help = 'difficulty mode for shapenet') 50 | args = parser.parse_args() 51 | 52 | if args.test and args.resume: 53 | raise ValueError( 54 | '--test and --resume cannot be both activate') 55 | 56 | if args.resume and args.start_ckpts is not None: 57 | raise ValueError( 58 | '--resume and --start_ckpts cannot be both activate') 59 | 60 | if args.test and args.ckpts is None: 61 | raise ValueError( 62 | 'ckpts shouldnt be None while test mode') 63 | 64 | if 'LOCAL_RANK' not in os.environ: 65 | os.environ['LOCAL_RANK'] = str(args.local_rank) 66 | 67 | if args.test: 68 | args.exp_name = 'test_' + args.exp_name 69 | if args.mode is not None: 70 | args.exp_name = args.exp_name + '_' +args.mode 71 | args.experiment_path = os.path.join('./experiments', Path(args.config).stem, Path(args.config).parent.stem, args.exp_name) 72 | args.tfboard_path = os.path.join('./experiments', Path(args.config).stem, Path(args.config).parent.stem,'TFBoard' ,args.exp_name) 73 | args.log_name = Path(args.config).stem 74 | create_experiment_dir(args) 75 | return args 76 | 77 | def create_experiment_dir(args): 78 | if not os.path.exists(args.experiment_path): 79 | os.makedirs(args.experiment_path) 80 | print('Create experiment path successfully at %s' % args.experiment_path) 81 | if not os.path.exists(args.tfboard_path): 82 | os.makedirs(args.tfboard_path) 83 | print('Create TFBoard path successfully at %s' % args.tfboard_path) 84 | 85 | -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/build/lib.linux-x86_64-3.7/pointnet2_ops/__init__.py: -------------------------------------------------------------------------------- 1 | import pointnet2_ops.pointnet2_modules 2 | import pointnet2_ops.pointnet2_utils 3 | from pointnet2_ops._version import __version__ 4 | -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/build/lib.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/ball_query.cpp: -------------------------------------------------------------------------------- 1 | #include "ball_query.h" 2 | #include "utils.h" 3 | 4 | void query_ball_point_kernel_wrapper(int b, int n, int m, float radius, 5 | int nsample, const float *new_xyz, 6 | const float *xyz, int *idx); 7 | 8 | at::Tensor ball_query(at::Tensor new_xyz, at::Tensor xyz, const float radius, 9 | const int nsample) { 10 | CHECK_CONTIGUOUS(new_xyz); 11 | CHECK_CONTIGUOUS(xyz); 12 | CHECK_IS_FLOAT(new_xyz); 13 | CHECK_IS_FLOAT(xyz); 14 | 15 | if (new_xyz.is_cuda()) { 16 | CHECK_CUDA(xyz); 17 | } 18 | 19 | at::Tensor idx = 20 | torch::zeros({new_xyz.size(0), new_xyz.size(1), nsample}, 21 | at::device(new_xyz.device()).dtype(at::ScalarType::Int)); 22 | 23 | if (new_xyz.is_cuda()) { 24 | query_ball_point_kernel_wrapper(xyz.size(0), xyz.size(1), new_xyz.size(1), 25 | radius, nsample, new_xyz.data_ptr(), 26 | xyz.data_ptr(), idx.data_ptr()); 27 | } else { 28 | AT_ASSERT(false, "CPU not supported"); 29 | } 30 | 31 | return idx; 32 | } 33 | -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/build/lib.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/ball_query_gpu.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | 5 | #include "cuda_utils.h" 6 | 7 | // input: new_xyz(b, m, 3) xyz(b, n, 3) 8 | // output: idx(b, m, nsample) 9 | __global__ void query_ball_point_kernel(int b, int n, int m, float radius, 10 | int nsample, 11 | const float *__restrict__ new_xyz, 12 | const float *__restrict__ xyz, 13 | int *__restrict__ idx) { 14 | int batch_index = blockIdx.x; 15 | xyz += batch_index * n * 3; 16 | new_xyz += batch_index * m * 3; 17 | idx += m * nsample * batch_index; 18 | 19 | int index = threadIdx.x; 20 | int stride = blockDim.x; 21 | 22 | float radius2 = radius * radius; 23 | for (int j = index; j < m; j += stride) { 24 | float new_x = new_xyz[j * 3 + 0]; 25 | float new_y = new_xyz[j * 3 + 1]; 26 | float new_z = new_xyz[j * 3 + 2]; 27 | for (int k = 0, cnt = 0; k < n && cnt < nsample; ++k) { 28 | float x = xyz[k * 3 + 0]; 29 | float y = xyz[k * 3 + 1]; 30 | float z = xyz[k * 3 + 2]; 31 | float d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) + 32 | (new_z - z) * (new_z - z); 33 | if (d2 < radius2) { 34 | if (cnt == 0) { 35 | for (int l = 0; l < nsample; ++l) { 36 | idx[j * nsample + l] = k; 37 | } 38 | } 39 | idx[j * nsample + cnt] = k; 40 | ++cnt; 41 | } 42 | } 43 | } 44 | } 45 | 46 | void query_ball_point_kernel_wrapper(int b, int n, int m, float radius, 47 | int nsample, const float *new_xyz, 48 | const float *xyz, int *idx) { 49 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 50 | query_ball_point_kernel<<>>( 51 | b, n, m, radius, nsample, new_xyz, xyz, idx); 52 | 53 | CUDA_CHECK_ERRORS(); 54 | } 55 | -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/build/lib.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/bindings.cpp: -------------------------------------------------------------------------------- 1 | #include "ball_query.h" 2 | #include "group_points.h" 3 | #include "interpolate.h" 4 | #include "sampling.h" 5 | 6 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 7 | m.def("gather_points", &gather_points); 8 | m.def("gather_points_grad", &gather_points_grad); 9 | m.def("furthest_point_sampling", &furthest_point_sampling); 10 | 11 | m.def("three_nn", &three_nn); 12 | m.def("three_interpolate", &three_interpolate); 13 | m.def("three_interpolate_grad", &three_interpolate_grad); 14 | 15 | m.def("ball_query", &ball_query); 16 | 17 | m.def("group_points", &group_points); 18 | m.def("group_points_grad", &group_points_grad); 19 | } 20 | -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/build/lib.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/group_points.cpp: -------------------------------------------------------------------------------- 1 | #include "group_points.h" 2 | #include "utils.h" 3 | 4 | void group_points_kernel_wrapper(int b, int c, int n, int npoints, int nsample, 5 | const float *points, const int *idx, 6 | float *out); 7 | 8 | void group_points_grad_kernel_wrapper(int b, int c, int n, int npoints, 9 | int nsample, const float *grad_out, 10 | const int *idx, float *grad_points); 11 | 12 | at::Tensor group_points(at::Tensor points, at::Tensor idx) { 13 | CHECK_CONTIGUOUS(points); 14 | CHECK_CONTIGUOUS(idx); 15 | CHECK_IS_FLOAT(points); 16 | CHECK_IS_INT(idx); 17 | 18 | if (points.is_cuda()) { 19 | CHECK_CUDA(idx); 20 | } 21 | 22 | at::Tensor output = 23 | torch::zeros({points.size(0), points.size(1), idx.size(1), idx.size(2)}, 24 | at::device(points.device()).dtype(at::ScalarType::Float)); 25 | 26 | if (points.is_cuda()) { 27 | group_points_kernel_wrapper(points.size(0), points.size(1), points.size(2), 28 | idx.size(1), idx.size(2), 29 | points.data_ptr(), idx.data_ptr(), 30 | output.data_ptr()); 31 | } else { 32 | AT_ASSERT(false, "CPU not supported"); 33 | } 34 | 35 | return output; 36 | } 37 | 38 | at::Tensor group_points_grad(at::Tensor grad_out, at::Tensor idx, const int n) { 39 | CHECK_CONTIGUOUS(grad_out); 40 | CHECK_CONTIGUOUS(idx); 41 | CHECK_IS_FLOAT(grad_out); 42 | CHECK_IS_INT(idx); 43 | 44 | if (grad_out.is_cuda()) { 45 | CHECK_CUDA(idx); 46 | } 47 | 48 | at::Tensor output = 49 | torch::zeros({grad_out.size(0), grad_out.size(1), n}, 50 | at::device(grad_out.device()).dtype(at::ScalarType::Float)); 51 | 52 | if (grad_out.is_cuda()) { 53 | group_points_grad_kernel_wrapper( 54 | grad_out.size(0), grad_out.size(1), n, idx.size(1), idx.size(2), 55 | grad_out.data_ptr(), idx.data_ptr(), 56 | output.data_ptr()); 57 | } else { 58 | AT_ASSERT(false, "CPU not supported"); 59 | } 60 | 61 | return output; 62 | } 63 | -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/build/lib.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/group_points_gpu.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #include "cuda_utils.h" 5 | 6 | // input: points(b, c, n) idx(b, npoints, nsample) 7 | // output: out(b, c, npoints, nsample) 8 | __global__ void group_points_kernel(int b, int c, int n, int npoints, 9 | int nsample, 10 | const float *__restrict__ points, 11 | const int *__restrict__ idx, 12 | float *__restrict__ out) { 13 | int batch_index = blockIdx.x; 14 | points += batch_index * n * c; 15 | idx += batch_index * npoints * nsample; 16 | out += batch_index * npoints * nsample * c; 17 | 18 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 19 | const int stride = blockDim.y * blockDim.x; 20 | for (int i = index; i < c * npoints; i += stride) { 21 | const int l = i / npoints; 22 | const int j = i % npoints; 23 | for (int k = 0; k < nsample; ++k) { 24 | int ii = idx[j * nsample + k]; 25 | out[(l * npoints + j) * nsample + k] = points[l * n + ii]; 26 | } 27 | } 28 | } 29 | 30 | void group_points_kernel_wrapper(int b, int c, int n, int npoints, int nsample, 31 | const float *points, const int *idx, 32 | float *out) { 33 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 34 | 35 | group_points_kernel<<>>( 36 | b, c, n, npoints, nsample, points, idx, out); 37 | 38 | CUDA_CHECK_ERRORS(); 39 | } 40 | 41 | // input: grad_out(b, c, npoints, nsample), idx(b, npoints, nsample) 42 | // output: grad_points(b, c, n) 43 | __global__ void group_points_grad_kernel(int b, int c, int n, int npoints, 44 | int nsample, 45 | const float *__restrict__ grad_out, 46 | const int *__restrict__ idx, 47 | float *__restrict__ grad_points) { 48 | int batch_index = blockIdx.x; 49 | grad_out += batch_index * npoints * nsample * c; 50 | idx += batch_index * npoints * nsample; 51 | grad_points += batch_index * n * c; 52 | 53 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 54 | const int stride = blockDim.y * blockDim.x; 55 | for (int i = index; i < c * npoints; i += stride) { 56 | const int l = i / npoints; 57 | const int j = i % npoints; 58 | for (int k = 0; k < nsample; ++k) { 59 | int ii = idx[j * nsample + k]; 60 | atomicAdd(grad_points + l * n + ii, 61 | grad_out[(l * npoints + j) * nsample + k]); 62 | } 63 | } 64 | } 65 | 66 | void group_points_grad_kernel_wrapper(int b, int c, int n, int npoints, 67 | int nsample, const float *grad_out, 68 | const int *idx, float *grad_points) { 69 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 70 | 71 | group_points_grad_kernel<<>>( 72 | b, c, n, npoints, nsample, grad_out, idx, grad_points); 73 | 74 | CUDA_CHECK_ERRORS(); 75 | } 76 | -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/build/lib.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/interpolate.cpp: -------------------------------------------------------------------------------- 1 | #include "interpolate.h" 2 | #include "utils.h" 3 | 4 | void three_nn_kernel_wrapper(int b, int n, int m, const float *unknown, 5 | const float *known, float *dist2, int *idx); 6 | void three_interpolate_kernel_wrapper(int b, int c, int m, int n, 7 | const float *points, const int *idx, 8 | const float *weight, float *out); 9 | void three_interpolate_grad_kernel_wrapper(int b, int c, int n, int m, 10 | const float *grad_out, 11 | const int *idx, const float *weight, 12 | float *grad_points); 13 | 14 | std::vector three_nn(at::Tensor unknowns, at::Tensor knows) { 15 | CHECK_CONTIGUOUS(unknowns); 16 | CHECK_CONTIGUOUS(knows); 17 | CHECK_IS_FLOAT(unknowns); 18 | CHECK_IS_FLOAT(knows); 19 | 20 | if (unknowns.is_cuda()) { 21 | CHECK_CUDA(knows); 22 | } 23 | 24 | at::Tensor idx = 25 | torch::zeros({unknowns.size(0), unknowns.size(1), 3}, 26 | at::device(unknowns.device()).dtype(at::ScalarType::Int)); 27 | at::Tensor dist2 = 28 | torch::zeros({unknowns.size(0), unknowns.size(1), 3}, 29 | at::device(unknowns.device()).dtype(at::ScalarType::Float)); 30 | 31 | if (unknowns.is_cuda()) { 32 | three_nn_kernel_wrapper(unknowns.size(0), unknowns.size(1), knows.size(1), 33 | unknowns.data_ptr(), knows.data_ptr(), 34 | dist2.data_ptr(), idx.data_ptr()); 35 | } else { 36 | AT_ASSERT(false, "CPU not supported"); 37 | } 38 | 39 | return {dist2, idx}; 40 | } 41 | 42 | at::Tensor three_interpolate(at::Tensor points, at::Tensor idx, 43 | at::Tensor weight) { 44 | CHECK_CONTIGUOUS(points); 45 | CHECK_CONTIGUOUS(idx); 46 | CHECK_CONTIGUOUS(weight); 47 | CHECK_IS_FLOAT(points); 48 | CHECK_IS_INT(idx); 49 | CHECK_IS_FLOAT(weight); 50 | 51 | if (points.is_cuda()) { 52 | CHECK_CUDA(idx); 53 | CHECK_CUDA(weight); 54 | } 55 | 56 | at::Tensor output = 57 | torch::zeros({points.size(0), points.size(1), idx.size(1)}, 58 | at::device(points.device()).dtype(at::ScalarType::Float)); 59 | 60 | if (points.is_cuda()) { 61 | three_interpolate_kernel_wrapper( 62 | points.size(0), points.size(1), points.size(2), idx.size(1), 63 | points.data_ptr(), idx.data_ptr(), weight.data_ptr(), 64 | output.data_ptr()); 65 | } else { 66 | AT_ASSERT(false, "CPU not supported"); 67 | } 68 | 69 | return output; 70 | } 71 | at::Tensor three_interpolate_grad(at::Tensor grad_out, at::Tensor idx, 72 | at::Tensor weight, const int m) { 73 | CHECK_CONTIGUOUS(grad_out); 74 | CHECK_CONTIGUOUS(idx); 75 | CHECK_CONTIGUOUS(weight); 76 | CHECK_IS_FLOAT(grad_out); 77 | CHECK_IS_INT(idx); 78 | CHECK_IS_FLOAT(weight); 79 | 80 | if (grad_out.is_cuda()) { 81 | CHECK_CUDA(idx); 82 | CHECK_CUDA(weight); 83 | } 84 | 85 | at::Tensor output = 86 | torch::zeros({grad_out.size(0), grad_out.size(1), m}, 87 | at::device(grad_out.device()).dtype(at::ScalarType::Float)); 88 | 89 | if (grad_out.is_cuda()) { 90 | three_interpolate_grad_kernel_wrapper( 91 | grad_out.size(0), grad_out.size(1), grad_out.size(2), m, 92 | grad_out.data_ptr(), idx.data_ptr(), 93 | weight.data_ptr(), output.data_ptr()); 94 | } else { 95 | AT_ASSERT(false, "CPU not supported"); 96 | } 97 | 98 | return output; 99 | } 100 | -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/build/lib.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/sampling.cpp: -------------------------------------------------------------------------------- 1 | #include "sampling.h" 2 | #include "utils.h" 3 | 4 | void gather_points_kernel_wrapper(int b, int c, int n, int npoints, 5 | const float *points, const int *idx, 6 | float *out); 7 | void gather_points_grad_kernel_wrapper(int b, int c, int n, int npoints, 8 | const float *grad_out, const int *idx, 9 | float *grad_points); 10 | 11 | void furthest_point_sampling_kernel_wrapper(int b, int n, int m, 12 | const float *dataset, float *temp, 13 | int *idxs); 14 | 15 | at::Tensor gather_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)}, 27 | at::device(points.device()).dtype(at::ScalarType::Float)); 28 | 29 | if (points.is_cuda()) { 30 | gather_points_kernel_wrapper(points.size(0), points.size(1), points.size(2), 31 | idx.size(1), points.data_ptr(), 32 | idx.data_ptr(), output.data_ptr()); 33 | } else { 34 | AT_ASSERT(false, "CPU not supported"); 35 | } 36 | 37 | return output; 38 | } 39 | 40 | at::Tensor gather_points_grad(at::Tensor grad_out, at::Tensor idx, 41 | const int n) { 42 | CHECK_CONTIGUOUS(grad_out); 43 | CHECK_CONTIGUOUS(idx); 44 | CHECK_IS_FLOAT(grad_out); 45 | CHECK_IS_INT(idx); 46 | 47 | if (grad_out.is_cuda()) { 48 | CHECK_CUDA(idx); 49 | } 50 | 51 | at::Tensor output = 52 | torch::zeros({grad_out.size(0), grad_out.size(1), n}, 53 | at::device(grad_out.device()).dtype(at::ScalarType::Float)); 54 | 55 | if (grad_out.is_cuda()) { 56 | gather_points_grad_kernel_wrapper(grad_out.size(0), grad_out.size(1), n, 57 | idx.size(1), grad_out.data_ptr(), 58 | idx.data_ptr(), 59 | output.data_ptr()); 60 | } else { 61 | AT_ASSERT(false, "CPU not supported"); 62 | } 63 | 64 | return output; 65 | } 66 | at::Tensor furthest_point_sampling(at::Tensor points, const int nsamples) { 67 | CHECK_CONTIGUOUS(points); 68 | CHECK_IS_FLOAT(points); 69 | 70 | at::Tensor output = 71 | torch::zeros({points.size(0), nsamples}, 72 | at::device(points.device()).dtype(at::ScalarType::Int)); 73 | 74 | at::Tensor tmp = 75 | torch::full({points.size(0), points.size(1)}, 1e10, 76 | at::device(points.device()).dtype(at::ScalarType::Float)); 77 | 78 | if (points.is_cuda()) { 79 | furthest_point_sampling_kernel_wrapper( 80 | points.size(0), points.size(1), nsamples, points.data_ptr(), 81 | tmp.data_ptr(), output.data_ptr()); 82 | } else { 83 | AT_ASSERT(false, "CPU not supported"); 84 | } 85 | 86 | return output; 87 | } 88 | -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/build/lib.linux-x86_64-3.7/pointnet2_ops/_ext.cpython-37m-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/utils/pointnet2_ops_lib/build/lib.linux-x86_64-3.7/pointnet2_ops/_ext.cpython-37m-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/build/lib.linux-x86_64-3.7/pointnet2_ops/_version.py: -------------------------------------------------------------------------------- 1 | __version__ = "3.0.0" 2 | -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/.ninja_deps: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/.ninja_deps -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/.ninja_log: -------------------------------------------------------------------------------- 1 | # ninja log v5 2 | 0 15064 1658931498929661702 /home/lss/Project/PoinTr/utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/ball_query.o ad15ea1f5172fcea 3 | 1 15324 1658931499185657976 /home/lss/Project/PoinTr/utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/bindings.o ceb5428682193f8b 4 | 1 17635 1658931501501624238 /home/lss/Project/PoinTr/utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/group_points.o 22866c2a26cc3124 5 | 1 19735 1658931503601593654 /home/lss/Project/PoinTr/utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/interpolate.o f1ef9afa379e0933 6 | 1 19834 1658931503701592200 /home/lss/Project/PoinTr/utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/sampling.o 8c73b19b37af6c49 7 | 1 34132 1658931517997384080 /home/lss/Project/PoinTr/utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/group_points_gpu.o bc58575b3e1973dc 8 | 0 34815 1658931518681374128 /home/lss/Project/PoinTr/utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/ball_query_gpu.o 5826afdb949ddd7d 9 | 1 35556 1658931519429363243 /home/lss/Project/PoinTr/utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/interpolate_gpu.o 617ac13796844030 10 | 1 37491 1658931521361335132 /home/lss/Project/PoinTr/utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/sampling_gpu.o c26ce1e92f674c40 11 | -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/build.ninja: -------------------------------------------------------------------------------- 1 | ninja_required_version = 1.3 2 | cxx = c++ 3 | nvcc = /usr/local/cuda-11.3/bin/nvcc 4 | 5 | cflags = -pthread -B /home/lss/anaconda3/envs/PoinTr/compiler_compat -Wl,--sysroot=/ -Wsign-compare -DNDEBUG -g -fwrapv -O3 -Wall -Wstrict-prototypes -fPIC -I/home/lss/Project/PoinTr/utils/pointnet2_ops_lib/pointnet2_ops/_ext-src/include -I/home/lss/anaconda3/envs/PoinTr/lib/python3.7/site-packages/torch/include -I/home/lss/anaconda3/envs/PoinTr/lib/python3.7/site-packages/torch/include/torch/csrc/api/include -I/home/lss/anaconda3/envs/PoinTr/lib/python3.7/site-packages/torch/include/TH -I/home/lss/anaconda3/envs/PoinTr/lib/python3.7/site-packages/torch/include/THC -I/usr/local/cuda-11.3/include -I/home/lss/anaconda3/envs/PoinTr/include/python3.7m -c 6 | post_cflags = -O3 -DTORCH_API_INCLUDE_EXTENSION_H '-DPYBIND11_COMPILER_TYPE="_gcc"' '-DPYBIND11_STDLIB="_libstdcpp"' '-DPYBIND11_BUILD_ABI="_cxxabi1011"' -DTORCH_EXTENSION_NAME=_ext -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++14 7 | cuda_cflags = -I/home/lss/Project/PoinTr/utils/pointnet2_ops_lib/pointnet2_ops/_ext-src/include -I/home/lss/anaconda3/envs/PoinTr/lib/python3.7/site-packages/torch/include -I/home/lss/anaconda3/envs/PoinTr/lib/python3.7/site-packages/torch/include/torch/csrc/api/include -I/home/lss/anaconda3/envs/PoinTr/lib/python3.7/site-packages/torch/include/TH -I/home/lss/anaconda3/envs/PoinTr/lib/python3.7/site-packages/torch/include/THC -I/usr/local/cuda-11.3/include -I/home/lss/anaconda3/envs/PoinTr/include/python3.7m -c 8 | cuda_post_cflags = -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr --compiler-options ''"'"'-fPIC'"'"'' -O3 -Xfatbin -compress-all -DTORCH_API_INCLUDE_EXTENSION_H '-DPYBIND11_COMPILER_TYPE="_gcc"' '-DPYBIND11_STDLIB="_libstdcpp"' '-DPYBIND11_BUILD_ABI="_cxxabi1011"' -DTORCH_EXTENSION_NAME=_ext -D_GLIBCXX_USE_CXX11_ABI=0 -gencode=arch=compute_37,code=compute_37 -gencode=arch=compute_37,code=sm_37 -gencode=arch=compute_50,code=sm_50 -gencode=arch=compute_60,code=sm_60 -gencode=arch=compute_61,code=sm_61 -gencode=arch=compute_62,code=sm_62 -gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_75,code=sm_75 -std=c++14 9 | ldflags = 10 | 11 | rule compile 12 | command = $cxx -MMD -MF $out.d $cflags -c $in -o $out $post_cflags 13 | depfile = $out.d 14 | deps = gcc 15 | 16 | rule cuda_compile 17 | depfile = $out.d 18 | deps = gcc 19 | command = $nvcc $cuda_cflags -c $in -o $out $cuda_post_cflags 20 | 21 | 22 | 23 | build /home/lss/Project/PoinTr/utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/ball_query.o: compile /home/lss/Project/PoinTr/utils/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/ball_query.cpp 24 | build /home/lss/Project/PoinTr/utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/ball_query_gpu.o: cuda_compile /home/lss/Project/PoinTr/utils/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/ball_query_gpu.cu 25 | build /home/lss/Project/PoinTr/utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/bindings.o: compile /home/lss/Project/PoinTr/utils/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/bindings.cpp 26 | build /home/lss/Project/PoinTr/utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/group_points.o: compile /home/lss/Project/PoinTr/utils/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/group_points.cpp 27 | build /home/lss/Project/PoinTr/utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/group_points_gpu.o: cuda_compile /home/lss/Project/PoinTr/utils/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/group_points_gpu.cu 28 | build /home/lss/Project/PoinTr/utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/interpolate.o: compile /home/lss/Project/PoinTr/utils/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/interpolate.cpp 29 | build /home/lss/Project/PoinTr/utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/interpolate_gpu.o: cuda_compile /home/lss/Project/PoinTr/utils/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/interpolate_gpu.cu 30 | build /home/lss/Project/PoinTr/utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/sampling.o: compile /home/lss/Project/PoinTr/utils/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/sampling.cpp 31 | build /home/lss/Project/PoinTr/utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/sampling_gpu.o: cuda_compile /home/lss/Project/PoinTr/utils/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/sampling_gpu.cu 32 | 33 | 34 | 35 | 36 | 37 | -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/ball_query.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/ball_query.o -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/ball_query_gpu.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/ball_query_gpu.o -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/bindings.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/bindings.o -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/group_points.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/group_points.o -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/group_points_gpu.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/group_points_gpu.o -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/interpolate.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/interpolate.o -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/interpolate_gpu.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/interpolate_gpu.o -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/sampling.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/sampling.o -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/sampling_gpu.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/utils/pointnet2_ops_lib/build/temp.linux-x86_64-3.7/pointnet2_ops/_ext-src/src/sampling_gpu.o -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/dist/pointnet2_ops-3.0.0-py3.7-linux-x86_64.egg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/I2-Multimedia-Lab/ProxyFormer/52e4d6dfe869b438e0a85bc128254edb94fd75a5/utils/pointnet2_ops_lib/dist/pointnet2_ops-3.0.0-py3.7-linux-x86_64.egg -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/pointnet2_ops.egg-info/PKG-INFO: -------------------------------------------------------------------------------- 1 | Metadata-Version: 2.1 2 | Name: pointnet2-ops 3 | Version: 3.0.0 4 | Summary: UNKNOWN 5 | Author: Erik Wijmans 6 | License: UNKNOWN 7 | Platform: UNKNOWN 8 | 9 | UNKNOWN 10 | 11 | -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/pointnet2_ops.egg-info/SOURCES.txt: -------------------------------------------------------------------------------- 1 | setup.py 2 | pointnet2_ops/__init__.py 3 | pointnet2_ops/_version.py 4 | pointnet2_ops/pointnet2_modules.py 5 | pointnet2_ops/pointnet2_utils.py 6 | pointnet2_ops.egg-info/PKG-INFO 7 | pointnet2_ops.egg-info/SOURCES.txt 8 | pointnet2_ops.egg-info/dependency_links.txt 9 | pointnet2_ops.egg-info/requires.txt 10 | pointnet2_ops.egg-info/top_level.txt 11 | pointnet2_ops/_ext-src/src/ball_query.cpp 12 | pointnet2_ops/_ext-src/src/ball_query_gpu.cu 13 | pointnet2_ops/_ext-src/src/bindings.cpp 14 | pointnet2_ops/_ext-src/src/group_points.cpp 15 | pointnet2_ops/_ext-src/src/group_points_gpu.cu 16 | pointnet2_ops/_ext-src/src/interpolate.cpp 17 | pointnet2_ops/_ext-src/src/interpolate_gpu.cu 18 | pointnet2_ops/_ext-src/src/sampling.cpp 19 | pointnet2_ops/_ext-src/src/sampling_gpu.cu -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/pointnet2_ops.egg-info/dependency_links.txt: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/pointnet2_ops.egg-info/requires.txt: -------------------------------------------------------------------------------- 1 | torch>=1.4 2 | -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/pointnet2_ops.egg-info/top_level.txt: -------------------------------------------------------------------------------- 1 | pointnet2_ops 2 | -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/pointnet2_ops/__init__.py: -------------------------------------------------------------------------------- 1 | import pointnet2_ops.pointnet2_modules 2 | import pointnet2_ops.pointnet2_utils 3 | from pointnet2_ops._version import __version__ 4 | -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/pointnet2_ops/_ext-src/include/ball_query.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include 3 | 4 | at::Tensor ball_query(at::Tensor new_xyz, at::Tensor xyz, const float radius, 5 | const int nsample); 6 | -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/pointnet2_ops/_ext-src/include/cuda_utils.h: -------------------------------------------------------------------------------- 1 | #ifndef _CUDA_UTILS_H 2 | #define _CUDA_UTILS_H 3 | 4 | #include 5 | #include 6 | #include 7 | 8 | #include 9 | #include 10 | 11 | #include 12 | 13 | #define TOTAL_THREADS 512 14 | 15 | inline int opt_n_threads(int work_size) { 16 | const int pow_2 = std::log(static_cast(work_size)) / std::log(2.0); 17 | 18 | return max(min(1 << pow_2, TOTAL_THREADS), 1); 19 | } 20 | 21 | inline dim3 opt_block_config(int x, int y) { 22 | const int x_threads = opt_n_threads(x); 23 | const int y_threads = 24 | max(min(opt_n_threads(y), TOTAL_THREADS / x_threads), 1); 25 | dim3 block_config(x_threads, y_threads, 1); 26 | 27 | return block_config; 28 | } 29 | 30 | #define CUDA_CHECK_ERRORS() \ 31 | do { \ 32 | cudaError_t err = cudaGetLastError(); \ 33 | if (cudaSuccess != err) { \ 34 | fprintf(stderr, "CUDA kernel failed : %s\n%s at L:%d in %s\n", \ 35 | cudaGetErrorString(err), __PRETTY_FUNCTION__, __LINE__, \ 36 | __FILE__); \ 37 | exit(-1); \ 38 | } \ 39 | } while (0) 40 | 41 | #endif 42 | -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/pointnet2_ops/_ext-src/include/group_points.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include 3 | 4 | at::Tensor group_points(at::Tensor points, at::Tensor idx); 5 | at::Tensor group_points_grad(at::Tensor grad_out, at::Tensor idx, const int n); 6 | -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/pointnet2_ops/_ext-src/include/interpolate.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | 6 | std::vector three_nn(at::Tensor unknowns, at::Tensor knows); 7 | at::Tensor three_interpolate(at::Tensor points, at::Tensor idx, 8 | at::Tensor weight); 9 | at::Tensor three_interpolate_grad(at::Tensor grad_out, at::Tensor idx, 10 | at::Tensor weight, const int m); 11 | -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/pointnet2_ops/_ext-src/include/sampling.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include 3 | 4 | at::Tensor gather_points(at::Tensor points, at::Tensor idx); 5 | at::Tensor gather_points_grad(at::Tensor grad_out, at::Tensor idx, const int n); 6 | at::Tensor furthest_point_sampling(at::Tensor points, const int nsamples); 7 | -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/pointnet2_ops/_ext-src/include/utils.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include 3 | #include 4 | 5 | #define CHECK_CUDA(x) \ 6 | do { \ 7 | AT_ASSERT(x.is_cuda(), #x " must be a CUDA tensor"); \ 8 | } while (0) 9 | 10 | #define CHECK_CONTIGUOUS(x) \ 11 | do { \ 12 | AT_ASSERT(x.is_contiguous(), #x " must be a contiguous tensor"); \ 13 | } while (0) 14 | 15 | #define CHECK_IS_INT(x) \ 16 | do { \ 17 | AT_ASSERT(x.scalar_type() == at::ScalarType::Int, \ 18 | #x " must be an int tensor"); \ 19 | } while (0) 20 | 21 | #define CHECK_IS_FLOAT(x) \ 22 | do { \ 23 | AT_ASSERT(x.scalar_type() == at::ScalarType::Float, \ 24 | #x " must be a float tensor"); \ 25 | } while (0) 26 | -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/ball_query.cpp: -------------------------------------------------------------------------------- 1 | #include "ball_query.h" 2 | #include "utils.h" 3 | 4 | void query_ball_point_kernel_wrapper(int b, int n, int m, float radius, 5 | int nsample, const float *new_xyz, 6 | const float *xyz, int *idx); 7 | 8 | at::Tensor ball_query(at::Tensor new_xyz, at::Tensor xyz, const float radius, 9 | const int nsample) { 10 | CHECK_CONTIGUOUS(new_xyz); 11 | CHECK_CONTIGUOUS(xyz); 12 | CHECK_IS_FLOAT(new_xyz); 13 | CHECK_IS_FLOAT(xyz); 14 | 15 | if (new_xyz.is_cuda()) { 16 | CHECK_CUDA(xyz); 17 | } 18 | 19 | at::Tensor idx = 20 | torch::zeros({new_xyz.size(0), new_xyz.size(1), nsample}, 21 | at::device(new_xyz.device()).dtype(at::ScalarType::Int)); 22 | 23 | if (new_xyz.is_cuda()) { 24 | query_ball_point_kernel_wrapper(xyz.size(0), xyz.size(1), new_xyz.size(1), 25 | radius, nsample, new_xyz.data_ptr(), 26 | xyz.data_ptr(), idx.data_ptr()); 27 | } else { 28 | AT_ASSERT(false, "CPU not supported"); 29 | } 30 | 31 | return idx; 32 | } 33 | -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/ball_query_gpu.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | 5 | #include "cuda_utils.h" 6 | 7 | // input: new_xyz(b, m, 3) xyz(b, n, 3) 8 | // output: idx(b, m, nsample) 9 | __global__ void query_ball_point_kernel(int b, int n, int m, float radius, 10 | int nsample, 11 | const float *__restrict__ new_xyz, 12 | const float *__restrict__ xyz, 13 | int *__restrict__ idx) { 14 | int batch_index = blockIdx.x; 15 | xyz += batch_index * n * 3; 16 | new_xyz += batch_index * m * 3; 17 | idx += m * nsample * batch_index; 18 | 19 | int index = threadIdx.x; 20 | int stride = blockDim.x; 21 | 22 | float radius2 = radius * radius; 23 | for (int j = index; j < m; j += stride) { 24 | float new_x = new_xyz[j * 3 + 0]; 25 | float new_y = new_xyz[j * 3 + 1]; 26 | float new_z = new_xyz[j * 3 + 2]; 27 | for (int k = 0, cnt = 0; k < n && cnt < nsample; ++k) { 28 | float x = xyz[k * 3 + 0]; 29 | float y = xyz[k * 3 + 1]; 30 | float z = xyz[k * 3 + 2]; 31 | float d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) + 32 | (new_z - z) * (new_z - z); 33 | if (d2 < radius2) { 34 | if (cnt == 0) { 35 | for (int l = 0; l < nsample; ++l) { 36 | idx[j * nsample + l] = k; 37 | } 38 | } 39 | idx[j * nsample + cnt] = k; 40 | ++cnt; 41 | } 42 | } 43 | } 44 | } 45 | 46 | void query_ball_point_kernel_wrapper(int b, int n, int m, float radius, 47 | int nsample, const float *new_xyz, 48 | const float *xyz, int *idx) { 49 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 50 | query_ball_point_kernel<<>>( 51 | b, n, m, radius, nsample, new_xyz, xyz, idx); 52 | 53 | CUDA_CHECK_ERRORS(); 54 | } 55 | -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/bindings.cpp: -------------------------------------------------------------------------------- 1 | #include "ball_query.h" 2 | #include "group_points.h" 3 | #include "interpolate.h" 4 | #include "sampling.h" 5 | 6 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 7 | m.def("gather_points", &gather_points); 8 | m.def("gather_points_grad", &gather_points_grad); 9 | m.def("furthest_point_sampling", &furthest_point_sampling); 10 | 11 | m.def("three_nn", &three_nn); 12 | m.def("three_interpolate", &three_interpolate); 13 | m.def("three_interpolate_grad", &three_interpolate_grad); 14 | 15 | m.def("ball_query", &ball_query); 16 | 17 | m.def("group_points", &group_points); 18 | m.def("group_points_grad", &group_points_grad); 19 | } 20 | -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/group_points.cpp: -------------------------------------------------------------------------------- 1 | #include "group_points.h" 2 | #include "utils.h" 3 | 4 | void group_points_kernel_wrapper(int b, int c, int n, int npoints, int nsample, 5 | const float *points, const int *idx, 6 | float *out); 7 | 8 | void group_points_grad_kernel_wrapper(int b, int c, int n, int npoints, 9 | int nsample, const float *grad_out, 10 | const int *idx, float *grad_points); 11 | 12 | at::Tensor group_points(at::Tensor points, at::Tensor idx) { 13 | CHECK_CONTIGUOUS(points); 14 | CHECK_CONTIGUOUS(idx); 15 | CHECK_IS_FLOAT(points); 16 | CHECK_IS_INT(idx); 17 | 18 | if (points.is_cuda()) { 19 | CHECK_CUDA(idx); 20 | } 21 | 22 | at::Tensor output = 23 | torch::zeros({points.size(0), points.size(1), idx.size(1), idx.size(2)}, 24 | at::device(points.device()).dtype(at::ScalarType::Float)); 25 | 26 | if (points.is_cuda()) { 27 | group_points_kernel_wrapper(points.size(0), points.size(1), points.size(2), 28 | idx.size(1), idx.size(2), 29 | points.data_ptr(), idx.data_ptr(), 30 | output.data_ptr()); 31 | } else { 32 | AT_ASSERT(false, "CPU not supported"); 33 | } 34 | 35 | return output; 36 | } 37 | 38 | at::Tensor group_points_grad(at::Tensor grad_out, at::Tensor idx, const int n) { 39 | CHECK_CONTIGUOUS(grad_out); 40 | CHECK_CONTIGUOUS(idx); 41 | CHECK_IS_FLOAT(grad_out); 42 | CHECK_IS_INT(idx); 43 | 44 | if (grad_out.is_cuda()) { 45 | CHECK_CUDA(idx); 46 | } 47 | 48 | at::Tensor output = 49 | torch::zeros({grad_out.size(0), grad_out.size(1), n}, 50 | at::device(grad_out.device()).dtype(at::ScalarType::Float)); 51 | 52 | if (grad_out.is_cuda()) { 53 | group_points_grad_kernel_wrapper( 54 | grad_out.size(0), grad_out.size(1), n, idx.size(1), idx.size(2), 55 | grad_out.data_ptr(), idx.data_ptr(), 56 | output.data_ptr()); 57 | } else { 58 | AT_ASSERT(false, "CPU not supported"); 59 | } 60 | 61 | return output; 62 | } 63 | -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/group_points_gpu.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #include "cuda_utils.h" 5 | 6 | // input: points(b, c, n) idx(b, npoints, nsample) 7 | // output: out(b, c, npoints, nsample) 8 | __global__ void group_points_kernel(int b, int c, int n, int npoints, 9 | int nsample, 10 | const float *__restrict__ points, 11 | const int *__restrict__ idx, 12 | float *__restrict__ out) { 13 | int batch_index = blockIdx.x; 14 | points += batch_index * n * c; 15 | idx += batch_index * npoints * nsample; 16 | out += batch_index * npoints * nsample * c; 17 | 18 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 19 | const int stride = blockDim.y * blockDim.x; 20 | for (int i = index; i < c * npoints; i += stride) { 21 | const int l = i / npoints; 22 | const int j = i % npoints; 23 | for (int k = 0; k < nsample; ++k) { 24 | int ii = idx[j * nsample + k]; 25 | out[(l * npoints + j) * nsample + k] = points[l * n + ii]; 26 | } 27 | } 28 | } 29 | 30 | void group_points_kernel_wrapper(int b, int c, int n, int npoints, int nsample, 31 | const float *points, const int *idx, 32 | float *out) { 33 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 34 | 35 | group_points_kernel<<>>( 36 | b, c, n, npoints, nsample, points, idx, out); 37 | 38 | CUDA_CHECK_ERRORS(); 39 | } 40 | 41 | // input: grad_out(b, c, npoints, nsample), idx(b, npoints, nsample) 42 | // output: grad_points(b, c, n) 43 | __global__ void group_points_grad_kernel(int b, int c, int n, int npoints, 44 | int nsample, 45 | const float *__restrict__ grad_out, 46 | const int *__restrict__ idx, 47 | float *__restrict__ grad_points) { 48 | int batch_index = blockIdx.x; 49 | grad_out += batch_index * npoints * nsample * c; 50 | idx += batch_index * npoints * nsample; 51 | grad_points += batch_index * n * c; 52 | 53 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 54 | const int stride = blockDim.y * blockDim.x; 55 | for (int i = index; i < c * npoints; i += stride) { 56 | const int l = i / npoints; 57 | const int j = i % npoints; 58 | for (int k = 0; k < nsample; ++k) { 59 | int ii = idx[j * nsample + k]; 60 | atomicAdd(grad_points + l * n + ii, 61 | grad_out[(l * npoints + j) * nsample + k]); 62 | } 63 | } 64 | } 65 | 66 | void group_points_grad_kernel_wrapper(int b, int c, int n, int npoints, 67 | int nsample, const float *grad_out, 68 | const int *idx, float *grad_points) { 69 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 70 | 71 | group_points_grad_kernel<<>>( 72 | b, c, n, npoints, nsample, grad_out, idx, grad_points); 73 | 74 | CUDA_CHECK_ERRORS(); 75 | } 76 | -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/interpolate.cpp: -------------------------------------------------------------------------------- 1 | #include "interpolate.h" 2 | #include "utils.h" 3 | 4 | void three_nn_kernel_wrapper(int b, int n, int m, const float *unknown, 5 | const float *known, float *dist2, int *idx); 6 | void three_interpolate_kernel_wrapper(int b, int c, int m, int n, 7 | const float *points, const int *idx, 8 | const float *weight, float *out); 9 | void three_interpolate_grad_kernel_wrapper(int b, int c, int n, int m, 10 | const float *grad_out, 11 | const int *idx, const float *weight, 12 | float *grad_points); 13 | 14 | std::vector three_nn(at::Tensor unknowns, at::Tensor knows) { 15 | CHECK_CONTIGUOUS(unknowns); 16 | CHECK_CONTIGUOUS(knows); 17 | CHECK_IS_FLOAT(unknowns); 18 | CHECK_IS_FLOAT(knows); 19 | 20 | if (unknowns.is_cuda()) { 21 | CHECK_CUDA(knows); 22 | } 23 | 24 | at::Tensor idx = 25 | torch::zeros({unknowns.size(0), unknowns.size(1), 3}, 26 | at::device(unknowns.device()).dtype(at::ScalarType::Int)); 27 | at::Tensor dist2 = 28 | torch::zeros({unknowns.size(0), unknowns.size(1), 3}, 29 | at::device(unknowns.device()).dtype(at::ScalarType::Float)); 30 | 31 | if (unknowns.is_cuda()) { 32 | three_nn_kernel_wrapper(unknowns.size(0), unknowns.size(1), knows.size(1), 33 | unknowns.data_ptr(), knows.data_ptr(), 34 | dist2.data_ptr(), idx.data_ptr()); 35 | } else { 36 | AT_ASSERT(false, "CPU not supported"); 37 | } 38 | 39 | return {dist2, idx}; 40 | } 41 | 42 | at::Tensor three_interpolate(at::Tensor points, at::Tensor idx, 43 | at::Tensor weight) { 44 | CHECK_CONTIGUOUS(points); 45 | CHECK_CONTIGUOUS(idx); 46 | CHECK_CONTIGUOUS(weight); 47 | CHECK_IS_FLOAT(points); 48 | CHECK_IS_INT(idx); 49 | CHECK_IS_FLOAT(weight); 50 | 51 | if (points.is_cuda()) { 52 | CHECK_CUDA(idx); 53 | CHECK_CUDA(weight); 54 | } 55 | 56 | at::Tensor output = 57 | torch::zeros({points.size(0), points.size(1), idx.size(1)}, 58 | at::device(points.device()).dtype(at::ScalarType::Float)); 59 | 60 | if (points.is_cuda()) { 61 | three_interpolate_kernel_wrapper( 62 | points.size(0), points.size(1), points.size(2), idx.size(1), 63 | points.data_ptr(), idx.data_ptr(), weight.data_ptr(), 64 | output.data_ptr()); 65 | } else { 66 | AT_ASSERT(false, "CPU not supported"); 67 | } 68 | 69 | return output; 70 | } 71 | at::Tensor three_interpolate_grad(at::Tensor grad_out, at::Tensor idx, 72 | at::Tensor weight, const int m) { 73 | CHECK_CONTIGUOUS(grad_out); 74 | CHECK_CONTIGUOUS(idx); 75 | CHECK_CONTIGUOUS(weight); 76 | CHECK_IS_FLOAT(grad_out); 77 | CHECK_IS_INT(idx); 78 | CHECK_IS_FLOAT(weight); 79 | 80 | if (grad_out.is_cuda()) { 81 | CHECK_CUDA(idx); 82 | CHECK_CUDA(weight); 83 | } 84 | 85 | at::Tensor output = 86 | torch::zeros({grad_out.size(0), grad_out.size(1), m}, 87 | at::device(grad_out.device()).dtype(at::ScalarType::Float)); 88 | 89 | if (grad_out.is_cuda()) { 90 | three_interpolate_grad_kernel_wrapper( 91 | grad_out.size(0), grad_out.size(1), grad_out.size(2), m, 92 | grad_out.data_ptr(), idx.data_ptr(), 93 | weight.data_ptr(), output.data_ptr()); 94 | } else { 95 | AT_ASSERT(false, "CPU not supported"); 96 | } 97 | 98 | return output; 99 | } 100 | -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/interpolate_gpu.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | 5 | #include "cuda_utils.h" 6 | 7 | // input: unknown(b, n, 3) known(b, m, 3) 8 | // output: dist2(b, n, 3), idx(b, n, 3) 9 | __global__ void three_nn_kernel(int b, int n, int m, 10 | const float *__restrict__ unknown, 11 | const float *__restrict__ known, 12 | float *__restrict__ dist2, 13 | int *__restrict__ idx) { 14 | int batch_index = blockIdx.x; 15 | unknown += batch_index * n * 3; 16 | known += batch_index * m * 3; 17 | dist2 += batch_index * n * 3; 18 | idx += batch_index * n * 3; 19 | 20 | int index = threadIdx.x; 21 | int stride = blockDim.x; 22 | for (int j = index; j < n; j += stride) { 23 | float ux = unknown[j * 3 + 0]; 24 | float uy = unknown[j * 3 + 1]; 25 | float uz = unknown[j * 3 + 2]; 26 | 27 | double best1 = 1e40, best2 = 1e40, best3 = 1e40; 28 | int besti1 = 0, besti2 = 0, besti3 = 0; 29 | for (int k = 0; k < m; ++k) { 30 | float x = known[k * 3 + 0]; 31 | float y = known[k * 3 + 1]; 32 | float z = known[k * 3 + 2]; 33 | float d = (ux - x) * (ux - x) + (uy - y) * (uy - y) + (uz - z) * (uz - z); 34 | if (d < best1) { 35 | best3 = best2; 36 | besti3 = besti2; 37 | best2 = best1; 38 | besti2 = besti1; 39 | best1 = d; 40 | besti1 = k; 41 | } else if (d < best2) { 42 | best3 = best2; 43 | besti3 = besti2; 44 | best2 = d; 45 | besti2 = k; 46 | } else if (d < best3) { 47 | best3 = d; 48 | besti3 = k; 49 | } 50 | } 51 | dist2[j * 3 + 0] = best1; 52 | dist2[j * 3 + 1] = best2; 53 | dist2[j * 3 + 2] = best3; 54 | 55 | idx[j * 3 + 0] = besti1; 56 | idx[j * 3 + 1] = besti2; 57 | idx[j * 3 + 2] = besti3; 58 | } 59 | } 60 | 61 | void three_nn_kernel_wrapper(int b, int n, int m, const float *unknown, 62 | const float *known, float *dist2, int *idx) { 63 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 64 | three_nn_kernel<<>>(b, n, m, unknown, known, 65 | dist2, idx); 66 | 67 | CUDA_CHECK_ERRORS(); 68 | } 69 | 70 | // input: points(b, c, m), idx(b, n, 3), weight(b, n, 3) 71 | // output: out(b, c, n) 72 | __global__ void three_interpolate_kernel(int b, int c, int m, int n, 73 | const float *__restrict__ points, 74 | const int *__restrict__ idx, 75 | const float *__restrict__ weight, 76 | float *__restrict__ out) { 77 | int batch_index = blockIdx.x; 78 | points += batch_index * m * c; 79 | 80 | idx += batch_index * n * 3; 81 | weight += batch_index * n * 3; 82 | 83 | out += batch_index * n * c; 84 | 85 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 86 | const int stride = blockDim.y * blockDim.x; 87 | for (int i = index; i < c * n; i += stride) { 88 | const int l = i / n; 89 | const int j = i % n; 90 | float w1 = weight[j * 3 + 0]; 91 | float w2 = weight[j * 3 + 1]; 92 | float w3 = weight[j * 3 + 2]; 93 | 94 | int i1 = idx[j * 3 + 0]; 95 | int i2 = idx[j * 3 + 1]; 96 | int i3 = idx[j * 3 + 2]; 97 | 98 | out[i] = points[l * m + i1] * w1 + points[l * m + i2] * w2 + 99 | points[l * m + i3] * w3; 100 | } 101 | } 102 | 103 | void three_interpolate_kernel_wrapper(int b, int c, int m, int n, 104 | const float *points, const int *idx, 105 | const float *weight, float *out) { 106 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 107 | three_interpolate_kernel<<>>( 108 | b, c, m, n, points, idx, weight, out); 109 | 110 | CUDA_CHECK_ERRORS(); 111 | } 112 | 113 | // input: grad_out(b, c, n), idx(b, n, 3), weight(b, n, 3) 114 | // output: grad_points(b, c, m) 115 | 116 | __global__ void three_interpolate_grad_kernel( 117 | int b, int c, int n, int m, const float *__restrict__ grad_out, 118 | const int *__restrict__ idx, const float *__restrict__ weight, 119 | float *__restrict__ grad_points) { 120 | int batch_index = blockIdx.x; 121 | grad_out += batch_index * n * c; 122 | idx += batch_index * n * 3; 123 | weight += batch_index * n * 3; 124 | grad_points += batch_index * m * c; 125 | 126 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 127 | const int stride = blockDim.y * blockDim.x; 128 | for (int i = index; i < c * n; i += stride) { 129 | const int l = i / n; 130 | const int j = i % n; 131 | float w1 = weight[j * 3 + 0]; 132 | float w2 = weight[j * 3 + 1]; 133 | float w3 = weight[j * 3 + 2]; 134 | 135 | int i1 = idx[j * 3 + 0]; 136 | int i2 = idx[j * 3 + 1]; 137 | int i3 = idx[j * 3 + 2]; 138 | 139 | atomicAdd(grad_points + l * m + i1, grad_out[i] * w1); 140 | atomicAdd(grad_points + l * m + i2, grad_out[i] * w2); 141 | atomicAdd(grad_points + l * m + i3, grad_out[i] * w3); 142 | } 143 | } 144 | 145 | void three_interpolate_grad_kernel_wrapper(int b, int c, int n, int m, 146 | const float *grad_out, 147 | const int *idx, const float *weight, 148 | float *grad_points) { 149 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 150 | three_interpolate_grad_kernel<<>>( 151 | b, c, n, m, grad_out, idx, weight, grad_points); 152 | 153 | CUDA_CHECK_ERRORS(); 154 | } 155 | -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/sampling.cpp: -------------------------------------------------------------------------------- 1 | #include "sampling.h" 2 | #include "utils.h" 3 | 4 | void gather_points_kernel_wrapper(int b, int c, int n, int npoints, 5 | const float *points, const int *idx, 6 | float *out); 7 | void gather_points_grad_kernel_wrapper(int b, int c, int n, int npoints, 8 | const float *grad_out, const int *idx, 9 | float *grad_points); 10 | 11 | void furthest_point_sampling_kernel_wrapper(int b, int n, int m, 12 | const float *dataset, float *temp, 13 | int *idxs); 14 | 15 | at::Tensor gather_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)}, 27 | at::device(points.device()).dtype(at::ScalarType::Float)); 28 | 29 | if (points.is_cuda()) { 30 | gather_points_kernel_wrapper(points.size(0), points.size(1), points.size(2), 31 | idx.size(1), points.data_ptr(), 32 | idx.data_ptr(), output.data_ptr()); 33 | } else { 34 | AT_ASSERT(false, "CPU not supported"); 35 | } 36 | 37 | return output; 38 | } 39 | 40 | at::Tensor gather_points_grad(at::Tensor grad_out, at::Tensor idx, 41 | const int n) { 42 | CHECK_CONTIGUOUS(grad_out); 43 | CHECK_CONTIGUOUS(idx); 44 | CHECK_IS_FLOAT(grad_out); 45 | CHECK_IS_INT(idx); 46 | 47 | if (grad_out.is_cuda()) { 48 | CHECK_CUDA(idx); 49 | } 50 | 51 | at::Tensor output = 52 | torch::zeros({grad_out.size(0), grad_out.size(1), n}, 53 | at::device(grad_out.device()).dtype(at::ScalarType::Float)); 54 | 55 | if (grad_out.is_cuda()) { 56 | gather_points_grad_kernel_wrapper(grad_out.size(0), grad_out.size(1), n, 57 | idx.size(1), grad_out.data_ptr(), 58 | idx.data_ptr(), 59 | output.data_ptr()); 60 | } else { 61 | AT_ASSERT(false, "CPU not supported"); 62 | } 63 | 64 | return output; 65 | } 66 | at::Tensor furthest_point_sampling(at::Tensor points, const int nsamples) { 67 | CHECK_CONTIGUOUS(points); 68 | CHECK_IS_FLOAT(points); 69 | 70 | at::Tensor output = 71 | torch::zeros({points.size(0), nsamples}, 72 | at::device(points.device()).dtype(at::ScalarType::Int)); 73 | 74 | at::Tensor tmp = 75 | torch::full({points.size(0), points.size(1)}, 1e10, 76 | at::device(points.device()).dtype(at::ScalarType::Float)); 77 | 78 | if (points.is_cuda()) { 79 | furthest_point_sampling_kernel_wrapper( 80 | points.size(0), points.size(1), nsamples, points.data_ptr(), 81 | tmp.data_ptr(), output.data_ptr()); 82 | } else { 83 | AT_ASSERT(false, "CPU not supported"); 84 | } 85 | 86 | return output; 87 | } 88 | -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/pointnet2_ops/_version.py: -------------------------------------------------------------------------------- 1 | __version__ = "3.0.0" 2 | -------------------------------------------------------------------------------- /utils/pointnet2_ops_lib/setup.py: -------------------------------------------------------------------------------- 1 | import glob 2 | import os 3 | import os.path as osp 4 | 5 | from setuptools import find_packages, setup 6 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 7 | 8 | this_dir = osp.dirname(osp.abspath(__file__)) 9 | _ext_src_root = osp.join("pointnet2_ops", "_ext-src") 10 | _ext_sources = glob.glob(osp.join(_ext_src_root, "src", "*.cpp")) + glob.glob( 11 | osp.join(_ext_src_root, "src", "*.cu") 12 | ) 13 | _ext_headers = glob.glob(osp.join(_ext_src_root, "include", "*")) 14 | 15 | requirements = ["torch>=1.4"] 16 | 17 | exec(open(osp.join("pointnet2_ops", "_version.py")).read()) 18 | 19 | os.environ["TORCH_CUDA_ARCH_LIST"] = "3.7+PTX;5.0;6.0;6.1;6.2;7.0;7.5" 20 | setup( 21 | name="pointnet2_ops", 22 | version=__version__, 23 | author="Erik Wijmans", 24 | packages=find_packages(), 25 | install_requires=requirements, 26 | ext_modules=[ 27 | CUDAExtension( 28 | name="pointnet2_ops._ext", 29 | sources=_ext_sources, 30 | extra_compile_args={ 31 | "cxx": ["-O3"], 32 | "nvcc": ["-O3", "-Xfatbin", "-compress-all"], 33 | }, 34 | include_dirs=[osp.join(this_dir, _ext_src_root, "include")], 35 | ) 36 | ], 37 | cmdclass={"build_ext": BuildExtension}, 38 | include_package_data=True, 39 | ) 40 | --------------------------------------------------------------------------------