├── .gitignore ├── LICENSE ├── README.md ├── auxiliary └── ChamferDistancePytorch │ ├── chamfer3D │ ├── .gitignore │ ├── chamfer3D.cu │ ├── chamfer_cuda.cpp │ ├── dist_chamfer_3D.py │ └── setup.py │ ├── chamfer_python.py │ ├── fscore.py │ └── unit_test.py ├── compile_chamfer_distance_op.sh ├── datasets ├── KITTI_mapping.txt ├── flyingthings3d_flownet3d.py ├── flyingthings3d_hplflownet.py ├── generic.py ├── kitti_flownet3d.py └── kitti_hplflownet.py ├── doc ├── poster.png ├── scoop_result.gif └── slides.pdf ├── download_pretrained_models.sh ├── install_environment.sh ├── models ├── gconv.py ├── graph.py ├── refiner.py └── scoop.py ├── scripts ├── evaluate_on_ft3d_s.sh ├── evaluate_on_kitti_o.sh ├── evaluate_on_kitti_o_all_points.sh ├── evaluate_on_kitti_s.sh ├── evaluate_on_kitti_t.sh ├── evaluate_on_kitti_t_all_points.sh ├── evaluate_scoop.py ├── train_on_ft3d_o.sh ├── train_on_ft3d_s.sh ├── train_on_kitti_v.sh ├── train_scoop.py └── visualize_scoop.py └── tools ├── losses.py ├── ot.py ├── reconstruction.py ├── seed.py ├── utils.py └── vis_utils.py /.gitignore: -------------------------------------------------------------------------------- 1 | data/ 2 | experiments/ 3 | pretrained_models/ 4 | .idea/ 5 | __pycache__/ 6 | *.egg-info 7 | *.pyc 8 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | SCOOP: Self-Supervised Correspondence and Optimization-Based Scene Flow 2 | 3 | MIT License 4 | 5 | Copyright (c) 2022 Itai Lang 6 | 7 | Permission is hereby granted, free of charge, to any person obtaining a copy 8 | of this software and associated documentation files (the "Software"), to deal 9 | in the Software without restriction, including without limitation the rights 10 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 11 | copies of the Software, and to permit persons to whom the Software is 12 | furnished to do so, subject to the following conditions: 13 | 14 | The above copyright notice and this permission notice shall be included in all 15 | copies or substantial portions of the Software. 16 | 17 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 18 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 19 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 20 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 21 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 22 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 23 | SOFTWARE. 24 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # SCOOP: Self-Supervised Correspondence and Optimization-Based Scene Flow 2 | [[Project Page]](https://itailang.github.io/SCOOP/) [[Paper]](https://openaccess.thecvf.com/content/CVPR2023/html/Lang_SCOOP_Self-Supervised_Correspondence_and_Optimization-Based_Scene_Flow_CVPR_2023_paper.html) [[Video]](https://www.youtube.com/watch?v=b8MVWGU7V4E) [[Slides]](./doc/slides.pdf) [[Poster]](./doc/poster.png) 3 | 4 | Created by [Itai Lang](https://itailang.github.io/)^1,2, [Dror Aiger](https://research.google/people/DrorAiger/)², [Forrester Cole](http://people.csail.mit.edu/fcole/)², [Shai Avidan](http://www.eng.tau.ac.il/~avidan/)¹, and [Michael Rubinstein](http://people.csail.mit.edu/mrub/)².
5 | ¹Tel Aviv University ²Google Research 6 | 7 | ![scoop_result](./doc/scoop_result.gif) 8 | 9 | ## Abstract 10 | Scene flow estimation is a long-standing problem in computer vision, where the goal is to find the 3D motion of a scene from its consecutive observations. 11 | Recently, there have been efforts to compute the scene flow from 3D point clouds. 12 | A common approach is to train a regression model that consumes source and target point clouds and outputs the per-point translation vector. 13 | An alternative is to learn point matches between the point clouds concurrently with regressing a refinement of the initial correspondence flow. 14 | In both cases, the learning task is very challenging since the flow regression is done in the free 3D space, and a typical solution is to resort to a large annotated synthetic dataset. 15 | 16 | We introduce SCOOP, a new method for scene flow estimation that can be learned on a small amount of data without employing ground-truth flow supervision. 17 | In contrast to previous work, we train a pure correspondence model focused on learning point feature representation and initialize the flow as the difference between a source point and its softly corresponding target point. 18 | Then, in the run-time phase, we directly optimize a flow refinement component with a self-supervised objective, which leads to a coherent and accurate flow field between the point clouds. 19 | Experiments on widespread datasets demonstrate the performance gains achieved by our method compared to existing leading techniques while using a fraction of the training data. 20 | 21 | ## Citation 22 | If you find our work useful in your research, please consider citing: 23 | 24 | @InProceedings{lang2023scoop, 25 | author = {Lang, Itai and Aiger, Dror and Cole, Forrester and Avidan, Shai and Rubinstein, Michael}, 26 | title = {{SCOOP: Self-Supervised Correspondence and Optimization-Based Scene Flow}}, 27 | booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)}, 28 | pages = {5281--5290}, 29 | year = {2023} 30 | } 31 | 32 | ## Installation 33 | The code has been tested with Python 3.6.13, PyTorch 1.6.0, CUDA 10.1, and cuDNN 7.6.5 on Ubuntu 16.04. 34 | 35 | Clone this repository: 36 | ```bash 37 | git clone https://github.com/itailang/SCOOP.git 38 | cd SCOOP/ 39 | ``` 40 | 41 | Create a conda environment: 42 | ```bash 43 | # create and activate a conda environment 44 | conda create -n scoop python=3.6.13 --yes 45 | conda activate scoop 46 | ``` 47 | 48 | Install required packages: 49 | ```bash 50 | sh install_environment.sh 51 | ``` 52 | 53 | Compile the Chamfer Distance op, implemented by [Groueix _et al._](https://github.com/ThibaultGROUEIX/ChamferDistancePytorch) The op is located under `auxiliary/ChamferDistancePytorch/chamfer3D` folder. The following compilation script uses a CUDA 10.1 path. If needed, modify script to point to your CUDA path. Then, use: 54 | ```bash 55 | sh compile_chamfer_distance_op.sh 56 | ``` 57 | 58 | The compilation results should be created under `auxiliary/ChamferDistancePytorch/chamfer3D/build` folder. 59 | 60 | ## Usage 61 | 62 | ### Data 63 | Create folders for the data: 64 | ```bash 65 | mkdir ./data/ 66 | mkdir ./data/FlowNet3D/ 67 | mkdir ./data/HPLFlowNet/ 68 | ``` 69 | 70 | We use the point cloud data version prepared Liu _et al._ from their work [FlowNet3D](https://openaccess.thecvf.com/content_CVPR_2019/html/Liu_FlowNet3D_Learning_Scene_Flow_in_3D_Point_Clouds_CVPR_2019_paper.html). Please follow their [code](https://github.com/xingyul/flownet3d) to acquire the data. 71 | * Put the preprocessed FlyingThings3D dataset at `./data/FlowNet3D/data_processed_maxcut_35_20k_2k_8192/`. This data is denoted as FT3D_o. 72 | * Put the preprocessed KITTI dataset at `./data/flownet3d/kitti_rm_ground/`. This dataset is denoted as KITTI_o and its subsets are denoted as KITTI_v and KITTI_t. 73 | 74 | We also use the point cloud data version prepared Gu _et al._ from their work [HPLFlowNet](https://openaccess.thecvf.com/content_CVPR_2019/html/Gu_HPLFlowNet_Hierarchical_Permutohedral_Lattice_FlowNet_for_Scene_Flow_Estimation_on_CVPR_2019_paper.html). Please follow their [code](https://github.com/laoreja/HPLFlowNet) to acquire the data. 75 | * Put the preprocessed FlyingThings3D dataset at `./data/HPLFlowNet/FlyingThings3D_subset_processed_35m/`. This dataset is denoted as FT3D_s. 76 | * Put the preprocessed KITTI dataset at `./data/flownet3d/KITTI_processed_occ_final/`. This dataset is denoted as KITTI_s. 77 | 78 | Note that you may put the data elsewhere and create a symbolic link to the actual location. For example: 79 | ```bash 80 | ln -s /path/to/the/actual/dataset/location ./data/FlowNet3D/data_processed_maxcut_35_20k_2k_8192 81 | ``` 82 | 83 | ### Training and Evaluation 84 | Switch to the `scripts` folder: 85 | ```bash 86 | cd ./scripts 87 | ``` 88 | 89 | #### FT3D_o / KITTI_o 90 | To train a model on 1,800 examples from the train set of FT3D_o, run the following command: 91 | ```bash 92 | sh train_on_ft3d_o.sh 93 | ``` 94 | 95 | Evaluate this model on KITTI_o with 2,048 point per point cloud using the following command: 96 | ```bash 97 | sh evaluate_on_kitti_o.sh 98 | ``` 99 | 100 | The results will be saved to the file `./experiments/ft3d_o_1800_examples/log_evaluation_kitti_o.txt`. 101 | 102 | Evaluate this model on KITTI_o with all the points in the point clouds using the following command: 103 | ```bash 104 | sh evaluate_on_kitti_o_all_point.sh 105 | ``` 106 | 107 | The results will be saved to the file `./experiments/ft3d_o_1800_examples/log_evaluation_kitti_o_all_points.txt`. 108 | 109 | #### KITTI_v / KITTI_t 110 | To train a model on the 100 examples of KITTI_v, run the following command: 111 | ```bash 112 | sh train_on_kitti_v.sh 113 | ``` 114 | 115 | Evaluate this model on KITTI_t with 2,048 point per point cloud using the following command: 116 | ```bash 117 | sh evaluate_on_kitti_t.sh 118 | ``` 119 | 120 | The results will be saved to the file `./experiments/kitti_v_100_examples/log_evaluation_kitti_t.txt`. 121 | 122 | Evaluate this model on KITTI_t with all the points in the point clouds using the following command: 123 | ```bash 124 | sh evaluate_on_kitti_t_all_points.sh 125 | ``` 126 | 127 | The results will be saved to the file `./experiments/kitti_v_100_examples/log_evaluation_kitti_t_all_points.txt`. 128 | 129 | #### FT3D_s / KITTI_s, FT3D_s / FT3D_s 130 | To train a model on 1,800 examples from the train set FT3D_s, run the following command: 131 | ```bash 132 | sh train_on_ft3d_s.sh 133 | ``` 134 | 135 | Evaluate this model on KITTI_s with 8,192 point per point cloud using the following command: 136 | ```bash 137 | sh evaluate_on_kitti_s.sh 138 | ``` 139 | 140 | The results will be saved to the file `./experiments/ft3d_s_1800_examples/log_evaluation_kitti_s.txt`. 141 | 142 | Evaluate this model on the test set of FT3D_s with 8,192 point per point cloud using the following command: 143 | ```bash 144 | sh evaluate_on_ft3d_s.sh 145 | ``` 146 | 147 | The results will be saved to the file `./experiments/ft3d_s_1800_examples/log_evaluation_ft3d_s.txt`. 148 | 149 | #### Visualization 150 | First, save results for visualization by adding the flag `--save_pc_res 1` when running the evaluation script. For example, the [script for evaluating on KITTI_t](./scripts/evaluate_on_kitti_t.sh). 151 | The results will be saved to the folder `./experiments/kitti_v_100_examples/pc_res/`. 152 | 153 | Then, select the scene index that you would like to visualize and run the visualization script. For example, visualizing scene index #1 from KITTI_t: 154 | ```python 155 | python visualize_scoop.py --res_dir ./../experiments/kitti_v_100_examples/pc_res --res_idx 1 156 | ``` 157 | 158 | The visualizations will be saved to the folder `./experiments/kitti_v_100_examples/pc_res/vis/`. 159 | 160 | ### Evaluation with Pretrained Models 161 | 162 | First, download our pretrained models with the following command: 163 | ```bash 164 | bash download_pretrained_models.sh 165 | ``` 166 | 167 | The models (about 2MB) will be saved under `pretrained_models` folder in the following structure: 168 | ``` 169 | pretrained_models/ 170 | ├── ft3d_o_1800_examples/model_e100.tar 171 | ├── ft3d_s_1800_examples/model_e060.tar 172 | ├── kitti_v_100_examples/model_e400.tar 173 | ``` 174 | 175 | Then, use the evaluation commands mentioned in section [Training and Evaluation](#training-and-Evaluation), 176 | after changing the `experiments` folder in the evaluation scripts to the `pretrained_models` folder. 177 | 178 | ## License 179 | This project is licensed under the terms of the MIT license (see the [LICENSE](./LICENSE) file for more details). 180 | 181 | ## Acknowledgment 182 | Our code builds upon the code provided by [Puy _et al._](https://github.com/valeoai/FLOT), [Groueix _et al._](https://github.com/ThibaultGROUEIX/ChamferDistancePytorch), [Liu _et al._](https://github.com/xingyul/flownet3d), and [Gu _et al._](https://github.com/laoreja/HPLFlowNet) We thank the authors for sharing their code. 183 | -------------------------------------------------------------------------------- /auxiliary/ChamferDistancePytorch/chamfer3D/.gitignore: -------------------------------------------------------------------------------- 1 | build/ 2 | chamfer_3D.egg-info/ 3 | dist/ 4 | -------------------------------------------------------------------------------- /auxiliary/ChamferDistancePytorch/chamfer3D/chamfer3D.cu: -------------------------------------------------------------------------------- 1 | 2 | #include 3 | #include 4 | 5 | #include 6 | #include 7 | 8 | #include 9 | 10 | 11 | 12 | __global__ void NmDistanceKernel(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i){ 13 | const int batch=512; 14 | __shared__ float buf[batch*3]; 15 | for (int i=blockIdx.x;ibest){ 127 | result[(i*n+j)]=best; 128 | result_i[(i*n+j)]=best_i; 129 | } 130 | } 131 | __syncthreads(); 132 | } 133 | } 134 | } 135 | // int chamfer_cuda_forward(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i,float * result2,int * result2_i, cudaStream_t stream){ 136 | int chamfer_cuda_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2, at::Tensor idx1, at::Tensor idx2){ 137 | 138 | const auto batch_size = xyz1.size(0); 139 | const auto n = xyz1.size(1); //num_points point cloud A 140 | const auto m = xyz2.size(1); //num_points point cloud B 141 | 142 | NmDistanceKernel<<>>(batch_size, n, xyz1.data(), m, xyz2.data(), dist1.data(), idx1.data()); 143 | NmDistanceKernel<<>>(batch_size, m, xyz2.data(), n, xyz1.data(), dist2.data(), idx2.data()); 144 | 145 | cudaError_t err = cudaGetLastError(); 146 | if (err != cudaSuccess) { 147 | printf("error in nnd updateOutput: %s\n", cudaGetErrorString(err)); 148 | //THError("aborting"); 149 | return 0; 150 | } 151 | return 1; 152 | 153 | 154 | } 155 | __global__ void NmDistanceGradKernel(int b,int n,const float * xyz1,int m,const float * xyz2,const float * grad_dist1,const int * idx1,float * grad_xyz1,float * grad_xyz2){ 156 | for (int i=blockIdx.x;i>>(batch_size,n,xyz1.data(),m,xyz2.data(),graddist1.data(),idx1.data(),gradxyz1.data(),gradxyz2.data()); 185 | NmDistanceGradKernel<<>>(batch_size,m,xyz2.data(),n,xyz1.data(),graddist2.data(),idx2.data(),gradxyz2.data(),gradxyz1.data()); 186 | 187 | cudaError_t err = cudaGetLastError(); 188 | if (err != cudaSuccess) { 189 | printf("error in nnd get grad: %s\n", cudaGetErrorString(err)); 190 | //THError("aborting"); 191 | return 0; 192 | } 193 | return 1; 194 | 195 | } 196 | 197 | -------------------------------------------------------------------------------- /auxiliary/ChamferDistancePytorch/chamfer3D/chamfer_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | ///TMP 5 | //#include "common.h" 6 | /// NOT TMP 7 | 8 | 9 | int chamfer_cuda_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2, at::Tensor idx1, at::Tensor idx2); 10 | 11 | 12 | int chamfer_cuda_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz1, at::Tensor gradxyz2, at::Tensor graddist1, at::Tensor graddist2, at::Tensor idx1, at::Tensor idx2); 13 | 14 | 15 | 16 | 17 | int chamfer_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2, at::Tensor idx1, at::Tensor idx2) { 18 | return chamfer_cuda_forward(xyz1, xyz2, dist1, dist2, idx1, idx2); 19 | } 20 | 21 | 22 | int chamfer_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz1, at::Tensor gradxyz2, at::Tensor graddist1, 23 | at::Tensor graddist2, at::Tensor idx1, at::Tensor idx2) { 24 | 25 | return chamfer_cuda_backward(xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2); 26 | } 27 | 28 | 29 | 30 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 31 | m.def("forward", &chamfer_forward, "chamfer forward (CUDA)"); 32 | m.def("backward", &chamfer_backward, "chamfer backward (CUDA)"); 33 | } -------------------------------------------------------------------------------- /auxiliary/ChamferDistancePytorch/chamfer3D/dist_chamfer_3D.py: -------------------------------------------------------------------------------- 1 | from torch import nn 2 | from torch.autograd import Function 3 | import torch 4 | import importlib 5 | import os 6 | chamfer_found = importlib.find_loader("chamfer_3D") is not None 7 | if not chamfer_found: 8 | ## Cool trick from https://github.com/chrdiller 9 | print("Jitting Chamfer 3D") 10 | 11 | from torch.utils.cpp_extension import load 12 | chamfer_3D = load(name="chamfer_3D", 13 | sources=[ 14 | "/".join(os.path.abspath(__file__).split('/')[:-1] + ["chamfer_cuda.cpp"]), 15 | "/".join(os.path.abspath(__file__).split('/')[:-1] + ["chamfer3D.cu"]), 16 | ]) 17 | print("Loaded JIT 3D CUDA chamfer distance") 18 | 19 | else: 20 | import chamfer_3D 21 | print("Loaded compiled 3D CUDA chamfer distance") 22 | 23 | 24 | # Chamfer's distance module @thibaultgroueix 25 | # GPU tensors only 26 | class chamfer_3DFunction(Function): 27 | @staticmethod 28 | def forward(ctx, xyz1, xyz2): 29 | batchsize, n, _ = xyz1.size() 30 | _, m, _ = xyz2.size() 31 | device = xyz1.device 32 | 33 | dist1 = torch.zeros(batchsize, n) 34 | dist2 = torch.zeros(batchsize, m) 35 | 36 | idx1 = torch.zeros(batchsize, n).type(torch.IntTensor) 37 | idx2 = torch.zeros(batchsize, m).type(torch.IntTensor) 38 | 39 | dist1 = dist1.to(device) 40 | dist2 = dist2.to(device) 41 | idx1 = idx1.to(device) 42 | idx2 = idx2.to(device) 43 | torch.cuda.set_device(device) 44 | 45 | chamfer_3D.forward(xyz1, xyz2, dist1, dist2, idx1, idx2) 46 | ctx.save_for_backward(xyz1, xyz2, idx1, idx2) 47 | return dist1, dist2, idx1, idx2 48 | 49 | @staticmethod 50 | def backward(ctx, graddist1, graddist2, gradidx1, gradidx2): 51 | xyz1, xyz2, idx1, idx2 = ctx.saved_tensors 52 | graddist1 = graddist1.contiguous() 53 | graddist2 = graddist2.contiguous() 54 | device = graddist1.device 55 | 56 | gradxyz1 = torch.zeros(xyz1.size()) 57 | gradxyz2 = torch.zeros(xyz2.size()) 58 | 59 | gradxyz1 = gradxyz1.to(device) 60 | gradxyz2 = gradxyz2.to(device) 61 | chamfer_3D.backward( 62 | xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2 63 | ) 64 | return gradxyz1, gradxyz2 65 | 66 | 67 | class chamfer_3DDist(nn.Module): 68 | def __init__(self): 69 | super(chamfer_3DDist, self).__init__() 70 | 71 | def forward(self, input1, input2): 72 | input1 = input1.contiguous() 73 | input2 = input2.contiguous() 74 | return chamfer_3DFunction.apply(input1, input2) 75 | 76 | -------------------------------------------------------------------------------- /auxiliary/ChamferDistancePytorch/chamfer3D/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 3 | 4 | setup( 5 | name='chamfer_3D', 6 | ext_modules=[ 7 | CUDAExtension('chamfer_3D', [ 8 | "/".join(__file__.split('/')[:-1] + ['chamfer_cuda.cpp']), 9 | "/".join(__file__.split('/')[:-1] + ['chamfer3D.cu']), 10 | ]), 11 | ], 12 | cmdclass={ 13 | 'build_ext': BuildExtension 14 | }) 15 | -------------------------------------------------------------------------------- /auxiliary/ChamferDistancePytorch/chamfer_python.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | def pairwise_dist(x, y): 5 | xx, yy, zz = torch.mm(x, x.t()), torch.mm(y, y.t()), torch.mm(x, y.t()) 6 | rx = xx.diag().unsqueeze(0).expand_as(xx) 7 | ry = yy.diag().unsqueeze(0).expand_as(yy) 8 | P = rx.t() + ry - 2 * zz 9 | return P 10 | 11 | 12 | def NN_loss(x, y, dim=0): 13 | dist = pairwise_dist(x, y) 14 | values, indices = dist.min(dim=dim) 15 | return values.mean() 16 | 17 | 18 | def distChamfer(a, b): 19 | """ 20 | :param a: Pointclouds Batch x nul_points x dim 21 | :param b: Pointclouds Batch x nul_points x dim 22 | :return: 23 | -closest point on b of points from a 24 | -closest point on a of points from b 25 | -idx of closest point on b of points from a 26 | -idx of closest point on a of points from b 27 | Works for pointcloud of any dimension 28 | """ 29 | x, y = a.double(), b.double() 30 | bs, num_points_x, points_dim = x.size() 31 | bs, num_points_y, points_dim = y.size() 32 | 33 | xx = torch.pow(x, 2).sum(2) 34 | yy = torch.pow(y, 2).sum(2) 35 | zz = torch.bmm(x, y.transpose(2, 1)) 36 | rx = xx.unsqueeze(1).expand(bs, num_points_y, num_points_x) # Diagonal elements xx 37 | ry = yy.unsqueeze(1).expand(bs, num_points_x, num_points_y) # Diagonal elements yy 38 | P = rx.transpose(2, 1) + ry - 2 * zz 39 | return torch.min(P, 2)[0].float(), torch.min(P, 1)[0].float(), torch.min(P, 2)[1].int(), torch.min(P, 1)[1].int() 40 | 41 | -------------------------------------------------------------------------------- /auxiliary/ChamferDistancePytorch/fscore.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | def fscore(dist1, dist2, threshold=0.001): 4 | """ 5 | Calculates the F-score between two point clouds with the corresponding threshold value. 6 | :param dist1: Batch, N-Points 7 | :param dist2: Batch, N-Points 8 | :param th: float 9 | :return: fscore, precision, recall 10 | """ 11 | 12 | precision_1 = torch.mean((dist1 < threshold).float(), dim=1) 13 | precision_2 = torch.mean((dist2 < threshold).float(), dim=1) 14 | fscore = 2 * precision_1 * precision_2 / (precision_1 + precision_2) 15 | fscore[torch.isnan(fscore)] = 0 16 | return fscore, precision_1, precision_2 17 | 18 | -------------------------------------------------------------------------------- /auxiliary/ChamferDistancePytorch/unit_test.py: -------------------------------------------------------------------------------- 1 | import torch, time 2 | import chamfer2D.dist_chamfer_2D 3 | import chamfer3D.dist_chamfer_3D 4 | import chamfer5D.dist_chamfer_5D 5 | import chamfer_python 6 | 7 | cham2D = chamfer2D.dist_chamfer_2D.chamfer_2DDist() 8 | cham3D = chamfer3D.dist_chamfer_3D.chamfer_3DDist() 9 | cham5D = chamfer5D.dist_chamfer_5D.chamfer_5DDist() 10 | 11 | from torch.autograd import Variable 12 | from fscore import fscore 13 | 14 | def test_chamfer(distChamfer, dim): 15 | points1 = torch.rand(4, 100, dim).cuda() 16 | points2 = torch.rand(4, 200, dim, requires_grad=True).cuda() 17 | dist1, dist2, idx1, idx2= distChamfer(points1, points2) 18 | 19 | loss = torch.sum(dist1) 20 | loss.backward() 21 | 22 | mydist1, mydist2, myidx1, myidx2 = chamfer_python.distChamfer(points1, points2) 23 | d1 = (dist1 - mydist1) ** 2 24 | d2 = (dist2 - mydist2) ** 2 25 | assert ( 26 | torch.mean(d1) + torch.mean(d2) < 0.00000001 27 | ), "chamfer cuda and chamfer normal are not giving the same results" 28 | 29 | xd1 = idx1 - myidx1 30 | xd2 = idx2 - myidx2 31 | assert ( 32 | torch.norm(xd1.float()) + torch.norm(xd2.float()) == 0 33 | ), "chamfer cuda and chamfer normal are not giving the same results" 34 | print(f"fscore :", fscore(dist1, dist2)) 35 | print("Unit test passed") 36 | 37 | 38 | def timings(distChamfer, dim): 39 | p1 = torch.rand(32, 2000, dim).cuda() 40 | p2 = torch.rand(32, 1000, dim).cuda() 41 | print("Timings : Start CUDA version") 42 | start = time.time() 43 | num_it = 100 44 | for i in range(num_it): 45 | points1 = Variable(p1, requires_grad=True) 46 | points2 = Variable(p2) 47 | mydist1, mydist2, idx1, idx2 = distChamfer(points1, points2) 48 | loss = torch.sum(mydist1) 49 | loss.backward() 50 | print(f"Ellapsed time forward backward is {(time.time() - start)/num_it} seconds.") 51 | 52 | 53 | print("Timings : Start Pythonic version") 54 | start = time.time() 55 | for i in range(num_it): 56 | points1 = Variable(p1, requires_grad=True) 57 | points2 = Variable(p2) 58 | mydist1, mydist2, idx1, idx2 = chamfer_python.distChamfer(points1, points2) 59 | loss = torch.sum(mydist1) 60 | loss.backward() 61 | print(f"Ellapsed time forward backward is {(time.time() - start)/num_it} seconds.") 62 | 63 | 64 | 65 | dims = [2,3,5] 66 | for i,cham in enumerate([cham2D, cham3D, cham5D]): 67 | print(f"testing Chamfer {dims[i]}D") 68 | test_chamfer(cham, dims[i]) 69 | timings(cham, dims[i]) 70 | -------------------------------------------------------------------------------- /compile_chamfer_distance_op.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | cd ./auxiliary/ChamferDistancePytorch/chamfer3D 4 | 5 | # update paths 6 | export CUDA_HOME=/usr/local/cuda-10.1 7 | export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda-10.1/lib64:/usr/local/cuda-10.1/extras/CUPTI/lib64 8 | export PATH=$PATH:$CUDA_HOME/bin 9 | 10 | # compile 3D-Chamfer Distance op 11 | python ./setup.py install 12 | -------------------------------------------------------------------------------- /datasets/KITTI_mapping.txt: -------------------------------------------------------------------------------- 1 | 2 | 3 | 2011_09_26 2011_09_26_drive_0005_sync 0000000010 4 | 2011_09_26 2011_09_26_drive_0005_sync 0000000059 5 | 6 | 7 | 8 | 2011_09_26 2011_09_26_drive_0009_sync 0000000354 9 | 2011_09_26 2011_09_26_drive_0009_sync 0000000364 10 | 2011_09_26 2011_09_26_drive_0009_sync 0000000374 11 | 2011_09_26 2011_09_26_drive_0009_sync 0000000384 12 | 2011_09_26 2011_09_26_drive_0009_sync 0000000394 13 | 2011_09_26 2011_09_26_drive_0009_sync 0000000414 14 | 2011_09_26 2011_09_26_drive_0011_sync 0000000111 15 | 2011_09_26 2011_09_26_drive_0011_sync 0000000127 16 | 2011_09_26 2011_09_26_drive_0011_sync 0000000147 17 | 2011_09_26 2011_09_26_drive_0011_sync 0000000157 18 | 2011_09_26 2011_09_26_drive_0011_sync 0000000167 19 | 2011_09_26 2011_09_26_drive_0013_sync 0000000010 20 | 2011_09_26 2011_09_26_drive_0013_sync 0000000020 21 | 2011_09_26 2011_09_26_drive_0013_sync 0000000040 22 | 2011_09_26 2011_09_26_drive_0013_sync 0000000070 23 | 2011_09_26 2011_09_26_drive_0014_sync 0000000010 24 | 2011_09_26 2011_09_26_drive_0014_sync 0000000020 25 | 2011_09_26 2011_09_26_drive_0014_sync 0000000030 26 | 2011_09_26 2011_09_26_drive_0014_sync 0000000050 27 | 2011_09_26 2011_09_26_drive_0014_sync 0000000060 28 | 2011_09_26 2011_09_26_drive_0014_sync 0000000129 29 | 2011_09_26 2011_09_26_drive_0014_sync 0000000141 30 | 2011_09_26 2011_09_26_drive_0014_sync 0000000152 31 | 2011_09_26 2011_09_26_drive_0014_sync 0000000172 32 | 2011_09_26 2011_09_26_drive_0014_sync 0000000192 33 | 2011_09_26 2011_09_26_drive_0014_sync 0000000213 34 | 2011_09_26 2011_09_26_drive_0014_sync 0000000240 35 | 2011_09_26 2011_09_26_drive_0015_sync 0000000187 36 | 2011_09_26 2011_09_26_drive_0015_sync 0000000197 37 | 2011_09_26 2011_09_26_drive_0015_sync 0000000209 38 | 2011_09_26 2011_09_26_drive_0015_sync 0000000219 39 | 2011_09_26 2011_09_26_drive_0015_sync 0000000229 40 | 2011_09_26 2011_09_26_drive_0015_sync 0000000239 41 | 2011_09_26 2011_09_26_drive_0015_sync 0000000264 42 | 2011_09_26 2011_09_26_drive_0015_sync 0000000273 43 | 2011_09_26 2011_09_26_drive_0015_sync 0000000286 44 | 2011_09_26 2011_09_26_drive_0017_sync 0000000010 45 | 2011_09_26 2011_09_26_drive_0017_sync 0000000030 46 | 2011_09_26 2011_09_26_drive_0017_sync 0000000040 47 | 2011_09_26 2011_09_26_drive_0017_sync 0000000050 48 | 2011_09_26 2011_09_26_drive_0018_sync 0000000046 49 | 2011_09_26 2011_09_26_drive_0018_sync 0000000066 50 | 2011_09_26 2011_09_26_drive_0018_sync 0000000076 51 | 2011_09_26 2011_09_26_drive_0018_sync 0000000086 52 | 2011_09_26 2011_09_26_drive_0018_sync 0000000096 53 | 2011_09_26 2011_09_26_drive_0018_sync 0000000106 54 | 2011_09_26 2011_09_26_drive_0018_sync 0000000133 55 | 2011_09_26 2011_09_26_drive_0019_sync 0000000030 56 | 2011_09_26 2011_09_26_drive_0019_sync 0000000087 57 | 2011_09_26 2011_09_26_drive_0019_sync 0000000097 58 | 2011_09_26 2011_09_26_drive_0022_sync 0000000634 59 | 2011_09_26 2011_09_26_drive_0022_sync 0000000644 60 | 2011_09_26 2011_09_26_drive_0022_sync 0000000654 61 | 2011_09_26 2011_09_26_drive_0027_sync 0000000053 62 | 2011_09_26 2011_09_26_drive_0027_sync 0000000103 63 | 2011_09_26 2011_09_26_drive_0028_sync 0000000071 64 | 2011_09_26 2011_09_26_drive_0028_sync 0000000118 65 | 2011_09_26 2011_09_26_drive_0028_sync 0000000228 66 | 2011_09_26 2011_09_26_drive_0028_sync 0000000269 67 | 2011_09_26 2011_09_26_drive_0028_sync 0000000284 68 | 2011_09_26 2011_09_26_drive_0028_sync 0000000303 69 | 2011_09_26 2011_09_26_drive_0028_sync 0000000313 70 | 2011_09_26 2011_09_26_drive_0028_sync 0000000378 71 | 2011_09_26 2011_09_26_drive_0029_sync 0000000016 72 | 2011_09_26 2011_09_26_drive_0029_sync 0000000123 73 | 2011_09_26 2011_09_26_drive_0032_sync 0000000095 74 | 2011_09_26 2011_09_26_drive_0032_sync 0000000114 75 | 2011_09_26 2011_09_26_drive_0032_sync 0000000125 76 | 2011_09_26 2011_09_26_drive_0032_sync 0000000207 77 | 2011_09_26 2011_09_26_drive_0032_sync 0000000218 78 | 2011_09_26 2011_09_26_drive_0032_sync 0000000330 79 | 2011_09_26 2011_09_26_drive_0032_sync 0000000340 80 | 2011_09_26 2011_09_26_drive_0032_sync 0000000350 81 | 2011_09_26 2011_09_26_drive_0032_sync 0000000360 82 | 2011_09_26 2011_09_26_drive_0032_sync 0000000378 83 | 84 | 2011_09_26 2011_09_26_drive_0036_sync 0000000054 85 | 2011_09_26 2011_09_26_drive_0036_sync 0000000402 86 | 2011_09_26 2011_09_26_drive_0046_sync 0000000052 87 | 2011_09_26 2011_09_26_drive_0046_sync 0000000062 88 | 89 | 2011_09_26 2011_09_26_drive_0051_sync 0000000023 90 | 2011_09_26 2011_09_26_drive_0051_sync 0000000218 91 | 2011_09_26 2011_09_26_drive_0051_sync 0000000230 92 | 2011_09_26 2011_09_26_drive_0051_sync 0000000282 93 | 2011_09_26 2011_09_26_drive_0051_sync 0000000292 94 | 2011_09_26 2011_09_26_drive_0051_sync 0000000302 95 | 2011_09_26 2011_09_26_drive_0051_sync 0000000312 96 | 2011_09_26 2011_09_26_drive_0051_sync 0000000322 97 | 2011_09_26 2011_09_26_drive_0051_sync 0000000342 98 | 2011_09_26 2011_09_26_drive_0051_sync 0000000356 99 | 2011_09_26 2011_09_26_drive_0051_sync 0000000379 100 | 101 | 102 | 103 | 104 | 105 | 106 | 2011_09_26 2011_09_26_drive_0056_sync 0000000010 107 | 2011_09_26 2011_09_26_drive_0056_sync 0000000082 108 | 2011_09_26 2011_09_26_drive_0056_sync 0000000122 109 | 2011_09_26 2011_09_26_drive_0056_sync 0000000132 110 | 2011_09_26 2011_09_26_drive_0056_sync 0000000191 111 | 2011_09_26 2011_09_26_drive_0056_sync 0000000201 112 | 2011_09_26 2011_09_26_drive_0056_sync 0000000282 113 | 2011_09_26 2011_09_26_drive_0057_sync 0000000125 114 | 2011_09_26 2011_09_26_drive_0057_sync 0000000140 115 | 2011_09_26 2011_09_26_drive_0057_sync 0000000176 116 | 2011_09_26 2011_09_26_drive_0057_sync 0000000299 117 | 2011_09_26 2011_09_26_drive_0057_sync 0000000319 118 | 2011_09_26 2011_09_26_drive_0057_sync 0000000339 119 | 2011_09_26 2011_09_26_drive_0059_sync 0000000026 120 | 2011_09_26 2011_09_26_drive_0059_sync 0000000046 121 | 2011_09_26 2011_09_26_drive_0059_sync 0000000137 122 | 2011_09_26 2011_09_26_drive_0059_sync 0000000150 123 | 2011_09_26 2011_09_26_drive_0059_sync 0000000260 124 | 2011_09_26 2011_09_26_drive_0059_sync 0000000280 125 | 2011_09_26 2011_09_26_drive_0059_sync 0000000290 126 | 2011_09_26 2011_09_26_drive_0059_sync 0000000300 127 | 2011_09_26 2011_09_26_drive_0059_sync 0000000310 128 | 2011_09_26 2011_09_26_drive_0059_sync 0000000320 129 | 2011_09_26 2011_09_26_drive_0070_sync 0000000069 130 | 2011_09_26 2011_09_26_drive_0070_sync 0000000224 131 | 2011_09_26 2011_09_26_drive_0084_sync 0000000084 132 | 2011_09_26 2011_09_26_drive_0084_sync 0000000179 133 | 2011_09_26 2011_09_26_drive_0084_sync 0000000238 134 | 135 | 136 | 137 | 138 | 139 | 140 | 141 | 142 | 2011_09_26 2011_09_26_drive_0096_sync 0000000020 143 | 2011_09_26 2011_09_26_drive_0096_sync 0000000278 144 | 2011_09_26 2011_09_26_drive_0096_sync 0000000381 145 | 2011_09_26 2011_09_26_drive_0101_sync 0000000109 146 | 2011_09_26 2011_09_26_drive_0101_sync 0000000175 147 | 2011_09_26 2011_09_26_drive_0101_sync 0000000447 148 | 2011_09_26 2011_09_26_drive_0101_sync 0000000457 149 | 2011_09_26 2011_09_26_drive_0101_sync 0000000809 150 | 2011_09_26 2011_09_26_drive_0104_sync 0000000015 151 | 2011_09_26 2011_09_26_drive_0104_sync 0000000035 152 | 153 | 154 | 155 | 156 | 2011_09_28 2011_09_28_drive_0002_sync 0000000343 157 | 158 | 2011_09_29 2011_09_29_drive_0004_sync 0000000036 159 | 2011_09_29 2011_09_29_drive_0004_sync 0000000079 160 | 2011_09_29 2011_09_29_drive_0004_sync 0000000094 161 | 2011_09_29 2011_09_29_drive_0004_sync 0000000105 162 | 2011_09_29 2011_09_29_drive_0004_sync 0000000162 163 | 2011_09_29 2011_09_29_drive_0004_sync 0000000258 164 | 2011_09_29 2011_09_29_drive_0004_sync 0000000285 165 | 2011_09_29 2011_09_29_drive_0004_sync 0000000308 166 | 167 | 168 | 169 | 2011_09_29 2011_09_29_drive_0071_sync 0000000059 170 | 2011_09_29 2011_09_29_drive_0071_sync 0000000943 171 | 172 | 173 | 174 | 175 | 176 | 177 | 178 | 179 | 180 | 181 | 182 | 183 | 184 | 185 | 186 | 187 | 188 | 189 | 190 | 191 | 192 | 193 | 194 | 195 | 196 | 197 | 198 | 199 | 200 | 2011_10_03 2011_10_03_drive_0047_sync 0000000556 201 | -------------------------------------------------------------------------------- /datasets/flyingthings3d_flownet3d.py: -------------------------------------------------------------------------------- 1 | import os 2 | import glob 3 | import numpy as np 4 | from .generic import SceneFlowDataset 5 | 6 | 7 | class FT3D(SceneFlowDataset): 8 | def __init__(self, root_dir, nb_points, all_points, mode, nb_examples): 9 | """ 10 | Construct the FlyingThing3D datatset as in: 11 | Liu, X., Qi, C.R., Guibas, L.J.: FlowNet3D: Learning scene ﬂow in 3D 12 | point clouds. IEEE Conf. Computer Vision and Pattern Recognition 13 | (CVPR). pp. 529–537 (2019) 14 | 15 | Parameters 16 | ---------- 17 | root_dir : str 18 | Path to root directory containing the datasets. 19 | nb_points : int 20 | Maximum number of points in point clouds. 21 | all_points : bool 22 | Whether to use all point in the point cloud (in chucks of nb_points) or only nb_points. 23 | mode : str 24 | 'train': training dataset. 25 | 26 | 'val': validation dataset. 27 | 28 | 'test': test dataset 29 | nb_examples: int 30 | Number of examples for the dataset. Active if > 0. 31 | 32 | """ 33 | 34 | super(FT3D, self).__init__(nb_points, all_points) 35 | self.mode = mode 36 | self.nb_examples = nb_examples 37 | self.root_dir = root_dir 38 | self.filenames = self.get_file_list() 39 | self.filename_curr = "" 40 | 41 | def __len__(self): 42 | 43 | return len(self.filenames) 44 | 45 | def get_file_list(self): 46 | """ 47 | Find and filter out paths to all examples in the dataset. 48 | 49 | """ 50 | 51 | # 52 | if self.mode == "train" or self.mode == "val": 53 | pattern = "TRAIN_*.npz" 54 | elif self.mode == "test": 55 | pattern = "TEST_*.npz" 56 | else: 57 | raise ValueError("Mode " + str(self.mode) + "unknown.") 58 | filenames = glob.glob(os.path.join(self.root_dir, pattern)) 59 | 60 | # Remove one sample containing a nan value in train set 61 | scan_with_nan_value = os.path.join( 62 | self.root_dir, "TRAIN_C_0140_left_0006-0.npz" 63 | ) 64 | if scan_with_nan_value in filenames: 65 | filenames.remove(scan_with_nan_value) 66 | 67 | # Remove samples with all points occluded in train set 68 | scan_with_points_all_occluded = [ 69 | "TRAIN_A_0364_left_0008-0.npz", 70 | "TRAIN_A_0364_left_0009-0.npz", 71 | "TRAIN_A_0658_left_0014-0.npz", 72 | "TRAIN_B_0053_left_0009-0.npz", 73 | "TRAIN_B_0053_left_0011-0.npz", 74 | "TRAIN_B_0424_left_0011-0.npz", 75 | "TRAIN_B_0609_right_0010-0.npz", 76 | ] 77 | for f in scan_with_points_all_occluded: 78 | if os.path.join(self.root_dir, f) in filenames: 79 | filenames.remove(os.path.join(self.root_dir, f)) 80 | 81 | # Remove samples with all points occluded in test set 82 | scan_with_points_all_occluded = [ 83 | "TEST_A_0149_right_0013-0.npz", 84 | "TEST_A_0149_right_0012-0.npz", 85 | "TEST_A_0123_right_0009-0.npz", 86 | "TEST_A_0123_right_0008-0.npz", 87 | ] 88 | for f in scan_with_points_all_occluded: 89 | if os.path.join(self.root_dir, f) in filenames: 90 | filenames.remove(os.path.join(self.root_dir, f)) 91 | 92 | # Train / val / test split 93 | if self.mode == "train" or self.mode == "val": 94 | ind_val = set(np.linspace(0, len(filenames) - 1, 2000).astype("int")) 95 | ind_all = set(np.arange(len(filenames)).astype("int")) 96 | ind_train = ind_all - ind_val 97 | assert ( 98 | len(ind_train.intersection(ind_val)) == 0 99 | ), "Train / Val not split properly" 100 | filenames = np.sort(filenames) 101 | if self.mode == "train": 102 | filenames = filenames[list(ind_train)] 103 | elif self.mode == "val": 104 | filenames = filenames[list(ind_val)] 105 | 106 | if 0 < self.nb_examples < len(filenames): 107 | idx_perm = np.random.permutation(len(filenames)) 108 | idx_sel = idx_perm[:self.nb_examples] 109 | filenames = filenames[idx_sel] 110 | 111 | return filenames 112 | 113 | def load_sequence(self, idx): 114 | """ 115 | Load a sequence of point clouds. 116 | 117 | Parameters 118 | ---------- 119 | idx : int 120 | Index of the sequence to load. 121 | 122 | Returns 123 | ------- 124 | sequence : list(np.array, np.array) 125 | List [pc1, pc2] of point clouds between which to estimate scene 126 | flow. pc1 has size n x 3 and pc2 has size m x 3. 127 | 128 | ground_truth : list(np.array, np.array) 129 | List [mask, flow]. mask has size n x 1 and pc1 has size n x 3. 130 | flow is the ground truth scene flow between pc1 and pc2. mask is 131 | binary with zeros indicating where the flow is not valid/occluded. 132 | 133 | """ 134 | 135 | # Load data 136 | self.filename_curr = self.filenames[idx] 137 | with np.load(self.filename_curr) as data: 138 | sequence = [data["points1"], data["points2"]] 139 | ground_truth = [data["valid_mask1"].reshape(-1, 1), data["flow"]] 140 | 141 | return sequence, ground_truth 142 | -------------------------------------------------------------------------------- /datasets/flyingthings3d_hplflownet.py: -------------------------------------------------------------------------------- 1 | import os 2 | import glob 3 | import numpy as np 4 | from .generic import SceneFlowDataset 5 | 6 | 7 | class FT3D(SceneFlowDataset): 8 | def __init__(self, root_dir, nb_points, all_points, mode, nb_examples): 9 | """ 10 | Construct the FlyingThing3D datatset as in: 11 | Gu, X., Wang, Y., Wu, C., Lee, Y.J., Wang, P., HPLFlowNet: Hierarchical 12 | Permutohedral Lattice FlowNet for scene ﬂow estimation on large-scale 13 | point clouds. IEEE Conf. Computer Vision and Pattern Recognition 14 | (CVPR). pp. 3254–3263 (2019) 15 | 16 | Parameters 17 | ---------- 18 | root_dir : str 19 | Path to root directory containing the datasets. 20 | nb_points : int 21 | Maximum number of points in point clouds. 22 | all_points : bool 23 | Whether to use all point in the point cloud (in chucks of nb_points) or only nb_points. 24 | mode : str 25 | 'train': training dataset. 26 | 27 | 'val': validation dataset. 28 | 29 | 'test': test dataset 30 | nb_examples: int 31 | Number of examples for the dataset. Active if > 0. 32 | 33 | """ 34 | 35 | super(FT3D, self).__init__(nb_points, all_points) 36 | 37 | self.mode = mode 38 | self.nb_examples = nb_examples 39 | self.root_dir = root_dir 40 | self.filenames = self.get_file_list() 41 | self.filename_curr = "" 42 | 43 | def __len__(self): 44 | 45 | return len(self.filenames) 46 | 47 | def get_file_list(self): 48 | """ 49 | Find and filter out paths to all examples in the dataset. 50 | 51 | """ 52 | 53 | # Get list of filenames / directories 54 | if self.mode == "train" or self.mode == "val": 55 | pattern = "train/0*" 56 | elif self.mode == "test": 57 | pattern = "val/0*" 58 | else: 59 | raise ValueError("Mode " + str(self.mode) + " unknown.") 60 | filenames = glob.glob(os.path.join(self.root_dir, pattern)) 61 | 62 | # Train / val / test split 63 | if self.mode == "train" or self.mode == "val": 64 | assert len(filenames) == 19640, "Problem with size of training set" 65 | ind_val = set(np.linspace(0, 19639, 2000).astype("int")) 66 | ind_all = set(np.arange(19640).astype("int")) 67 | ind_train = ind_all - ind_val 68 | assert ( 69 | len(ind_train.intersection(ind_val)) == 0 70 | ), "Train / Val not split properly" 71 | filenames = np.sort(filenames) 72 | if self.mode == "train": 73 | filenames = filenames[list(ind_train)] 74 | elif self.mode == "val": 75 | filenames = filenames[list(ind_val)] 76 | else: 77 | assert len(filenames) == 3824, "Problem with size of test set" 78 | 79 | if 0 < self.nb_examples < len(filenames): 80 | idx_perm = np.random.permutation(len(filenames)) 81 | idx_sel = idx_perm[:self.nb_examples] 82 | filenames = filenames[idx_sel] 83 | 84 | return filenames 85 | 86 | def load_sequence(self, idx): 87 | """ 88 | Load a sequence of point clouds. 89 | 90 | Parameters 91 | ---------- 92 | idx : int 93 | Index of the sequence to load. 94 | 95 | Returns 96 | ------- 97 | sequence : list(np.array, np.array) 98 | List [pc1, pc2] of point clouds between which to estimate scene 99 | flow. pc1 has size n x 3 and pc2 has size m x 3. 100 | 101 | ground_truth : list(np.array, np.array) 102 | List [mask, flow]. mask has size n x 1 and pc1 has size n x 3. 103 | flow is the ground truth scene flow between pc1 and pc2. mask is 104 | binary with zeros indicating where the flow is not valid/occluded. 105 | 106 | """ 107 | 108 | # Load data 109 | self.filename_curr = self.filenames[idx] 110 | sequence = [] # [Point cloud 1, Point cloud 2] 111 | for fname in ["pc1.npy", "pc2.npy"]: 112 | pc = np.load(os.path.join(self.filename_curr, fname)) 113 | pc[..., 0] *= -1 114 | pc[..., -1] *= -1 115 | sequence.append(pc) 116 | ground_truth = [ 117 | np.ones_like(sequence[0][:, 0:1]), 118 | sequence[1] - sequence[0], 119 | ] # [Occlusion mask, flow] 120 | 121 | return sequence, ground_truth 122 | -------------------------------------------------------------------------------- /datasets/generic.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | from torch.utils.data import Dataset 4 | 5 | 6 | class Batch: 7 | def __init__(self, batch): 8 | """ 9 | Concatenate list of dataset.generic.SceneFlowDataset's item in batch 10 | dimension. 11 | 12 | Parameters 13 | ---------- 14 | batch : list 15 | list of dataset.generic.SceneFlowDataset's item. 16 | 17 | """ 18 | 19 | self.data = {} 20 | batch_size = len(batch) 21 | for key in ["sequence", "ground_truth", "orig_size"]: 22 | self.data[key] = [] 23 | for ind_seq in range(len(batch[0][key])): 24 | tmp = [] 25 | for ind_batch in range(batch_size): 26 | item = batch[ind_batch][key][ind_seq] 27 | if len(item.shape) > 3: 28 | tmp.append(item.reshape([-1, item.shape[-2], item.shape[-1]])) 29 | else: 30 | tmp.append(item) 31 | self.data[key].append(torch.cat(tmp, 0)) 32 | 33 | def __getitem__(self, item): 34 | """ 35 | Get 'sequence' or 'ground_truth' from the batch. 36 | 37 | Parameters 38 | ---------- 39 | item : str 40 | Accept two keys 'sequence' or 'ground_truth'. 41 | 42 | Returns 43 | ------- 44 | list(torch.Tensor, torch.Tensor) 45 | item='sequence': returns a list [pc1, pc2] of point clouds between 46 | which to estimate scene flow. pc1 has size B x n x 3 and pc2 has 47 | size B x m x 3. 48 | 49 | item='ground_truth': returns a list [mask, flow]. mask has size 50 | B x n x 1 and flow has size B x n x 3. flow is the ground truth 51 | scene flow between pc1 and pc2. flow is the ground truth scene 52 | flow. mask is binary with zeros indicating where the flow is not 53 | valid or occluded. 54 | 55 | """ 56 | return self.data[item] 57 | 58 | def to(self, *args, **kwargs): 59 | 60 | for key in self.data.keys(): 61 | self.data[key] = [d.to(*args, **kwargs) for d in self.data[key]] 62 | 63 | return self 64 | 65 | def pin_memory(self): 66 | 67 | for key in self.data.keys(): 68 | self.data[key] = [d.pin_memory() for d in self.data[key]] 69 | 70 | return self 71 | 72 | 73 | class SceneFlowDataset(Dataset): 74 | def __init__(self, nb_points, all_points=False): 75 | """ 76 | Abstract constructor for scene flow datasets. 77 | 78 | Each item of the dataset is returned in a dictionary with two keys: 79 | (key = 'sequence', value=list(torch.Tensor, torch.Tensor)): 80 | list [pc1, pc2] of point clouds between which to estimate scene 81 | flow. pc1 has size 1 x n x 3 and pc2 has size 1 x m x 3. 82 | 83 | (key = 'ground_truth', value = list(torch.Tensor, torch.Tensor)): 84 | list [mask, flow]. mask has size 1 x n x 1 and flow has size 85 | 1 x n x 3. flow is the ground truth scene flow between pc1 and pc2. 86 | mask is binary with zeros indicating where the flow is not 87 | valid/occluded. 88 | 89 | Parameters 90 | ---------- 91 | nb_points : int 92 | Maximum number of points in point clouds: self.nb_points <= m, n. 93 | all_points : bool 94 | Whether to use all point in the point cloud (in chucks of nb_points) or only nb_points. 95 | 96 | """ 97 | 98 | super(SceneFlowDataset, self).__init__() 99 | self.nb_points = nb_points 100 | self.all_points = all_points 101 | 102 | def __getitem__(self, idx): 103 | sequence, ground_truth, orig_size = self.to_torch(*self.subsample_points_rnd(*self.load_sequence(idx))) 104 | data = {"sequence": sequence, "ground_truth": ground_truth, "orig_size": orig_size} 105 | 106 | return data 107 | 108 | def to_torch(self, sequence, ground_truth, orig_size): 109 | """ 110 | Convert numpy array and torch.Tensor. 111 | 112 | Parameters 113 | ---------- 114 | sequence : list(np.array, np.array) 115 | List [pc1, pc2] of point clouds between which to estimate scene 116 | flow. pc1 has size n x 3 and pc2 has size m x 3. 117 | 118 | ground_truth : list(np.array, np.array) 119 | List [mask, flow]. mask has size n x 1 and pc1 has size n x 3. 120 | flow is the ground truth scene flow between pc1 and pc2. mask is 121 | binary with zeros indicating where the flow is not valid/occluded. 122 | 123 | orig_size : list(np.array, np.array) 124 | List [n1, n2]. Original size of the point clouds. 125 | 126 | Returns 127 | ------- 128 | sequence : list(torch.Tensor, torch.Tensor) 129 | List [pc1, pc2] of point clouds between which to estimate scene 130 | flow. pc1 has size 1 x n x 3 and pc2 has size 1 x m x 3. 131 | 132 | ground_truth : list(torch.Tensor, torch.Tensor) 133 | List [mask, flow]. mask has size 1 x n x 1 and pc1 has size 134 | 1 x n x 3. flow is the ground truth scene flow between pc1 and pc2. 135 | mask is binary with zeros indicating where the flow is not 136 | valid/occluded. 137 | 138 | """ 139 | 140 | sequence = [torch.unsqueeze(torch.from_numpy(s), 0).float() for s in sequence] 141 | ground_truth = [torch.unsqueeze(torch.from_numpy(gt), 0).float() for gt in ground_truth] 142 | orig_size = [torch.unsqueeze(torch.from_numpy(os), 0) for os in orig_size] 143 | 144 | return sequence, ground_truth, orig_size 145 | 146 | def subsample_points_rnd(self, sequence, ground_truth): 147 | """ 148 | Subsample point clouds randomly. 149 | 150 | Parameters 151 | ---------- 152 | sequence : list(np.array, np.array) 153 | List [pc1, pc2] of point clouds between which to estimate scene 154 | flow. pc1 has size 1 x N x 3 and pc2 has size 1 x M x 3. 155 | 156 | ground_truth : list(np.array, np.array) 157 | List [mask, flow]. mask has size 1 x N x 1 and flow has size 158 | 1 x N x 3. flow is the ground truth scene flow between pc1 and pc2. 159 | mask is binary with zeros indicating where the flow is not 160 | valid/occluded. 161 | 162 | Returns 163 | ------- 164 | sequence : list(np.array, np.array) 165 | List [pc1, pc2] of point clouds between which to estimate scene 166 | flow. pc1 has size 1 x n x 3 and pc2 has size 1 x m x 3. The n 167 | points are chosen randomly among the N available ones. The m points 168 | are chosen randomly among the M available ones. If N, M >= 169 | self.nb_point then n, m = self.nb_points. If N, M < 170 | self.nb_point then n, m = N, M. 171 | 172 | ground_truth : list(np.array, np.array) 173 | List [mask, flow]. mask has size 1 x n x 1 and flow has size 174 | 1 x n x 3. flow is the ground truth scene flow between pc1 and pc2. 175 | mask is binary with zeros indicating where the flow is not 176 | valid/occluded. 177 | 178 | """ 179 | 180 | # Permute indices 181 | n1 = sequence[0].shape[0] 182 | idx1_perm = np.random.permutation(n1) 183 | 184 | n2 = sequence[1].shape[0] 185 | idx2_perm = np.random.permutation(n2) 186 | 187 | # Prepare indices for sampling 188 | if self.all_points: 189 | if self.nb_points == 1: 190 | idx1 = np.arange(n1) 191 | idx2 = np.arange(n2) 192 | else: 193 | n1_div_ceil = n1 // self.nb_points + int((n1 % self.nb_points) > 0) 194 | n1_ceil = n1_div_ceil * self.nb_points 195 | 196 | n2_div_ceil = n2 // self.nb_points + int((n2 % self.nb_points) > 0) 197 | n2_ceil = n2_div_ceil * self.nb_points 198 | 199 | # Take larger point cloud size, in order to have all point from both point clouds for evaluation 200 | n_ceil = n1_ceil > n2_ceil and n1_ceil or n2_ceil 201 | 202 | idx1 = np.concatenate([idx1_perm, idx1_perm[:(n_ceil - n1)]]) 203 | idx2 = np.concatenate([idx2_perm, idx2_perm[:(n_ceil - n2)]]) 204 | else: 205 | idx1 = idx1_perm[:self.nb_points] 206 | idx2 = idx2_perm[:self.nb_points] 207 | 208 | # Sample points in the first scan 209 | sequence[0] = sequence[0][idx1] 210 | ground_truth = [g[idx1] for g in ground_truth] 211 | 212 | # Sample point in the second scan 213 | sequence[1] = sequence[1][idx2] 214 | 215 | # Reshape data 216 | if self.all_points: 217 | sequence[0] = sequence[0].reshape([-1, self.nb_points, 3]) 218 | ground_truth = [g.reshape([-1, self.nb_points, g.shape[1]]) for g in ground_truth] 219 | 220 | sequence[1] = sequence[1].reshape([-1, self.nb_points, 3]) 221 | 222 | if self.nb_points == 1: 223 | sequence = [s.transpose(1, 0, 2) for s in sequence] 224 | ground_truth = [g.transpose(1, 0, 2) for g in ground_truth] 225 | 226 | orig_size = [np.array([n1], dtype=np.int32), np.array([n2], dtype=np.int32)] 227 | 228 | return sequence, ground_truth, orig_size 229 | 230 | def load_sequence(self, idx): 231 | """ 232 | Abstract function to be implemented to load a sequence of point clouds. 233 | 234 | Parameters 235 | ---------- 236 | idx : int 237 | Index of the sequence to load. 238 | 239 | Must return: 240 | ------- 241 | sequence : list(np.array, np.array) 242 | List [pc1, pc2] of point clouds between which to estimate scene 243 | flow. pc1 has size N x 3 and pc2 has size M x 3. 244 | 245 | ground_truth : list(np.array, np.array) 246 | List [mask, flow]. mask has size N x 1 and flow has size N x 3. 247 | flow is the ground truth scene flow between pc1 and pc2. mask is 248 | binary with zeros indicating where the flow is not valid/occluded. 249 | 250 | """ 251 | 252 | raise NotImplementedError 253 | -------------------------------------------------------------------------------- /datasets/kitti_flownet3d.py: -------------------------------------------------------------------------------- 1 | import os 2 | import glob 3 | import numpy as np 4 | from .generic import SceneFlowDataset 5 | 6 | 7 | class Kitti(SceneFlowDataset): 8 | def __init__(self, root_dir, nb_points, all_points, same_v_t_split, mode): 9 | """ 10 | Construct the KITTI scene flow datatset as in: 11 | Liu, X., Qi, C.R., Guibas, L.J.: FlowNet3D: Learning scene ﬂow in 3D 12 | point clouds. IEEE Conf. Computer Vision and Pattern Recognition 13 | (CVPR). pp. 529–537 (2019) 14 | 15 | Parameters 16 | ---------- 17 | root_dir : str 18 | Path to root directory containing the datasets. 19 | nb_points : int 20 | Maximum number of points in point clouds. 21 | all_points : bool 22 | Whether to use all point in the point cloud (in chucks of nb_points) or only nb_points. 23 | same_v_t_split: bool 24 | Whether to use the same validation and test split. 25 | mode : str 26 | 'train': training dataset. 27 | 28 | 'val': validation dataset. 29 | 30 | 'test': test dataset 31 | 32 | 'all': all dataset 33 | 34 | """ 35 | 36 | super(Kitti, self).__init__(nb_points, all_points) 37 | self.mode = mode 38 | self.root_dir = root_dir 39 | self.same_v_t_split = same_v_t_split 40 | self.filenames = self.make_dataset() 41 | self.filename_curr = "" 42 | 43 | def __len__(self): 44 | 45 | return len(self.filenames) 46 | 47 | def make_dataset(self): 48 | """ 49 | Find and filter out paths to all examples in the dataset. 50 | 51 | """ 52 | len_dataset = 150 53 | filenames_all = glob.glob(os.path.join(self.root_dir, "*.npz")) 54 | 55 | test_list = [1, 5, 7, 8, 10, 12, 15, 17, 20, 21, 24, 25, 29, 30, 31, 32, 34, 35, 36, 39, 40, 44, 45, 47, 48, 56 | 49, 50, 51, 53, 55, 56, 58, 59, 60, 70, 71, 72, 74, 76, 77, 78, 79, 81, 82, 88, 91, 93, 94, 95, 98] 57 | val_list = [4, 54, 73, 101, 102, 104, 115, 130, 136, 147] 58 | 59 | if self.same_v_t_split: 60 | val_list = test_list 61 | 62 | train_list = [i for i in range(len_dataset) if i not in test_list and i not in val_list] 63 | 64 | if self.mode == "train": 65 | filenames_train = [fn for fn in filenames_all if int(os.path.split(fn)[1].split(".")[0]) in train_list] 66 | train_size = 100 if self.same_v_t_split else 90 67 | assert len(filenames_train) == train_size, "Problem with size of kitti train dataset" 68 | filenames = filenames_train 69 | 70 | elif self.mode == "val": 71 | filenames_val = [fn for fn in filenames_all if int(os.path.split(fn)[1].split(".")[0]) in val_list] 72 | val_size = 50 if self.same_v_t_split else 10 73 | assert len(filenames_val) == val_size, "Problem with size of kitti validation dataset" 74 | filenames = filenames_val 75 | 76 | elif self.mode == "test": 77 | filenames_test = [fn for fn in filenames_all if int(os.path.split(fn)[1].split(".")[0]) in test_list] 78 | assert len(filenames_test) == 50, "Problem with size of kitti test dataset" 79 | filenames = filenames_test 80 | 81 | elif self.mode == "all": 82 | assert len(filenames_all) == 150, "Problem with size of kitti dataset" 83 | filenames = filenames_all 84 | else: 85 | raise ValueError("Mode " + str(self.mode) + "unknown.") 86 | 87 | return filenames 88 | 89 | def load_sequence(self, idx): 90 | """ 91 | Load a sequence of point clouds. 92 | 93 | Parameters 94 | ---------- 95 | idx : int 96 | Index of the sequence to load. 97 | 98 | Returns 99 | ------- 100 | sequence : list(np.array, np.array) 101 | List [pc1, pc2] of point clouds between which to estimate scene 102 | flow. pc1 has size n x 3 and pc2 has size m x 3. 103 | 104 | ground_truth : list(np.array, np.array) 105 | List [mask, flow]. mask has size n x 1 and pc1 has size n x 3. 106 | flow is the ground truth scene flow between pc1 and pc2. mask is 107 | binary with zeros indicating where the flow is not valid/occluded. 108 | 109 | """ 110 | 111 | # Load data 112 | self.filename_curr = self.filenames[idx] 113 | with np.load(self.filename_curr) as data: 114 | sequence = [data["pos1"][:, (1, 2, 0)], data["pos2"][:, (1, 2, 0)]] 115 | ground_truth = [ 116 | np.ones_like(data["pos1"][:, 0:1]), 117 | data["gt"][:, (1, 2, 0)], 118 | ] 119 | 120 | # Restrict to 35m 121 | loc = sequence[0][:, 2] < 35 122 | sequence[0] = sequence[0][loc] 123 | ground_truth[0] = ground_truth[0][loc] 124 | ground_truth[1] = ground_truth[1][loc] 125 | loc = sequence[1][:, 2] < 35 126 | sequence[1] = sequence[1][loc] 127 | 128 | return sequence, ground_truth 129 | -------------------------------------------------------------------------------- /datasets/kitti_hplflownet.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | from .generic import SceneFlowDataset 4 | 5 | 6 | class Kitti(SceneFlowDataset): 7 | def __init__(self, root_dir, nb_points, all_points, mode): 8 | """ 9 | Construct the KITTI scene flow datatset as in: 10 | Gu, X., Wang, Y., Wu, C., Lee, Y.J., Wang, P., HPLFlowNet: Hierarchical 11 | Permutohedral Lattice FlowNet for scene ﬂow estimation on large-scale 12 | point clouds. IEEE Conf. Computer Vision and Pattern Recognition 13 | (CVPR). pp. 3254–3263 (2019) 14 | 15 | Parameters 16 | ---------- 17 | root_dir : str 18 | Path to root directory containing the datasets. 19 | nb_points : int 20 | Maximum number of points in point clouds. 21 | all_points : bool 22 | Whether to use all point in the point cloud (in chucks of nb_points) or only nb_points. 23 | mode : str 24 | 'train': training dataset. 25 | 26 | 'val': validation dataset. 27 | 28 | 'test': test dataset 29 | 30 | 'all': all dataset 31 | 32 | """ 33 | 34 | super(Kitti, self).__init__(nb_points, all_points) 35 | self.mode = mode 36 | self.root_dir = root_dir 37 | self.paths = self.make_dataset() 38 | self.filename_curr = "" 39 | 40 | def __len__(self): 41 | 42 | return len(self.paths) 43 | 44 | def make_dataset(self): 45 | """ 46 | Find and filter out paths to all examples in the dataset. 47 | 48 | """ 49 | 50 | # 51 | root = os.path.realpath(os.path.expanduser(self.root_dir)) 52 | all_paths = sorted(os.walk(root)) 53 | useful_paths = [item[0] for item in all_paths if len(item[1]) == 0] 54 | assert len(useful_paths) == 200, "Problem with size of kitti dataset" 55 | 56 | # Mapping / Filtering of scans as in HPLFlowNet code 57 | mapping_path = os.path.join(os.path.dirname(__file__), "KITTI_mapping.txt") 58 | with open(mapping_path) as fd: 59 | lines = fd.readlines() 60 | lines = [line.strip() for line in lines] 61 | useful_paths = [ 62 | path for path in useful_paths if lines[int(os.path.split(path)[-1])] != "" 63 | ] 64 | 65 | useful_paths = np.array(useful_paths) 66 | len_dataset = len(useful_paths) 67 | 68 | # the train indices was randomly selected by: 69 | # train_idx = self.make_subset_idx(total_examples=len_dataset, nb_examples=92, seed=42) 70 | train_idx = [0, 3, 4, 5, 7, 8, 9, 10, 11, 12, 13, 15, 16, 17, 18, 19, 22, 23, 24, 25, 26, 27, 28, 30, 31, 33, 71 | 34, 35, 36, 39, 40, 42, 43, 44, 45, 47, 49, 51, 53, 55, 56, 60, 62, 64, 65, 66, 67, 68, 69, 70, 72 | 73, 76, 77, 78, 80, 81, 82, 83, 84, 85, 86, 89, 93, 94, 95, 96, 97, 98, 101, 105, 108, 109, 110, 73 | 111, 112, 113, 114, 115, 117, 118, 122, 123, 124, 126, 128, 130, 131, 133, 135, 137, 138, 140] 74 | 75 | test_idx = [i for i in range(len_dataset) if i not in train_idx] 76 | val_idx = test_idx 77 | 78 | if self.mode == "train": 79 | useful_paths_train = useful_paths[train_idx] 80 | train_size = 92 81 | assert len(useful_paths_train) == train_size, "Problem with size of kitti train dataset" 82 | useful_paths = useful_paths_train 83 | 84 | elif self.mode == "val": 85 | useful_paths_val = useful_paths[val_idx] 86 | val_size = 50 87 | assert len(useful_paths_val) == val_size, "Problem with size of kitti validation dataset" 88 | useful_paths = useful_paths_val 89 | 90 | elif self.mode == "test": 91 | useful_paths_test = useful_paths[test_idx] 92 | assert len(useful_paths_test) == 50, "Problem with size of kitti test dataset" 93 | useful_paths = useful_paths_test 94 | 95 | elif self.mode == "all": 96 | assert len(useful_paths) == 142, "Problem with size of kitti dataset" 97 | else: 98 | raise ValueError("Mode " + str(self.mode) + "unknown.") 99 | 100 | useful_paths = list(useful_paths) 101 | 102 | return useful_paths 103 | 104 | def make_subset_idx(self, total_examples, nb_examples, seed=42): 105 | np.random.seed(seed) 106 | idx_perm = np.random.permutation(total_examples) 107 | idx_sel = np.sort(idx_perm[:nb_examples]) 108 | 109 | return idx_sel 110 | 111 | def load_sequence(self, idx): 112 | """ 113 | Load a sequence of point clouds. 114 | 115 | Parameters 116 | ---------- 117 | idx : int 118 | Index of the sequence to load. 119 | 120 | Returns 121 | ------- 122 | sequence : list(np.array, np.array) 123 | List [pc1, pc2] of point clouds between which to estimate scene 124 | flow. pc1 has size n x 3 and pc2 has size m x 3. 125 | 126 | ground_truth : list(np.array, np.array) 127 | List [mask, flow]. mask has size n x 1 and pc1 has size n x 3. 128 | flow is the ground truth scene flow between pc1 and pc2. mask is 129 | binary with zeros indicating where the flow is not valid/occluded. 130 | 131 | """ 132 | 133 | # Load data 134 | self.filename_curr = self.paths[idx] 135 | sequence = [np.load(os.path.join(self.paths[idx], "pc1.npy"))] 136 | sequence.append(np.load(os.path.join(self.paths[idx], "pc2.npy"))) 137 | 138 | # Remove ground points 139 | is_ground = np.logical_and(sequence[0][:, 1] < -1.4, sequence[1][:, 1] < -1.4) 140 | not_ground = np.logical_not(is_ground) 141 | sequence = [sequence[i][not_ground] for i in range(2)] 142 | 143 | # Remove points further than 35 meter away as in HPLFlowNet code 144 | is_close = np.logical_and(sequence[0][:, 2] < 35, sequence[1][:, 2] < 35) 145 | sequence = [sequence[i][is_close] for i in range(2)] 146 | 147 | # Scene flow 148 | ground_truth = [ 149 | np.ones_like(sequence[0][:, 0:1]), 150 | sequence[1] - sequence[0], 151 | ] # [Occlusion mask, scene flow] 152 | 153 | return sequence, ground_truth 154 | -------------------------------------------------------------------------------- /doc/poster.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/itailang/SCOOP/9f41e8baafed8689be867b8ab7b10159371de37d/doc/poster.png -------------------------------------------------------------------------------- /doc/scoop_result.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/itailang/SCOOP/9f41e8baafed8689be867b8ab7b10159371de37d/doc/scoop_result.gif -------------------------------------------------------------------------------- /doc/slides.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/itailang/SCOOP/9f41e8baafed8689be867b8ab7b10159371de37d/doc/slides.pdf -------------------------------------------------------------------------------- /download_pretrained_models.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | # define a download function 4 | function google_drive_download() 5 | { 6 | CONFIRM=$(wget --quiet --save-cookies /tmp/cookies.txt --keep-session-cookies --no-check-certificate "https://docs.google.com/uc?export=download&id=$1" -O- | sed -rn 's/.*confirm=([0-9A-Za-z_]+).*/\1\n/p') 7 | wget --load-cookies /tmp/cookies.txt "https://docs.google.com/uc?export=download&confirm=$CONFIRM&id=$1" -O $2 8 | rm -rf /tmp/cookies.txt 9 | } 10 | 11 | # download models 12 | google_drive_download 14M0R0AMJhsyGZYHDIp3p8QStCkKCy_bO pretrained_models.zip 13 | unzip pretrained_models.zip 14 | rm pretrained_models.zip 15 | -------------------------------------------------------------------------------- /install_environment.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | # install dependencies 4 | conda install pytorch==1.6.0 torchvision==0.7.0 cudatoolkit=10.1 -c pytorch --yes 5 | pip install tensorboard==2.4.1 --no-cache-dir 6 | pip install tqdm==4.60.0 --no-cache-dir 7 | 8 | # install mayavi for visualization 9 | pip install vtk==9.1.0 --no-cache-dir 10 | pip install pyQt5==5.15.2 --no-cache-dir 11 | pip install mayavi==4.7.4 --no-cache-dir 12 | -------------------------------------------------------------------------------- /models/gconv.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | class SetConv(torch.nn.Module): 5 | def __init__(self, nb_feat_in, nb_feat_out): 6 | """ 7 | Module that performs PointNet++-like convolution on point clouds. 8 | 9 | Parameters 10 | ---------- 11 | nb_feat_in : int 12 | Number of input channels. 13 | nb_feat_out : int 14 | Number of ouput channels. 15 | 16 | Returns 17 | ------- 18 | None. 19 | 20 | """ 21 | 22 | super(SetConv, self).__init__() 23 | 24 | self.fc1 = torch.nn.Conv2d(nb_feat_in + 3, nb_feat_out, 1, bias=False) 25 | self.bn1 = torch.nn.InstanceNorm2d(nb_feat_out, affine=True) 26 | 27 | self.fc2 = torch.nn.Conv2d(nb_feat_out, nb_feat_out, 1, bias=False) 28 | self.bn2 = torch.nn.InstanceNorm2d(nb_feat_out, affine=True) 29 | 30 | self.fc3 = torch.nn.Conv2d(nb_feat_out, nb_feat_out, 1, bias=False) 31 | self.bn3 = torch.nn.InstanceNorm2d(nb_feat_out, affine=True) 32 | 33 | self.pool = lambda x: torch.max(x, 2)[0] 34 | self.lrelu = torch.nn.LeakyReLU(negative_slope=0.1) 35 | 36 | def forward(self, signal, graph): 37 | """ 38 | Performs PointNet++-like convolution 39 | 40 | Parameters 41 | ---------- 42 | signal : torch.Tensor 43 | Input features of size B x N x nb_feat_in. 44 | graph : scoop.models.graph.Graph 45 | Graph build on the input point cloud on with the input features 46 | live. The graph contains the list of nearest neighbors (NN) for 47 | each point and all edge features (relative point coordinates with 48 | NN). 49 | 50 | Returns 51 | ------- 52 | torch.Tensor 53 | Ouput features of size B x N x nb_feat_out. 54 | 55 | """ 56 | 57 | # Input features dimension 58 | b, n, c = signal.shape 59 | n_out = graph.size[0] // b 60 | 61 | # Concatenate input features with edge features 62 | signal = signal.reshape(b * n, c) 63 | signal = torch.cat((signal[graph.edges], graph.edge_feats), -1) 64 | signal = signal.view(b, n_out, graph.k_neighbors, c + 3) 65 | signal = signal.transpose(1, -1) 66 | 67 | # Pointnet++-like convolution 68 | for func in [ 69 | self.fc1, 70 | self.bn1, 71 | self.lrelu, 72 | self.fc2, 73 | self.bn2, 74 | self.lrelu, 75 | self.fc3, 76 | self.bn3, 77 | self.lrelu, 78 | self.pool, 79 | ]: 80 | signal = func(signal) 81 | 82 | return signal.transpose(1, -1) 83 | -------------------------------------------------------------------------------- /models/graph.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import torch 4 | import numpy as np 5 | 6 | # add path 7 | project_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) 8 | if project_dir not in sys.path: 9 | sys.path.append(project_dir) 10 | 11 | from tools.utils import iterate_in_chunks 12 | 13 | class Graph: 14 | def __init__(self, edges, edge_feats, k_neighbors, size): 15 | """ 16 | Directed nearest neighbor graph constructed on a point cloud. 17 | 18 | Parameters 19 | ---------- 20 | edges : torch.Tensor 21 | Contains list with nearest neighbor indices. 22 | edge_feats : torch.Tensor 23 | Contains edge features: relative point coordinates. 24 | k_neighbors : int 25 | Number of nearest neighbors. 26 | size : tuple(int, int) 27 | Number of points. 28 | 29 | """ 30 | 31 | self.edges = edges 32 | self.size = tuple(size) 33 | self.edge_feats = edge_feats 34 | self.k_neighbors = k_neighbors 35 | 36 | @staticmethod 37 | def construct_graph(pcloud, nb_neighbors): 38 | """ 39 | Construct a directed nearest neighbor graph on the input point cloud. 40 | 41 | Parameters 42 | ---------- 43 | pcloud : torch.Tensor 44 | Input point cloud. Size B x N x 3. 45 | nb_neighbors : int 46 | Number of nearest neighbors per point. 47 | 48 | Returns 49 | ------- 50 | graph : scoop.models.graph.Graph 51 | Graph build on input point cloud containing the list of nearest 52 | neighbors (NN) for each point and all edge features (relative 53 | coordinates with NN). 54 | 55 | """ 56 | 57 | # Size 58 | nb_points = pcloud.shape[1] 59 | size_batch = pcloud.shape[0] 60 | 61 | # Distance between points 62 | distance_matrix = torch.sum(pcloud ** 2, -1, keepdim=True) 63 | distance_matrix = distance_matrix + distance_matrix.transpose(1, 2) 64 | distance_matrix = distance_matrix - 2 * torch.bmm( 65 | pcloud, pcloud.transpose(1, 2) 66 | ) 67 | 68 | # Find nearest neighbors 69 | neighbors = torch.argsort(distance_matrix, -1)[..., :nb_neighbors] 70 | effective_nb_neighbors = neighbors.shape[-1] 71 | neighbors = neighbors.reshape(size_batch, -1) 72 | 73 | # Edge origin 74 | idx = torch.arange(nb_points, device=distance_matrix.device).long() 75 | idx = torch.repeat_interleave(idx, effective_nb_neighbors) 76 | 77 | # Edge features 78 | edge_feats = [] 79 | for ind_batch in range(size_batch): 80 | edge_feats.append( 81 | pcloud[ind_batch, neighbors[ind_batch]] - pcloud[ind_batch, idx] 82 | ) 83 | edge_feats = torch.cat(edge_feats, 0) 84 | 85 | # Handle batch dimension to get indices of nearest neighbors 86 | for ind_batch in range(1, size_batch): 87 | neighbors[ind_batch] = neighbors[ind_batch] + ind_batch * nb_points 88 | neighbors = neighbors.view(-1) 89 | 90 | # Create graph 91 | graph = Graph( 92 | neighbors, 93 | edge_feats, 94 | effective_nb_neighbors, 95 | [size_batch * nb_points, size_batch * nb_points], 96 | ) 97 | 98 | return graph 99 | 100 | @staticmethod 101 | def construct_graph_in_chunks(pcloud, nb_neighbors, chunk_size): 102 | # Size 103 | size_batch, nb_points, _ = pcloud.shape 104 | assert size_batch == 1, "For construction of graph in chucks, the batch size should be 1, got %d." % size_batch 105 | 106 | # Find nearest neighbors 107 | #distance_matrix = -1 * torch.ones([nb_points, nb_points], dtype=torch.float32, device=pcloud.device) 108 | neighbors = -1 * torch.ones([nb_points, nb_neighbors], dtype=torch.int64, device=pcloud.device) 109 | idx = np.arange(nb_points) 110 | for b in iterate_in_chunks(idx, chunk_size): 111 | # Distance between points 112 | points_curr = torch.transpose(pcloud[:, b], 0, 1) 113 | #distance_matrix_curr = torch.sum(torch.pow(points_curr - pcloud, 2), -1) 114 | distance_matrix_curr = torch.sum(points_curr ** 2, -1) - 2 * torch.sum(points_curr * pcloud, dim=-1) + torch.sum(pcloud ** 2, -1) 115 | #distance_matrix[b, :] = distance_matrix_curr 116 | 117 | # Find nearest neighbors 118 | neighbors_curr = torch.argsort(distance_matrix_curr, -1)[..., :nb_neighbors] 119 | neighbors[b, :] = neighbors_curr 120 | 121 | assert torch.all(neighbors >= 0), "Problem with nearest neighbors computation. Not all indices filled correctly." 122 | #assert torch.all(distance_matrix >= 0), "Problem with distance matrix computation. Not all distances filled correctly." 123 | 124 | effective_nb_neighbors = neighbors.shape[-1] 125 | edge_feats = torch.empty(0) 126 | 127 | neighbors = neighbors.view(-1) 128 | 129 | # Create graph 130 | graph = Graph( 131 | neighbors, 132 | edge_feats, 133 | effective_nb_neighbors, 134 | [size_batch * nb_points, size_batch * nb_points], 135 | ) 136 | 137 | return graph 138 | -------------------------------------------------------------------------------- /models/refiner.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import numpy as np 4 | import time 5 | import torch 6 | 7 | # add path 8 | project_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) 9 | if project_dir not in sys.path: 10 | sys.path.append(project_dir) 11 | 12 | from models.graph import Graph 13 | from tools.losses import chamfer_loss, smooth_loss 14 | 15 | 16 | class Refiner(torch.nn.Module): 17 | def __init__(self, shape, device): 18 | """ 19 | Construct a model for refining scene flow between two point clouds. 20 | 21 | Parameters 22 | ---------- 23 | shape: 24 | Shape of the refinement tensor. 25 | device: 26 | Device of the refinement tensor. 27 | """ 28 | 29 | super(Refiner, self).__init__() 30 | 31 | self.refinement = torch.nn.Parameter(torch.zeros(shape, dtype=torch.float32, device=device, requires_grad=True)) 32 | 33 | def forward(self, flow): 34 | refined_flow = flow + self.refinement 35 | 36 | return refined_flow 37 | 38 | def refine_flow(self, batch, flow, corr_conf, optimizer, args): 39 | pc_0, pc_1 = batch["sequence"][0], batch["sequence"][1] 40 | gt_flow = batch["ground_truth"][1] 41 | 42 | n0 = int(batch["orig_size"][0].cpu().numpy()) 43 | n1 = int(batch["orig_size"][1].cpu().numpy()) 44 | 45 | b, nb_points0, c = pc_0.shape 46 | b, nb_points1, c = pc_1.shape 47 | 48 | pc_0_orig = torch.unsqueeze(torch.reshape(pc_0, (b * nb_points0, 3))[:n0], dim=0) 49 | pc_1_orig = torch.unsqueeze(torch.reshape(pc_1, (b * nb_points1, 3))[:n1], dim=0) 50 | gt_flow_orig = torch.unsqueeze(torch.reshape(gt_flow, (b * nb_points0, 3))[:n0], dim=0) 51 | 52 | start_time = time.time() 53 | 54 | graph = Graph.construct_graph_in_chunks(pc_0_orig, 32, 2048) 55 | 56 | # results aggregation 57 | target_recon_loss_all = np.zeros(args.test_time_num_step + 1, dtype=np.float32) 58 | smooth_flow_loss_all = np.zeros(args.test_time_num_step + 1, dtype=np.float32) 59 | epe_all = np.zeros(args.test_time_num_step + 1, dtype=np.float32) 60 | 61 | for step in range(args.test_time_num_step): 62 | refined_flow = self(flow) 63 | target_pc_recon = pc_0_orig + refined_flow 64 | 65 | target_recon_loss = chamfer_loss(target_pc_recon, pc_1_orig, corr_conf, 66 | backward_dist_weight=args.backward_dist_weight, mask=None, use_chamfer_cuda=bool(args.use_chamfer_cuda)) 67 | loss = args.target_recon_loss_weight * target_recon_loss 68 | 69 | if args.use_smooth_flow and args.smooth_flow_loss_weight > 0: 70 | smooth_flow_loss, _ = smooth_loss(refined_flow, graph, args.nb_neigh_smooth_flow, loss_norm=1, mask=None) 71 | loss = loss + (args.smooth_flow_loss_weight * smooth_flow_loss) 72 | else: 73 | smooth_flow_loss = 0 74 | 75 | # Gradient step 76 | optimizer.zero_grad() 77 | loss.backward() 78 | optimizer.step() 79 | 80 | # Loss evolution 81 | loss_curr = loss.item() 82 | target_recon_loss_curr = target_recon_loss.item() 83 | smooth_flow_loss_curr = smooth_flow_loss.item() if args.use_smooth_flow and args.smooth_flow_loss_weight > 0 else smooth_flow_loss 84 | 85 | # EPE 86 | error = refined_flow - gt_flow_orig 87 | epe_per_point = torch.sqrt(torch.sum(torch.pow(error, 2.0), -1)) 88 | epe = epe_per_point.mean() 89 | epe_curr = epe.item() 90 | 91 | if args.test_time_verbose: 92 | print("Refinement step %04d/%04d: loss: %.6f, target_recon_loss: %.6f, smooth_flow_loss: %.6f, epe: %.3f" % 93 | ((step + 1), args.test_time_num_step, loss_curr, target_recon_loss_curr, smooth_flow_loss_curr, epe_curr)) 94 | 95 | # aggregate results 96 | target_recon_loss_all[step] = target_recon_loss_curr 97 | smooth_flow_loss_all[step] = smooth_flow_loss_curr 98 | epe_all[step] = epe_curr 99 | 100 | refined_flow = self(flow) 101 | 102 | duration = time.time() - start_time 103 | 104 | # EPE last 105 | error = refined_flow - gt_flow_orig 106 | epe_per_point = torch.sqrt(torch.sum(torch.pow(error, 2.0), -1)) 107 | epe = epe_per_point.mean() 108 | epe_curr = epe.item() 109 | 110 | epe_all[-1] = epe_curr 111 | 112 | refine_metrics = {"target_recon_loss_all": target_recon_loss_all, "smooth_flow_loss_all": smooth_flow_loss_all, "epe_all": epe_all} 113 | 114 | return refined_flow, refine_metrics, duration 115 | -------------------------------------------------------------------------------- /models/scoop.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from tools import ot, reconstruction as R 3 | from models.graph import Graph 4 | from models.gconv import SetConv 5 | 6 | 7 | class SCOOP(torch.nn.Module): 8 | def __init__(self, args): 9 | """ 10 | Construct a model that, once trained, estimate the scene flow between 11 | two point clouds. 12 | 13 | Parameters 14 | ---------- 15 | args.nb_iter : int 16 | Number of iterations to unroll in the Sinkhorn algorithm. 17 | 18 | """ 19 | 20 | super(SCOOP, self).__init__() 21 | 22 | # Hand-chosen parameters. Define the number of channels. 23 | n = 32 24 | 25 | # OT parameters 26 | # Number of unrolled iterations in the Sinkhorn algorithm 27 | self.nb_iter = args.nb_iter 28 | # Mass regularisation 29 | self.gamma = torch.nn.Parameter(torch.zeros(1)) 30 | # Entropic regularisation 31 | self.epsilon = torch.nn.Parameter(torch.zeros(1)) 32 | 33 | # Architecture parameters 34 | self.nb_neigh_cross_recon = args.nb_neigh_cross_recon 35 | 36 | try: 37 | self.linear_corr_conf = args.linear_corr_conf 38 | except AttributeError: 39 | self.linear_corr_conf = 0 40 | 41 | # Feature extraction 42 | self.feat_conv1 = SetConv(3, n) 43 | self.feat_conv2 = SetConv(n, 2 * n) 44 | self.feat_conv3 = SetConv(2 * n, 4 * n) 45 | 46 | def get_features(self, pcloud, nb_neighbors=32): 47 | """ 48 | Compute deep features for each point of the input point cloud. These 49 | features are used to compute the transport cost matrix between two 50 | point clouds. 51 | 52 | Parameters 53 | ---------- 54 | pcloud : torch.Tensor 55 | Input point cloud of size B x N x 3 56 | nb_neighbors : int 57 | Number of nearest neighbors for each point. 58 | 59 | Returns 60 | ------- 61 | x : torch.Tensor 62 | Deep features for each point. Size B x N x 128 63 | graph : scoop.models.graph.Graph 64 | Graph build on input point cloud containing list of nearest 65 | neighbors (NN) and edge features (relative coordinates with NN). 66 | 67 | """ 68 | 69 | graph = Graph.construct_graph(pcloud, nb_neighbors) 70 | x = self.feat_conv1(pcloud, graph) 71 | x = self.feat_conv2(x, graph) 72 | x = self.feat_conv3(x, graph) 73 | 74 | return x, graph 75 | 76 | def get_recon_flow(self, pclouds, feats): 77 | feats_0, feats_1 = feats[0], feats[1] 78 | 79 | # Reconstructed target point cloud 80 | transport_cross, similarity_cross = ot.sinkhorn( 81 | feats_0, 82 | feats_1, 83 | pclouds[0], 84 | pclouds[1], 85 | epsilon=torch.exp(self.epsilon) + 0.03, 86 | gamma=torch.exp(self.gamma), 87 | max_iter=self.nb_iter, 88 | ) 89 | 90 | if self.nb_neigh_cross_recon > 0: 91 | source_cross_nn_weight, _, source_cross_nn_idx, _, _, _ = \ 92 | R.get_s_t_neighbors(self.nb_neigh_cross_recon, transport_cross, sim_normalization="none", s_only=True) 93 | 94 | # Target point cloud cross reconstruction 95 | cross_weight_sum = source_cross_nn_weight.sum(-1, keepdim=True) 96 | source_cross_nn_weight_normalized = source_cross_nn_weight / (cross_weight_sum + 1e-8) 97 | target_cross_recon = R.reconstruct(pclouds[1], source_cross_nn_idx, source_cross_nn_weight_normalized, self.nb_neigh_cross_recon) 98 | 99 | # Matching probability 100 | cross_nn_sim, _, _, _ = R.get_s_t_topk(similarity_cross, self.nb_neigh_cross_recon, s_only=True, nn_idx=source_cross_nn_idx) 101 | nn_sim_weighted = cross_nn_sim * source_cross_nn_weight_normalized 102 | nn_sim_weighted = torch.sum(nn_sim_weighted, dim=2) 103 | if self.linear_corr_conf: 104 | corr_conf = (nn_sim_weighted + 1) / 2 105 | else: 106 | corr_conf = torch.clamp_min(nn_sim_weighted, 0.0) 107 | else: 108 | row_sum = transport_cross.sum(-1, keepdim=True) 109 | target_cross_recon = (transport_cross @ pclouds[1]) / (row_sum + 1e-8) 110 | corr_conf = None 111 | 112 | # Estimate flow from target cross reconstruction 113 | recon_flow = target_cross_recon - pclouds[0] 114 | 115 | return recon_flow, corr_conf, target_cross_recon 116 | 117 | def forward(self, pclouds): 118 | """ 119 | Estimate scene flow between two input point clouds. 120 | 121 | Parameters 122 | ---------- 123 | pclouds : (torch.Tensor, torch.Tensor) 124 | List of input point clouds (pc1, pc2). pc1 has size B x N x 3. 125 | pc2 has size B x M x 3. 126 | 127 | Returns 128 | ------- 129 | est_flow : torch.Tensor 130 | Estimated scene flow of size B x N x 3. 131 | 132 | """ 133 | 134 | # Extract features 135 | feats_0, graph = self.get_features(pclouds[0]) 136 | feats_1, _ = self.get_features(pclouds[1]) 137 | feats = [feats_0, feats_1] 138 | 139 | # Get reconstruction-based flow 140 | recon_flow, corr_conf, target_cross_recon = self.get_recon_flow(pclouds, feats) 141 | 142 | return recon_flow, corr_conf, target_cross_recon, graph 143 | -------------------------------------------------------------------------------- /scripts/evaluate_on_ft3d_s.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | # evaluate on ft3d_s 4 | python evaluate_scoop.py --dataset_name HPLFlowNet_ft3d --mode test --nb_points 8192 --all_points 0 --all_candidates 0 \ 5 | --path2ckpt ./../experiments/ft3d_s_1800_examples/model_e060.tar --backward_dist_weight 1.0 \ 6 | --use_test_time_refinement 1 --test_time_num_step 1000 --test_time_update_rate 0.1 --nb_neigh_smooth_flow 16 \ 7 | --log_fname log_evaluation_ft3d_s.txt 8 | -------------------------------------------------------------------------------- /scripts/evaluate_on_kitti_o.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | # evaluate on kitti_o 4 | python evaluate_scoop.py --dataset_name FlowNet3D_kitti --mode all --nb_points 2048 --all_points 0 --all_candidates 0 \ 5 | --path2ckpt ./../experiments/ft3d_o_1800_examples/model_e100.tar \ 6 | --use_test_time_refinement 1 --test_time_num_step 1000 --test_time_update_rate 0.05 \ 7 | --log_fname log_evaluation_kitti_o.txt 8 | -------------------------------------------------------------------------------- /scripts/evaluate_on_kitti_o_all_points.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | # evaluate on kitti_o with all the points 4 | python evaluate_scoop.py --dataset_name FlowNet3D_kitti --mode all --nb_points 2048 \ 5 | --path2ckpt ./../experiments/ft3d_o_1800_examples/model_e100.tar \ 6 | --use_test_time_refinement 1 --test_time_num_step 150 --test_time_update_rate 0.2 \ 7 | --log_fname log_evaluation_kitti_o_all_points.txt 8 | -------------------------------------------------------------------------------- /scripts/evaluate_on_kitti_s.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | # evaluate on kitti_s 4 | python evaluate_scoop.py --dataset_name HPLFlowNet_kitti --mode all --nb_points 8192 --all_points 0 --all_candidates 0 \ 5 | --path2ckpt ./../experiments/ft3d_s_1800_examples/model_e060.tar --backward_dist_weight 1.0 \ 6 | --use_test_time_refinement 1 --test_time_num_step 1000 --test_time_update_rate 0.05 \ 7 | --log_fname log_evaluation_kitti_s.txt -------------------------------------------------------------------------------- /scripts/evaluate_on_kitti_t.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | # evaluate on kitti_t 4 | python evaluate_scoop.py --dataset_name FlowNet3D_kitti --mode test --nb_points 2048 --all_points 0 --all_candidates 0 \ 5 | --path2ckpt ./../experiments/kitti_v_100_examples/model_e400.tar \ 6 | --use_test_time_refinement 1 --test_time_num_step 1000 --test_time_update_rate 0.05 \ 7 | --log_fname log_evaluation_kitti_t.txt 8 | -------------------------------------------------------------------------------- /scripts/evaluate_on_kitti_t_all_points.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | # evaluate on kitti_t with all the points 4 | python evaluate_scoop.py --dataset_name FlowNet3D_kitti --mode test --nb_points 2048 \ 5 | --path2ckpt ./../experiments/kitti_v_100_examples/model_e400.tar \ 6 | --use_test_time_refinement 1 --test_time_num_step 150 --test_time_update_rate 0.2 \ 7 | --log_fname log_evaluation_kitti_t_all_points.txt 8 | -------------------------------------------------------------------------------- /scripts/evaluate_scoop.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import torch 4 | import argparse 5 | import numpy as np 6 | import time 7 | from tqdm import tqdm 8 | 9 | # add path 10 | project_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) 11 | if project_dir not in sys.path: 12 | sys.path.append(project_dir) 13 | 14 | from models.scoop import SCOOP 15 | from models.refiner import Refiner 16 | from datasets.generic import Batch 17 | from tools.seed import seed_everything 18 | from tools.utils import log_string, iterate_in_chunks 19 | from torch.utils.data import DataLoader 20 | 21 | 22 | def compute_flow(scoop, batch, args): 23 | pc_0, pc_1 = batch["sequence"][0], batch["sequence"][1] 24 | 25 | n0 = int(batch["orig_size"][0].cpu().numpy()) 26 | n1 = int(batch["orig_size"][1].cpu().numpy()) 27 | 28 | with torch.no_grad(): 29 | est_flow = torch.zeros([1, n0, 3], dtype=torch.float32, device=pc_0.device) 30 | corr_conf = torch.zeros([1, n0], dtype=torch.float32, device=pc_0.device) 31 | 32 | feats_0, graph = scoop.get_features(pc_0) 33 | feats_1, _ = scoop.get_features(pc_1) 34 | 35 | b, nb_points0, c = feats_0.shape 36 | b, nb_points1, c = feats_1.shape 37 | 38 | pc_0_orig = torch.unsqueeze(torch.reshape(pc_0, (b * nb_points0, 3))[:n0], dim=0) 39 | pc_1_orig = torch.unsqueeze(torch.reshape(pc_1, (b * nb_points1, 3))[:n1], dim=0) 40 | feats_0_orig = torch.unsqueeze(torch.reshape(feats_0, (b * nb_points0, c))[:n0], dim=0) 41 | feats_1_orig = torch.unsqueeze(torch.reshape(feats_1, (b * nb_points1, c))[:n1], dim=0) 42 | idx = np.arange(n0) 43 | for b in iterate_in_chunks(idx, args.nb_points_chunk): 44 | points = pc_0_orig[:, b] 45 | feats = feats_0_orig[:, b] 46 | points_flow, points_conf, _ = scoop.get_recon_flow([points, pc_1_orig], [feats, feats_1_orig]) 47 | est_flow[:, b] = points_flow 48 | corr_conf[:, b] = points_conf 49 | 50 | return est_flow, corr_conf, graph 51 | 52 | 53 | def compute_epe_test(est_flow, batch, args): 54 | """ 55 | Compute EPE, accuracy and number of outliers. 56 | 57 | Parameters 58 | ---------- 59 | est_flow : torch.Tensor 60 | Estimated flow. 61 | batch : scoop.datasets.generic.Batch 62 | Contains ground truth flow and mask. 63 | args : Namespace 64 | Arguments for evaluation. 65 | 66 | Returns 67 | ------- 68 | EPE3D : float 69 | End point error. 70 | acc3d_strict : float 71 | Strict accuracy. 72 | acc3d_relax : float 73 | Relax accuracy. 74 | outlier : float 75 | Percentage of outliers. 76 | 77 | """ 78 | 79 | # Extract occlusion mask 80 | mask = batch["ground_truth"][0].cpu().numpy()[..., 0] 81 | 82 | # Flow 83 | sf_gt_before_mask = batch["ground_truth"][1].cpu().numpy() 84 | sf_pred_before_mask = est_flow.cpu().numpy() 85 | 86 | # In all_points evaluation mode, take only the original points of the source point cloud 87 | n1 = batch["orig_size"][0].cpu().numpy() 88 | if args.all_points: 89 | assert len(n1) == 1, "If evaluating with all points, the batch size should be equal 1 (got %d)" % len(n1) 90 | 91 | mask = mask.reshape(-1)[:int(n1)] 92 | 93 | sf_gt_before_mask = sf_gt_before_mask.reshape([-1, 3])[:int(n1)] 94 | sf_pred_before_mask = sf_pred_before_mask.reshape([-1, 3])[:int(n1)] 95 | 96 | # Flow 97 | sf_gt = sf_gt_before_mask[mask > 0] 98 | sf_pred = sf_pred_before_mask[mask > 0] 99 | 100 | # EPE 101 | epe3d_per_point = np.linalg.norm(sf_gt - sf_pred, axis=-1) 102 | epe3d = epe3d_per_point.mean() 103 | 104 | # 105 | sf_norm = np.linalg.norm(sf_gt, axis=-1) 106 | relative_err_per_point = epe3d_per_point / (sf_norm + 1e-4) 107 | acc3d_strict_per_point = (np.logical_or(epe3d_per_point < 0.05, relative_err_per_point < 0.05)).astype(np.float32) 108 | acc3d_strict = acc3d_strict_per_point.mean() 109 | acc3d_relax_per_point = (np.logical_or(epe3d_per_point < 0.1, relative_err_per_point < 0.1)).astype(np.float32) 110 | acc3d_relax = acc3d_relax_per_point.mean() 111 | outlier_per_point = (np.logical_or(epe3d_per_point > 0.3, relative_err_per_point > 0.1)).astype(np.float32) 112 | outlier = outlier_per_point.mean() 113 | 114 | return epe3d, acc3d_strict, acc3d_relax, outlier, epe3d_per_point, acc3d_strict_per_point, acc3d_relax_per_point, outlier_per_point 115 | 116 | 117 | def eval_model(scoop, testloader, log_file, log_dir, res_dir, args): 118 | """ 119 | Compute performance metrics on test / validation set. 120 | 121 | Parameters 122 | ---------- 123 | scoop : scoop.models.SCOOP 124 | SCOOP model to evaluate. 125 | testloader : scoop.datasets.generic.SceneFlowDataset 126 | Dataset loader. 127 | log_file: file 128 | Evaluation log file. 129 | log_dir: 130 | Directory for saving results. 131 | res_dir: srt 132 | Directory for saving point cloud results (active if not None). 133 | args : Namespace 134 | Arguments for evaluation. 135 | 136 | Returns 137 | ------- 138 | mean_epe : float 139 | Average EPE on dataset. 140 | mean_outlier : float 141 | Average percentage of outliers. 142 | mean_acc3d_relax : float 143 | Average relaxed accuracy. 144 | mean_acc3d_strict : TYPE 145 | Average strict accuracy. 146 | 147 | """ 148 | 149 | # Init. 150 | start_time_eval = time.time() 151 | num_batch = len(testloader) 152 | fname_list = [None] * num_batch 153 | epe_per_scene = np.zeros(num_batch, dtype=np.float32) 154 | acc3d_strict_per_scene = np.zeros(num_batch, dtype=np.float32) 155 | acc3d_relax_per_scene = np.zeros(num_batch, dtype=np.float32) 156 | outlier_per_scene = np.zeros(num_batch, dtype=np.float32) 157 | duration_per_scene = np.zeros(num_batch, dtype=np.float32) 158 | epe_per_point_list = [None] * num_batch 159 | acc3d_strict_per_point_list = [None] * num_batch 160 | acc3d_relax_per_point_list = [None] * num_batch 161 | outlier_per_point_list = [None] * num_batch 162 | running_epe = 0 163 | running_acc3d_strict = 0 164 | running_acc3d_relax = 0 165 | running_outlier = 0 166 | 167 | if args.use_test_time_refinement: 168 | target_recon_loss_refinement = np.zeros([num_batch, args.test_time_num_step + 1], dtype=np.float32) 169 | smooth_flow_loss_refinement = np.zeros([num_batch, args.test_time_num_step + 1], dtype=np.float32) 170 | epe_refinement = np.zeros([num_batch, args.test_time_num_step + 1], dtype=np.float32) 171 | duration_refinement = np.zeros([num_batch], dtype=np.float32) 172 | 173 | save_pc_res = res_dir is not None 174 | 175 | # 176 | scoop = scoop.eval() 177 | device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") 178 | for it, batch in enumerate(tqdm(testloader)): 179 | 180 | # Current file name 181 | fname = os.path.split(testloader.dataset.filename_curr)[-1] 182 | fname_list[it] = fname 183 | 184 | # Send data to GPU 185 | batch = batch.to(device, non_blocking=True) 186 | 187 | # Estimate flow 188 | start_time = time.time() 189 | if args.all_candidates: 190 | est_flow, corr_conf, graph = compute_flow(scoop, batch, args) 191 | else: 192 | with torch.no_grad(): 193 | est_flow, corr_conf, _, graph = scoop(batch["sequence"]) 194 | duration_corr = time.time() - start_time 195 | 196 | duration_total = duration_corr 197 | if args.use_test_time_refinement: 198 | refiner = Refiner(est_flow.shape, est_flow.device) 199 | test_time_optimizer = torch.optim.Adam(refiner.parameters(), lr=args.test_time_update_rate) 200 | refined_flow, refine_metrics_curr, duration_curr = refiner.refine_flow(batch, est_flow.detach(), corr_conf.detach(), test_time_optimizer, args) 201 | est_flow = refined_flow.detach() 202 | 203 | target_recon_loss_refinement[it] = refine_metrics_curr["target_recon_loss_all"] 204 | smooth_flow_loss_refinement[it] = refine_metrics_curr["smooth_flow_loss_all"] 205 | epe_refinement[it] = refine_metrics_curr["epe_all"] 206 | duration_refinement[it] = duration_curr 207 | 208 | duration_total = duration_total + duration_curr 209 | 210 | # Performance metrics 211 | epe3d, acc3d_strict, acc3d_relax, outlier, epe3d_pp, acc3d_strict_pp, acc3d_relax_pp, outlier_pp =\ 212 | compute_epe_test(est_flow, batch, args) 213 | epe_per_scene[it] = epe3d 214 | acc3d_strict_per_scene[it] = acc3d_strict 215 | acc3d_relax_per_scene[it] = acc3d_relax 216 | outlier_per_scene[it] = outlier 217 | duration_per_scene[it] = duration_total 218 | epe_per_point_list[it] = epe3d_pp 219 | acc3d_strict_per_point_list[it] = acc3d_strict_pp 220 | acc3d_relax_per_point_list[it] = acc3d_relax_pp 221 | outlier_per_point_list[it] = outlier_pp 222 | running_epe += epe3d 223 | running_outlier += outlier 224 | running_acc3d_relax += acc3d_relax 225 | running_acc3d_strict += acc3d_strict 226 | 227 | # Save point cloud results 228 | if save_pc_res: 229 | n1 = int(batch["orig_size"][0]) 230 | n1_save = n1 if args.all_points else int(batch["sequence"][0].shape[1]) 231 | pc1_save = batch["sequence"][0].cpu().numpy().reshape([-1, 3])[:n1_save] 232 | est_flow_save = est_flow.cpu().numpy().reshape([-1, 3])[:n1_save] 233 | corr_conf_save = corr_conf.cpu().numpy().reshape([-1])[:n1_save] 234 | mask_save = batch["ground_truth"][0].cpu().numpy().reshape([-1, 1])[:n1_save] == 1 235 | gt_save = batch["ground_truth"][1].cpu().numpy().reshape([-1, 3])[:n1_save] 236 | 237 | n2 = int(batch["orig_size"][1]) 238 | n2_save = n2 if args.all_points else int(batch["sequence"][1].shape[1]) 239 | pc2_save = batch["sequence"][1].cpu().numpy().reshape([-1, 3])[:n2_save] 240 | 241 | fname_split = fname.split(".") 242 | fname_pc_save = ".".join([fname_split[0] + "_res", fname_split[1]]) 243 | path_pc_save = os.path.join(res_dir, fname_pc_save) 244 | np.savez(path_pc_save, pc1=pc1_save, pc2=pc2_save, gt_mask_for_pc1=mask_save, gt_flow_for_pc1=gt_save, est_flow_for_pc1=est_flow_save, corr_conf_for_pc1=corr_conf_save) 245 | 246 | mean_epe = running_epe / num_batch 247 | mean_outlier = running_outlier / num_batch 248 | mean_acc3d_relax = running_acc3d_relax / num_batch 249 | mean_acc3d_strict = running_acc3d_strict / num_batch 250 | 251 | log_string(log_file, "EPE: %.4f, ACC3DS: %.4f, ACC3DR: %.4f, Outlier: %.4f, Dataset Size: %d" % 252 | (mean_epe, mean_acc3d_strict, mean_acc3d_relax, mean_outlier, num_batch)) 253 | 254 | duration_eval = time.time() - start_time_eval 255 | log_string(log_file, "Evaluation duration: %.2f minutes (time per example: %.2f seconds)" % 256 | (duration_eval/60, duration_eval/num_batch)) 257 | 258 | fnames = np.array(fname_list) 259 | if args.save_metrics: 260 | path_metrics_save = os.path.join(log_dir, args.metrics_fname) 261 | data_for_save = {"fnames": fnames, 262 | "epe_per_scene": epe_per_scene, 263 | "acc3d_strict_per_scene": acc3d_strict_per_scene, 264 | "acc3d_relax_per_scene": acc3d_relax_per_scene, 265 | "outlier_per_scene": outlier_per_scene, 266 | "duration_per_scene": duration_per_scene, 267 | "epe_per_point": epe_per_point_list 268 | } 269 | 270 | if args.use_test_time_refinement: 271 | data_for_save["target_recon_loss_refinement"] = target_recon_loss_refinement 272 | data_for_save["smooth_flow_loss_refinement"] = smooth_flow_loss_refinement 273 | data_for_save["epe_refinement"] = epe_refinement 274 | data_for_save["duration_refinement"] = duration_refinement 275 | 276 | np.savez(path_metrics_save, **data_for_save) 277 | 278 | return mean_epe, mean_outlier, mean_acc3d_relax, mean_acc3d_strict 279 | 280 | 281 | def my_main(args): 282 | """ 283 | Entry point of the script. 284 | 285 | Parameters 286 | ---------- 287 | args.dataset_name : str 288 | Dataset for evaluation. Either FlowNet3D_kitti or FlowNet3D_ft3d or HPLFlowNet_kitti or HPLFlowNet_ft3d. 289 | args.batch_size: int 290 | Batch size for evaluation. 291 | args.nb_points : int 292 | Number of points in point clouds. 293 | args.path2ckpt : str 294 | Path to saved model. 295 | args.mode : str 296 | Whether to use test set of validation set or all set. 297 | 298 | Raises 299 | ------ 300 | ValueError 301 | Unknown dataset. 302 | 303 | """ 304 | 305 | # Set seed 306 | seed = seed_everything(seed=42) 307 | 308 | # Path to current file 309 | pathroot = os.path.dirname(__file__) 310 | 311 | # Select dataset 312 | if args.dataset_name.split("_")[0].lower() == "HPLFlowNet".lower(): 313 | 314 | # HPLFlowNet version of the datasets 315 | path2data = os.path.join(pathroot, "..", "data", "HPLFlowNet") 316 | 317 | # KITTI 318 | if args.dataset_name.split("_")[1].lower() == "kitti".lower(): 319 | path2data = os.path.join(path2data, "KITTI_processed_occ_final") 320 | from datasets.kitti_hplflownet import Kitti 321 | 322 | assert args.mode == "val" or args.mode == "test" or args.mode == "all", "Problem with mode " + args.mode 323 | dataset = Kitti(root_dir=path2data, nb_points=args.nb_points, all_points=args.all_points, mode=args.mode) 324 | 325 | # FlyingThing3D 326 | elif args.dataset_name.split("_")[1].lower() == "ft3d".lower(): 327 | path2data = os.path.join(path2data, "FlyingThings3D_subset_processed_35m") 328 | from datasets.flyingthings3d_hplflownet import FT3D 329 | 330 | assert args.mode == "val" or args.mode == "test", "Problem with mode " + args.mode 331 | dataset = FT3D(root_dir=path2data, nb_points=args.nb_points, all_points=args.all_points, mode=args.mode, nb_examples=-1) 332 | 333 | else: 334 | raise ValueError("Unknown dataset " + args.dataset_name) 335 | 336 | elif args.dataset_name.split("_")[0].lower() == "FlowNet3D".lower(): 337 | 338 | # FlowNet3D version of the datasets 339 | path2data = os.path.join(pathroot, "..", "data", "FlowNet3D") 340 | 341 | # KITTI 342 | if args.dataset_name.split("_")[1].lower() == "kitti".lower(): 343 | path2data = os.path.join(path2data, "kitti_rm_ground") 344 | from datasets.kitti_flownet3d import Kitti 345 | 346 | assert args.mode == "val" or args.mode == "test" or args.mode == "all", "Problem with mode " + args.mode 347 | dataset = Kitti(root_dir=path2data, nb_points=args.nb_points, all_points=args.all_points, 348 | same_v_t_split=True, mode=args.mode) 349 | 350 | # FlyingThing3D 351 | elif args.dataset_name.split("_")[1].lower() == "ft3d".lower(): 352 | path2data = os.path.join(path2data, "data_processed_maxcut_35_20k_2k_8192") 353 | from datasets.flyingthings3d_flownet3d import FT3D 354 | 355 | assert args.mode == "val" or args.mode == "test", "Problem with mode " + args.mode 356 | dataset = FT3D(root_dir=path2data, nb_points=args.nb_points, all_points=args.all_points, mode=args.mode, nb_examples=-1) 357 | 358 | else: 359 | raise ValueError("Unknown dataset" + args.dataset_name) 360 | 361 | else: 362 | raise ValueError("Unknown dataset " + args.dataset_name) 363 | print("\n\nDataset: " + path2data + " " + args.mode) 364 | 365 | # Dataloader 366 | testloader = DataLoader( 367 | dataset, 368 | batch_size=args.batch_size, 369 | pin_memory=True, 370 | shuffle=False, 371 | num_workers=args.nb_workers, 372 | collate_fn=Batch, 373 | drop_last=False, 374 | ) 375 | 376 | # Load Checkpoint 377 | file = torch.load(args.path2ckpt) 378 | 379 | # Load parameters 380 | saved_args = file["args"] 381 | 382 | # Load model 383 | scoop = SCOOP(saved_args) 384 | device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") 385 | scoop = scoop.to(device, non_blocking=True) 386 | scoop.load_state_dict(file["model"]) 387 | scoop = scoop.eval() 388 | 389 | # Log file 390 | log_dir = os.path.split(args.path2ckpt)[0] 391 | log_file = open(os.path.join(log_dir, args.log_fname), 'w') 392 | log_string(log_file, 'Evaluation arguments:') 393 | log_file.write(str(args) + '\n') 394 | log_string(log_file, "Seed: %d" % seed) 395 | 396 | # Point cloud results directory 397 | res_dir = None 398 | if args.save_pc_res: 399 | res_dir = os.path.join(log_dir, args.res_folder) 400 | if not os.path.exists(res_dir): 401 | os.mkdir(res_dir) 402 | 403 | # Evaluation 404 | epsilon = 0.03 + torch.exp(scoop.epsilon).item() 405 | gamma = torch.exp(scoop.gamma).item() 406 | power = gamma / (gamma + epsilon) 407 | log_string(log_file, "Epsilon: %.4f, Power: %.4f" % (epsilon, power)) 408 | eval_model(scoop, testloader, log_file, log_dir, res_dir, args) 409 | 410 | log_string(log_file, "Finished Evaluation.") 411 | log_file.close() 412 | 413 | 414 | if __name__ == "__main__": 415 | 416 | # Args 417 | parser = argparse.ArgumentParser(description="Evaluate SCOOP.") 418 | parser.add_argument("--dataset_name", type=str, default="FlowNet3D_kitti", help="Dataset. FlowNet3D_kitti or FlowNet3D_FT3D or Either HPLFlowNet_kitti or HPLFlowNet_FT3D.") 419 | parser.add_argument("--batch_size", type=int, default=1, help="Batch size for evaluation.") 420 | parser.add_argument("--mode", type=str, default="test", help="Test or validation or all dataset (options: [val, test, all]).") 421 | parser.add_argument("--use_test_time_refinement", type=int, default=1, help="1: Use test time refinement, 0: Do not use test time refinement.") 422 | parser.add_argument("--test_time_num_step", type=int, default=150, help="1: Number of steps for test time refinement.") 423 | parser.add_argument("--test_time_update_rate", type=float, default=0.2, help="1: Update rate for test time refinement.") 424 | parser.add_argument("--backward_dist_weight", type=float, default=0.0, help="Backward distance weight for target reconstruction loss in test time refinement.") 425 | parser.add_argument("--target_recon_loss_weight", type=float, default=1.0, help="Weight for target reconstruction loss in test time refinement.") 426 | parser.add_argument("--use_smooth_flow", type=int, default=1, help="1: Use self smooth flow loss in test time refinement, 0: Do not use smooth flow loss.") 427 | parser.add_argument("--nb_neigh_smooth_flow", type=int, default=32, help="Number of neighbor points for smooth flow loss in test time refinement.") 428 | parser.add_argument("--smooth_flow_loss_weight", type=float, default=1.0, help="Weight for smooth flow loss in test time refinement. Active if > 0.") 429 | parser.add_argument("--test_time_verbose", type=int, default=0, help="1: Print test time results during optimization, 0: Do not print.") 430 | parser.add_argument("--use_chamfer_cuda", type=int, default=1, help="1: Use chamfer distance cuda implementation in test time refinement, 0: Use chamfer distance pytorch implementation in test time refinement.") 431 | parser.add_argument("--nb_points", type=int, default=2048, help="Maximum number of points in point cloud.") 432 | parser.add_argument("--all_points", type=int, default=1, help="1: use all point in the source point cloud for evaluation in chunks of nb_points, 0: use only nb_points.") 433 | parser.add_argument("--all_candidates", type=int, default=1, help="1: use all points in the target point cloud as candidates concurrently, 0: use chunks of nb_points from the target point cloud each time.") 434 | parser.add_argument("--nb_points_chunk", type=int, default=2048, help="Number of source points chuck for evaluation with all candidate target points.") 435 | parser.add_argument("--nb_workers", type=int, default=0, help="Number of workers for the dataloader.") 436 | parser.add_argument("--path2ckpt", type=str, default="./../pretrained_models/kitti_v_100_examples/model_e400.tar", help="Path to saved checkpoint.") 437 | parser.add_argument("--log_fname", type=str, default="log_evaluation.txt", help="Evaluation log file name.") 438 | parser.add_argument("--save_pc_res", type=int, default=0, help="1: save point cloud results, 0: do not save point cloud results [default: 0]") 439 | parser.add_argument("--res_folder", type=str, default="pc_res", help="Folder name for saving results.") 440 | parser.add_argument("--save_metrics", type=int, default=0, help="1: save evaluation metrics results, 0: do not save evaluation metrics results [default: 0]") 441 | parser.add_argument("--metrics_fname", type=str, default="metrics_results.npz", help="Name for metrics file.") 442 | args = parser.parse_args() 443 | 444 | # Check arguments 445 | if args.all_points: 446 | assert args.batch_size == 1, "For evaluation with all source points, the batch_size should be equal to 1 (got %d)" % args.batch_size 447 | 448 | if args.all_candidates: 449 | assert args.batch_size == 1, "For evaluation with all candidate target points, the batch_size should be equal to 1 (got %d)" % args.batch_size 450 | 451 | if args.save_pc_res: 452 | assert args.batch_size == 1, "For evaluation with saving point cloud results, the batch_size should be equal to 1 (got %d)" % args.batch_size 453 | 454 | # Launch evaluation 455 | my_main(args) 456 | -------------------------------------------------------------------------------- /scripts/train_on_ft3d_o.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | # train on ft3d_o 4 | python train_scoop.py --dataset_name FlowNet3D_ft3d --nb_train_examples 1800 --nb_points 2048 \ 5 | --batch_size_train 4 --batch_size_val 10 --nb_epochs 100 --nb_workers 8 \ 6 | --backward_dist_weight 0.0 --use_corr_conf 1 --corr_conf_loss_weight 0.1 \ 7 | --add_model_suff 1 --save_model_epoch 25 --log_dir ft3d_o_1800_examples 8 | -------------------------------------------------------------------------------- /scripts/train_on_ft3d_s.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | # train on ft3d_s 4 | python train_scoop.py --dataset_name HPLFlowNet_ft3d --nb_train_examples 1800 --nb_val_examples 200 --nb_points 8192 \ 5 | --batch_size_train 1 --batch_size_val 1 --nb_epochs 60 --nb_workers 8 \ 6 | --backward_dist_weight 1.0 --use_corr_conf 0 \ 7 | --add_model_suff 1 --save_model_epoch 15 --log_dir ft3d_s_1800_examples 8 | -------------------------------------------------------------------------------- /scripts/train_on_kitti_v.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | # train on kitti_v 4 | python train_scoop.py --dataset_name FlowNet3D_kitti --nb_points 2048 \ 5 | --batch_size_train 4 --batch_size_val 10 --nb_epochs 400 --nb_workers 8 \ 6 | --backward_dist_weight 0.0 --use_corr_conf 1 --corr_conf_loss_weight 0.1 \ 7 | --add_model_suff 1 --save_model_epoch 100 --log_dir kitti_v_100_examples 8 | -------------------------------------------------------------------------------- /scripts/train_scoop.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import time 4 | import torch 5 | import argparse 6 | from tqdm import tqdm 7 | from torch.utils.data import DataLoader 8 | from torch.utils.tensorboard import SummaryWriter 9 | 10 | # add path 11 | project_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) 12 | if project_dir not in sys.path: 13 | sys.path.append(project_dir) 14 | 15 | from datasets.generic import Batch 16 | from models.scoop import SCOOP 17 | from tools.seed import seed_everything 18 | from tools.losses import compute_loss_unsupervised 19 | from tools.utils import log_string 20 | 21 | 22 | def compute_epe_train(recon_flow, batch): 23 | """ 24 | Compute EPE during training. 25 | 26 | Parameters 27 | ---------- 28 | recon_flow: torch.Tensor 29 | Flow from reconstruction of the target point cloud by the source point cloud. 30 | batch : scoop.datasets.generic.Batch 31 | Contains ground truth flow and mask. 32 | 33 | Returns 34 | ------- 35 | epe : torch.Tensor 36 | Mean EPE for current batch. 37 | 38 | """ 39 | 40 | mask = batch["ground_truth"][0][..., 0] 41 | true_flow = batch["ground_truth"][1] 42 | error = recon_flow - true_flow 43 | error = error[mask > 0] 44 | epe_per_point = torch.sqrt(torch.sum(torch.pow(error, 2.0), -1)) 45 | epe = epe_per_point.mean() 46 | 47 | return epe 48 | 49 | 50 | def train(scoop, train_dataloader, val_dataloader, delta, optimizer, scheduler, path2log, args): 51 | """ 52 | Train scene flow model. 53 | 54 | Parameters 55 | ---------- 56 | scoop : scoop.models.SCOOP 57 | SCOOP model 58 | train_dataloader : scoop.datasets.generic.SceneFlowDataset 59 | Training Dataset loader. 60 | val_dataloader : scoop.datasets.generic.SceneFlowDataset 61 | Validation Dataset loader. 62 | delta : int 63 | Frequency of logs in number of iterations. 64 | optimizer : torch.optim.Optimizer 65 | Optimiser. 66 | scheduler : 67 | Scheduler. 68 | path2log : str 69 | Where to save logs / model. 70 | args : Namespace 71 | Arguments for training. 72 | 73 | """ 74 | 75 | # Set seed 76 | seed = seed_everything(seed=42) 77 | 78 | # Log directory 79 | if not os.path.exists(path2log): 80 | os.makedirs(path2log) 81 | writer = SummaryWriter(path2log) 82 | 83 | # Log file 84 | log_file = open(os.path.join(path2log, args.log_fname), 'w') 85 | log_file.write('Training arguments:') 86 | log_file.write(str(args) + '\n') 87 | log_string(log_file, "Seed: %d" % seed) 88 | 89 | # Train 90 | device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") 91 | scoop = scoop.to(device, non_blocking=True) 92 | 93 | total_it = 0 94 | epoch_start = 0 95 | for epoch in range(epoch_start, args.nb_epochs): 96 | 97 | # Initialization 98 | running_loss = 0 99 | training_loss_tot = 0 100 | running_target_recon_loss = 0 101 | training_target_recon_loss_tot = 0 102 | if args.use_corr_conf: 103 | running_corr_conf_loss = 0 104 | training_corr_conf_loss_tot = 0 105 | if args.use_smooth_flow: 106 | running_smooth_flow_loss = 0 107 | training_smooth_flow_loss_tot = 0 108 | running_epe = 0 109 | training_epe_tot = 0 110 | epoch_it = 0 111 | 112 | # Train for one epoch 113 | start = time.time() 114 | scoop = scoop.train() 115 | log_string(log_file, "Training epoch %d out of %d" % (epoch + 1, args.nb_epochs)) 116 | for it, batch in enumerate(tqdm(train_dataloader)): 117 | 118 | # Send data to GPU 119 | batch = batch.to(device, non_blocking=True) 120 | 121 | # Run model 122 | recon_flow, corr_conf, target_recon, graph = scoop(batch["sequence"]) 123 | 124 | # Compute loss 125 | loss, target_recon_loss, corr_conf_loss, smooth_flow_loss =\ 126 | compute_loss_unsupervised(recon_flow, corr_conf, target_recon, graph, batch, args) 127 | 128 | # Gradient step 129 | optimizer.zero_grad() 130 | loss.backward() 131 | optimizer.step() 132 | 133 | # Loss evolution 134 | loss_curr = loss.item() 135 | running_loss += loss_curr 136 | training_loss_tot += loss_curr 137 | target_recon_loss_curr = target_recon_loss.item() 138 | running_target_recon_loss += target_recon_loss_curr 139 | training_target_recon_loss_tot += target_recon_loss_curr 140 | if args.use_corr_conf: 141 | corr_conf_loss_curr = corr_conf_loss.item() 142 | running_corr_conf_loss += corr_conf_loss_curr 143 | training_corr_conf_loss_tot += corr_conf_loss_curr 144 | if args.use_smooth_flow: 145 | smooth_flow_loss_curr = smooth_flow_loss.item() 146 | running_smooth_flow_loss += smooth_flow_loss_curr 147 | training_smooth_flow_loss_tot += smooth_flow_loss_curr 148 | 149 | epe_curr = compute_epe_train(recon_flow, batch).item() 150 | running_epe += epe_curr 151 | training_epe_tot += epe_curr 152 | 153 | # Logs 154 | if it % delta == delta - 1: 155 | # Print / save logs 156 | writer.add_scalar("Loss/training_loss", running_loss / delta, total_it) 157 | writer.add_scalar("Loss/training_target_recon_loss", running_target_recon_loss / delta, total_it) 158 | if args.use_corr_conf: 159 | writer.add_scalar("Loss/training_corr_conf_loss", running_corr_conf_loss / delta, total_it) 160 | if args.use_smooth_flow: 161 | writer.add_scalar("Loss/training_smooth_flow_loss", running_smooth_flow_loss / delta, total_it) 162 | writer.add_scalar("Loss/training_epe", running_epe / delta, total_it) 163 | print("Epoch {0:d} - It. {1:d}: loss = {2:e}".format(epoch + 1, total_it, running_loss / delta)) 164 | print((time.time() - start) / 60, "minutes") 165 | # Re-init. 166 | running_loss = 0 167 | running_target_recon_loss = 0 168 | if args.use_corr_conf: 169 | running_corr_conf_loss = 0 170 | if args.use_smooth_flow: 171 | running_smooth_flow_loss = 0 172 | running_epe = 0 173 | start = time.time() 174 | 175 | epoch_it += 1 176 | total_it += 1 177 | 178 | # Training loss 179 | training_loss = training_loss_tot / epoch_it 180 | writer.add_scalar("Loss/training_loss", training_loss, total_it) 181 | training_target_recon_loss = training_target_recon_loss_tot / epoch_it 182 | writer.add_scalar("Loss/training_target_recon_loss", training_target_recon_loss, total_it) 183 | if args.use_corr_conf: 184 | training_corr_conf_loss = training_corr_conf_loss_tot / epoch_it 185 | writer.add_scalar("Loss/training_corr_conf_loss", training_corr_conf_loss, total_it) 186 | if args.use_smooth_flow: 187 | training_smooth_flow_loss = training_smooth_flow_loss_tot / epoch_it 188 | writer.add_scalar("Loss/training_smooth_flow_loss", training_smooth_flow_loss, total_it) 189 | training_epe = training_epe_tot / epoch_it 190 | writer.add_scalar("Loss/training_epe", training_epe, total_it) 191 | log_string(log_file, "Training: loss: %.6f, epe: %.3f" % (training_loss, training_epe)) 192 | 193 | # Scheduler 194 | scheduler.step() 195 | 196 | # Validation 197 | scoop = scoop.eval() 198 | val_loss_tot = 0 199 | val_target_recon_loss_tot = 0 200 | if args.use_corr_conf: 201 | val_corr_conf_loss_tot = 0 202 | if args.use_smooth_flow: 203 | val_smooth_flow_loss_tot = 0 204 | val_epe_tot = 0 205 | val_it = 0 206 | with torch.no_grad(): 207 | for it, batch in enumerate(tqdm(val_dataloader)): 208 | 209 | # Send data to GPU 210 | batch = batch.to(device, non_blocking=True) 211 | 212 | # Run model 213 | recon_flow, corr_conf, target_recon, graph = scoop(batch["sequence"]) 214 | 215 | # Compute loss 216 | loss, target_recon_loss, corr_conf_loss, smooth_flow_loss =\ 217 | compute_loss_unsupervised(recon_flow, corr_conf, target_recon, graph, batch, args) 218 | 219 | # Validation loss 220 | val_loss_tot += loss.item() 221 | val_target_recon_loss_tot += target_recon_loss.item() 222 | if args.use_corr_conf: 223 | val_corr_conf_loss_tot += corr_conf_loss.item() 224 | if args.use_smooth_flow: 225 | val_smooth_flow_loss_tot += smooth_flow_loss.item() 226 | val_epe_tot += compute_epe_train(recon_flow, batch).item() 227 | val_it += 1 228 | 229 | val_loss = val_loss_tot / val_it 230 | writer.add_scalar("Loss/validation_loss", val_loss, total_it) 231 | val_target_recon_loss = val_target_recon_loss_tot / val_it 232 | writer.add_scalar("Loss/validation_target_recon_loss", val_target_recon_loss, total_it) 233 | if args.use_corr_conf: 234 | val_corr_conf_loss = val_corr_conf_loss_tot / val_it 235 | writer.add_scalar("Loss/validation_corr_conf_loss", val_corr_conf_loss, total_it) 236 | if args.use_smooth_flow: 237 | val_smooth_flow_loss = val_smooth_flow_loss_tot / val_it 238 | writer.add_scalar("Loss/validation_smooth_flow_loss", val_smooth_flow_loss, total_it) 239 | val_epe = val_epe_tot / val_it 240 | writer.add_scalar("Loss/validation_epe", val_epe, total_it) 241 | log_string(log_file, "Validation: loss: %.6f, epe: %.3f" % (val_loss, val_epe)) 242 | 243 | # Save model after each epoch 244 | state = { 245 | "args": args, 246 | "model": scoop.state_dict(), 247 | "optimizer": optimizer.state_dict(), 248 | "scheduler": scheduler.state_dict(), 249 | } 250 | if epoch == 0 or (epoch + 1) % args.save_model_epoch == 0: 251 | suff = args.add_model_suff and "_e%03d" % (epoch + 1) or "" 252 | torch.save(state, os.path.join(path2log, "model%s.tar" % suff)) 253 | 254 | log_string(log_file, "Finished Training.") 255 | log_file.close() 256 | 257 | return None 258 | 259 | 260 | def my_main(args): 261 | """ 262 | Entry point of the script. 263 | 264 | Parameters 265 | ---------- 266 | args.dataset_name : str 267 | Dataset for training. Either FlowNet3D_kitti or FlowNet3D_ft3d or HPLFlowNet_kitti or HPLFlowNet_ft3d. 268 | args.nb_iter : int 269 | Number of unrolled iteration of Sinkhorn algorithm in SCOOP. 270 | args.batch_size_train : int 271 | Batch size fot training dataset. 272 | args.batch_size_val : int 273 | Batch size fot validation dataset. 274 | args.nb_points : int 275 | Number of points in point clouds. 276 | args.nb_epochs : int 277 | Number of epochs. 278 | args.nb_workers: int 279 | Number of workers for the dataloader. 280 | args.log_dir: 281 | Logging directory. 282 | 283 | Raises 284 | ------ 285 | ValueError 286 | If dataset_name is an unknown dataset. 287 | 288 | """ 289 | 290 | # Path to current file 291 | pathroot = os.path.dirname(__file__) 292 | 293 | # Select dataset 294 | if args.dataset_name.split("_")[0].lower() == "HPLFlowNet".lower(): 295 | # HPLFlowNet version of the datasets 296 | path2data = os.path.join(pathroot, "..", "data", "HPLFlowNet") 297 | 298 | # KITTI 299 | if args.dataset_name.split("_")[1].lower() == "kitti".lower(): 300 | path2data = os.path.join(path2data, "KITTI_processed_occ_final") 301 | from datasets.kitti_hplflownet import Kitti 302 | 303 | # datasets 304 | train_dataset = Kitti(root_dir=path2data, nb_points=args.nb_points, all_points=False, mode="train") 305 | val_dataset = Kitti(root_dir=path2data, nb_points=args.nb_points, all_points=False, mode="val") 306 | 307 | # learning rate schedule 308 | lr_lambda = lambda epoch: 1.0 if epoch < 50 else 1.0 309 | 310 | # FlyingThing3D 311 | elif args.dataset_name.split("_")[1].lower() == "ft3d".lower(): 312 | path2data = os.path.join(path2data, "FlyingThings3D_subset_processed_35m") 313 | from datasets.flyingthings3d_hplflownet import FT3D 314 | 315 | # datasets 316 | train_dataset = FT3D(root_dir=path2data, nb_points=args.nb_points, all_points=False, mode="train", nb_examples=args.nb_train_examples) 317 | val_dataset = FT3D(root_dir=path2data, nb_points=args.nb_points, all_points=False, mode="val", nb_examples=args.nb_val_examples) 318 | 319 | # learning rate schedule 320 | lr_lambda = lambda epoch: 1.0 if epoch < 50 else 0.1 321 | 322 | else: 323 | raise ValueError("Unknown dataset " + args.dataset_name) 324 | 325 | elif args.dataset_name.split("_")[0].lower() == "FlowNet3D".lower(): 326 | # FlowNet3D version of the datasets 327 | path2data = os.path.join(pathroot, "..", "data", "FlowNet3D") 328 | 329 | # KITTI 330 | if args.dataset_name.split("_")[1].lower() == "kitti".lower(): 331 | path2data = os.path.join(path2data, "kitti_rm_ground") 332 | from datasets.kitti_flownet3d import Kitti 333 | 334 | # datasets 335 | train_dataset = Kitti(root_dir=path2data, nb_points=args.nb_points, all_points=False, same_v_t_split=args.same_val_test_kitti, mode="train") 336 | val_dataset = Kitti(root_dir=path2data, nb_points=args.nb_points, all_points=False, same_v_t_split=args.same_val_test_kitti, mode="val") 337 | 338 | # learning rate schedule 339 | lr_lambda = lambda epoch: 1.0 if epoch < 340 else 0.1 340 | 341 | # FlyingThing3D 342 | elif args.dataset_name.split("_")[1].lower() == "ft3d".lower(): 343 | path2data = os.path.join(path2data, "data_processed_maxcut_35_20k_2k_8192") 344 | from datasets.flyingthings3d_flownet3d import FT3D 345 | 346 | # datasets 347 | train_dataset = FT3D(root_dir=path2data, nb_points=args.nb_points, all_points=False, mode="train", nb_examples=args.nb_train_examples) 348 | val_dataset = FT3D(root_dir=path2data, nb_points=args.nb_points, all_points=False, mode="val", nb_examples=args.nb_val_examples) 349 | 350 | # learning rate schedule 351 | lr_lambda = lambda epoch: 1.0 if epoch < 340 else 1.0 352 | 353 | else: 354 | raise ValueError("Unknown dataset" + args.dataset_name) 355 | 356 | else: 357 | raise ValueError("Invalid dataset name: " + args.dataset_name) 358 | 359 | # Training data 360 | train_dataloader = DataLoader( 361 | train_dataset, 362 | batch_size=args.batch_size_train, 363 | pin_memory=True, 364 | shuffle=True, 365 | num_workers=args.nb_workers, 366 | collate_fn=Batch, 367 | drop_last=True, 368 | ) 369 | 370 | # Validation data 371 | val_dataloader = DataLoader( 372 | val_dataset, 373 | batch_size=args.batch_size_val, 374 | pin_memory=True, 375 | shuffle=False, 376 | num_workers=args.nb_workers, 377 | collate_fn=Batch, 378 | drop_last=False, 379 | ) 380 | 381 | # Model 382 | scoop = SCOOP(args) 383 | 384 | # Optimizer 385 | optimizer = torch.optim.Adam(scoop.parameters(), lr=args.learning_rate) 386 | 387 | # Scheduler 388 | scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda) 389 | 390 | # Log directory 391 | path2log = os.path.join(pathroot, "..", "experiments", args.log_dir) 392 | 393 | 394 | # Train 395 | print("Training started. Logs in " + path2log) 396 | train(scoop, train_dataloader, val_dataloader, 500, optimizer, scheduler, path2log, args) 397 | 398 | return None 399 | 400 | 401 | if __name__ == "__main__": 402 | 403 | # Args 404 | parser = argparse.ArgumentParser(description="Train SCOOP.") 405 | parser.add_argument("--dataset_name", type=str, default="FlowNet3D_kitti", help="Dataset name for train and validation. Either FlowNet3D_kitti or FlowNet3D_ft3d or HPLFlowNet_kitti or HPLFlowNet_ft3d.") 406 | parser.add_argument("--nb_train_examples", type=int, default=-1, help="Number of examples for the training dataset. active if > 0.") 407 | parser.add_argument("--nb_val_examples", type=int, default=-1, help="Number of examples for the validation dataset. active if > 0.") 408 | parser.add_argument("--same_val_test_kitti", type=int, default=1, help="1: Use the same validation and test split for KITTI dataset, 0: Do not use the same validation and test split.") 409 | parser.add_argument("--batch_size_train", type=int, default=4, help="Batch size fot training dataset.") 410 | parser.add_argument("--batch_size_val", type=int, default=10, help="Batch size for validation dataset.") 411 | parser.add_argument("--nb_epochs", type=int, default=400, help="Number of epochs.") 412 | parser.add_argument("--learning_rate", type=float, default=0.001, help="Learning rate.") 413 | parser.add_argument("--use_corr_conf", type=int, default=1, help="1: Use correspondence confidence for training, 0: Do not use correspondence confidence.") 414 | parser.add_argument("--linear_corr_conf", type=int, default=0, help="1: Use linear normalization for correspondence confidence, 0: Use ReLU normalization.") 415 | parser.add_argument("--corr_conf_loss_weight", type=float, default=0.1, help="Weight for correspondence confidence loss.") 416 | parser.add_argument("--nb_points", type=int, default=2048, help="Maximum number of points in point cloud.") 417 | parser.add_argument("--nb_iter", type=int, default=1, help="Number of unrolled iterations of the Sinkhorn algorithm.") 418 | parser.add_argument("--nb_neigh_cross_recon", type=int, default=64, help="Number of neighbor points for cross reconstruction. Active if > 0") 419 | parser.add_argument("--use_smooth_flow", type=int, default=1, help="1: Use smooth flow loss, 0: Do not use smooth flow loss.") 420 | parser.add_argument("--nb_neigh_smooth_flow", type=int, default=32, help="Number of neighbor points for smooth flow loss.") 421 | parser.add_argument("--smooth_flow_loss_weight", type=float, default=10.0, help="Weight for smooth flow loss. Active if > 0.") 422 | parser.add_argument("--backward_dist_weight", type=float, default=0.0, help="Backward distance weight in Chamfer Distance loss.") 423 | parser.add_argument("--nb_workers", type=int, default=0, help="Number of workers for the dataloader.") 424 | parser.add_argument("--log_dir", type=str, default="log_scoop", help="Logging directory.") 425 | parser.add_argument("--log_fname", type=str, default="log_train.txt", help="Evaluation log file name.") 426 | parser.add_argument("--save_model_epoch", type=int, default=1, help="Number of epochs difference for saving the model.") 427 | parser.add_argument("--add_model_suff", type=int, default=0, help="1: Add suffix to model name, 0: do not add suffix.") 428 | args = parser.parse_args() 429 | 430 | args.use_corr_conf = bool(args.use_corr_conf) 431 | args.use_smooth_flow = bool(args.use_smooth_flow) 432 | args.add_model_suff = bool(args.add_model_suff) 433 | 434 | # Check arguments 435 | if args.use_smooth_flow: 436 | assert args.nb_neigh_smooth_flow > 0, "If use_smooth_flow is on, nb_neigh_smooth_flow should be > 0 (got %d)." % args.nb_neigh_smooth_flow 437 | assert args.smooth_flow_loss_weight >= 0, "If use_smooth_flow is on, smooth_flow_loss_weight should be >= 0 (got %f)." % args.smooth_flow_loss_weight 438 | 439 | # Launch training 440 | my_main(args) 441 | -------------------------------------------------------------------------------- /scripts/visualize_scoop.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import argparse 4 | 5 | # add path 6 | project_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) 7 | if project_dir not in sys.path: 8 | sys.path.append(project_dir) 9 | 10 | from tools.utils import create_dir, load_data 11 | from tools.vis_utils import plot_pc_list, plot_pcs, save_pc_plot 12 | 13 | COLOR_RED = (1, 0, 0) 14 | COLOR_GREEN = (0, 1, 0) 15 | COLOR_BLUE = (0, 0, 1) 16 | COLOR_PURPLE = (127/255., 0, 1) 17 | 18 | 19 | def vis_result(res_dir, res_srt, view_dict, save_dir, save_plot=True, show_plot=False, size=(1600, 800)): 20 | res_path = os.path.join(res_dir, "%s_res.npz" % res_srt) 21 | pc1, pc2, gt_mask, gt_flow, est_flow, corr_conf, pc1_warped_gt, pc1_warped_est = load_data(res_path) 22 | 23 | # pc1, pc2 24 | pc_list = [pc1, pc2] 25 | color_list = [COLOR_RED, COLOR_GREEN] 26 | plot_pc_list(pc_list, color_list, view_dict=view_dict, get_view=False, fig=None, size=size) 27 | 28 | save_path = os.path.join(save_dir, "%s_01_pc1_pc2.png" % res_srt) 29 | save_pc_plot(save_path, save=save_plot, show=show_plot) 30 | 31 | # pc1, pc1_warped_est 32 | pc_list = [pc1, pc1_warped_est] 33 | color_list = [COLOR_RED, COLOR_BLUE] 34 | plot_pc_list(pc_list, color_list, view_dict=view_dict, get_view=False, fig=None, size=size) 35 | 36 | save_path = os.path.join(save_dir, "%s_02_pc1_pc1_warped_est.png" % res_srt) 37 | save_pc_plot(save_path, save=save_plot, show=show_plot) 38 | 39 | # pc2, pc1_warped_est 40 | pc_list = [pc2, pc1_warped_est] 41 | color_list = [COLOR_GREEN, COLOR_BLUE] 42 | plot_pc_list(pc_list, color_list, view_dict=view_dict, get_view=False, fig=None, size=size) 43 | 44 | save_path = os.path.join(save_dir, "%s_03_pc2_pc1_warped_est.png" % res_srt) 45 | save_pc_plot(save_path, save=save_plot, show=show_plot) 46 | 47 | # pc1, pc2, pc1_warped_est 48 | pc_list = [pc1, pc2, pc1_warped_est] 49 | color_list = [COLOR_RED, COLOR_GREEN, COLOR_BLUE] 50 | plot_pc_list(pc_list, color_list, view_dict=view_dict, get_view=False, fig=None, size=size) 51 | 52 | save_path = os.path.join(save_dir, "%s_04_pc1_pc2_pc1_warped_est.png" % res_srt) 53 | save_pc_plot(save_path, save=save_plot, show=show_plot) 54 | 55 | # pc1, pc1_warped_gt 56 | pc_list = [pc1, pc1_warped_gt] 57 | color_list = [COLOR_RED, COLOR_PURPLE] 58 | plot_pc_list(pc_list, color_list, view_dict=view_dict, get_view=False, fig=None, size=size) 59 | 60 | save_path = os.path.join(save_dir, "%s_05_pc1_pc1_warped_gt.png" % res_srt) 61 | save_pc_plot(save_path, save=save_plot, show=show_plot) 62 | 63 | # pc2, pc1_warped_gt 64 | pc_list = [pc2, pc1_warped_gt] 65 | color_list = [COLOR_GREEN, COLOR_PURPLE] 66 | plot_pc_list(pc_list, color_list, view_dict=view_dict, get_view=False, fig=None, size=size) 67 | 68 | save_path = os.path.join(save_dir, "%s_06_pc2_pc1_warped_gt.png" % res_srt) 69 | save_pc_plot(save_path, save=save_plot, show=show_plot) 70 | 71 | # pc1, pc2, pc1_warped_gt 72 | pc_list = [pc1, pc2, pc1_warped_gt] 73 | color_list = [COLOR_RED, COLOR_GREEN, COLOR_PURPLE] 74 | plot_pc_list(pc_list, color_list, view_dict=view_dict, get_view=False, fig=None, size=size) 75 | 76 | save_path = os.path.join(save_dir, "%s_07_pc1_pc2_pc1_warped_gt.png" % res_srt) 77 | save_pc_plot(save_path, save=save_plot, show=show_plot) 78 | 79 | # pc1, est_flow (quiver), pc2 80 | color_dict = {"pc1_c": COLOR_RED, "pc2_c": COLOR_GREEN, "flow_c": COLOR_BLUE} 81 | plot_pcs(pc1, pc2, flow=est_flow, color_dict=color_dict, view_dict=view_dict, get_view=False, size=size) 82 | 83 | save_path = os.path.join(save_dir, "%s_08_pc1_flow_est_pc2.png" % res_srt) 84 | save_pc_plot(save_path, save=save_plot, show=show_plot) 85 | 86 | # pc1, gt_flow (quiver), pc2 87 | color_dict = {"pc1_c": COLOR_RED, "pc2_c": COLOR_GREEN, "flow_c": COLOR_PURPLE} 88 | plot_pcs(pc1, pc2, flow=gt_flow, color_dict=color_dict, view_dict=view_dict, get_view=False, size=size) 89 | 90 | save_path = os.path.join(save_dir, "%s_09_pc1_flow_gt_pc2.png" % res_srt) 91 | save_pc_plot(save_path, save=save_plot, show=show_plot) 92 | 93 | 94 | if __name__ == "__main__": 95 | # Args 96 | parser = argparse.ArgumentParser(description="Visualize SCOOP.") 97 | parser.add_argument("--res_dir", type=str, default="./../pretrained_models/kitti_v_100_examples/pc_res", help="Point cloud results directory.") 98 | parser.add_argument("--res_idx", type=int, default=1, help="Index of the result to visualize.") 99 | args = parser.parse_args() 100 | 101 | view_dict = {'azimuth': 21.245700205624324, 'elevation': 99.80932242859983, 'distance': 13.089033739603677, 102 | 'focalpoint': [0.63038374, -1.79234603, 10.63751715], 'roll': 0.7910293664830116} 103 | 104 | save_dir = create_dir(os.path.join(args.res_dir, "vis")) 105 | vis_result(args.res_dir, "%06d" % args.res_idx, view_dict, save_dir, save_plot=True, show_plot=False, size=(1600, 800)) 106 | -------------------------------------------------------------------------------- /tools/losses.py: -------------------------------------------------------------------------------- 1 | """ Loss functions. """ 2 | import os 3 | import sys 4 | import torch 5 | import numpy as np 6 | 7 | # add path 8 | project_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) 9 | if project_dir not in sys.path: 10 | sys.path.append(project_dir) 11 | 12 | try: 13 | from auxiliary.ChamferDistancePytorch.chamfer3D import dist_chamfer_3D 14 | chamfer_dist_3d_cu = dist_chamfer_3D.chamfer_3DDist() 15 | except: 16 | print("Could not load compiled 3D CUDA chamfer distance") 17 | 18 | from tools.utils import iterate_in_chunks 19 | 20 | 21 | def compute_loss_unsupervised(recon_flow, corr_conf, target_pc_recon, graph, batch, args): 22 | """ 23 | Compute unsupervised training loss. 24 | 25 | Parameters 26 | ---------- 27 | recon_flow: torch.Tensor 28 | Flow from reconstruction of the target point cloud by the source point cloud. 29 | corr_conf: torch.Tensor 30 | Correspondence confidence. 31 | target_pc_recon: torch.Tensor 32 | Cross reconstructed target point cloud. 33 | graph: scoop.models.graph.Graph 34 | Nearest neighbor graph for the source point cloud. 35 | batch: scoop.datasets.generic.Batch 36 | Contains ground truth flow and mask. 37 | args: dictionary. 38 | Arguments for loss terms. 39 | 40 | Returns 41 | ------- 42 | loss : torch.Tensor 43 | Training loss for current batch. 44 | 45 | """ 46 | mask = None 47 | 48 | if args.use_corr_conf: 49 | point_weight = corr_conf 50 | else: 51 | point_weight = None 52 | 53 | target_pc_input = batch["sequence"][1] 54 | target_recon_loss = chamfer_loss(target_pc_recon, target_pc_input, point_weight, args.backward_dist_weight, mask) 55 | 56 | loss = target_recon_loss 57 | 58 | if args.use_corr_conf and args.corr_conf_loss_weight > 0: 59 | if mask is not None: 60 | corr_conf_masked = corr_conf[mask > 0] 61 | else: 62 | corr_conf_masked = corr_conf 63 | 64 | corr_conf_loss = 1 - torch.mean(corr_conf_masked) 65 | loss = loss + (args.corr_conf_loss_weight * corr_conf_loss) 66 | else: 67 | corr_conf_loss = 0 68 | 69 | if args.use_smooth_flow and args.smooth_flow_loss_weight > 0: 70 | smooth_flow_loss, _ = smooth_loss(recon_flow, graph, args.nb_neigh_smooth_flow, loss_norm=1, mask=mask) 71 | loss = loss + (args.smooth_flow_loss_weight * smooth_flow_loss) 72 | else: 73 | smooth_flow_loss = 0 74 | 75 | return loss, target_recon_loss, corr_conf_loss, smooth_flow_loss 76 | 77 | 78 | def chamfer_dist_3d_pt(pc1, pc2, backward_dist_weight=0.0, chunk_size=2048): 79 | """ 80 | Compute Chamfer Distance between two point clouds. 81 | Input: 82 | pc1: (b, n, 3) torch.Tensor, first point cloud xyz coordinates. 83 | pc2: (b, m, 3) torch.Tensor, second point cloud xyz coordinates. 84 | backward_dist_weight: float, weight for backward distance 85 | chunk_size: int, chunk size for distance computation. 86 | 87 | Output: 88 | dist1: (b, n) torch.Tensor float32, for each point in pc1, the distance to the closest point in pc2. 89 | dist2: (b, m) torch.Tensor float32, for each point in pc2, the distance to the closest point in pc1. 90 | idx1: (b, n) torch.Tensor int32, for each point in pc1, the index of the closest point in pc2 (values are in the range [0, ..., m-1]). 91 | idx2: (b, m) torch.Tensor int32, for each point in pc2, the index of the closest point in pc1 (values are in the range [0, ..., n-1]). 92 | """ 93 | 94 | b = pc1.shape[0] 95 | n = pc1.shape[1] 96 | m = pc2.shape[1] 97 | device = pc1.device 98 | 99 | dist1 = torch.zeros([b, n], dtype=torch.float32, device=device) 100 | idx1 = torch.zeros([b, n], dtype=torch.int32, device=device) 101 | 102 | rng1 = np.arange(n) 103 | for chunk in iterate_in_chunks(rng1, chunk_size): 104 | pc1_curr = torch.unsqueeze(pc1[:, chunk], dim=2).repeat(1, 1, m, 1) 105 | pc2_curr = torch.unsqueeze(pc2, dim=1).repeat(1, len(chunk), 1, 1) 106 | diff = pc1_curr - pc2_curr # shape (b, cs, m, 3) 107 | dist = torch.sum(diff ** 2, dim=-1) # shape (b, cs, m) 108 | 109 | min1 = torch.min(dist, dim=2) 110 | dist1_curr = min1.values 111 | idx1_curr = min1.indices.type(torch.IntTensor) 112 | idx1_curr = idx1_curr.to(dist.device) 113 | 114 | dist1[:, chunk] = dist1_curr 115 | idx1[:, chunk] = idx1_curr 116 | 117 | if backward_dist_weight == 0.0: 118 | dist2 = None 119 | idx2 = None 120 | else: 121 | dist2 = torch.zeros([b, m], dtype=torch.float32, device=device) 122 | idx2 = torch.zeros([b, m], dtype=torch.int32, device=device) 123 | 124 | rng2 = np.arange(m) 125 | for chunk in iterate_in_chunks(rng2, chunk_size): 126 | pc1_curr = torch.unsqueeze(pc1, dim=2).repeat(1, 1, len(chunk), 1) 127 | pc2_curr = torch.unsqueeze(pc2[:, chunk], dim=1).repeat(1, n, 1, 1) 128 | diff = pc1_curr - pc2_curr # shape (b, n, cs, 3) 129 | dist = torch.sum(diff ** 2, dim=-1) # shape (b, n, cs) 130 | 131 | min2 = torch.min(dist, dim=1) 132 | dist2_curr = min2.values 133 | idx2_curr = min2.indices.type(torch.IntTensor) 134 | idx2_curr = idx2_curr.to(dist.device) 135 | 136 | dist2[:, chunk] = dist2_curr 137 | idx2[:, chunk] = idx2_curr 138 | 139 | return dist1, dist2, idx1, idx2 140 | 141 | 142 | def chamfer_loss(pc1, pc2, point_weight=None, backward_dist_weight=1.0, mask=None, use_chamfer_cuda=True): 143 | if not pc1.is_cuda: 144 | pc1 = pc1.cuda() 145 | 146 | if not pc2.is_cuda: 147 | pc2 = pc2.cuda() 148 | 149 | if use_chamfer_cuda: 150 | dist1, dist2, idx1, idx2 = chamfer_dist_3d_cu(pc1, pc2) 151 | else: 152 | dist1, dist2, idx1, idx2 = chamfer_dist_3d_pt(pc1, pc2, backward_dist_weight) 153 | 154 | if point_weight is not None: 155 | dist1_weighted = dist1 * point_weight 156 | else: 157 | dist1_weighted = dist1 158 | 159 | if mask is not None: 160 | dist1_masked = dist1_weighted[mask > 0] 161 | dist1_mean = torch.mean(dist1_masked) 162 | else: 163 | dist1_mean = torch.mean(dist1_weighted) 164 | 165 | if backward_dist_weight == 1.0: 166 | loss = dist1_mean + torch.mean(dist2) 167 | elif backward_dist_weight == 0.0: 168 | loss = dist1_mean 169 | else: 170 | loss = dist1_mean + backward_dist_weight * torch.mean(dist2) 171 | 172 | return loss 173 | 174 | 175 | def smooth_loss(est_flow, graph, nb_neigh, loss_norm=1, mask=None): 176 | b, n, c = est_flow.shape 177 | est_flow_neigh = est_flow.reshape(b * n, c) 178 | est_flow_neigh = est_flow_neigh[graph.edges] 179 | est_flow_neigh = est_flow_neigh.view(b, n, graph.k_neighbors, c) 180 | est_flow_neigh = est_flow_neigh[:, :, 1:(nb_neigh + 1)] 181 | flow_diff = (est_flow.unsqueeze(2) - est_flow_neigh).norm(p=loss_norm, dim=-1) 182 | 183 | if mask is not None: 184 | mask_neigh = mask.reshape(b * n) 185 | mask_neigh = mask_neigh[graph.edges] 186 | mask_neigh = mask_neigh.view(b, n, graph.k_neighbors) 187 | mask_neigh = mask_neigh[:, :, 1:(nb_neigh + 1)] 188 | mask_neigh_sum = mask_neigh.sum(dim=-1) 189 | 190 | flow_diff_masked = flow_diff * mask_neigh 191 | flow_diff_masked_sum = flow_diff_masked.sum(dim=-1) 192 | smooth_flow_per_point = flow_diff_masked_sum / (mask_neigh_sum + 1e-8) 193 | smooth_flow_per_point = smooth_flow_per_point[mask > 0] 194 | else: 195 | smooth_flow_per_point = flow_diff.mean(dim=-1) 196 | 197 | smooth_flow_loss = smooth_flow_per_point.mean() 198 | 199 | return smooth_flow_loss, smooth_flow_per_point 200 | -------------------------------------------------------------------------------- /tools/ot.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | def sinkhorn(feature1, feature2, pcloud1, pcloud2, epsilon, gamma, max_iter): 5 | """ 6 | Sinkhorn algorithm 7 | 8 | Parameters 9 | ---------- 10 | feature1 : torch.Tensor 11 | Feature for points cloud 1. Used to computed transport cost. 12 | Size B x N x C. 13 | feature2 : torch.Tensor 14 | Feature for points cloud 2. Used to computed transport cost. 15 | Size B x M x C. 16 | pcloud1 : torch.Tensor 17 | Point cloud 1. Size B x N x 3. 18 | pcloud2 : torch.Tensor 19 | Point cloud 2. Size B x M x 3. 20 | epsilon : torch.Tensor 21 | Entropic regularisation. Scalar. 22 | gamma : torch.Tensor 23 | Mass regularisation. Scalar. 24 | max_iter : int 25 | Number of unrolled iteration of the Sinkhorn algorithm. 26 | 27 | Returns 28 | ------- 29 | torch.Tensor 30 | Transport plan between point cloud 1 and 2. Size B x N x M. 31 | """ 32 | 33 | # Squared l2 distance between points points of both point clouds 34 | distance_matrix = torch.sum(pcloud1 ** 2, -1, keepdim=True) 35 | distance_matrix = distance_matrix + torch.sum( 36 | pcloud2 ** 2, -1, keepdim=True 37 | ).transpose(1, 2) 38 | distance_matrix = distance_matrix - 2 * torch.bmm(pcloud1, pcloud2.transpose(1, 2)) 39 | # Force transport to be zero for points further than 10 m apart 40 | support = (distance_matrix < 10 ** 2).float() 41 | 42 | # Transport cost matrix 43 | feature1 = feature1 / torch.sqrt(torch.sum(feature1 ** 2, -1, keepdim=True) + 1e-8) 44 | feature2 = feature2 / torch.sqrt(torch.sum(feature2 ** 2, -1, keepdim=True) + 1e-8) 45 | S = torch.bmm(feature1, feature2.transpose(1, 2)) 46 | C = 1.0 - S 47 | 48 | # Entropic regularisation 49 | K = torch.exp(-C / epsilon) * support 50 | 51 | # Early return if no iterations 52 | if max_iter == 0: 53 | return K, S 54 | 55 | # Init. of Sinkhorn algorithm 56 | power = gamma / (gamma + epsilon) 57 | a = ( 58 | torch.ones( 59 | (K.shape[0], K.shape[1], 1), device=feature1.device, dtype=feature1.dtype 60 | ) 61 | / K.shape[1] 62 | ) 63 | prob1 = ( 64 | torch.ones( 65 | (K.shape[0], K.shape[1], 1), device=feature1.device, dtype=feature1.dtype 66 | ) 67 | / K.shape[1] 68 | ) 69 | prob2 = ( 70 | torch.ones( 71 | (K.shape[0], K.shape[2], 1), device=feature2.device, dtype=feature2.dtype 72 | ) 73 | / K.shape[2] 74 | ) 75 | 76 | # Sinkhorn algorithm 77 | for _ in range(max_iter): 78 | # Update b 79 | KTa = torch.bmm(K.transpose(1, 2), a) 80 | b = torch.pow(prob2 / (KTa + 1e-8), power) 81 | # Update a 82 | Kb = torch.bmm(K, b) 83 | a = torch.pow(prob1 / (Kb + 1e-8), power) 84 | 85 | # Transportation map 86 | T = torch.mul(torch.mul(a, K), b.transpose(1, 2)) 87 | 88 | return T, S 89 | -------------------------------------------------------------------------------- /tools/reconstruction.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | 4 | 5 | def get_support_matrix(pcloud1, pcloud2, dist_thresh=10): 6 | """ 7 | Compute support matrix for two point clouds. 8 | The matrix indicates if the distance between two points is less than a threshold. 9 | 10 | Parameters 11 | ---------- 12 | pcloud1 : torch.Tensor 13 | Point cloud 1. Size B x N x 3. 14 | pcloud2 : torch.Tensor 15 | Point cloud 2. Size B x M x 3. 16 | dist_thresh: 17 | Threshold on the Euclidean distance between points. 18 | 19 | Returns 20 | ------- 21 | torch.Tensor 22 | Support matrix. Size B x N x M. 23 | 24 | """ 25 | # Squared l2 distance between points points of both point clouds 26 | distance_matrix = torch.sum(pcloud1 ** 2, -1, keepdim=True) 27 | distance_matrix = distance_matrix + torch.sum(pcloud2 ** 2, -1, keepdim=True).transpose(1, 2) 28 | distance_matrix = distance_matrix - 2 * torch.bmm(pcloud1, pcloud2.transpose(1, 2)) 29 | 30 | # Force transport to be zero for points further than 10 m apart 31 | support_mat = (distance_matrix < dist_thresh ** 2).float() 32 | 33 | return support_mat 34 | 35 | 36 | def get_similarity_matrix(feature1, feature2): 37 | """ 38 | Cosine similarity between point cloud features 39 | 40 | Parameters 41 | ---------- 42 | feature1 : torch.Tensor 43 | Feature for points cloud 1. Used to computed transport cost. Size B x N x C. 44 | feature2 : torch.Tensor 45 | Feature for points cloud 2. Used to computed transport cost. Size B x M x C. 46 | 47 | Returns 48 | ------- 49 | torch.Tensor 50 | Feature similarity matrix. Size B x N x M. 51 | """ 52 | # Normalize features 53 | feature1_normalized = feature1 / (feature1.norm(dim=-1, p=2, keepdim=True) + 1e-8) 54 | feature2_normalized = feature2 / (feature2.norm(dim=-1, p=2, keepdim=True) + 1e-8) 55 | 56 | # Compute feature similarity 57 | sim_mat = torch.bmm(feature1_normalized, feature2_normalized.transpose(1, 2)) 58 | 59 | return sim_mat 60 | 61 | 62 | def normalize_mat(mat, mat_normalization, dim=None): 63 | """ 64 | The method to normalize the input matrix to be like a statistical matrix. 65 | """ 66 | if dim is None: 67 | dim = 1 if len(mat.shape) == 3 else 0 68 | 69 | if mat_normalization == "none": 70 | return mat 71 | if mat_normalization == "softmax": 72 | return F.softmax(mat, dim=dim) 73 | raise NameError 74 | 75 | 76 | def get_s_t_topk(mat, k, s_only=False, nn_idx=None): 77 | """ 78 | Get nearest neighbors per point (similarity value and index) for source and target shapes 79 | 80 | Args: 81 | mat (BxNsxNb Tensor): Similarity matrix 82 | k: Number of neighbors per point 83 | s_only: Whether to get neighbors only for the source point cloud or also for the target point cloud. 84 | nn_idx: An optional pre-computed nearest neighbor indices. 85 | """ 86 | if nn_idx is not None: 87 | assert s_only, "Pre-computed nearest neighbor indices is allowed for the source point cloud only." 88 | s_nn_idx = nn_idx 89 | s_nn_val = mat.gather(dim=2, index=nn_idx) 90 | t_nn_val = t_nn_idx = None 91 | else: 92 | s_nn_val, s_nn_idx = mat.topk(k=min(k, mat.shape[2]), dim=2) 93 | 94 | if not s_only: 95 | t_nn_val, t_nn_idx = mat.topk(k=k, dim=1) 96 | 97 | t_nn_val = t_nn_val.transpose(2, 1) 98 | t_nn_idx = t_nn_idx.transpose(2, 1) 99 | else: 100 | t_nn_val = None 101 | t_nn_idx = None 102 | 103 | return s_nn_val, s_nn_idx, t_nn_val, t_nn_idx 104 | 105 | 106 | def get_s_t_neighbors(k, mat, sim_normalization, s_only=False, ignore_first=False, nn_idx=None): 107 | s_nn_sim, s_nn_idx, t_nn_sim, t_nn_idx = get_s_t_topk(mat, k, s_only=s_only, nn_idx=nn_idx) 108 | if ignore_first: 109 | s_nn_sim, s_nn_idx = s_nn_sim[:, :, 1:], s_nn_idx[:, :, 1:] 110 | 111 | s_nn_weight = normalize_mat(s_nn_sim, sim_normalization, dim=2) 112 | 113 | if not s_only: 114 | if ignore_first: 115 | t_nn_sim, t_nn_idx = t_nn_sim[:, :, 1:], t_nn_idx[:, :, 1:] 116 | 117 | t_nn_weight = normalize_mat(t_nn_sim, sim_normalization, dim=2) 118 | else: 119 | t_nn_weight = None 120 | 121 | return s_nn_weight, s_nn_sim, s_nn_idx, t_nn_weight, t_nn_sim, t_nn_idx 122 | 123 | 124 | def knn(x, k): 125 | inner = -2 * torch.matmul(x.transpose(2, 1), x) 126 | xx = torch.sum(x ** 2, dim=1, keepdim=True) 127 | pairwise_distance = -xx - inner - xx.transpose(2, 1) 128 | 129 | idx = pairwise_distance.topk(k=k, dim=-1)[1] # (batch_size, num_points, k) 130 | 131 | return idx 132 | 133 | 134 | def get_graph_feature(x, k, idx=None, only_intrinsic='neighs', permute_feature=True): 135 | batch_size = x.size(0) 136 | num_points = x.size(2) 137 | x = x.view(batch_size, -1, num_points) 138 | if idx is None: 139 | idx = knn(x, k=k) # (batch_size, num_points, k) 140 | else: 141 | if len(idx.shape) == 2: 142 | idx = idx.unsqueeze(0).repeat(batch_size, 1, 1) 143 | idx = idx[:, :, :k] 144 | k = min(k, idx.shape[-1]) 145 | 146 | num_idx = idx.shape[1] 147 | 148 | idx_base = torch.arange(0, batch_size, device=idx.device).view(-1, 1, 1) * num_points 149 | 150 | idx = idx + idx_base 151 | 152 | idx = idx.contiguous() 153 | idx = idx.view(-1) 154 | 155 | _, num_dims, _ = x.size() 156 | 157 | x = x.transpose(2, 1).contiguous() # (batch_size, num_points, num_dims) -> (batch_size*num_points, num_dims) # batch_size * num_points * k + range(0, batch_size*num_points) 158 | feature = x.view(batch_size * num_points, -1)[idx, :] 159 | feature = feature.view(batch_size, num_idx, k, num_dims) 160 | x = x.view(batch_size, num_points, 1, num_dims).repeat(1, 1, k, 1) 161 | 162 | if only_intrinsic == 'true': 163 | feature = feature - x 164 | elif only_intrinsic == 'neighs': 165 | feature = feature 166 | elif only_intrinsic == 'concat': 167 | feature = torch.cat((feature, x), dim=3) 168 | else: 169 | feature = torch.cat((feature - x, x), dim=3) 170 | 171 | if permute_feature: 172 | feature = feature.permute(0, 3, 1, 2).contiguous() 173 | 174 | return feature 175 | 176 | 177 | def reconstruct(pos, nn_idx, nn_weight, k): 178 | nn_pos = get_graph_feature(pos.transpose(1, 2), k=k, idx=nn_idx, only_intrinsic='neighs', permute_feature=False) 179 | nn_weighted = nn_pos * nn_weight.unsqueeze(dim=3) 180 | recon = torch.sum(nn_weighted, dim=2) 181 | 182 | return recon 183 | -------------------------------------------------------------------------------- /tools/seed.py: -------------------------------------------------------------------------------- 1 | """Helper functions to help with reproducibility of models. """ 2 | 3 | import logging 4 | import os 5 | import random 6 | from typing import Optional 7 | 8 | import numpy as np 9 | import torch 10 | 11 | log = logging.getLogger(__name__) 12 | 13 | 14 | def seed_everything(seed: Optional[int] = None) -> int: 15 | """ 16 | Function that sets seed for pseudo-random number generators in: 17 | pytorch, numpy, python.random 18 | In addition, sets the env variable `PL_GLOBAL_SEED` which will be passed to 19 | spawned subprocesses (e.g. ddp_spawn backend). 20 | 21 | Args: 22 | seed: the integer value seed for global random state in Lightning. 23 | If `None`, will read seed from `PL_GLOBAL_SEED` env variable 24 | or select it randomly. 25 | """ 26 | max_seed_value = np.iinfo(np.uint32).max 27 | min_seed_value = np.iinfo(np.uint32).min 28 | 29 | try: 30 | if seed is None: 31 | seed = os.environ.get("PL_GLOBAL_SEED") 32 | seed = int(seed) 33 | except (TypeError, ValueError): 34 | seed = _select_seed_randomly(min_seed_value, max_seed_value) 35 | print(f"No correct seed found, seed set to {seed}") 36 | 37 | if not (min_seed_value <= seed <= max_seed_value): 38 | print(f"{seed} is not in bounds, numpy accepts from {min_seed_value} to {max_seed_value}") 39 | seed = _select_seed_randomly(min_seed_value, max_seed_value) 40 | 41 | log.info(f"Global seed set to {seed}") 42 | os.environ["PL_GLOBAL_SEED"] = str(seed) 43 | random.seed(seed) 44 | np.random.seed(seed) 45 | torch.manual_seed(seed) 46 | torch.cuda.manual_seed_all(seed) 47 | return seed 48 | 49 | 50 | def _select_seed_randomly(min_seed_value: int = 0, max_seed_value: int = 255) -> int: 51 | return random.randint(min_seed_value, max_seed_value) 52 | -------------------------------------------------------------------------------- /tools/utils.py: -------------------------------------------------------------------------------- 1 | """ Helper utility functions. """ 2 | import os 3 | import numpy as np 4 | 5 | 6 | def create_dir(directory): 7 | if not os.path.exists(directory): 8 | os.mkdir(directory) 9 | return directory 10 | 11 | 12 | def log_string(log_file, log_str): 13 | log_file.write(log_str + '\n') 14 | log_file.flush() 15 | print(log_str) 16 | 17 | 18 | def iterate_in_chunks(l, n): 19 | """ Yield successive 'n'-sized chunks from iterable 'l'. 20 | Note: last chunk will be smaller than l if n doesn't divide l perfectly. 21 | """ 22 | for i in range(0, len(l), n): 23 | yield l[i: i + n] 24 | 25 | 26 | def load_data(load_path): 27 | pc_res = np.load(load_path) 28 | 29 | pc1, pc2, gt_flow, est_flow = pc_res['pc1'], pc_res['pc2'], pc_res['gt_flow_for_pc1'], pc_res['est_flow_for_pc1'] 30 | try: 31 | gt_mask = pc_res['gt_mask_for_pc1'] 32 | gt_mask = gt_mask.reshape(len(pc1)) 33 | except KeyError: 34 | gt_mask = np.ones(len(pc1)) == 1 35 | 36 | try: 37 | corr_conf = pc_res['corr_conf_for_pc1'] 38 | except KeyError: 39 | corr_conf = np.ones(len(pc1), dtype=np.float32) 40 | 41 | pc1_warped_gt = pc1 + gt_flow 42 | pc1_warped_est = pc1 + est_flow 43 | 44 | return pc1, pc2, gt_mask, gt_flow, est_flow, corr_conf, pc1_warped_gt, pc1_warped_est 45 | 46 | -------------------------------------------------------------------------------- /tools/vis_utils.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | from mayavi import mlab 4 | from PIL import Image 5 | 6 | 7 | def get_view_params(mlab_fig): 8 | azimuth, elevation, distance, focalpoint = mlab_fig.view() 9 | roll = mlab_fig.roll() 10 | 11 | view_dict = {"azimuth": azimuth, "elevation": elevation, "distance": distance, "focalpoint": focalpoint, "roll": roll} 12 | return view_dict 13 | 14 | 15 | def plot_pc_list(pc_list, color_list, view_dict=None, get_view=False, fig=None, size=(1600, 800)): 16 | if fig is None: 17 | fig = mlab.figure(figure=None, bgcolor=(1, 1, 1), fgcolor=None, engine=None, size=size) 18 | 19 | for pc, color in zip(pc_list, color_list): 20 | if pc is not None: 21 | mlab.points3d(pc[:, 0], pc[:, 1], pc[:, 2], scale_factor=0.05, color=color, figure=fig) 22 | 23 | if view_dict is not None: 24 | mlab.view(**view_dict) 25 | 26 | if get_view: 27 | view_dict = get_view_params(mlab) 28 | print(view_dict) 29 | 30 | return fig 31 | 32 | 33 | def plot_pcs(pc1, pc2, mask=None, flow=None, color_dict=None, view_dict=None, get_view=False, fig=None, size=(1600, 800)): 34 | if color_dict is None: 35 | color_dict = {"pc1_c": (1, 0, 0), "pc2_c": (0, 1, 0), "flow_c": (0, 0, 1)} 36 | 37 | if fig is None: 38 | fig = mlab.figure(figure=None, bgcolor=(1, 1, 1), fgcolor=None, engine=None, size=size) 39 | 40 | num_points = len(pc1) 41 | if len(color_dict["pc1_c"]) == num_points: # color per-point: 42 | p3d = mlab.points3d(pc1[:, 0], pc1[:, 1], pc1[:, 2], scale_factor=0.05, figure=fig) 43 | p3d.glyph.scale_mode = 'scale_by_vector' 44 | p3d.mlab_source.dataset.point_data.scalars = color_dict["pc1_c"] 45 | else: 46 | mlab.points3d(pc1[:, 0], pc1[:, 1], pc1[:, 2], scale_factor=0.05, color=color_dict["pc1_c"], figure=fig) 47 | 48 | if flow is not None: 49 | if mask is not None: 50 | pc1_plot = pc1[mask] 51 | flow_plot = flow[mask] 52 | else: 53 | pc1_plot = pc1 54 | flow_plot = flow 55 | mlab.quiver3d(pc1_plot[:, 0], pc1_plot[:, 1], pc1_plot[:, 2], flow_plot[:, 0], flow_plot[:, 1], flow_plot[:, 2], scale_factor=1, color=color_dict["flow_c"], line_width=2.0, figure=fig) 56 | if pc2 is not None: 57 | mlab.points3d(pc2[:, 0], pc2[:, 1], pc2[:, 2], scale_factor=0.05, color=color_dict["pc2_c"], figure=fig) 58 | 59 | if view_dict is not None: 60 | mlab.view(**view_dict) 61 | 62 | if get_view: 63 | view_dict = get_view_params(mlab) 64 | print(view_dict) 65 | 66 | return fig 67 | 68 | 69 | def save_pc_plot(save_path, save=True, show=False): 70 | try: 71 | screenshot = mlab.screenshot() 72 | except: 73 | screenshot = mlab.screenshot() 74 | img = Image.fromarray(screenshot) 75 | 76 | if save: 77 | img.save(save_path) 78 | 79 | if show: 80 | mlab.show() 81 | else: 82 | mlab.close() 83 | --------------------------------------------------------------------------------