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)