diff --git a/README.md b/README.md
index b491fb381a6c0a59148ef671e80c5567fcc5a2ef..a365c0d5b13adc1b64d60574127b45b860aca844 100644
--- a/README.md
+++ b/README.md
@@ -36,14 +36,6 @@ This repository is the implementation code of the paper "DenseFusion: 6D Object
* CUDA 7.5/8.0/9.0 (Required. CPU-only will lead to extreme slow training speed because of the loss calculation of the symmetry objects (pixel-wise nearest neighbour loss).)
## Code Structure
-* **lib**
- * **lib/loss.py**: Loss calculation for DenseFusion model.
- * **lib/loss_refiner.py**: Loss calculation for iterative refinement model.
- * **lib/transformations.py**: [Transformation Function Library](https://www.lfd.uci.edu/~gohlke/code/transformations.py.html).
- * **lib/network.py**: Network architecture.
- * **lib/extractors.py**: Encoder network architecture.
- * **lib/pspnet.py**: Decoder network architecture.
- * **lib/utils.py**: Logger code.
* **datasets**
* **datasets/ycb**
* **datasets/ycb/dataset.py**: Data loader for YCB_Video dataset.
@@ -55,33 +47,50 @@ This repository is the implementation code of the paper "DenseFusion: 6D Object
* **datasets/linemod/dataset.py**: Data loader for LineMOD dataset.
* **datasets/linemod/dataset_config**:
* **datasets/linemod/dataset_config/models_info.yml**: Object model info of LineMOD dataset.
-* **knn**: CUDA K-nearest neighbours library adapted from [pytorch_knn_cuda](https://github.com/chrischoy/pytorch_knn_cuda).
+* **replace_ycb_toolbox**: Replacement codes for the evaluation with [YCB_Video_toolbox](https://github.com/yuxng/YCB_Video_toolbox).
* **trained_models**
* **trained_models/ycb**: Checkpoints of YCB_Video dataset.
* **trained_models/linemod**: Checkpoints of LineMOD dataset.
-* **eval_result**
- * **eval_result/ycb**
- * **eval_result/ycb/Densefusion_wo_refine_result**: Evaluation result on YCB_Video dataset without refinement.
- * **eval_result/ycb/Densefusion_iterative_result**: Evaluation result on YCB_Video dataset with iterative refinement.
- * **eval_result/linemod**: Evaluation results on LineMOD dataset with iterative refinement.
-* **replace_ycb_toolbox**: Replacement codes for the evaluation with [YCB_Video_toolbox](https://github.com/yuxng/YCB_Video_toolbox).
-* **eval_ycb.py**: Evaluation script for YCB_Video dataset.
-* **eval_linemod.py**: Evaluation script for LineMOD dataset.
-* **train.py**: Training script for YCB_Video dataset and LineMOD dataset.
+* **lib**
+ * **lib/loss.py**: Loss calculation for DenseFusion model.
+ * **lib/loss_refiner.py**: Loss calculation for iterative refinement model.
+ * **lib/transformations.py**: [Transformation Function Library](https://www.lfd.uci.edu/~gohlke/code/transformations.py.html).
+ * **lib/network.py**: Network architecture.
+ * **lib/extractors.py**: Encoder network architecture adapted from [pspnet-pytorch](https://github.com/Lextal/pspnet-pytorch)
+ * **lib/pspnet.py**: Decoder network architecture.
+ * **lib/utils.py**: Logger code.
+ * **lib/knn/**: CUDA K-nearest neighbours library adapted from [pytorch_knn_cuda](https://github.com/chrischoy/pytorch_knn_cuda).
+* **tools**
+ * **tools/_init_paths.py**: Add local path.
+ * **tools/eval_ycb.py**: Evaluation code for YCB_Video dataset.
+ * **tools/eval_linemod.py**: Evaluation code for LineMOD dataset.
+ * **tools/train.py**: Training code for YCB_Video dataset and LineMOD dataset.
+* **experiments**
+ * **experiments/eval_result**
+ * **experiments/eval_result/ycb**
+ * **experiments/eval_result/ycb/Densefusion_wo_refine_result**: Evaluation result on YCB_Video dataset without refinement.
+ * **experiments/eval_result/ycb/Densefusion_iterative_result**: Evaluation result on YCB_Video dataset with iterative refinement.
+ * **experiments/eval_result/linemod**: Evaluation results on LineMOD dataset with iterative refinement.
+ * **experiments/logs/**: Training log files.
+ * **experiments/scripts**
+ * **experiments/scripts/train_ycb.sh**: Training script on the YCB_Video dataset.
+ * **experiments/scripts/train_linemod.sh**: Training script on the LineMOD dataset.
+ * **experiments/scripts/eval_ycb.sh**: Evaluation script on the YCB_Video dataset.
+ * **experiments/scripts/eval_linemod.sh**: Evaluation script on the LineMOD dataset.
+* **download.sh**: Script for downloading YCB_Video Dataset, preprocessed LineMOD dataset and the trained checkpoints.
+
## Datasets
This work is tested on two 6D object pose estimation datasets:
-* [YCB_Video Dataset](https://rse-lab.cs.washington.edu/projects/posecnn/): Download from google drive:
-```
-wget --load-cookies /tmp/cookies.txt "https://docs.google.com/uc?export=download&confirm=$(wget --quiet --save-cookies /tmp/cookies.txt --keep-session-cookies --no-check-certificate 'https://docs.google.com/uc?export=download&id=1if4VoEXNx9W3XCn0Y7Fp15B4GpcYbyYi' -O- | sed -rn 's/.*confirm=([0-9A-Za-z_]+).*/\1\n/p')&id=1if4VoEXNx9W3XCn0Y7Fp15B4GpcYbyYi" -O YCB_Video_Dataset.zip && rm -rf /tmp/cookies.txt
-```
-Training and Testing sets follow [PoseCNN](https://arxiv.org/abs/1711.00199). The training set includes 80 training videos 0000-0047 & 0060-0091 (choosen by 7 frame as a gap in our training) and synthetic data 000000-079999. The testing set includes 2949 keyframes from 10 testing videos 0048-0059.
+* [YCB_Video Dataset](https://rse-lab.cs.washington.edu/projects/posecnn/): Training and Testing sets follow [PoseCNN](https://arxiv.org/abs/1711.00199). The training set includes 80 training videos 0000-0047 & 0060-0091 (choosen by 7 frame as a gap in our training) and synthetic data 000000-079999. The testing set includes 2949 keyframes from 10 testing videos 0048-0059.
-* [LineMOD](http://campar.in.tum.de/Main/StefanHinterstoisser): Download the [preprocessed LineMOD dataset](https://drive.google.com/drive/folders/19ivHpaKm9dOrr12fzC8IDFczWRPFxho7) (including the testing results outputted by the trained vanilla SegNet used for evaluation):
-```
-wget --load-cookies /tmp/cookies.txt "https://docs.google.com/uc?export=download&confirm=$(wget --quiet --save-cookies /tmp/cookies.txt --keep-session-cookies --no-check-certificate 'https://docs.google.com/uc?export=download&id=1YFUra533pxS_IHsb9tB87lLoxbcHYXt8' -O- | sed -rn 's/.*confirm=([0-9A-Za-z_]+).*/\1\n/p')&id=1YFUra533pxS_IHsb9tB87lLoxbcHYXt8" -O Linemod_preprocessed.zip && rm -rf /tmp/cookies.txt
+* [LineMOD](http://campar.in.tum.de/Main/StefanHinterstoisser): Download the [preprocessed LineMOD dataset](https://drive.google.com/drive/folders/19ivHpaKm9dOrr12fzC8IDFczWRPFxho7) (including the testing results outputted by the trained vanilla SegNet used for evaluation).
+
+Download YCB_Video Dataset, preprocessed LineMOD dataset and the trained checkpoints:
+```
+./download.sh
```
## Training
@@ -89,12 +98,12 @@ wget --load-cookies /tmp/cookies.txt "https://docs.google.com/uc?export=download
* YCB_Video Dataset:
After you have downloaded and unzipped the YCB_Video_Dataset.zip and installed all the dependency packages, please run:
```
-python train.py --dataset=ycb --dataset_root=PATH_TO_YCB_DATASET/
+./experiments/scripts/train_ycb.sh
```
* LineMOD Dataset:
After you have downloaded and unzipped the Linemod_preprocessed.zip, please run:
```
-python train.py --dataset=linemod --dataset_root=PATH_TO_LINEMOD_DATASET/
+./experiments/scripts/train_linemod.sh
```
**Training Process**: The training process contains two components: (i) Training of the DenseFusion model. (ii) Training of the Iterative Refinement model. In this code, a DenseFusion model will be trained first. When the average testing distance result (ADD for non-symmetry objects, ADD-S for symmetry objects) is smaller than a certain margin, the training of the Iterative Refinement model will start automatically and the DenseFusion model will then be fixed. You can change this margin to have better DenseFusion result without refinement but it's inferior than the final result after the iterative refinement.
@@ -107,7 +116,7 @@ python train.py --dataset=linemod --dataset_root=PATH_TO_LINEMOD_DATASET/
Just run:
```
cd vanilla_segmentation/
-python train.py --dataset_root=PATH_TO_YCB_DATASET/
+python train.py --dataset_root=./datasets/ycb/YCB_Video_Dataset
```
To make the best use of the training set, several data augementation techniques are used in this code:
@@ -120,29 +129,22 @@ To make the best use of the training set, several data augementation techniques
## Evaluation
### Evaluation on YCB_Video Dataset
-For fair comparsion, we use the same segmentation results of [PoseCNN](https://rse-lab.cs.washington.edu/projects/posecnn/) and compare with their results after ICP refinement. Please first download the `YCB_Video_toolbox` to the root folder of this repo:
-```
-cd Densefusion
-git clone https://github.com/yuxng/YCB_Video_toolbox.git
-cd YCB_Video_toolbox
-unzip results_PoseCNN_RSS2018.zip
-cd ..
-```
-Then please run:
+For fair comparsion, we use the same segmentation results of [PoseCNN](https://rse-lab.cs.washington.edu/projects/posecnn/) and compare with their results after ICP refinement.
+Please run:
```
-python eval_ycb.py --dataset_root=PATH_TO_YCB_DATASET/ --model=NAME_OF_MODEL -- refine_model=NAME_OF_REFINE_MODEL
+./experiments/scripts/eval_ycb.sh
```
-This script will test the selected DenseFusion and Iterative Refinement models on the 2949 keyframes of the 10 testing video in YCB_Video Dataset with the same segmentation result of PoseCNN. The result without refinement is stored in `eval_result/ycb/Densefusion_wo_refine_result` and the refined result is in `eval_result/ycb/Densefusion_iterative_result`.
+This script will first download the `YCB_Video_toolbox` to the root folder of this repo and test the selected DenseFusion and Iterative Refinement models on the 2949 keyframes of the 10 testing video in YCB_Video Dataset with the same segmentation result of PoseCNN. The result without refinement is stored in `eval_result/ycb/Densefusion_wo_refine_result` and the refined result is in `eval_result/ycb/Densefusion_iterative_result`.
-After that, you can add the path of `eval_result/ycb/Densefusion_wo_refine_result/` and `eval_result/ycb/Densefusion_iterative_result/` to the code `YCB_Video_toolbox/evaluate_poses_keyframe.m` and run it with [MATLAB](https://www.mathworks.com/products/matlab.html). The code `YCB_Video_toolbox/plot_accuracy_keyframe.m` can show you the comparsion plot result. You can easily make it by copying the adapted codes from the `replace_ycb_toolbox/` folder and replace them in the `YCB_Video_toolbox/` folder. But you might still need to change the path of your `YCB_Video Dataset/` in the `globals.m` and copy two result folders(`Densefusion_wo_refine_result/` and `Densefusion_iterative_result/`) to the `YCB_Video_toolbox/` folder.
+After that, you can add the path of `experiments/eval_result/ycb/Densefusion_wo_refine_result/` and `experiments/eval_result/ycb/Densefusion_iterative_result/` to the code `YCB_Video_toolbox/evaluate_poses_keyframe.m` and run it with [MATLAB](https://www.mathworks.com/products/matlab.html). The code `YCB_Video_toolbox/plot_accuracy_keyframe.m` can show you the comparsion plot result. You can easily make it by copying the adapted codes from the `replace_ycb_toolbox/` folder and replace them in the `YCB_Video_toolbox/` folder. But you might still need to change the path of your `YCB_Video Dataset/` in the `globals.m` and copy two result folders(`Densefusion_wo_refine_result/` and `Densefusion_iterative_result/`) to the `YCB_Video_toolbox/` folder.
### Evaluation on LineMOD Dataset
Just run:
```
-python eval_linemod.py --dataset_root=PATH_TO_YCB_DATASET/ --model=NAME_OF_MODEL --refine_model=NAME_OF_REFINE_MODEL
+./experiments/scripts/eval_linemod.sh
```
-This script will test the models on the testing set of the LineMOD dataset with the masks outputted by the trained vanilla SegNet model. The result will be printed at the end of the execution and saved as a log in `eval_result/linemod/`.
+This script will test the models on the testing set of the LineMOD dataset with the masks outputted by the trained vanilla SegNet model. The result will be printed at the end of the execution and saved as a log in `experiments/eval_result/linemod/`.
## Results
diff --git a/download.sh b/download.sh
new file mode 100644
index 0000000000000000000000000000000000000000..5dec465d5c7207965a44c0d822fa975b0c715f30
--- /dev/null
+++ b/download.sh
@@ -0,0 +1,27 @@
+# Download the datasets and checkpoints
+
+echo 'Downloading the YCB-Video Dataset'
+
+wget --load-cookies /tmp/cookies.txt "https://docs.google.com/uc?export=download&confirm=$(wget --quiet --save-cookies /tmp/cookies.txt --keep-session-cookies --no-check-certificate 'https://docs.google.com/uc?export=download&id=1if4VoEXNx9W3XCn0Y7Fp15B4GpcYbyYi' -O- | sed -rn 's/.*confirm=([0-9A-Za-z_]+).*/\1\n/p')&id=1if4VoEXNx9W3XCn0Y7Fp15B4GpcYbyYi" -O YCB_Video_Dataset.zip && rm -rf /tmp/cookies.txt
+unzip YCB_Video_Dataset.zip
+mv YCB_Video_Dataset/ datasets/ycb/
+rm YCB_Video_Dataset.zip
+
+echo 'Downloading the preprocessed LineMOD dataset'
+
+wget --load-cookies /tmp/cookies.txt "https://docs.google.com/uc?export=download&confirm=$(wget --quiet --save-cookies /tmp/cookies.txt --keep-session-cookies --no-check-certificate 'https://docs.google.com/uc?export=download&id=1YFUra533pxS_IHsb9tB87lLoxbcHYXt8' -O- | sed -rn 's/.*confirm=([0-9A-Za-z_]+).*/\1\n/p')&id=1YFUra533pxS_IHsb9tB87lLoxbcHYXt8" -O Linemod_preprocessed.zip && rm -rf /tmp/cookies.txt
+unzip Linemod_preprocessed.zip
+mv Linemod_preprocessed/ datasets/linemod/
+rm Linemod_preprocessed.zip
+
+echo 'Downloading the trained checkpoints...'
+
+wget --load-cookies /tmp/cookies.txt "https://docs.google.com/uc?export=download&confirm=$(wget --quiet --save-cookies /tmp/cookies.txt --keep-session-cookies --no-check-certificate 'https://docs.google.com/uc?export=download&id=1bQ9H-fyZplQoNt1qRwdIUX5_3_1pj6US' -O- | sed -rn 's/.*confirm=([0-9A-Za-z_]+).*/\1\n/p')&id=1bQ9H-fyZplQoNt1qRwdIUX5_3_1pj6US" -O trained_checkpoints.zip && rm -rf /tmp/cookies.txt
+unzip trained_checkpoints.zip
+echo 'Copying trained checkpoints to trained_models/'
+cp trained_checkpoints/ycb/*.pth trained_models/ycb/
+cp trained_checkpoints/linemod/*.pth trained_models/linemod/
+rm trained_checkpoints.zip
+rm -r trained_checkpoints/
+
+echo 'done'
\ No newline at end of file
diff --git a/eval_result/linemod/.gitignore b/experiments/eval_result/linemod/.gitignore
similarity index 100%
rename from eval_result/linemod/.gitignore
rename to experiments/eval_result/linemod/.gitignore
diff --git a/eval_result/ycb/Densefusion_iterative_result/.gitignore b/experiments/eval_result/ycb/Densefusion_iterative_result/.gitignore
similarity index 100%
rename from eval_result/ycb/Densefusion_iterative_result/.gitignore
rename to experiments/eval_result/ycb/Densefusion_iterative_result/.gitignore
diff --git a/eval_result/ycb/Densefusion_wo_refine_result/.gitignore b/experiments/eval_result/ycb/Densefusion_wo_refine_result/.gitignore
similarity index 100%
rename from eval_result/ycb/Densefusion_wo_refine_result/.gitignore
rename to experiments/eval_result/ycb/Densefusion_wo_refine_result/.gitignore
diff --git a/logs/linemod/.gitignore b/experiments/logs/linemod/.gitignore
similarity index 100%
rename from logs/linemod/.gitignore
rename to experiments/logs/linemod/.gitignore
diff --git a/logs/ycb/.gitignore b/experiments/logs/ycb/.gitignore
similarity index 100%
rename from logs/ycb/.gitignore
rename to experiments/logs/ycb/.gitignore
diff --git a/experiments/scripts/eval_linemod.sh b/experiments/scripts/eval_linemod.sh
new file mode 100644
index 0000000000000000000000000000000000000000..fc806a1242507ec6f79a8b7cab897567954edf62
--- /dev/null
+++ b/experiments/scripts/eval_linemod.sh
@@ -0,0 +1,11 @@
+#!/bin/bash
+
+set -x
+set -e
+
+export PYTHONUNBUFFERED="True"
+export CUDA_VISIBLE_DEVICES=0
+
+python ./tools/eval_linemod.py --dataset_root ./datasets/linemod/Linemod_preprocessed\
+ --model pose_model_9_0.01310166542980859.pth\
+ --refine_model pose_refine_model_493_0.006761023565178073.pth
\ No newline at end of file
diff --git a/experiments/scripts/eval_ycb.sh b/experiments/scripts/eval_ycb.sh
new file mode 100644
index 0000000000000000000000000000000000000000..a4c12a010c05df816729bd655833258317feebcb
--- /dev/null
+++ b/experiments/scripts/eval_ycb.sh
@@ -0,0 +1,20 @@
+#!/bin/bash
+
+set -x
+set -e
+
+export PYTHONUNBUFFERED="True"
+export CUDA_VISIBLE_DEVICES=0
+
+if [ ! -d YCB_Video_toolbox ];then
+ echo 'Downloading the YCB_Video_toolbox...'
+ git clone https://github.com/yuxng/YCB_Video_toolbox.git
+ cd YCB_Video_toolbox
+ unzip results_PoseCNN_RSS2018.zip
+ cd ..
+ cp replace_ycb_toolbox/*.m YCB_Video_toolbox/
+fi
+
+python ./tools/eval_ycb.py --dataset_root ./datasets/ycb/YCB_Video_Dataset\
+ --model pose_model_26_0.012863246640872631.pth\
+ --refine_model pose_refine_model_69_0.009449292959118935.pth
\ No newline at end of file
diff --git a/experiments/scripts/train_linemod.sh b/experiments/scripts/train_linemod.sh
new file mode 100644
index 0000000000000000000000000000000000000000..5697159db5bca594a44543c99672469f9b4f6465
--- /dev/null
+++ b/experiments/scripts/train_linemod.sh
@@ -0,0 +1,10 @@
+#!/bin/bash
+
+set -x
+set -e
+
+export PYTHONUNBUFFERED="True"
+export CUDA_VISIBLE_DEVICES=0
+
+python ./tools/train.py --dataset linemod\
+ --dataset_root ./datasets/linemod/Linemod_preprocessed
\ No newline at end of file
diff --git a/experiments/scripts/train_ycb.sh b/experiments/scripts/train_ycb.sh
new file mode 100644
index 0000000000000000000000000000000000000000..449da49e6b10bd4ad19cedacdf94ba8f361ba98e
--- /dev/null
+++ b/experiments/scripts/train_ycb.sh
@@ -0,0 +1,10 @@
+#!/bin/bash
+
+set -x
+set -e
+
+export PYTHONUNBUFFERED="True"
+export CUDA_VISIBLE_DEVICES=0
+
+python ./tools/train.py --dataset ycb\
+ --dataset_root ./datasets/ycb/YCB_Video_Dataset
\ No newline at end of file
diff --git a/knn/Makefile b/lib/knn/Makefile
similarity index 100%
rename from knn/Makefile
rename to lib/knn/Makefile
diff --git a/knn/__init__.py b/lib/knn/__init__.py
similarity index 95%
rename from knn/__init__.py
rename to lib/knn/__init__.py
index afef478b331748da85ddfc775c9f1a7ad03eff3a..ed16fe3b8f0b58e1a1ec49cb55d21ad7c7b3a53a 100755
--- a/knn/__init__.py
+++ b/lib/knn/__init__.py
@@ -4,7 +4,7 @@ import operator as op
import functools
import torch
from torch.autograd import Variable, Function
-import knn.knn_pytorch as knn_pytorch
+from lib.knn import knn_pytorch as knn_pytorch
class KNearestNeighbor(Function):
""" Compute k nearest neighbors for each query point.
diff --git a/knn/build/knn_cuda_kernel.so b/lib/knn/build/knn_cuda_kernel.so
similarity index 100%
rename from knn/build/knn_cuda_kernel.so
rename to lib/knn/build/knn_cuda_kernel.so
diff --git a/knn/build_ffi.py b/lib/knn/build_ffi.py
similarity index 100%
rename from knn/build_ffi.py
rename to lib/knn/build_ffi.py
diff --git a/knn/knn_pytorch/__init__.py b/lib/knn/knn_pytorch/__init__.py
similarity index 100%
rename from knn/knn_pytorch/__init__.py
rename to lib/knn/knn_pytorch/__init__.py
diff --git a/knn/knn_pytorch/__pycache__/__init__.cpython-35.pyc b/lib/knn/knn_pytorch/__pycache__/__init__.cpython-35.pyc
similarity index 100%
rename from knn/knn_pytorch/__pycache__/__init__.cpython-35.pyc
rename to lib/knn/knn_pytorch/__pycache__/__init__.cpython-35.pyc
diff --git a/knn/knn_pytorch/__pycache__/__init__.cpython-36.pyc b/lib/knn/knn_pytorch/__pycache__/__init__.cpython-36.pyc
similarity index 100%
rename from knn/knn_pytorch/__pycache__/__init__.cpython-36.pyc
rename to lib/knn/knn_pytorch/__pycache__/__init__.cpython-36.pyc
diff --git a/knn/knn_pytorch/_knn_pytorch.so b/lib/knn/knn_pytorch/_knn_pytorch.so
similarity index 100%
rename from knn/knn_pytorch/_knn_pytorch.so
rename to lib/knn/knn_pytorch/_knn_pytorch.so
diff --git a/knn/src/knn_cuda_kernel.cu b/lib/knn/src/knn_cuda_kernel.cu
similarity index 100%
rename from knn/src/knn_cuda_kernel.cu
rename to lib/knn/src/knn_cuda_kernel.cu
diff --git a/knn/src/knn_cuda_kernel.h b/lib/knn/src/knn_cuda_kernel.h
similarity index 100%
rename from knn/src/knn_cuda_kernel.h
rename to lib/knn/src/knn_cuda_kernel.h
diff --git a/knn/src/knn_pytorch.c b/lib/knn/src/knn_pytorch.c
similarity index 100%
rename from knn/src/knn_pytorch.c
rename to lib/knn/src/knn_pytorch.c
diff --git a/knn/src/knn_pytorch.h b/lib/knn/src/knn_pytorch.h
similarity index 100%
rename from knn/src/knn_pytorch.h
rename to lib/knn/src/knn_pytorch.h
diff --git a/lib/loss.py b/lib/loss.py
index 657bc869f0d116ca8061f6d93b89b079bcb54119..4ad1064741d09f1048169bf162a07ff33790a1dc 100755
--- a/lib/loss.py
+++ b/lib/loss.py
@@ -6,7 +6,7 @@ import numpy as np
import torch.nn as nn
import random
import torch.backends.cudnn as cudnn
-from knn.__init__ import KNearestNeighbor
+from lib.knn.__init__ import KNearestNeighbor
knn = KNearestNeighbor(1)
diff --git a/lib/loss_refiner.py b/lib/loss_refiner.py
index b2a9aae7216899139eb7b6a2d8fe747d9013e0cc..04430dbe0bdfbb49a60b1fd8554301c24203bfa9 100755
--- a/lib/loss_refiner.py
+++ b/lib/loss_refiner.py
@@ -6,7 +6,7 @@ import numpy as np
import torch.nn as nn
import random
import torch.backends.cudnn as cudnn
-from knn.__init__ import KNearestNeighbor
+from lib.knn.__init__ import KNearestNeighbor
knn = KNearestNeighbor(1)
diff --git a/tools/_init_paths.py b/tools/_init_paths.py
new file mode 100644
index 0000000000000000000000000000000000000000..3a37635fbf949ab7e20948a296c776caa7943348
--- /dev/null
+++ b/tools/_init_paths.py
@@ -0,0 +1,3 @@
+import os
+import sys
+sys.path.insert(0, os.getcwd())
\ No newline at end of file
diff --git a/eval_linemod.py b/tools/eval_linemod.py
similarity index 90%
rename from eval_linemod.py
rename to tools/eval_linemod.py
index d3ab7249d486cb85eba7d5ce6cfdfbaafbf770cc..3bd949b45e3e1d13086923825d54a58081707707 100644
--- a/eval_linemod.py
+++ b/tools/eval_linemod.py
@@ -1,3 +1,4 @@
+import _init_paths
import argparse
import os
import random
@@ -19,9 +20,9 @@ from lib.loss import Loss
from lib.loss_refiner import Loss_refine
parser = argparse.ArgumentParser()
-parser.add_argument('--dataset_root', type=str, default = '/home/data1/jeremy/Linemod_preprocessed', help='dataset root dir')
-parser.add_argument('--model', type=str, default = 'pose_model_9_0.01310166542980859.pth', help='resume PoseNet model')
-parser.add_argument('--refine_model', type=str, default = 'pose_refine_model_493_0.006761023565178073.pth', help='resume PoseRefineNet model')
+parser.add_argument('--dataset_root', type=str, default = '', help='dataset root dir')
+parser.add_argument('--model', type=str, default = '', help='resume PoseNet model')
+parser.add_argument('--refine_model', type=str, default = '', help='resume PoseRefineNet model')
opt = parser.parse_args()
num_objects = 13
@@ -30,7 +31,7 @@ num_points = 500
iteration = 2
dataset_config_dir = 'datasets/linemod/dataset_config'
trained_models_dir = 'trained_models/linemod'
-output_result_dir = 'eval_result/linemod'
+output_result_dir = 'experiments/eval_result/linemod'
estimator = PoseNet(num_points = num_points, num_obj = num_objects)
diff --git a/eval_ycb.py b/tools/eval_ycb.py
similarity index 94%
rename from eval_ycb.py
rename to tools/eval_ycb.py
index ac949a3f717db303493c9f3aba824ade4401828e..39833471b5a87809100609ee8cf3124980a8e2f1 100644
--- a/eval_ycb.py
+++ b/tools/eval_ycb.py
@@ -1,3 +1,4 @@
+import _init_paths
import argparse
import os
import copy
@@ -24,9 +25,9 @@ from lib.network import PoseNet, PoseRefineNet
from lib.transformations import euler_matrix, quaternion_matrix, quaternion_from_matrix
parser = argparse.ArgumentParser()
-parser.add_argument('--dataset_root', type=str, default = '/home/data1/jeremy/YCB_Video_Dataset', help='dataset root dir')
-parser.add_argument('--model', type=str, default = 'pose_model_23_0.012863246640872631.pth', help='resume PoseNet model')
-parser.add_argument('--refine_model', type=str, default = 'pose_refine_model_49_0.009449292959118935.pth', help='resume PoseRefineNet model')
+parser.add_argument('--dataset_root', type=str, default = '', help='dataset root dir')
+parser.add_argument('--model', type=str, default = '', help='resume PoseNet model')
+parser.add_argument('--refine_model', type=str, default = '', help='resume PoseRefineNet model')
opt = parser.parse_args()
norm = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
@@ -47,8 +48,8 @@ iteration = 2
bs = 1
dataset_config_dir = 'datasets/ycb/dataset_config'
ycb_toolbox_dir = 'YCB_Video_toolbox'
-result_wo_refine_dir = 'eval_result/ycb/Densefusion_wo_refine_result'
-result_refine_dir = 'eval_result/ycb/Densefusion_iterative_result'
+result_wo_refine_dir = 'experiments/eval_result/ycb/Densefusion_wo_refine_result'
+result_refine_dir = 'experiments/eval_result/ycb/Densefusion_iterative_result'
trained_models_dir = 'trained_models/ycb'
def get_bbox(posecnn_rois):
@@ -235,7 +236,7 @@ for now in range(0, 2949):
my_t = my_t_final
my_result.append(my_pred.tolist())
- except:
+ except ZeroDivisionError:
print("PoseCNN Detector Lost {0} at No.{1} keyframe".format(itemid, now))
my_result_wo_refine.append([0.0 for i in range(7)])
my_result.append([0.0 for i in range(7)])
diff --git a/train.py b/tools/train.py
similarity index 96%
rename from train.py
rename to tools/train.py
index 69d6b7fec2fd1bc80646ea72f086b1cf538d2c87..78b8fc82f728010a47537f6a8c7b811ba33b219e 100755
--- a/train.py
+++ b/tools/train.py
@@ -1,3 +1,10 @@
+# --------------------------------------------------------
+# DenseFusion 6D Object Pose Estimation by Iterative Dense Fusion
+# Licensed under The MIT License [see LICENSE for details]
+# Written by Chen
+# --------------------------------------------------------
+
+import _init_paths
import argparse
import os
import random
@@ -49,15 +56,13 @@ def main():
opt.num_objects = 21 #number of object classes in the dataset
opt.num_points = 1000 #number of points on the input pointcloud
opt.outf = 'trained_models/ycb' #folder to save trained models
- opt.log_dir = 'logs/ycb' #folder to save logs
- opt.dataset_root = '/home/data1/jeremy/YCB_Video_Dataset' #dataset root dir
+ opt.log_dir = 'experiments/logs/ycb' #folder to save logs
opt.repeat_epoch = 1 #number of repeat times for one epoch training
elif opt.dataset == 'linemod':
opt.num_objects = 13
opt.num_points = 500
opt.outf = 'trained_models/linemod'
- opt.log_dir = 'logs/linemod'
- opt.dataset_root = '/home/data1/jeremy/Linemod_preprocessed'
+ opt.log_dir = 'experiments/logs/linemod'
opt.repeat_epoch = 20
else:
print('Unknown dataset')