add scripts

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)
 

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)
 

tools/_init_paths.py

+import os
+import sys
+sys.path.insert(0, os.getcwd())
\ No newline at end of file

eval_linemod.py → tools/eval_linemod.py

+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)

eval_ycb.py → tools/eval_ycb.py

+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)])

train.py → tools/train.py

+# --------------------------------------------------------
+# 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')