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
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(
'dataset/eval_list/eval_list_{0}.txt'.format(choose_cate), 'r')
while 1:
input_line = eval_list_file.readline()
if not input_line:
break
if input_line[-1:] == '\n':
input_line = input_line[:-1]
_, choose_obj, choose_video = input_line.split(' ')
try:
current_r, current_t = test_dataset.getfirst(
choose_obj, choose_video)
rad_t = np.array([random.uniform(-0.02, 0.02)
knn = KNearestNeighbor(1)
estimator = poseNet(numPoints, numObjects)
estimator.cuda()
estimator.load_state_dict(torch.load(opt.model))
estimator.eval()
testDataset = PoseDataset('eval', numPoints, False, opt.datasetRoot, 0.0, True)
testDataLoader = torch.utils.data.DataLoader(testDataset,
batch_size=1,
shuffle=False,
num_workers=10)
symList = testDataset.get_sym_list()
numPointsMesh = testDataset.get_num_points_mesh()
poseNetLoss = Loss(numPointsMesh, symList)
diameter = []
meta_file = open('{0}/models_info.yml'.format(datasetConfigDir), 'r')
meta = yaml.load(meta_file)
kn = 0.1 # ADD 参数设置
for obj in objList:
diameter.append(meta[obj]['diameter'] / 1000.0 * kn)
print(diameter)
successCount = [0 for i in range(numObjects)]
numCount = [0 for i in range(numObjects)]
fw = open('{0}/eval_result_logs.txt'.format(output_result_dir), 'w')
for i, data in enumerate(testDataLoader, 0):
if len(data) != 7:
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,
batch_size=1,
shuffle=True,
num_workers=opt.workers)
criterion = Loss(opt.num_kp, opt.num_cates)
best_test = np.Inf
optimizer = optim.Adam(model.parameters(), lr=opt.lr)
for epoch in range(0, 500):
model.train()
train_dis_avg = 0.0
train_count = 0
optimizer.zero_grad()
for i, data in enumerate(dataloader, 0):
img_fr, choose_fr, cloud_fr, r_fr, t_fr, img_to, choose_to, cloud_to, r_to, t_to, mesh, anchor, scale, cate = data
img_fr, choose_fr, cloud_fr, r_fr, t_fr, img_to, choose_to, cloud_to, r_to, t_to, mesh, anchor, scale, cate = Variable(img_fr).cuda(), \
Variable(choose_fr).cuda(), \
def main():
opt.manualSeed = random.randint(1, 10000)
random.seed(opt.manualSeed)
torch.manual_seed(
opt.manualSeed) #torch.manual_seed可以保证在种子值不变的情况下，每次torch.rand产生的结果相同
# torch.cuda.manual_seed(opt.manualSeed) torch.cuda.manual_seed_all(seed)
# 对于可重复的实验，有必要为任何使用随机数生成的进行随机种子设置。注意，cuDNN使用非确定性算法，并且可以使用torch.backends.cudnn.enabled = False来进行禁用。
# 根据数据集配置物体种类，用于预测位姿的点数，每个epoch训练的次数，模型保存的位置，log保存的位置
if opt.dataset == 'linemod':
opt.numObjects = 13
opt.numPoints = 576
opt.repeatEpoch = 20
opt.modelFolder = 'trainedModels/'
opt.logFolder = 'experimentResult/logs/'
else:
print('Unknown dataset')
return
# 定义网络，如果有训练过的模型，可加载； 确定网络优化方法
estimator = poseNet(opt.numPoints, opt.numObjects)
estimator.cuda() # or estimator.to('cuda')
if opt.resumePosenet != '':
estimator.load_state_dict(
torch.load('{0}/{1}'.format(opt.modelFolder, opt.resumePosenet)))
optimizer = optim.Adam(estimator.parameters(), lr=opt.lr)
opt.decay_start = False
opt.refine_start = False
# 加载训练数据和测试数据
if opt.dataset == 'linemod':
dataset = PoseDataset('train', opt.numPoints, opt.addNoise,
opt.datasetRoot, opt.noiseTrans,
opt.refine_start)
test_dataset = PoseDataset('test', opt.numPoints, False,
opt.datasetRoot, 0.0, opt.refine_start)
dataLoader = torch.utils.data.DataLoader(dataset,
batch_size=1,
shuffle=True,
num_workers=opt.workers)
testdataLoader = torch.utils.data.DataLoader(test_dataset,
batch_size=1,
shuffle=False,
num_workers=opt.workers)
opt.symList = dataset.get_sym_list()
opt.numPointsMesh = dataset.get_num_points_mesh()
print(
'>>>>>>--------Dataset Loaded!---------<<<<<<<\nlength of the training set:{0}\nlength of the testing set:{1}\nnumber of sample points on mesh: {2}\nsymmetry object list: {3}'
.format(len(dataset), len(test_dataset), opt.numPointsMesh,
opt.symList))
# 定义网络训练的误差
poseNetLoss = Loss(opt.numPointsMesh, opt.symList)
best_distance = np.Inf
if opt.startEpoch == 1:
for log in os.listdir(opt.logFolder):
os.remove(os.path.join(opt.logFolder, log))
st_time = time.time()
for epoch in range(opt.startEpoch, opt.nepoch):
# ------------------模型训练阶段------------------
logger = setup_logger(
'epoch%d' % epoch,
os.path.join(opt.logFolder, 'epoch_%d_log.txt' % epoch))
logger.info('Train time {0}'.format(
time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)) +
',' + 'Training started'))
train_count = 0
train_dis_sum = 0.0
estimator.train()
optimizer.zero_grad()
for rep in range(opt.repeatEpoch):
for i, data in enumerate(dataLoader, 0):
img, cloud, pointIndex, tarPoints, modelPoints, idx, ori_img = data
img = Variable(img).cuda()
cloud = Variable(cloud).cuda()
pointIndex = Variable(pointIndex).cuda()
tarPoints = Variable(tarPoints).cuda()
modelPoints = Variable(modelPoints).cuda()
idx = Variable(idx).cuda()
ori_img = np.array(ori_img)
pred_r, pred_t, pred_c, colorEmb = estimator(
img, cloud, pointIndex, idx)
loss, dis, newCloud, newTarPoints = poseNetLoss(
pred_r, pred_t, pred_c, tarPoints, modelPoints, idx, cloud,
opt.w, opt.refine_start)
loss.backward()
train_dis_sum += dis.item()
train_count += 1
if train_count % opt.batchSize == 0:
logger.info(
'Train time {0} Epoch {1} Batch{2} Frame{3} Avg_dis:{4}'
.format(
time.strftime("%Hh %Mm %Ss",
time.gmtime(time.time() - st_time)),
epoch, int(train_count / opt.batchSize),
train_count, train_dis_sum / opt.batchSize))
optimizer.step()
optimizer.zero_grad()
train_dis_sum = 0
if train_count != 0 and train_count % 1000 == 0:
torch.save(
estimator.state_dict(),
'{0}/pose_model_current.pth'.format(opt.modelFolder))
print(
'>>>>>>>>----------epoch {0} train finish---------<<<<<<<<'.format(
epoch))
# ------------------模型测试阶段------------------
logger = setup_logger(
'epoch%d_test' % epoch,
os.path.join(opt.logFolder, 'epoch_%d_test_log.txt' % epoch))
logger.info('Test time {0}'.format(
time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)) +
', ' + 'Testing started'))
test_dis = 0.0
test_count = 0
estimator.eval()
for j, data in enumerate(testdataLoader, 0):
img, cloud, pointIndex, tarPoints, modelPoints, idx, ori_img = data
img = Variable(img).cuda()
cloud = Variable(cloud).cuda()
pointIndex = Variable(pointIndex).cuda()
tarPoints = Variable(tarPoints).cuda()
modelPoints = Variable(modelPoints).cuda()
idx = Variable(idx).cuda()
ori_img = np.array(ori_img)
pred_r, pred_t, pred_c, colorEmb = estimator(
img, cloud, pointIndex, idx)
loss, dis, newCloud, newTarPoints = poseNetLoss(
pred_r, pred_t, pred_c, tarPoints, modelPoints, idx, cloud,
opt.w, opt.refine_start)
logger.info('Test time {0} Test Frame No.{1} dis:{2}'.format(
time.strftime("%Hh %Mm %Ss",
time.gmtime(time.time() - st_time)), test_count,
dis))
test_dis += dis.item()
test_count += 1
test_dis = test_dis / test_count
logger.info('Test time {0} Epoch {1} TEST FINISH Avg dis: {2}'.format(
time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)),
epoch, test_dis))
# ------------------模型评估阶段-------------------
if test_dis < best_distance:
best_distance = test_dis
torch.save(
estimator.state_dict(),
'{0}/pose_model_{1}_{2}.pth'.format(opt.modelFolder, epoch,
test_dis))
print(epoch,
'>>>>>>>>----------BEST TEST MODEL SAVED---------<<<<<<<<')
if best_distance < opt.decayMargin and not opt.decay_start:
opt.decay_start = True
opt.lr *= opt.lr_dr
optimizer = optim.Adam(estimator.parameters(), lr=opt.lr)