This repository contains the PyTorch implementation for PT-FlowNet
Feel free to open issues and thank you for your attention!
- CUDA 11.3
- Python 3.8
- PyTorch 1.10
- torch-scatter, h5py, pyyaml, tqdm, tensorboard, scipy, imageio, png
- Install pointops lib:
cd lib/pointops
python3 setup.py install
cd ../..
We follow HPLFlowNet to prepare the datasets:
- FlyingThings3D:
Download and unzip the "Disparity", "Disparity Occlusions", "Disparity change", "Optical flow", "Flow Occlusions" for DispNet/FlowNet2.0 dataset subsets from the FlyingThings3D website (we used the paths from this file, now they added torrent downloads)
. They will be upzipped into the same directory,
RAW_DATA_PATH. Then run the following script for 3D reconstruction:
python3 data_preprocess/process_flyingthings3d_subset.py --raw_data_path RAW_DATA_PATH --save_path SAVE_PATH/FlyingThings3D_subset_processed_35m --only_save_near_pts- KITTI Scene Flow 2015
Download and unzip KITTI Scene Flow Evaluation 2015 to directory
RAW_DATA_PATH. Run the following script for 3D reconstruction:
python data_preprocess/process_kitti.py --raw_data_path=RAW_DATA_PATH --save_path=SAVE_PATH/KITTI_processed_occ_final --calib_path=util/calib_cam_to_camsh train.shft3d_dataset_dir and kitti_dataset_dir are the preprocessed dataset paths, please specify the storage location of the dataset on disk. Relevant data file paths can also be predefined in ./tools/parser.py.
exp_path is the experiment folder name and root is the project root path. The memory requirement for network training is at least 9573 MiB.
sh train_refine.shPlease add the path of the pre-trained model through --weights parameter.
Afer training, you can test the model on both FT3Dand KITTI datasets as follows:
sh test.shsh test_refine.sh--weights is the absolute path of checkpoint file. The memory requirement for inference is 5249 MiB.
Our code is based on PV-RAFT and point-transformer. We also refer to FLOT and HPLFlowNet.