PlaneRecNet
This is an official implementation for PlaneRecNet: A multi-task convolutional neural network provides instance segmentation for piece-wise planes and monocular depth estimation, and focus on the cross-task consistency between two branches. 
Changing Logs
22th. Oct. 2021: Initial update, some trained models and data annotation will be uploaded very soon.
29th. Oct. 2021: Upload ResNet-50 based model.
3rd. Nov. 2021: Nice to know that "prn" or "PRN" is a forbiden name in Windows.
4th. Nov. 2021: For inference, input image will be resized to max(H, W) == cfg.max_size, and reserve the aspect ratio. Update enviroment.yml, so that newest GPU can run it as well.
Installation
Install environment:
- Clone this repository and enter it:
git clone https://github.com/EryiXie/PlaneRecNet.git
cd PlaneRecNet
- Set up the environment using one of the following methods:
Download trained model:
Here are our models (released on Oct 22th, 2021), which can reproduce the results in the paper:
All models below are trained with batch_size=8 and a single RTX3090 or a single RTXA6000 on the plane annotation for ScanNet dataset:
| Image Size | Backbone | FPS | Weights |
|---|---|---|---|
| 480x640 | Resnet50-DCN | 19.1 | PlaneRecNet_50 |
| 480x640 | Resnet101-DCN | 14.4 | PlaneRecNet_101 |
Simple Inference
Inference with an single image(*.jpg or *.png format):
python3 simple_inference.py --config=PlaneRecNet_101_config --trained_model=weights/PlaneRecNet_101_9_125000.pth --image=data/example_nyu.jpg
Inference with images in a folder:
python3 simple_inference.py --config=PlaneRecNet_101_config --trained_model=weights/PlaneRecNet_101_9_125000.pth --images=input_folder:output_folder
Inference with .mat files from iBims-1 Dataset:
python3 simple_inference.py --config=PlaneRecNet_101_config --trained_model=weights/PlaneRecNet_101_9_125000.pth --ibims1=input_folder:output_folder
Then you will get segmentation and depth estimation results like these:
Training
PlaneRecNet is trained on ScanNet with 100k samples on one single RTX 3090 with batch_size=8, it takes approximate 37 hours. Here are the data annotations(about 1.0 GB) for training of ScanNet datasets, which is based on the annotation given by PlaneRCNN and converted into json file. Please not that, our training sample is not same as PlaneRCNN, because we don't have their training split at hand.
Please notice, the pathing and naming rules in our data/dataset.py, is not compatable with the raw data extracted with the ScanNetv2 original code. Please refer to this issue for fixing tips, thanks uyoung-jeong for that. I will add the data preprocessing script to fix this, once I have time.
Of course, please download ScanNet too for rgb image, depth image and camera intrinsic etc.. The annotation file we provide only contains paths for images and camera intrinsic and the ground truth of piece-wise plane instance and its plane parameters.
- To train, grab an imagenet-pretrained model and put it in
./weights. - Run one of the training commands below.
- Press ctrl+c while training and it will save an
*_interrupt.pthfile at the current iteration. - All weights are saved in the
./weightsdirectory by default with the file name<config>_<epoch>_<iter>.pth.
- Press ctrl+c while training and it will save an
Trains PlaneRecNet_101_config with a batch_size of 8.
python3 train.py --config=PlaneRecNet_101_config --batch_size=8
Trains PlaneRecNet, without writing any logs to tensorboard.
python3 train.py --config=PlaneRecNet_101_config --batch_size=8 --no_tensorboard
Run Tensorboard on local dir "./logs" to check the visualization. So far we provide loss recording and image sample visualization, may consider to add more (22.Oct.2021).
tenosrborad --logdir /log/folder/
Resume training PlaneRecNet with a specific weight file and start from the iteration specified in the weight file's name.
python3 train.py --config=PlaneRecNet_101_config --resume=weights/PlaneRecNet_101_X_XXXXX.pth
Use the help option to see a description of all available command line arguments.
python3 train.py --help
Multi-GPU Support
We adapted the Multi-GPU support from YOLACT, as well as the introduction of how to use it as follow:
- Put
CUDA_VISIBLE_DEVICES=[gpus]on the beginning of the training command.- Where you should replace [gpus] with a comma separated list of the index of each GPU you want to use (e.g., 0,1,2,3).
- You should still do this if only using 1 GPU.
- You can check the indices of your GPUs with
nvidia-smi.
- Then, simply set the batch size to
8*num_gpuswith the training commands above. The training script will automatically scale the hyperparameters to the right values.- If you have memory to spare you can increase the batch size further, but keep it a multiple of the number of GPUs you're using.
- If you want to allocate the images per GPU specific for different GPUs, you can use
--batch_alloc=[alloc]where [alloc] is a comma seprated list containing the number of images on each GPU. This must sum tobatch_size.
Known Issues
- Userwarning of torch.max_pool2d. This has no real affect. It appears when using PyTorch 1.9. And it is claimed "fixed" for the nightly version of PyTorch.
UserWarning: Named tensors and all their associated APIs are an experimental feature and subject to change. Please do not use them for anything important until they are released as stable. (Triggered internally at /pytorch/c10/core/TensorImpl.h:1156.)
return torch.max_pool2d(input, kernel_size, stride, padding, dilation, ceil_mode)
- Userwarning of leaking Caffe2 while training. This issues related to dataloader in PyTorch1.9, to avoid showing this warning, set pin_memory=False for dataloader. But you don't necessarily need to do this.
[W pthreadpool-cpp.cc:90] Warning: Leaking Caffe2 thread-pool after fork. (function pthreadpool)
Citation
If you use PlaneRecNet or this code base in your work, please cite
@misc{xie2021planerecnet,
title={PlaneRecNet: Multi-Task Learning with Cross-Task Consistency for Piece-Wise Plane Detection and Reconstruction from a Single RGB Image},
author={Yaxu Xie and Fangwen Shu and Jason Rambach and Alain Pagani and Didier Stricker},
year={2021},
eprint={2110.11219},
archivePrefix={arXiv},
primaryClass={cs.CV}
}
Contact
For questions about our paper or code, please contact Yaxu Xie, or take a good use at the Issues section of this repository.
49 Nov 23, 2022
76 Nov 23, 2022
3.2k Dec 30, 2022
272 Dec 28, 2022
443 Dec 6, 2022
201 Dec 21, 2022
364 Dec 14, 2022
111 Dec 31, 2022
28 Nov 16, 2022
35 Dec 6, 2022
127 Dec 28, 2022
87 Dec 30, 2022
59 Dec 17, 2022
121 Dec 25, 2022
11.3k Dec 30, 2022
46 Nov 9, 2022
34 Jan 19, 2022
67 Dec 20, 2022