add scripts

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,46 +47,63 @@ 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
 
 * 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

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'

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

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

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

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

knn/__init__.py → lib/knn/__init__.py

@@ -4,7 +4,7 @@
 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.

lib/loss.py

@@ -6,7 +6,7 @@
 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)
 

lib/loss_refiner.py

@@ -6,7 +6,7 @@
 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)
 

tools/_init_paths.py

@@ -0,0 +1,3 @@
+import os
+import sys
+sys.path.insert(0, os.getcwd())

eval_linemod.py → tools/eval_linemod.py

@@ -1,3 +1,4 @@
+import _init_paths
 import argparse
 import os
 import random
@@ -19,9 +20,9 @@
 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 @@
 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)

eval_ycb.py → tools/eval_ycb.py

@@ -1,3 +1,4 @@
+import _init_paths
 import argparse
 import os
 import copy
@@ -24,9 +25,9 @@
 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 @@
 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 @@ def get_bbox(posecnn_rois):
                 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)])