def __init__(self, choose_cate):
self.resume_models = ['model_112_0.184814792342484_bottle.pth',
'model_120_0.10268432162888348_bowl.pth',
'model_118_0.2008235973417759_camera.pth',
'model_107_0.18291547849029302_can.pth',
'model_117_0.12762234719470145_laptop.pth',
'model_102_0.1468337191492319_mug.pth']
self.first_ite = 10
self.temp_gap = 3
self.num_kp = 8
self.num_points = 500
self.num_cates = 6
self.outf = 'models'
self.min_dis = 0.0005
self.model = KeyNet(num_points = self.num_points, num_key = self.num_kp, num_cates = self.num_cates)
self.model.cuda()
self.model.load_state_dict(torch.load('{0}/{1}'.format(self.outf, self.resume_models[choose_cate-1])))
self.model.eval()
self.dataprocess = Dataset(self.num_points)
self.criterion = Loss(self.num_kp, self.num_cates)
self.temp_dir = []
for temp_ite in range(self.temp_gap):
self.temp_dir.append(np.array([0.0, 0.0, 0.0]))
self.Kp_fr = Variable(torch.from_numpy(np.array([0, 0, 0]).astype(np.float32))).cuda()
self.Kp_fr.view(1, 1, 3).repeat(1, self.num_kp, 1)
class tracker():
def __init__(self, choose_cate):
self.resume_models = ['model_112_0.184814792342484_bottle.pth',
'model_120_0.10268432162888348_bowl.pth',
'model_118_0.2008235973417759_camera.pth',
'model_107_0.18291547849029302_can.pth',
'model_117_0.12762234719470145_laptop.pth',
'model_102_0.1468337191492319_mug.pth']
self.first_ite = 10
self.temp_gap = 3
self.num_kp = 8
self.num_points = 500
self.num_cates = 6
self.outf = 'models'
self.min_dis = 0.0005
self.model = KeyNet(num_points = self.num_points, num_key = self.num_kp, num_cates = self.num_cates)
self.model.cuda()
self.model.load_state_dict(torch.load('{0}/{1}'.format(self.outf, self.resume_models[choose_cate-1])))
self.model.eval()
self.dataprocess = Dataset(self.num_points)
self.criterion = Loss(self.num_kp, self.num_cates)
self.temp_dir = []
for temp_ite in range(self.temp_gap):
self.temp_dir.append(np.array([0.0, 0.0, 0.0]))
self.Kp_fr = Variable(torch.from_numpy(np.array([0, 0, 0]).astype(np.float32))).cuda()
self.Kp_fr.view(1, 1, 3).repeat(1, self.num_kp, 1)
def init_estimation(self, rgb_dir, depth_dir, current_r, current_t, bbox):
for temp_ite in range(self.temp_gap):
self.temp_dir.append(np.array([0.0, 0.0, 0.0]))
while len(self.temp_dir) > self.temp_gap:
del self.temp_dir[0]
# print(rgb_dir, depth_dir)
self.dataprocess.add_bbox(bbox)
if self.first_ite != 0:
min_dis = 1000.0
for iterative in range(self.first_ite):
img_fr, choose_fr, cloud_fr, anchor, scale = self.dataprocess.getone(rgb_dir, depth_dir, current_r, current_t)
img_fr, choose_fr, cloud_fr, anchor, scale = Variable(img_fr).cuda(), \
Variable(choose_fr).cuda(), \
Variable(cloud_fr).cuda(), \
Variable(anchor).cuda(), \
Variable(scale).cuda()
Kp_fr, _ = self.model.eval_forward(img_fr, choose_fr, cloud_fr, anchor, scale, 0.0, True)
new_t, kp_dis = self.criterion.inf_zero(Kp_fr[0])
if min_dis > kp_dis:
min_dis = kp_dis
best_current_r = copy.deepcopy(current_r)
best_current_t = copy.deepcopy(current_t)
print(min_dis)
current_t = current_t + np.dot(new_t, current_r.T)
current_r, current_t = best_current_r, best_current_t
img_fr, choose_fr, cloud_fr, anchor, scale = self.dataprocess.getone(rgb_dir, depth_dir, current_r, current_t)
img_fr, choose_fr, cloud_fr, anchor, scale = Variable(img_fr).cuda(), \
Variable(choose_fr).cuda(), \
Variable(cloud_fr).cuda(), \
Variable(anchor).cuda(), \
Variable(scale).cuda()
Kp_fr, _ = self.model.eval_forward(img_fr, choose_fr, cloud_fr, anchor, scale, 0.0, True)
self.Kp_fr = Kp_fr[0]
self.dataprocess.projection(rgb_dir, current_r, current_t)
print("NEXT!!!")
return current_r, current_t
def next_estimation(self, rgb_dir, depth_dir, current_r, current_t):
img_fr, choose_fr, cloud_fr, anchor, scale = self.dataprocess.getone(rgb_dir, depth_dir, current_r, current_t)
img_fr, choose_fr, cloud_fr, anchor, scale = Variable(img_fr).cuda(), \
Variable(choose_fr).cuda(), \
Variable(cloud_fr).cuda(), \
Variable(anchor).cuda(), \
Variable(scale).cuda()
Kp_to, _ = self.model.eval_forward(img_fr, choose_fr, cloud_fr, anchor, scale, 0.0, False)
self.min_dis = 1000.0
lenggth = len(Kp_to)
for idx in range(lenggth):
new_r, new_t, kp_dis = self.criterion.inf(self.Kp_fr, Kp_to[idx])
if self.min_dis > kp_dis:
self.min_dis = kp_dis
best_r = new_r
best_t = new_t
print(self.min_dis)
current_t = current_t + np.dot(best_t, current_r.T)
current_r = np.dot(current_r, best_r)
self.temp_dir.append(copy.deepcopy(best_t / 1000.0))
if len(self.temp_dir) > self.temp_gap:
del self.temp_dir[0]
self.dataprocess.projection(rgb_dir, current_r, current_t)
print("NEXT!!!")
return current_r, current_t
default=5,
help='number of data loading workers')
parser.add_argument('--num_kp', type=int, default=8, help='number of kp')
parser.add_argument('--outf', type=str, default='models/', help='save dir')
parser.add_argument('--lr', default=0.0001, help='learning rate')
opt = parser.parse_args()
# model = KeyNet(num_points=opt.num_points, num_key=opt.num_kp, num_cates=opt.num_cates)
# model.cuda()
# dataset = Dataset(opt.dataset_root, opt.num_points, opt.image_num)
# test_dataset = Dataset('val', opt.dataset_root, False, opt.num_points, opt.num_cates, 1000, opt.category)
# testdataloader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=True, num_workers=opt.workers)
model = KeyNet(num_points=opt.num_points,
num_key=opt.num_kp,
num_cates=opt.num_cates)
model.cuda()
model.eval()
# for direct forward
choose_cate = 1
choose_obj = "x_hammer1"
choose_video = "x_hammer1"
# choose_video = ""
# mode, root, add_noise, num_pt, num_cates count, cate_id
# test_dataset = Dataset('val', opt.dataset_root, False, opt.num_points, choose_cate, 1000)
# for reading from a txt file
# eval_list_file = open('dataset/eval_list/eval_list_{0}.txt'.format(choose_cate), 'r')
# while 1:
type=str,
default='models/',
help='load model dir')
opt = parser.parse_args()
if not os.path.exists('eval_results'):
os.makedirs('eval_results')
if not os.path.exists('eval_results/TEST_{0}'.format(opt.eval_id)):
os.makedirs('eval_results/TEST_{0}'.format(opt.eval_id))
for item in choose_cate_list:
os.makedirs('eval_results/TEST_{0}/temp_{1}'.format(opt.eval_id, item))
for choose_cate in choose_cate_list:
model = KeyNet(num_points=opt.num_points,
num_key=opt.num_kp,
num_cates=opt.num_cates)
model.cuda()
model.eval()
model.load_state_dict(
torch.load('{0}/{1}'.format(opt.outf, resume_models[choose_cate - 1])))
pconf = torch.ones(opt.num_kp) / opt.num_kp
pconf = Variable(pconf).cuda()
test_dataset = Dataset('val', opt.dataset_root, False, opt.num_points,
choose_cate, 1000)
criterion = Loss(opt.num_kp, opt.num_cates)
eval_list_file = open(
parser.add_argument('--num_points', type=int, default=500, help='points')
parser.add_argument('--num_cates',
type=int,
default=6,
help='number of categories')
parser.add_argument('--workers',
type=int,
default=5,
help='number of data loading workers')
parser.add_argument('--num_kp', type=int, default=8, help='number of kp')
parser.add_argument('--outf', type=str, default='models/', help='save dir')
parser.add_argument('--lr', default=0.0001, help='learning rate')
opt = parser.parse_args()
model = KeyNet(num_points=opt.num_points,
num_key=opt.num_kp,
num_cates=opt.num_cates)
model.cuda()
if opt.resume != '':
model.load_state_dict(torch.load('{0}/{1}'.format(opt.outf, opt.resume)))
dataset = Dataset('train', opt.dataset_root, True, opt.num_points,
opt.num_cates, 5000, opt.category)
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=1,
shuffle=True,
num_workers=opt.workers)
test_dataset = Dataset('val', opt.dataset_root, False, opt.num_points,
opt.num_cates, 1000, opt.category)
testdataloader = torch.utils.data.DataLoader(test_dataset,