Our work mainly based on the paper: Estimating and Exploiting the Aleatoric Uncertainty in Surface Normal Estimation. Based on this, we changed the training pipeline to fit efficientNet V2. We rewrote parts of validation that is not efficient and would potentially cause out of memory error, and added the functionality of resuming training if failed. In addition, we chose NYU v2 large dataset as our training data, for which we wrote a dataloader for .mat input files, and added support for ScanNet.
-
NYUv2 (official): The official train/test split contains 795/654 images. The dataset can be downloaded from this link. Unzip the file
nyu_dataset.zipunder./datasets, so that./datasets/nyu/trainand./datasets/nyu/test/exist. -
NYUv2 (big): Please visit GeoNet to download a larger training set consisting of 30907 images. This is the training set used to train our model.
-
ScanNet: Please visit FrameNet to download ScanNet with ground truth surface normals.
- Check the args in train.py and simply run
python train.py