def __init__(self, types_list=None, learn_std=False, activation_layer='ReLU', K=1, M=1): """ Initialize BaseMissVAE. Args: types_list (list of dictionaries): Each dictionary contains: name, type, dim, nclass, index; for every attribute. learn_std (boolean): Learn the :math:`\sigma` for the real and positive distributions. activation_layer (string): Choose "relu", "tanh" or "sigmoid". K: number of importance weights for IWAE model (see: https://arxiv.org/abs/1509.00519) M: number of Monte Carlo samples for ELBO estimation """ super(BaseMissVAE, self).__init__() # Heterogeneous vars assert types_list is not None self.types_list = utils.reindex_types_list(types_list) self.transform_idx = utils.get_idx_transform(self.types_list) self.device = set_device() self.learn_std = learn_std self.activation = utils.set_activation_layer(activation_layer) # Sampler self.sampler = samplers.Sampler() # Loss self.loss = Loss() self.K = K self.M = M
def __init__(self, cfg): """ Implementation of the CoMVC model. :param cfg: Model config. See `config.defaults.CoMVC` for documentation on the config object. """ super().__init__() self.cfg = cfg self.output = self.hidden = self.fused = self.backbone_outputs = self.projections = None # Define Backbones and Fusion modules self.backbones = Backbones(cfg.backbone_configs) self.fusion = get_fusion_module(cfg.fusion_config, self.backbones.output_sizes) bb_sizes = self.backbones.output_sizes assert all([bb_sizes[0] == s for s in bb_sizes]), f"CoMVC requires all backbones to have the same " \ f"output size. Got: {bb_sizes}" if cfg.projector_config is None: self.projector = nn.Identity() else: self.projector = MLP(cfg.projector_config, input_size=bb_sizes[0]) # Define clustering module self.ddc = DDC(input_dim=self.fusion.output_size, cfg=cfg.cm_config) # Define loss-module self.loss = Loss(cfg=cfg.loss_config) # Initialize weights. self.apply(helpers.he_init_weights) # Instantiate optimizer self.optimizer = Optimizer(cfg.optimizer_config, self.parameters())
def __init__(self, cfg): """ Full DDC model :param cfg: DDC model config :type cfg: config.defaults.DDCModel """ super().__init__() self.cfg = cfg self.backbone_output = self.output = self.hidden = None self.backbone = Backbones.create_backbone(cfg.backbone_config) self.ddc_input_size = np.prod(self.backbone.output_size) self.ddc = DDC([self.ddc_input_size], cfg.cm_config) self.loss = Loss(cfg.loss_config) # Initialize weights. self.apply(helpers.he_init_weights) # Instantiate optimizer self.optimizer = Optimizer(cfg.optimizer_config, self.parameters())
def __init__(self, cfg): """ Implementation of the SiMVC model. :param cfg: Model config. See `config.defaults.SiMVC` for documentation on the config object. """ super().__init__() self.cfg = cfg self.output = self.hidden = self.fused = self.backbone_outputs = None # Define Backbones and Fusion modules self.backbones = Backbones(cfg.backbone_configs) self.fusion = get_fusion_module(cfg.fusion_config, self.backbones.output_sizes) # Define clustering module self.ddc = DDC(input_dim=self.fusion.output_size, cfg=cfg.cm_config) # Define loss-module self.loss = Loss(cfg=cfg.loss_config) # Initialize weights. self.apply(helpers.he_init_weights) # Instantiate optimizer self.optimizer = Optimizer(cfg.optimizer_config, self.parameters())
def main(): opt.manualSeed = random.randint(1, 10000) random.seed(opt.manualSeed) torch.manual_seed(opt.manualSeed) if opt.dataset == 'ycb': opt.num_objects = 21 # number of object classes in the dataset opt.num_points = 1000 # number of points on the input pointcloud opt.outf = proj_dir+'trained_models/ycb' # folder to save trained models opt.log_dir = proj_dir+'experiments/logs/ycb' # folder to save logs opt.repeat_epoch = 1 # number of repeat times for one epoch training else: print('Unknown dataset') return estimator = SymNet(num_points = opt.num_points) estimator.cuda() if opt.resume_symnet != '': estimator.load_state_dict(torch.load('{0}/{1}'.format(opt.outf, opt.resume_symnet))) opt.refine_start = False opt.decay_start = False optimizer = optim.Adam(estimator.parameters(), lr=opt.lr) opt.w *= opt.w_rate if opt.dataset == 'ycb': dataset = SymDataset_ycb('train', opt.num_points, False, opt.dataset_root, proj_dir,opt.noise_trans, opt.refine_start) test_dataset = SymDataset_ycb('test', opt.num_points, False, opt.dataset_root, proj_dir,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.sym_list = dataset.get_sym_list() opt.num_points_mesh = 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.num_points_mesh, opt.sym_list)) criterion = Loss(opt.num_points_mesh) best_test = 0 st_time = time.time() for epoch in range(opt.start_epoch, opt.nepoch): logger = setup_logger('epoch%d' % epoch, os.path.join(opt.log_dir, '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_avg = 0.0 train_err_cent = 0.0 train_loss_ref = 0.0 train_err_ref = 0.0 train_err_num = 0.0 train_err_mode = 0.0 estimator.train() optimizer.zero_grad() for rep in range(opt.repeat_epoch): for i, data in enumerate(dataloader, 0): points, choose, img, idx, target_s, target_num, target_mode, pt_num = data # the original version if idx not in sym_list: continue points, choose, img, idx, target_s, target_num, target_mode = Variable(points).cuda(), \ Variable(choose).cuda(), \ Variable(img).cuda(), \ Variable(idx).cuda(),\ Variable(target_s).cuda(), \ Variable(target_num).cuda(),\ Variable(target_mode).cuda() pred_cent, pred_ref,pred_foot_ref,pred_rot, pred_num, pred_mode, emb = estimator(img, points, choose) loss, dis, error_cent, loss_ref, error_ref, error_num, error_mode = criterion( pred_cent, pred_ref,pred_foot_ref,pred_rot, pred_num, pred_mode, target_s, points, opt.w, target_mode) loss.backward() train_dis_avg += dis.item() train_err_cent += error_cent.item() train_loss_ref += loss_ref.item() train_err_ref += error_ref.item() train_err_num += error_num.item() train_err_mode += error_mode.item() train_count += 1 if train_count % opt.batch_size == 0: logger.info( 'Train time {0} Epoch {1} Batch {2} Frame {3} error_ref: {8} loss_ref:{9} loss_cent: {7} loss_num: {5} loss_mode:{6} Avg_loss:{4} cls_id: {10}'.format( time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)), epoch, int(train_count / opt.batch_size), train_count, train_dis_avg / opt.batch_size, train_err_num / opt.batch_size, train_err_mode / opt.batch_size, train_err_cent / opt.batch_size, train_err_ref / opt.batch_size, train_loss_ref / opt.batch_size, idx.data.cpu().numpy().reshape(-1)[0])) optimizer.step() optimizer.zero_grad() train_dis_avg = 0 train_err_cent = 0 train_err_num = 0 train_err_mode = 0 train_err_ref = 0 train_loss_ref = 0 if train_count != 0 and train_count % 1000 == 0: torch.save(estimator.state_dict(), '{0}/sym_model_current.pth'.format(opt.outf)) print('>>>>>>>>----------epoch {0} train finish---------<<<<<<<<'.format(epoch)) logger = setup_logger('epoch%d_test' % epoch, os.path.join(opt.log_dir, '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 # add symmetry dis test_err_num = 0.0 test_err_mode = 0.0 test_err_ref = 0.0 test_loss_ref = 0.0 test_err_cent = 0.0 test_count = 0 ang_tps = 0 estimator.eval() # refiner.eval() for j, data in enumerate(testdataloader, 0): points, choose, img, idx, target_s, target_num,target_mode,pt_num = data if idx not in sym_list: continue points, choose, img, idx, target_s, target_num, target_mode = \ Variable(points).cuda(), \ Variable(choose).cuda(), \ Variable(img).cuda(), \ Variable(idx).cuda(), \ Variable(target_s).cuda(), \ Variable(target_num).cuda(),\ Variable(target_mode).cuda() pred_cent, pred_ref, pred_foot_ref, pred_rot, pred_num, pred_mode, emb = estimator(img, points, choose) _, dis, error_cent, loss_ref, error_ref, error_num, error_mode = criterion( pred_cent, pred_ref, pred_foot_ref, pred_rot, pred_num, pred_mode, target_s, points, opt.w, target_mode) test_dis += dis.item() test_err_cent += error_cent.item() test_err_num += error_num.item() test_err_mode += error_mode.item() test_loss_ref += loss_ref.item() test_err_ref += error_ref.item() logger.info( 'Test time {0} Test Frame:{1} error_ref:{6} loss_ref:{7} loss_cent:{5} loss_num:{3} loss_mode:{4} total_loss:{2} cls_id{8}'.format( time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)), test_count, dis, error_num, error_mode, error_cent, error_ref, loss_ref, idx.data.cpu().numpy().reshape(-1)[0])) test_count += 1 if error_ref <= 20: ang_tps += 1 test_dis = test_dis / test_count test_err_num = test_err_num / test_count test_err_mode = test_err_mode / test_count test_err_ref = test_err_ref / test_count test_loss_ref = test_loss_ref / test_count test_err_cent = test_err_cent / test_count pect_ang_tps = ang_tps / test_count # angle_loss = math.cos(test_err_ref) logger.info('Test time {0} Epoch {1} TEST FINISH loss_ref:{7} angle_tps{8} Avg dis:{2} error_num:{3} error_mode:{4} error_cent:{5} error_ref:{6} '.format(time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)), epoch, test_dis, test_err_num, test_err_mode, test_err_cent,test_err_ref, test_loss_ref, pect_ang_tps)) if pect_ang_tps >= best_test: best_test = pect_ang_tps torch.save(estimator.state_dict(), '{0}/sym_model_{1}_{2}.pth'.format(opt.outf, epoch, test_dis)) print(epoch, '>>>>>>>>----------BEST TEST MODEL SAVED---------<<<<<<<<') if test_err_ref < opt.decay_margin and not opt.decay_start: opt.decay_start = True opt.lr *= opt.lr_rate opt.w *= opt.w_rate optimizer = optim.Adam(estimator.parameters(), lr=opt.lr)
def main(): # opt.manualSeed = random.randint(1, 10000) # # opt.manualSeed = 1 # random.seed(opt.manualSeed) # torch.manual_seed(opt.manualSeed) torch.set_printoptions(threshold=5000) # device_ids = [0,1] cudnn.benchmark = True if opt.dataset == 'ycb': 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 = 'experiments/logs/ycb' #folder to save logs opt.repeat_epoch = 3 #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 = 'experiments/logs/linemod' opt.repeat_epoch = 20 else: print('Unknown dataset') return estimator = PoseNet(num_points=opt.num_points, num_obj=opt.num_objects) estimator.cuda() refiner = PoseRefineNet(num_points=opt.num_points, num_obj=opt.num_objects) # refiner.cuda() # estimator = nn.DataParallel(estimator, device_ids=device_ids) if opt.resume_posenet != '': estimator.load_state_dict( torch.load('{0}/{1}'.format(opt.outf, opt.resume_posenet))) print('LOADED!!') if opt.resume_refinenet != '': refiner.load_state_dict( torch.load('{0}/{1}'.format(opt.outf, opt.resume_refinenet))) opt.refine_start = True opt.decay_start = True opt.lr *= opt.lr_rate opt.w *= opt.w_rate opt.batch_size = int(opt.batch_size / opt.iteration) optimizer = optim.Adam(refiner.parameters(), lr=opt.lr) else: print('no refinement') opt.refine_start = False opt.decay_start = False optimizer = optim.Adam(estimator.parameters(), lr=opt.lr) # optimizer = nn.DataParallel(optimizer, device_ids=device_ids) if opt.dataset == 'ycb': dataset = PoseDataset_ycb('train', opt.num_points, False, opt.dataset_root, opt.noise_trans, opt.refine_start) # print(dataset.list) elif opt.dataset == 'linemod': dataset = PoseDataset_linemod('train', opt.num_points, True, opt.dataset_root, opt.noise_trans, opt.refine_start) dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True, num_workers=opt.workers) if opt.dataset == 'ycb': test_dataset = PoseDataset_ycb('test', opt.num_points, False, opt.dataset_root, 0.0, opt.refine_start) elif opt.dataset == 'linemod': test_dataset = PoseDataset_linemod('test', opt.num_points, False, opt.dataset_root, 0.0, opt.refine_start) testdataloader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=opt.workers) opt.sym_list = dataset.get_sym_list() opt.num_points_mesh = 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.num_points_mesh, opt.sym_list)) criterion = Loss(opt.num_points_mesh, opt.sym_list) # criterion_refine = Loss_refine(opt.num_points_mesh, opt.sym_list) best_test = np.Inf best_epoch = 0 if opt.start_epoch == 1: for log in os.listdir(opt.log_dir): os.remove(os.path.join(opt.log_dir, log)) st_time = time.time() count_gen = 0 mode = 1 if mode == 1: for epoch in range(opt.start_epoch, opt.nepoch): logger = setup_logger( 'epoch%d' % epoch, os.path.join(opt.log_dir, '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_avg = 0.0 if opt.refine_start: estimator.eval() refiner.train() else: estimator.train() optimizer.zero_grad() for rep in range(opt.repeat_epoch): for i, data in enumerate(dataloader, 0): points, choose, img, target_sym, target_cen, idx, file_list_idx = data if idx is 9 or idx is 16: continue points, choose, img, target_sym, target_cen, idx = Variable(points).cuda(), \ Variable(choose).cuda(), \ Variable(img).cuda(), \ Variable(target_sym).cuda(), \ Variable(target_cen).cuda(), \ Variable(idx).cuda() pred_norm, pred_on_plane, emb = estimator( img, points, choose, idx) loss = criterion(pred_norm, pred_on_plane, target_sym, target_cen, idx, points, opt.w, opt.refine_start) # scene_idx = dataset.list[file_list_idx] loss.backward() # train_dis_avg += dis.item() train_count += 1 if train_count % opt.batch_size == 0: logger.info( 'Train time {0} Epoch {1} Batch {2} Frame {3}'. format( time.strftime( "%Hh %Mm %Ss", time.gmtime(time.time() - st_time)), epoch, int(train_count / opt.batch_size), train_count)) optimizer.step() # for param_lr in optimizer.module.param_groups: # param_lr['lr'] /= 2 optimizer.zero_grad() train_dis_avg = 0 if train_count % 8 == 0: print(pred_on_plane.max()) print(pred_on_plane.mean()) print(idx) if train_count != 0 and train_count % 1000 == 0: if opt.refine_start: torch.save( refiner.state_dict(), '{0}/pose_refine_model_current.pth'.format( opt.outf)) else: torch.save( estimator.state_dict(), '{0}/pose_model_current.pth'.format(opt.outf)) print('>>>>>>>>----------epoch {0} train finish---------<<<<<<<<'. format(epoch)) logger = setup_logger( 'epoch%d_test' % epoch, os.path.join(opt.log_dir, '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_loss = 0.0 test_count = 0 estimator.eval() 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_loss)) print(pred_on_plane.max()) print(pred_on_plane.mean()) bs, num_p, _ = pred_on_plane.size() # if epoch % 40 == 0: # import pdb;pdb.set_trace() best_test = test_loss best_epoch = epoch if opt.refine_start: torch.save( refiner.state_dict(), '{0}/pose_refine_model_{1}_{2}.pth'.format( opt.outf, epoch, test_loss)) else: torch.save( estimator.state_dict(), '{0}/pose_model_{1}_{2}.pth'.format( opt.outf, epoch, test_loss)) print(epoch, '>>>>>>>>----------BEST TEST MODEL SAVED---------<<<<<<<<') if best_test < opt.decay_margin and not opt.decay_start: opt.decay_start = True opt.lr *= opt.lr_rate # opt.w *= opt.w_rate optimizer = optim.Adam(estimator.parameters(), lr=opt.lr) estimator.load_state_dict( torch.load('{0}/pose_model_{1}_{2}.pth'.format( opt.outf, best_epoch, best_test))) else: estimator.load_state_dict( torch.load('{0}/pose_model_45_0.0.pth'.format(opt.outf), map_location='cpu'))
def train_net(): os.environ['CUDA_VISIBLE_DEVICES'] = opt.gpu physical_devices = tf.config.experimental.list_physical_devices('GPU') if len(physical_devices) > 0: for k in range(len(physical_devices)): tf.config.experimental.set_memory_growth(physical_devices[k], True) print( 'memory growth:', tf.config.experimental.get_memory_growth(physical_devices[k])) else: print("Not enough GPU hardware devices available") # set result directory if not os.path.exists(opt.result_dir): os.makedirs(opt.result_dir) tb_writer = tf.summary.create_file_writer(opt.result_dir) logger = setup_logger('train_log', os.path.join(opt.result_dir, 'log.txt')) logger.propagate = 0 for key, value in vars(opt).items(): logger.info(key + ': ' + str(value)) # model & loss estimator = DeformNet(opt.n_cat, opt.nv_prior) estimator.cuda() criterion = Loss(opt.corr_wt, opt.cd_wt, opt.entropy_wt, opt.deform_wt) if opt.resume_model != '': estimator.load_state_dict(torch.load(opt.resume_model)) # dataset train_dataset = PoseDataset(opt.dataset, 'train', opt.data_dir, opt.n_pts, opt.img_size, opt.points_process, vis=visflag) val_dataset = PoseDataset(opt.dataset, 'test', opt.data_dir, opt.n_pts, opt.img_size, opt.points_process, vis=visflag) # start training st_time = time.time() train_steps = 1500 global_step = train_steps * (opt.start_epoch - 1) n_decays = len(opt.decay_epoch) assert len(opt.decay_rate) == n_decays for i in range(n_decays): if opt.start_epoch > opt.decay_epoch[i]: decay_count = i current_lr = opt.lr * opt.decay_rate[decay_count] optimizer = torch.optim.Adam(estimator.parameters(), lr=current_lr) train_size = train_steps * opt.batch_size indices = [] page_start = -train_size for epoch in range(opt.start_epoch, opt.max_epoch + 1): # train one epoch logger.info('Time {0}'.format(time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)) + \ ', ' + 'Epoch %02d' % epoch + ', ' + 'Training started')) # # create optimizer and adjust learning rate if needed # if decay_count < len(opt.decay_rate): # if epoch > opt.decay_epoch[decay_count]: # current_lr = opt.lr * opt.decay_rate[decay_count] # optimizer = torch.optim.Adam(estimator.parameters(), lr=current_lr) # decay_count += 1 # sample train subset page_start += train_size len_last = len(indices) - page_start if len_last < train_size: indices = indices[page_start:] if opt.dataset == 'CAMERA+Real': # CAMERA : Real = 3 : 1 camera_len = train_dataset.subset_len[0] real_len = train_dataset.subset_len[1] real_indices = list(range(camera_len, camera_len + real_len)) camera_indices = list(range(camera_len)) n_repeat = (train_size - len_last) // (4 * real_len) + 1 data_list = random.sample(camera_indices, 3 * n_repeat * real_len) + real_indices * n_repeat random.shuffle(data_list) indices += data_list else: data_list = list(range(train_dataset.length)) for i in range((train_size - len_last) // train_dataset.length + 1): random.shuffle(data_list) indices += data_list page_start = 0 train_idx = indices[page_start:(page_start + train_size)] train_sampler = torch.utils.data.sampler.SubsetRandomSampler(train_idx) train_dataloader = torch.utils.data.DataLoader( train_dataset, batch_size=opt.batch_size, sampler=train_sampler, num_workers=opt.num_workers, pin_memory=True) estimator.train() for i, data in enumerate(train_dataloader, 1): points, points_pro, rgb, choose, cat_id, model, prior, sRT, nocs = data points_pro = points_pro.cuda() rgb = rgb.cuda() choose = choose.cuda() cat_id = cat_id.cuda() model = model.cuda() prior = prior.cuda() sRT = sRT.cuda() nocs = nocs.cuda() assign_mat, deltas = estimator(points_pro, rgb, choose, cat_id, prior) loss, corr_loss, cd_loss, entropy_loss, deform_loss = criterion( assign_mat, deltas, prior, nocs, model) optimizer.zero_grad() loss.backward() optimizer.step() global_step += 1 # write results to tensorboard with tb_writer.as_default(): tf.summary.scalar('learning_rate', current_lr, step=global_step) tf.summary.scalar('train_loss', loss.item(), step=global_step) tf.summary.scalar('corr_loss', corr_loss.item(), step=global_step) tf.summary.scalar('cd_loss', cd_loss.item(), step=global_step) tf.summary.scalar('entropy_loss', entropy_loss.item(), step=global_step) tf.summary.scalar('deform_loss', deform_loss.item(), step=global_step) tb_writer.flush() if i % 10 == 0: logger.info( 'Batch {0} Loss:{1:f}, corr_loss:{2:f}, cd_loss:{3:f}, entropy_loss:{4:f}, deform_loss:{5:f}' .format(i, loss.item(), corr_loss.item(), cd_loss.item(), entropy_loss.item(), deform_loss.item())) # adjust learning rate if needed if decay_count < len(opt.decay_rate): if epoch >= opt.decay_epoch[decay_count]: current_lr = opt.lr * opt.decay_rate[decay_count] optimizer = torch.optim.Adam(estimator.parameters(), lr=current_lr) decay_count += 1 logger.info( '>>>>>>>>----------Epoch {:02d} train finish---------<<<<<<<<'. format(epoch)) # evaluate one epoch logger.info('Time {0}'.format( time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)) + ', ' + 'Epoch %02d' % epoch + ', ' + 'Testing started')) val_loss = 0.0 total_count = np.zeros((opt.n_cat, ), dtype=int) strict_success = np.zeros((opt.n_cat, ), dtype=int) # 5 degree and 5 cm easy_success = np.zeros((opt.n_cat, ), dtype=int) # 10 degree and 5 cm iou_success = np.zeros((opt.n_cat, ), dtype=int) # relative scale error < 0.1 # sample validation subset # opt.val_size = 2500 val_idx = random.sample(list(range(val_dataset.length)), opt.val_size) val_sampler = torch.utils.data.sampler.SubsetRandomSampler(val_idx) val_dataloader = torch.utils.data.DataLoader( val_dataset, batch_size=1, sampler=val_sampler, num_workers=opt.num_workers, pin_memory=True) estimator.eval() for i, data in enumerate(val_dataloader, 1): points, points_pro, rgb, choose, cat_id, model, prior, sRT, nocs = data points_pro = points_pro.cuda() points = points.cuda() rgb = rgb.cuda() choose = choose.cuda() cat_id = cat_id.cuda() model = model.cuda() prior = prior.cuda() sRT = sRT.cuda() nocs = nocs.cuda() assign_mat, deltas = estimator(points_pro, rgb, choose, cat_id, prior) loss, _, _, _, _ = criterion(assign_mat, deltas, prior, nocs, model) # estimate pose and scale inst_shape = prior + deltas assign_mat = F.softmax(assign_mat, dim=2) nocs_coords = torch.bmm(assign_mat, inst_shape) nocs_coords = nocs_coords.detach().cpu().numpy()[0] points = points.cpu().numpy()[0] # use choose to remove repeated points choose = choose.cpu().numpy()[0] _, choose = np.unique(choose, return_index=True) nocs_coords = nocs_coords[choose, :] points = points[choose, :] _, _, _, pred_sRT = estimateSimilarityTransform( nocs_coords, points) # evaluate pose cat_id = cat_id.item() if pred_sRT is not None: sRT = sRT.detach().cpu().numpy()[0] R_error, T_error, IoU = compute_sRT_errors(pred_sRT, sRT) if R_error < 5 and T_error < 0.05: strict_success[cat_id] += 1 if R_error < 10 and T_error < 0.05: easy_success[cat_id] += 1 if IoU < 0.1: iou_success[cat_id] += 1 total_count[cat_id] += 1 val_loss += loss.item() if i % 100 == 0: logger.info('Batch {0} Loss:{1:f}'.format(i, loss.item())) # compute accuracy strict_acc = 100 * (strict_success / total_count) easy_acc = 100 * (easy_success / total_count) iou_acc = 100 * (iou_success / total_count) for i in range(opt.n_cat): logger.info('{} accuracies:'.format(val_dataset.cat_names[i])) logger.info('5^o 5cm: {:4f}'.format(strict_acc[i])) logger.info('10^o 5cm: {:4f}'.format(easy_acc[i])) logger.info('IoU < 0.1: {:4f}'.format(iou_acc[i])) strict_acc = np.mean(strict_acc) easy_acc = np.mean(easy_acc) iou_acc = np.mean(iou_acc) val_loss = val_loss / opt.val_size with tb_writer.as_default(): tf.summary.scalar('val_loss', val_loss, step=global_step) tf.summary.scalar('5^o5cm_acc', strict_acc, step=global_step) tf.summary.scalar('10^o5cm_acc', easy_acc, step=global_step) tf.summary.scalar('iou_acc', iou_acc, step=global_step) tb_writer.flush() logger.info('Epoch {0:02d} test average loss: {1:06f}'.format( epoch, val_loss)) logger.info('Overall accuracies:') logger.info('5^o 5cm: {:4f} 10^o 5cm: {:4f} IoU: {:4f}'.format( strict_acc, easy_acc, iou_acc)) logger.info( '>>>>>>>>----------Epoch {:02d} test finish---------<<<<<<<<'. format(epoch)) # save model after each epoch torch.save(estimator.state_dict(), '{0}/model_{1:02d}.pth'.format(opt.result_dir, epoch))
def main(): # g13: parameter setting ------------------- batch_id = 1 opt.dataset ='linemod' opt.dataset_root = './datasets/linemod/Linemod_preprocessed' estimator_path = 'trained_checkpoints/linemod/pose_model_9_0.01310166542980859.pth' refiner_path = 'trained_checkpoints/linemod/pose_refine_model_493_0.006761023565178073.pth' opt.resume_posenet = estimator_path opt.resume_posenet = refiner_path dataset_config_dir = 'datasets/linemod/dataset_config' output_result_dir = 'experiments/eval_result/linemod' bs = 1 #fixed because of the default setting in torch.utils.data.DataLoader opt.iteration = 2 #default is 4 in eval_linemod.py t1_idx = 0 t1_total_eval_num = 3 axis_range = 0.1 # the length of X, Y, and Z axis in 3D vimg_dir = 'verify_img' if not os.path.exists(vimg_dir): os.makedirs(vimg_dir) #------------------------------------------- if opt.dataset == 'ycb': 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 = '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 = 'experiments/logs/linemod' opt.repeat_epoch = 20 else: print('Unknown dataset') return estimator = PoseNet(num_points = opt.num_points, num_obj = opt.num_objects) estimator.cuda() refiner = PoseRefineNet(num_points = opt.num_points, num_obj = opt.num_objects) refiner.cuda() if opt.resume_posenet != '': estimator.load_state_dict(torch.load(estimator_path)) if opt.resume_refinenet != '': refiner.load_state_dict(torch.load(refiner_path)) opt.refine_start = True opt.decay_start = True opt.lr *= opt.lr_rate opt.w *= opt.w_rate opt.batch_size = int(opt.batch_size / opt.iteration) optimizer = optim.Adam(refiner.parameters(), lr=opt.lr) else: opt.refine_start = False opt.decay_start = False optimizer = optim.Adam(estimator.parameters(), lr=opt.lr) if opt.dataset == 'ycb': test_dataset = PoseDataset_ycb('test', opt.num_points, False, opt.dataset_root, 0.0, opt.refine_start) elif opt.dataset == 'linemod': test_dataset = PoseDataset_linemod('test', opt.num_points, False, opt.dataset_root, 0.0, opt.refine_start) testdataloader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=opt.workers) print('complete loading testing loader\n') opt.sym_list = test_dataset.get_sym_list() opt.num_points_mesh = test_dataset.get_num_points_mesh() print('>>>>>>>>----------Dataset loaded!---------<<<<<<<<\n\ length of the testing set: {0}\nnumber of sample points on mesh: {1}\n\ symmetry object list: {2}'\ .format( len(test_dataset), opt.num_points_mesh, opt.sym_list)) #load pytorch model estimator.eval() refiner.eval() criterion = Loss(opt.num_points_mesh, opt.sym_list) criterion_refine = Loss_refine(opt.num_points_mesh, opt.sym_list) fw = open('{0}/t1_eval_result_logs.txt'.format(output_result_dir), 'w') #Pose estimation for j, data in enumerate(testdataloader, 0): # g13: modify this part for evaluation target-------------------- if j == t1_total_eval_num: break #---------------------------------------------------------------- points, choose, img, target, model_points, idx = data if len(points.size()) == 2: print('No.{0} NOT Pass! Lost detection!'.format(j)) fw.write('No.{0} NOT Pass! Lost detection!\n'.format(j)) continue points, choose, img, target, model_points, idx = Variable(points).cuda(), \ Variable(choose).cuda(), \ Variable(img).cuda(), \ Variable(target).cuda(), \ Variable(model_points).cuda(), \ Variable(idx).cuda() pred_r, pred_t, pred_c, emb = estimator(img, points, choose, idx) _, dis, new_points, new_target = criterion(pred_r, pred_t, pred_c, target, model_points, idx, points, opt.w, opt.refine_start) #if opt.refine_start: #iterative poserefinement # for ite in range(0, opt.iteration): # pred_r, pred_t = refiner(new_points, emb, idx) # dis, new_points, new_target = criterion_refine(pred_r, pred_t, new_target, model_points, idx, new_points) pred_r = pred_r / torch.norm(pred_r, dim=2).view(1, opt.num_points, 1) pred_c = pred_c.view(bs, opt.num_points) how_max, which_max = torch.max(pred_c, 1) pred_t = pred_t.view(bs * opt.num_points, 1, 3) my_r = pred_r[0][which_max[0]].view(-1).cpu().data.numpy() my_t = (points.view(bs * opt.num_points, 1, 3) + pred_t)[which_max[0]].view(-1).cpu().data.numpy() my_pred = np.append(my_r, my_t) for ite in range(0, opt.iteration): T = Variable(torch.from_numpy(my_t.astype(np.float32))).cuda().view(1, 3).repeat(opt.num_points, 1).contiguous().view(1, opt.num_points, 3) my_mat = quaternion_matrix(my_r) R = Variable(torch.from_numpy(my_mat[:3, :3].astype(np.float32))).cuda().view(1, 3, 3) my_mat[0:3, 3] = my_t new_points = torch.bmm((points - T), R).contiguous() pred_r, pred_t = refiner(new_points, emb, idx) pred_r = pred_r.view(1, 1, -1) pred_r = pred_r / (torch.norm(pred_r, dim=2).view(1, 1, 1)) my_r_2 = pred_r.view(-1).cpu().data.numpy() my_t_2 = pred_t.view(-1).cpu().data.numpy() my_mat_2 = quaternion_matrix(my_r_2) my_mat_2[0:3, 3] = my_t_2 my_mat_final = np.dot(my_mat, my_mat_2) my_r_final = copy.deepcopy(my_mat_final) my_r_final[0:3, 3] = 0 my_r_final = quaternion_from_matrix(my_r_final, True) my_t_final = np.array([my_mat_final[0][3], my_mat_final[1][3], my_mat_final[2][3]]) my_pred = np.append(my_r_final, my_t_final) my_r = my_r_final my_t = my_t_final # g13: start drawing pose on image------------------------------------ # pick up image print("index {0}: {1}".format(j, test_dataset.list_rgb[j])) img = Image.open(test_dataset.list_rgb[j]) # pick up center position by bbox meta_file = open('{0}/data/{1}/gt.yml'.format(opt.dataset_root, '%02d' % test_dataset.list_obj[j]), 'r') meta = {} meta = yaml.load(meta_file) which_item = test_dataset.list_rank[j] bbx = meta[which_item][0]['obj_bb'] draw = ImageDraw.Draw(img) # draw box (ensure this is the right object) draw.line((bbx[0],bbx[1], bbx[0], bbx[1]+bbx[3]), fill=(255,0,0), width=5) draw.line((bbx[0],bbx[1], bbx[0]+bbx[2], bbx[1]), fill=(255,0,0), width=5) draw.line((bbx[0],bbx[1]+bbx[3], bbx[0]+bbx[2], bbx[1]+bbx[3]), fill=(255,0,0), width=5) draw.line((bbx[0]+bbx[2],bbx[1], bbx[0]+bbx[2], bbx[1]+bbx[3]), fill=(255,0,0), width=5) #get center c_x = bbx[0]+int(bbx[2]/2) c_y = bbx[1]+int(bbx[3]/2) draw.point((c_x,c_y), fill=(255,255,0)) #get the 3D position of center cam_intrinsic = np.zeros((3,3)) cam_intrinsic.itemset(0, test_dataset.cam_fx) cam_intrinsic.itemset(4, test_dataset.cam_fy) cam_intrinsic.itemset(2, test_dataset.cam_cx) cam_intrinsic.itemset(5, test_dataset.cam_cy) cam_intrinsic.itemset(8, 1) cam_extrinsic = my_mat_final[0:3, :] cam2d_3d = np.matmul(cam_intrinsic, cam_extrinsic) cen_3d = np.matmul(np.linalg.pinv(cam2d_3d), [[c_x],[c_y],[1]]) # replace img.show() with plt.imshow(img) #transpose three 3D axis point into 2D x_3d = cen_3d + [[axis_range],[0],[0],[0]] y_3d = cen_3d + [[0],[axis_range],[0],[0]] z_3d = cen_3d + [[0],[0],[axis_range],[0]] x_2d = np.matmul(cam2d_3d, x_3d) y_2d = np.matmul(cam2d_3d, y_3d) z_2d = np.matmul(cam2d_3d, z_3d) #draw the axis on 2D draw.line((c_x, c_y, x_2d[0], x_2d[1]), fill=(255,255,0), width=5) draw.line((c_x, c_y, y_2d[0], y_2d[1]), fill=(0,255,0), width=5) draw.line((c_x, c_y, z_2d[0], z_2d[1]), fill=(0,0,255), width=5) #g13: show image #img.show() #save file under file img_file_name = '{0}/pred_obj{1}_pic{2}.png'.format(vimg_dir, test_dataset.list_obj[j], which_item) img.save( img_file_name, "PNG" ) img.close()
def main(): class_id = 0 class_file = open('datasets/ycb/dataset_config/classes.txt') cld = {} while 1: class_input = class_file.readline() if not class_input: break input_file = open('{0}/models/{1}/points.xyz'.format( opt.dataset_root, class_input[:-1])) cld[class_id] = [] while 1: input_line = input_file.readline() if not input_line: break input_line = input_line[:-1].split(' ') cld[class_id].append([ float(input_line[0]), float(input_line[1]), float(input_line[2]) ]) cld[class_id] = np.array(cld[class_id]) input_file.close() class_id += 1 opt.manualSeed = random.randint(1, 10000) random.seed(opt.manualSeed) torch.manual_seed(opt.manualSeed) symmetry_obj_idx = [12, 15, 18, 19, 20] if opt.dataset == 'ycb': 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/' + opt.output_dir # folder to save trained models opt.test_output = 'experiments/output/ycb/' + opt.output_dir if not os.path.exists(opt.test_output): os.makedirs(opt.test_output, exist_ok=True) 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 = 'experiments/logs/linemod' opt.repeat_epoch = 20 else: print('Unknown dataset') return estimator = PoseNet(num_points=opt.num_points, num_obj=opt.num_objects, object_max=opt.object_max) estimator.cuda() if opt.resume_posenet != '': estimator.load_state_dict( torch.load('{0}/{1}'.format(opt.outf, opt.resume_posenet))) opt.refine_start = False opt.decay_start = False dataset = PoseDataset_ycb('train', opt.num_points, False, opt.dataset_root, opt.noise_trans, opt.seg_type, True) test_dataset = PoseDataset_ycb('test', opt.num_points, False, opt.dataset_root, 0.0, opt.seg_type, True) testdataloader = torch.utils.data.DataLoader(test_dataset, shuffle=False, num_workers=opt.workers) opt.sym_list = dataset.get_sym_list() opt.num_points_mesh = 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.num_points_mesh, opt.sym_list)) criterion = Loss(opt.num_points_mesh, opt.sym_list) logger = setup_logger( 'final_eval_tf_with_seg_square', os.path.join(opt.test_output, 'final_eval_tf_with_seg_square.txt')) object_max = opt.object_max total_test_dis = {key: [] for key in range(0, object_max)} total_test_count = {key: [] for key in range(0, object_max)} dir_test_dis = {key: [] for key in range(0, object_max)} dir_test_count = {key: [] for key in range(0, object_max)} # for add total_unseen_objects = {key: [] for key in range(0, object_max)} total_object_without_pose = {key: [] for key in range(0, object_max)} dir_add_count = {key: [] for key in range(0, object_max)} dir_add_count_unseen = {key: [] for key in range(0, object_max)} dir_add_02_count_unseen = {key: [] for key in range(0, object_max)} dir_add_pure_count = {key: [] for key in range(0, object_max)} dir_add_s_count = {key: [] for key in range(0, object_max)} dir_add_02_count = {key: [] for key in range(0, object_max)} dir_add_pure_02_count = {key: [] for key in range(0, object_max)} dir_add_s_02_count = {key: [] for key in range(0, object_max)} total_add_count = {key: [] for key in range(0, object_max)} total_add_count_unseen = {key: [] for key in range(0, object_max)} total_add_02_count_unseen = {key: [] for key in range(0, object_max)} total_add_pure_count = {key: [] for key in range(0, object_max)} total_add_s_count = {key: [] for key in range(0, object_max)} total_add_02_count = {key: [] for key in range(0, object_max)} total_add_pure_02_count = {key: [] for key in range(0, object_max)} total_add_s_02_count = {key: [] for key in range(0, object_max)} dir_dbd_count = {key: [] for key in range(0, object_max)} dir_drr_count = {key: [] for key in range(0, object_max)} dir_ada_count = {key: [] for key in range(0, object_max)} dir_distance_1_count = {key: [] for key in range(0, object_max)} total_dbd_count = {key: [] for key in range(0, object_max)} total_drr_count = {key: [] for key in range(0, object_max)} total_ada_count = {key: [] for key in range(0, object_max)} total_distance_1_count = {key: [] for key in range(0, object_max)} last_dis = {key: [] for key in range(0, object_max)} for i in range(object_max): total_unseen_objects[i] = 0 total_object_without_pose[i] = 0 total_test_dis[i] = 0. total_test_count[i] = 0 dir_test_dis[i] = 0. dir_test_count[i] = 0 # for add dir_add_count[i] = 0 dir_add_count_unseen[i] = 0 dir_add_02_count_unseen[i] = 0 dir_add_pure_count[i] = 0 dir_add_s_count[i] = 0 dir_add_02_count[i] = 0 total_add_count[i] = 0 total_add_count_unseen[i] = 0 total_add_02_count_unseen[i] = 0 total_add_pure_count[i] = 0 total_add_s_count[i] = 0 total_add_02_count[i] = 0 dir_add_pure_02_count[i] = 0 dir_add_s_02_count[i] = 0 total_add_pure_02_count[i] = 0 total_add_s_02_count[i] = 0 # for stable dir_dbd_count[i] = 0. dir_drr_count[i] = 0 dir_ada_count[i] = 0. dir_distance_1_count[i] = 0. total_dbd_count[i] = 0. total_drr_count[i] = 0 total_ada_count[i] = 0. total_distance_1_count[i] = 0. last_dis[i] = None st_time = time.time() isFirstInitLastDatafolder = True estimator.eval() with torch.no_grad(): for j, data in enumerate(testdataloader, 0): if opt.dataset == 'ycb': list_points, list_choose, list_img, list_target, list_model_points, list_idx, list_filename, \ list_full_img, list_focal_length, list_principal_point, list_motion = data output_image = Image.open('{0}/{1}-color-masked-square.png'.format( opt.dataset_root, list_filename[0][0])) OUTPUT_IMAGE_PATH = '{0}/{1}-color-seg-square-output-tf.png'.format( opt.dataset_root, list_filename[0][0]) for list_index in range(len(list_points)): points, choose, img, target, model_points, idx, filename, full_img, focal_length, principal_point, motion \ = list_points[list_index], list_choose[list_index], list_img[list_index], \ list_target[list_index], list_model_points[list_index], list_idx[list_index], \ list_filename[list_index], list_full_img[list_index], list_focal_length[list_index], \ list_principal_point[list_index], list_motion[list_index] # Temporal Clean when Changing datafolder datafolder = filename[0].split('/')[1] filehead = filename[0].split('/')[2] if isFirstInitLastDatafolder: lastdatafolder = datafolder isFirstInitLastDatafolder = False if datafolder != lastdatafolder: logger.info('changing folder from {0} to {1}'.format( lastdatafolder, datafolder)) estimator.temporalClear(opt.object_max) # handle dir output for i in range(0, object_max): if dir_test_count[i] != 0: logger.info( 'Dir {0} Object {1} dis:{2} with {3} samples'. format(lastdatafolder, i, dir_test_dis[i] / dir_test_count[i], dir_test_count[i])) if dir_add_count[i] != 0: logger.info( 'Dir {0} Object {1} add:{2} with 0.02: {3}' .format( lastdatafolder, i, dir_add_count[i] / dir_test_count[i], dir_add_02_count[i] / dir_add_count[i])) else: logger.info( 'Dir {0} Object {1} add:{2} with 0.02: {3}' .format( lastdatafolder, i, dir_add_count[i] / dir_test_count[i], 0)) if dir_add_pure_count[i] != -0: logger.info( 'Dir {0} Object {1} add_pure:{2} with 0.02: {3}' .format( lastdatafolder, i, dir_add_pure_count[i] / dir_test_count[i], dir_add_pure_02_count[i] / dir_add_pure_count[i])) else: logger.info( 'Dir {0} Object {1} add_pure:{2} with 0.02: {3}' .format( lastdatafolder, i, dir_add_pure_count[i] / dir_test_count[i], 0)) if dir_add_s_count[i] != 0: logger.info( 'Dir {0} Object {1} add_s:{2} with 0.02: {3}' .format( lastdatafolder, i, dir_add_s_count[i] / dir_test_count[i], dir_add_s_02_count[i] / dir_add_s_count[i])) else: logger.info( 'Dir {0} Object {1} add_s:{2} with 0.02: {3}' .format( lastdatafolder, i, dir_add_s_count[i] / dir_test_count[i], 0)) logger.info('Dir {0} Object {1} dbd:{2}'.format( lastdatafolder, i, dir_dbd_count[i] / dir_test_count[i])) logger.info('Dir {0} Object {1} drr:{2}'.format( lastdatafolder, i, dir_drr_count[i] / dir_test_count[i])) logger.info('Dir {0} Object {1} ada:{2}'.format( lastdatafolder, i, dir_ada_count[i] / dir_test_count[i])) logger.info( 'Dir {0} Object {1} distance_1:{2}'.format( lastdatafolder, i, dir_distance_1_count[i] / dir_test_count[i])) dir_dbd = 0. dir_drr = 0. dir_ada = 0. dir_distance_1 = 0. dir_dis = 0. dir_add = 0 dir_add_s = 0 dir_add_pure = 0 dir_add_02 = 0 dir_add_s_02 = 0 dir_add_pure_02 = 0 dir_count = 0 for i in range(object_max): if total_test_count[i] != 0: dir_count += dir_test_count[i] dir_dis += dir_test_dis[i] dir_add += dir_add_count[i] dir_add_pure += dir_add_pure_count[i] dir_add_s += dir_add_s_count[i] dir_add_02 += dir_add_02_count[i] dir_add_pure_02 += dir_add_pure_02_count[i] dir_add_s_02 += dir_add_s_02_count[i] dir_dbd += dir_dbd_count[i] dir_drr += dir_drr_count[i] dir_ada += dir_ada_count[i] dir_distance_1 += dir_distance_1_count[i] dir_test_dis[i] = 0 dir_test_count[i] = 0 dir_add_count[i] = 0 dir_add_pure_count[i] = 0 dir_add_s_count[i] = 0 dir_add_02_count[i] = 0 dir_add_pure_02_count[i] = 0 dir_add_s_02_count[i] = 0 dir_dbd_count[i] = 0 dir_drr_count[i] = 0 dir_ada_count[i] = 0 dir_distance_1_count[i] = 0 last_dis[i] = None logger.info( 'Dir {0} \'s total dis:{1} with {2} samples'.format( lastdatafolder, dir_dis / dir_count, dir_count)) logger.info( 'Dir {0} \'s total add:{1} with 0.02: {2}'.format( lastdatafolder, dir_add / dir_count, dir_add_02 / dir_add)) logger.info( 'Dir {0} \'s total add_s:{1} with 0.02: {2}'.format( lastdatafolder, dir_add_s / dir_count, dir_add_s_02 / dir_add_s)) logger.info( 'Dir {0} \'s total add_pure:{1} with 0.02: {2}'.format( lastdatafolder, dir_add_pure / dir_count, dir_add_pure_02 / dir_add_pure)) logger.info('Dir {0} \'s total dbd:{1}'.format( lastdatafolder, dir_dbd / dir_count)) logger.info('Dir {0} \'s total drr:{1}'.format( lastdatafolder, dir_drr / dir_count)) logger.info('Dir {0} \'s total ada:{1}'.format( lastdatafolder, dir_ada / dir_count)) logger.info('Dir {0} \'s total distance_1:{1}'.format( lastdatafolder, dir_distance_1 / dir_count)) # end of handle dir output lastdatafolder = datafolder points, choose, img, target, model_points, idx = points.cuda(), \ choose.cuda(), \ img.cuda(), \ target.cuda(), \ model_points.cuda(), \ idx.cuda() cloud_path = "experiments/clouds/ycb/{0}/{1}/{2}/{3}_{4}".format( opt.output_dir, 1, datafolder, filehead, int(idx)) # folder to save logs pred_r, pred_t, pred_c, x_return = estimator( img, points, choose, idx, focal_length, principal_point, motion, cloud_path) # count for unseen object if pred_r is None: last_dis[int(idx)] = None total_unseen_objects[int(idx)] += 1 total_object_without_pose[int(idx)] += 1 continue pred_r_ori = copy.deepcopy(pred_r) pred_t_ori = copy.deepcopy(pred_t) pred_c_ori = copy.deepcopy(pred_c) x_return_ori = copy.deepcopy(x_return) gt_r, gt_t = get_target(opt.dataset_root, filename, idx) if gt_r is None: print('gtr is None') is_sym = int(idx) in symmetry_obj_idx dis, dis_vector, pred_cloud = calDistance( pred_r_ori, pred_t_ori, pred_c_ori, x_return_ori, gt_r, gt_t, cld[int(idx)], is_sym) dis_s, dis_vector_s, _ = calDistance(pred_r_ori, pred_t_ori, pred_c_ori, x_return_ori, gt_r, gt_t, cld[int(idx)], True) dis_pure, dis_vector_pure, _ = calDistance( pred_r_ori, pred_t_ori, pred_c_ori, x_return_ori, gt_r, gt_t, cld[int(idx)], False) if last_dis[int(idx)] is not None: dir_dbd_count[int(idx)] += torch.norm(dis_vector - last_dis[int(idx)]) total_dbd_count[int(idx)] += torch.norm(dis_vector - last_dis[int(idx)]) dir_distance_1_count[int(idx)] += torch.norm( (dis_vector / torch.norm(dis_vector)) - (last_dis[int(idx)] / torch.norm(last_dis[int(idx)]))) total_distance_1_count[int(idx)] += torch.norm( (dis_vector / torch.norm(dis_vector)) - (last_dis[int(idx)] / torch.norm(last_dis[int(idx)]))) if torch.dot(last_dis[int(idx)], dis_vector) < 0: dir_drr_count[int(idx)] += 1 total_drr_count[int(idx)] += 1 dir_ada_count[int(idx)] += torch.acos( (torch.dot(last_dis[int(idx)], dis_vector)) / (torch.norm(last_dis[int(idx)]) * torch.norm(dis_vector))) total_ada_count[int(idx)] += torch.acos( (torch.dot(last_dis[int(idx)], dis_vector)) / (torch.norm(last_dis[int(idx)]) * torch.norm(dis_vector))) last_dis[int(idx)] = dis_vector # calc adds if img.shape[1] != 0: dir_test_dis[int(idx)] += dis.item() total_test_dis[int(idx)] += dis.item() dir_test_count[int(idx)] += 1 total_test_count[int(idx)] += 1 if dis < 0.1: dir_add_count[int(idx)] += 1 total_add_count[int(idx)] += 1 if dis < 0.02: dir_add_02_count[int(idx)] += 1 total_add_02_count[int(idx)] += 1 if dis_s < 0.1: dir_add_s_count[int(idx)] += 1 total_add_s_count[int(idx)] += 1 if dis_s < 0.02: dir_add_s_02_count[int(idx)] += 1 total_add_s_02_count[int(idx)] += 1 if dis_pure < 0.1: dir_add_pure_count[int(idx)] += 1 total_add_pure_count[int(idx)] += 1 if dis_pure < 0.02: dir_add_pure_02_count[int(idx)] += 1 total_add_pure_02_count[int(idx)] += 1 else: last_dis[int(idx)] = None if dis < 0.1: dir_add_count_unseen[int(idx)] += 1 total_add_count_unseen[int(idx)] += 1 total_unseen_objects[int(idx)] += 1 if dis < 0.02: dir_add_02_count_unseen[int(idx)] += 1 total_add_02_count_unseen[int(idx)] += 1 total_unseen_objects[int(idx)] += 1 output_image = output_transformed_image( OUTPUT_IMAGE_PATH, output_image, pred_cloud, focal_length, principal_point, int(idx)) logger.info('Test time {0} Test Frame {1} {2} dis:{3}'.format( time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)), filename, idx.item(), dis)) output_image.save(OUTPUT_IMAGE_PATH) # handle dir output for i in range(0, object_max): if dir_test_count[i] != 0: logger.info( 'Dir {0} Object {1} dis:{2} with {3} samples'.format( lastdatafolder, i, dir_test_dis[i] / dir_test_count[i], dir_test_count[i])) if dir_add_count[i] != 0: logger.info( 'Dir {0} Object {1} add:{2} with 0.02: {3}'.format( lastdatafolder, i, dir_add_count[i] / dir_test_count[i], dir_add_02_count[i] / dir_add_count[i])) else: logger.info( 'Dir {0} Object {1} add:{2} with 0.02: {3}'.format( lastdatafolder, i, dir_add_count[i] / dir_test_count[i], 0)) if dir_add_pure_count[i] != -0: logger.info( 'Dir {0} Object {1} add_pure:{2} with 0.02: {3}'. format( lastdatafolder, i, dir_add_pure_count[i] / dir_test_count[i], dir_add_pure_02_count[i] / dir_add_pure_count[i])) else: logger.info( 'Dir {0} Object {1} add_pure:{2} with 0.02: {3}'. format(lastdatafolder, i, dir_add_pure_count[i] / dir_test_count[i], 0)) if dir_add_s_count[i] != 0: logger.info( 'Dir {0} Object {1} add_s:{2} with 0.02: {3}'.format( lastdatafolder, i, dir_add_s_count[i] / dir_test_count[i], dir_add_s_02_count[i] / dir_add_s_count[i])) else: logger.info( 'Dir {0} Object {1} add_s:{2} with 0.02: {3}'.format( lastdatafolder, i, dir_add_s_count[i] / dir_test_count[i], 0)) logger.info('Dir {0} Object {1} dbd:{2}'.format( lastdatafolder, i, dir_dbd_count[i] / dir_test_count[i])) logger.info('Dir {0} Object {1} drr:{2}'.format( lastdatafolder, i, dir_drr_count[i] / dir_test_count[i])) logger.info('Dir {0} Object {1} ada:{2}'.format( lastdatafolder, i, dir_ada_count[i] / dir_test_count[i])) logger.info('Dir {0} Object {1} distance_1:{2}'.format( lastdatafolder, i, dir_distance_1_count[i] / dir_test_count[i])) dir_dbd = 0. dir_drr = 0. dir_ada = 0. dir_distance_1 = 0. dir_dis = 0. dir_add = 0 dir_add_s = 0 dir_add_pure = 0 dir_add_02 = 0 dir_add_s_02 = 0 dir_add_pure_02 = 0 dir_count = 0 for i in range(object_max): if total_test_count[i] != 0: dir_count += dir_test_count[i] dir_dis += dir_test_dis[i] dir_add += dir_add_count[i] dir_add_pure += dir_add_pure_count[i] dir_add_s += dir_add_s_count[i] dir_add_02 += dir_add_02_count[i] dir_add_pure_02 += dir_add_pure_02_count[i] dir_add_s_02 += dir_add_s_02_count[i] dir_dbd += dir_dbd_count[i] dir_drr += dir_drr_count[i] dir_ada += dir_ada_count[i] dir_distance_1 += dir_distance_1_count[i] dir_test_dis[i] = 0 dir_test_count[i] = 0 dir_add_count[i] = 0 dir_add_pure_count[i] = 0 dir_add_s_count[i] = 0 dir_add_02_count[i] = 0 dir_add_pure_02_count[i] = 0 dir_add_s_02_count[i] = 0 dir_dbd_count[i] = 0 dir_drr_count[i] = 0 dir_ada_count[i] = 0 dir_distance_1_count[i] = 0 logger.info('Dir {0} \'s total dis:{1} with {2} samples'.format( lastdatafolder, dir_dis / dir_count, dir_count)) logger.info('Dir {0} \'s total add:{1} with 0.02: {2}'.format( lastdatafolder, dir_add / dir_count, dir_add_02 / dir_add)) logger.info('Dir {0} \'s total add_s:{1} with 0.02: {2}'.format( lastdatafolder, dir_add_s / dir_count, dir_add_s_02 / dir_add_s)) logger.info('Dir {0} \'s total add_pure:{1} with 0.02: {2}'.format( lastdatafolder, dir_add_pure / dir_count, dir_add_pure_02 / dir_add_pure)) logger.info('Dir {0} \'s total dbd:{1}'.format(lastdatafolder, dir_dbd / dir_count)) logger.info('Dir {0} \'s total drr:{1}'.format(lastdatafolder, dir_drr / dir_count)) logger.info('Dir {0} \'s total ada:{1}'.format(lastdatafolder, dir_ada / dir_count)) logger.info('Dir {0} \'s total distance_1:{1}'.format( lastdatafolder, dir_distance_1 / dir_count)) # end of handle dir output # handle global output total_unseen_count = 0 total_without_pose_count = 0 total_add_count_unseen_count = 0 total_add_02_count_unseen_count = 0 total_drr = 0. total_dbd = 0. total_ada = 0. total_distance_1 = 0. total_dis = 0. total_add = 0 total_add_s = 0 total_add_pure = 0 total_add_02 = 0 total_add_s_02 = 0 total_add_pure_02 = 0 total_count = 0 for i in range(object_max): if total_test_count[i] != 0: logger.info( 'Total: Object {0} dis:{1} with {2} samples'.format( i, total_test_dis[i] / total_test_count[i], total_test_count[i])) logger.info('Total: Object {0} add:{1} with 0.02: {2}'.format( i, total_add_count[i] / total_test_count[i], total_add_02_count[i] / total_add_count[i])) logger.info('Total: Object {0} drr:{1}'.format( i, total_drr_count[i] / total_test_count[i])) logger.info('Total: Object {0} ada:{1}'.format( i, total_ada_count[i] / total_test_count[i])) logger.info('Total: Object {0} distance_1:{1}'.format( i, total_distance_1_count[i] / total_test_count[i])) if total_unseen_objects[i] != 0: if total_unseen_objects[i] - total_object_without_pose[ i] != 0: logger.info( 'Total: Unseen Object {0} add:{1} with 0.02: {2} with {3} samples ' .format( i, total_add_count_unseen[i] / (total_unseen_objects[i] - total_object_without_pose[i]), total_add_02_count_unseen[i] / total_add_count_unseen[i], (total_unseen_objects[i] - total_object_without_pose[i]))) logger.info( 'Total: Object {0} unseen :{1} times, {2} of them without poses, success rate:{3}' .format(i, total_unseen_objects[i], total_object_without_pose[i], (total_unseen_objects[i] - total_object_without_pose[i]) / total_unseen_objects[i])) total_unseen_count += total_unseen_objects[i] total_without_pose_count += total_object_without_pose[i] total_count += total_test_count[i] total_dis += total_test_dis[i] total_add += total_add_count[i] total_add_count_unseen_count += total_add_count_unseen[i] total_add_02_count_unseen_count += total_add_02_count_unseen[i] total_add_s += total_add_s_count[i] total_add_pure += total_add_pure_count[i] total_add_02 += total_add_02_count[i] total_add_s_02 += total_add_s_02_count[i] total_add_pure_02 += total_add_pure_02_count[i] total_dbd += total_dbd_count[i] total_drr += total_drr_count[i] total_ada += total_ada_count[i] total_distance_1 += total_distance_1_count[i] logger.info('total dis:{0} with {1} samples'.format( total_dis / total_count, total_count)) logger.info('total add:{0} with 0.02: {1}'.format( total_add / total_count, total_add_02 / total_add)) logger.info('total unseen add:{0} with 0.02: {1}'.format( total_add_count_unseen_count / (total_unseen_count - total_without_pose_count), total_add_02_count_unseen_count / total_add_count_unseen_count)) logger.info('total add_pure:{0} with 0.02: {1}'.format( total_add_pure / total_count, total_add_pure_02 / total_add_pure)) logger.info('total add_s:{0} with 0.02: {1}'.format( total_add_s / total_count, total_add_s_02 / total_add_s)) logger.info( 'detected unseen object :{0}, failed calculate {1} poses with success rate: {2}' .format(total_unseen_count, total_without_pose_count, (total_unseen_count - total_without_pose_count) / total_unseen_count)) logger.info('Total drr:{0}'.format(total_drr / total_count)) logger.info('Total ada:{0}'.format(total_ada / total_count)) logger.info('Total distance_1:{0}'.format(total_distance_1 / total_count))
def main(): if opt.dataset == 'linemod': opt.num_obj = 1 opt.list_obj = [1, 2, 4, 5, 6, 8, 9, 10, 11, 12, 13, 14, 15] opt.occ_list_obj = [1, 5, 6, 8, 9, 10, 11, 12] opt.list_name = ['ape', 'benchvise', 'cam', 'can', 'cat', 'driller', 'duck', 'eggbox', 'glue', 'holepuncher', 'iron', 'lamp', 'phone'] obj_name = opt.list_name[opt.list_obj.index(opt.obj_id)] opt.sym_list = [10, 11] opt.num_points = 500 meta_file = open('{0}/models/models_info.yml'.format(opt.dataset_root), 'r') meta = yaml.load(meta_file) diameter = meta[opt.obj_id]['diameter'] / 1000.0 * 0.1 if opt.render: opt.repeat_num = 1 elif opt.fuse: opt.repeat_num = 1 else: opt.repeat_num = 5 writer = SummaryWriter('experiments/runs/linemod/{}{}'.format(obj_name, opt.experiment_name)) opt.outf = 'trained_models/linemod/{}{}'.format(obj_name, opt.experiment_name) opt.log_dir = 'experiments/logs/linemod/{}{}'.format(obj_name, opt.experiment_name) if not os.path.exists(opt.outf): os.mkdir(opt.outf) if not os.path.exists(opt.log_dir): os.mkdir(opt.log_dir) else: print('Unknown dataset') return estimator = PoseNet(num_points = opt.num_points, num_vote = 9, num_obj = opt.num_obj) estimator.cuda() refiner = PoseRefineNet(num_points = opt.num_points, num_obj = opt.num_obj) refiner.cuda() if opt.resume_posenet != '': estimator.load_state_dict(torch.load('{0}/{1}'.format(opt.outf, opt.resume_posenet))) if opt.resume_refinenet != '': refiner.load_state_dict(torch.load('{0}/{1}'.format(opt.outf, opt.resume_refinenet))) opt.refine_start = True opt.lr = opt.lr_refine opt.batch_size = int(opt.batch_size / opt.iteration) optimizer = optim.Adam(refiner.parameters(), lr=opt.lr) else: opt.refine_start = False optimizer = optim.Adam(estimator.parameters(), lr=opt.lr) dataset = PoseDataset_linemod('train', opt.num_points, opt.dataset_root, opt.real, opt.render, opt.fuse, opt.obj_id) dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True, num_workers=opt.workers) test_dataset = PoseDataset_linemod('test', opt.num_points, opt.dataset_root, True, False, False, opt.obj_id) testdataloader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=opt.workers) print('>>>>>>>>----------Dataset loaded!---------<<<<<<<<\nlength of the training set: {0}\nlength of the testing set: {1}\nnumber of sample points on mesh: {2}'.format(len(dataset), len(test_dataset), opt.num_points)) if opt.obj_id in opt.occ_list_obj: occ_test_dataset = PoseDataset_occ('test', opt.num_points, opt.occ_dataset_root, opt.obj_id) occtestdataloader = torch.utils.data.DataLoader(occ_test_dataset, batch_size=1, shuffle=False, num_workers=opt.workers) print('length of the occ testing set: {}'.format(len(occ_test_dataset))) criterion = Loss(opt.num_points, opt.sym_list) criterion_refine = Loss_refine(opt.num_points, opt.sym_list) best_test = np.Inf if opt.start_epoch == 1: for log in os.listdir(opt.log_dir): os.remove(os.path.join(opt.log_dir, log)) st_time = time.time() train_scalar = 0 for epoch in range(opt.start_epoch, opt.nepoch): logger = setup_logger('epoch%d' % epoch, os.path.join(opt.log_dir, '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_loss_avg = 0.0 train_loss = 0.0 train_dis_avg = 0.0 train_dis = 0.0 if opt.refine_start: estimator.eval() refiner.train() else: estimator.train() optimizer.zero_grad() for rep in range(opt.repeat_num): for i, data in enumerate(dataloader, 0): points, choose, img, target, model_points, model_kp, vertex_gt, idx, target_r, target_t = data if len(points.size()) == 2: print('pass') continue points, choose, img, target, model_points, model_kp, vertex_gt, idx, target_r, target_t = points.cuda(), choose.cuda(), img.cuda(), target.cuda(), model_points.cuda(), model_kp.cuda(), vertex_gt.cuda(), idx.cuda(), target_r.cuda(), target_t.cuda() vertex_pred, c_pred, emb = estimator(img, points, choose, idx) vertex_loss, pose_loss, dis, new_points, new_target = criterion(vertex_pred, vertex_gt, c_pred, points, target, model_points, model_kp, opt.obj_id, target_r, target_t) loss = 10 * vertex_loss + pose_loss if opt.refine_start: for ite in range(0, opt.iteration): pred_r, pred_t = refiner(new_points, emb, idx) dis, new_points, new_target = criterion_refine(pred_r, pred_t, new_points, new_target, model_points, opt.obj_id) dis.backward() else: loss.backward() train_loss_avg += loss.item() train_loss += loss.item() train_dis_avg += dis.item() train_dis += dis.item() train_count += 1 train_scalar += 1 if train_count % opt.batch_size == 0: logger.info('Train time {0} Epoch {1} Batch {2} Frame {3} Avg_loss:{4} Avg_diss:{5}'.format(time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)), epoch, int(train_count / opt.batch_size), train_count, train_loss_avg / opt.batch_size, train_dis_avg / opt.batch_size)) writer.add_scalar('linemod training loss', train_loss_avg / opt.batch_size, train_scalar) writer.add_scalar('linemod training dis', train_dis_avg / opt.batch_size, train_scalar) optimizer.step() optimizer.zero_grad() train_loss_avg = 0 train_dis_avg = 0 if train_count != 0 and train_count % 1000 == 0: if opt.refine_start: torch.save(refiner.state_dict(), '{0}/pose_refine_model_current.pth'.format(opt.outf)) else: torch.save(estimator.state_dict(), '{0}/pose_model_current.pth'.format(opt.outf)) print('>>>>>>>>----------epoch {0} train finish---------<<<<<<<<'.format(epoch)) train_loss = train_loss / train_count train_dis = train_dis / train_count logger.info('Train time {0} Epoch {1} TRAIN FINISH Avg loss: {2} Avg dis: {3}'.format(time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)), epoch, train_loss, train_dis)) logger = setup_logger('epoch%d_test' % epoch, os.path.join(opt.log_dir, '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_loss = 0.0 test_vertex_loss = 0.0 test_pose_loss = 0.0 test_dis = 0.0 test_count = 0 success_count = 0 estimator.eval() refiner.eval() for j, data in enumerate(testdataloader, 0): points, choose, img, target, model_points, model_kp, vertex_gt, idx, target_r, target_t = data if len(points.size()) == 2: logger.info('Test time {0} Lost detection!'.format(time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)))) continue points, choose, img, target, model_points, model_kp, vertex_gt, idx, target_r, target_t = points.cuda(), choose.cuda(), img.cuda(), target.cuda(), model_points.cuda(), model_kp.cuda(), vertex_gt.cuda(), idx.cuda(), target_r.cuda(), target_t.cuda() vertex_pred, c_pred, emb = estimator(img, points, choose, idx) vertex_loss, pose_loss, dis, new_points, new_target = criterion(vertex_pred, vertex_gt, c_pred, points, target, model_points, model_kp, opt.obj_id, target_r, target_t) loss = 10 * vertex_loss + pose_loss if opt.refine_start: for ite in range(0, opt.iteration): pred_r, pred_t = refiner(new_points, emb, idx) dis, new_points, new_target = criterion_refine(pred_r, pred_t, new_points, new_target, model_points, opt.obj_id) test_loss += loss.item() test_vertex_loss += vertex_loss.item() test_pose_loss += pose_loss.item() test_dis += dis.item() logger.info('Test time {0} Test Frame No.{1} loss:{2} dis:{3}'.format(time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)), test_count, loss, dis)) if dis.item() < diameter: success_count += 1 test_count += 1 test_loss = test_loss / test_count test_vertex_loss = test_vertex_loss / test_count test_pose_loss = test_pose_loss / test_count test_dis = test_dis / test_count success_rate = float(success_count) / test_count logger.info('Test time {0} Epoch {1} TEST FINISH Avg loss: {2} Avg dis: {3} Success rate: {4}'.format(time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)), epoch, test_loss, test_dis, success_rate)) writer.add_scalar('linemod test loss', test_loss, epoch) writer.add_scalar('linemod test vertex loss', test_vertex_loss, epoch) writer.add_scalar('linemod test pose loss', test_pose_loss, epoch) writer.add_scalar('linemod test dis', test_dis, epoch) writer.add_scalar('linemod success rate', success_rate, epoch) writer.add_scalar('lr', optimizer.param_groups[0]['lr'], epoch) if test_dis <= best_test: best_test = test_dis if opt.refine_start: torch.save(refiner.state_dict(), '{0}/pose_refine_model_{1}_{2}.pth'.format(opt.outf, epoch, test_dis)) else: torch.save(estimator.state_dict(), '{0}/pose_model_{1}_{2}.pth'.format(opt.outf, epoch, test_dis)) print(epoch, '>>>>>>>>----------MODEL SAVED---------<<<<<<<<') if opt.obj_id in opt.occ_list_obj: logger = setup_logger('epoch%d_occ_test' % epoch, os.path.join(opt.log_dir, 'epoch_%d_occ_test_log.txt' % epoch)) logger.info('Occ test time {0}'.format(time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)) + ', ' + 'Testing started')) occ_test_dis = 0.0 occ_test_count = 0 occ_success_count = 0 estimator.eval() refiner.eval() for j, data in enumerate(occtestdataloader, 0): points, choose, img, target, model_points, model_kp, vertex_gt, idx, target_r, target_t = data if len(points.size()) == 2: logger.info('Occ test time {0} Lost detection!'.format(time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)))) continue points, choose, img, target, model_points, model_kp, vertex_gt, idx, target_r, target_t = points.cuda(), choose.cuda(), img.cuda(), target.cuda(), model_points.cuda(), model_kp.cuda(), vertex_gt.cuda(), idx.cuda(), target_r.cuda(), target_t.cuda() vertex_pred, c_pred, emb = estimator(img, points, choose, idx) vertex_loss, pose_loss, dis, new_points, new_target = criterion(vertex_pred, vertex_gt, c_pred, points, target, model_points, model_kp, opt.obj_id, target_r, target_t) if opt.refine_start: for ite in range(0, opt.iteration): pred_r, pred_t = refiner(new_points, emb, idx) dis, new_points, new_target = criterion_refine(pred_r, pred_t, new_points, new_target, model_points, opt.obj_id) occ_test_dis += dis.item() logger.info('Occ test time {0} Test Frame No.{1} dis:{2}'.format(time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)), occ_test_count, dis)) if dis.item() < diameter: occ_success_count += 1 occ_test_count += 1 occ_test_dis = occ_test_dis / occ_test_count occ_success_rate = float(occ_success_count) / occ_test_count logger.info('Occ test time {0} Epoch {1} TEST FINISH Avg dis: {2} Success rate: {3}'.format(time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)), epoch, occ_test_dis, occ_success_rate)) writer.add_scalar('occ test dis', occ_test_dis, epoch) writer.add_scalar('occ success rate', occ_success_rate, epoch) if best_test < opt.refine_margin and not opt.refine_start: opt.refine_start = True opt.lr = opt.lr_refine opt.batch_size = int(opt.batch_size / opt.iteration) optimizer = optim.Adam(refiner.parameters(), lr=opt.lr) print('>>>>>>>>----------Refine started---------<<<<<<<<') writer.close()
def main(): # g13: parameter setting ------------------- ''' posemodel is trained_checkpoints/linemod/pose_model_9_0.01310166542980859.pth refine model is trained_checkpoints/linemod/pose_refine_model_493_0.006761023565178073.pth ''' objlist = [1, 2, 4, 5, 6, 8, 9, 10, 11, 12, 13, 14, 15] knn = KNearestNeighbor(1) opt.dataset ='linemod' opt.dataset_root = './datasets/linemod/Linemod_preprocessed' estimator_path = 'trained_checkpoints/linemod/pose_model_9_0.01310166542980859.pth' refiner_path = 'trained_checkpoints/linemod/pose_refine_model_493_0.006761023565178073.pth' opt.model = estimator_path opt.refine_model = refiner_path dataset_config_dir = 'datasets/linemod/dataset_config' output_result_dir = 'experiments/eval_result/linemod' opt.refine_start = True bs = 1 #fixed because of the default setting in torch.utils.data.DataLoader opt.iteration = 2 #default is 4 in eval_linemod.py t1_start = True t1_idx = 0 t1_total_eval_num = 3 t2_start = False t2_target_list = [22, 30, 172, 187, 267, 363, 410, 471, 472, 605, 644, 712, 1046, 1116, 1129, 1135, 1263] #t2_target_list = [0, 1] axis_range = 0.1 # the length of X, Y, and Z axis in 3D vimg_dir = 'verify_img' diameter = [] meta_file = open('{0}/models_info.yml'.format(dataset_config_dir), 'r') meta_d = yaml.load(meta_file) for obj in objlist: diameter.append(meta_d[obj]['diameter'] / 1000.0 * 0.1) print(diameter) if not os.path.exists(vimg_dir): os.makedirs(vimg_dir) #------------------------------------------- if opt.dataset == 'ycb': 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 = '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 = 'experiments/logs/linemod' opt.repeat_epoch = 20 else: print('Unknown dataset') return estimator = PoseNet(num_points = opt.num_points, num_obj = opt.num_objects) estimator.cuda() refiner = PoseRefineNet(num_points = opt.num_points, num_obj = opt.num_objects) refiner.cuda() estimator.load_state_dict(torch.load(estimator_path)) refiner.load_state_dict(torch.load(refiner_path)) opt.refine_start = True test_dataset = PoseDataset_linemod('test', opt.num_points, False, opt.dataset_root, 0.0, opt.refine_start) testdataloader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=opt.workers) opt.sym_list = test_dataset.get_sym_list() opt.num_points_mesh = test_dataset.get_num_points_mesh() print('>>>>>>>>----------Dataset loaded!---------<<<<<<<<\n\ length of the testing set: {0}\nnumber of sample points on mesh: {1}\n\ symmetry object list: {2}'\ .format( len(test_dataset), opt.num_points_mesh, opt.sym_list)) #load pytorch model estimator.eval() refiner.eval() criterion = Loss(opt.num_points_mesh, opt.sym_list) criterion_refine = Loss_refine(opt.num_points_mesh, opt.sym_list) fw = open('{0}/t1_eval_result_logs.txt'.format(output_result_dir), 'w') #Pose estimation for j, data in enumerate(testdataloader, 0): # g13: modify this part for evaluation target-------------------- if t1_start and j == t1_total_eval_num: break if t2_start and not (j in t2_target_list): continue #---------------------------------------------------------------- points, choose, img, target, model_points, idx = data if len(points.size()) == 2: print('No.{0} NOT Pass! Lost detection!'.format(j)) fw.write('No.{0} NOT Pass! Lost detection!\n'.format(j)) continue points, choose, img, target, model_points, idx = Variable(points).cuda(), \ Variable(choose).cuda(), \ Variable(img).cuda(), \ Variable(target).cuda(), \ Variable(model_points).cuda(), \ Variable(idx).cuda() pred_r, pred_t, pred_c, emb = estimator(img, points, choose, idx) _, dis, new_points, new_target = criterion(pred_r, pred_t, pred_c, target, model_points, idx, points, opt.w, opt.refine_start) #if opt.refine_start: #iterative poserefinement # for ite in range(0, opt.iteration): # pred_r, pred_t = refiner(new_points, emb, idx) # dis, new_points, new_target = criterion_refine(pred_r, pred_t, new_target, model_points, idx, new_points) pred_r = pred_r / torch.norm(pred_r, dim=2).view(1, opt.num_points, 1) pred_c = pred_c.view(bs, opt.num_points) how_max, which_max = torch.max(pred_c, 1) pred_t = pred_t.view(bs * opt.num_points, 1, 3) my_r = pred_r[0][which_max[0]].view(-1).cpu().data.numpy() my_t = (points.view(bs * opt.num_points, 1, 3) + pred_t)[which_max[0]].view(-1).cpu().data.numpy() my_pred = np.append(my_r, my_t) for ite in range(0, opt.iteration): T = Variable(torch.from_numpy(my_t.astype(np.float32))).cuda().view(1, 3).repeat(opt.num_points, 1).contiguous().view(1, opt.num_points, 3) my_mat = quaternion_matrix(my_r) R = Variable(torch.from_numpy(my_mat[:3, :3].astype(np.float32))).cuda().view(1, 3, 3) my_mat[0:3, 3] = my_t new_points = torch.bmm((points - T), R).contiguous() pred_r, pred_t = refiner(new_points, emb, idx) pred_r = pred_r.view(1, 1, -1) pred_r = pred_r / (torch.norm(pred_r, dim=2).view(1, 1, 1)) my_r_2 = pred_r.view(-1).cpu().data.numpy() my_t_2 = pred_t.view(-1).cpu().data.numpy() my_mat_2 = quaternion_matrix(my_r_2) my_mat_2[0:3, 3] = my_t_2 my_mat_final = np.dot(my_mat, my_mat_2) my_r_final = copy.deepcopy(my_mat_final) my_r_final[0:3, 3] = 0 my_r_final = quaternion_from_matrix(my_r_final, True) my_t_final = np.array([my_mat_final[0][3], my_mat_final[1][3], my_mat_final[2][3]]) my_pred = np.append(my_r_final, my_t_final) my_r = my_r_final my_t = my_t_final # Here 'my_pred' is the final pose estimation result after refinement ('my_r': quaternion, 'my_t': translation) #g13: checking the dis value success_count = [0 for i in range(opt.num_objects)] num_count = [0 for i in range(opt.num_objects)] model_points = model_points[0].cpu().detach().numpy() my_r = quaternion_matrix(my_r)[:3, :3] pred = np.dot(model_points, my_r.T) + my_t target = target[0].cpu().detach().numpy() if idx[0].item() in opt.sym_list: pred = torch.from_numpy(pred.astype(np.float32)).cuda().transpose(1, 0).contiguous() target = torch.from_numpy(target.astype(np.float32)).cuda().transpose(1, 0).contiguous() inds = knn(target.unsqueeze(0), pred.unsqueeze(0)) target = torch.index_select(target, 1, inds.view(-1) - 1) dis = torch.mean(torch.norm((pred.transpose(1, 0) - target.transpose(1, 0)), dim=1), dim=0).item() else: dis = np.mean(np.linalg.norm(pred - target, axis=1)) if dis < diameter[idx[0].item()]: success_count[idx[0].item()] += 1 print('No.{0} Pass! Distance: {1}'.format(j, dis)) fw.write('No.{0} Pass! Distance: {1}\n'.format(j, dis)) else: print('No.{0} NOT Pass! Distance: {1}'.format(j, dis)) fw.write('No.{0} NOT Pass! Distance: {1}\n'.format(j, dis)) num_count[idx[0].item()] += 1 # g13: start drawing pose on image------------------------------------ # pick up image print('{0}:\nmy_r is {1}\nmy_t is {2}\ndis:{3}'.format(j, my_r, my_t, dis.item())) print("index {0}: {1}".format(j, test_dataset.list_rgb[j])) img = Image.open(test_dataset.list_rgb[j]) # pick up center position by bbox meta_file = open('{0}/data/{1}/gt.yml'.format(opt.dataset_root, '%02d' % test_dataset.list_obj[j]), 'r') meta = {} meta = yaml.load(meta_file) which_item = test_dataset.list_rank[j] which_obj = test_dataset.list_obj[j] which_dict = 0 dict_leng = len(meta[which_item]) #print('get meta[{0}][{1}][obj_bb]'.format(which_item, which_obj)) k_idx = 0 while 1: if meta[which_item][k_idx]['obj_id'] == which_obj: which_dict = k_idx break k_idx = k_idx+1 bbx = meta[which_item][which_dict]['obj_bb'] draw = ImageDraw.Draw(img) # draw box (ensure this is the right object) draw.line((bbx[0],bbx[1], bbx[0], bbx[1]+bbx[3]), fill=(255,0,0), width=5) draw.line((bbx[0],bbx[1], bbx[0]+bbx[2], bbx[1]), fill=(255,0,0), width=5) draw.line((bbx[0],bbx[1]+bbx[3], bbx[0]+bbx[2], bbx[1]+bbx[3]), fill=(255,0,0), width=5) draw.line((bbx[0]+bbx[2],bbx[1], bbx[0]+bbx[2], bbx[1]+bbx[3]), fill=(255,0,0), width=5) #get center c_x = bbx[0]+int(bbx[2]/2) c_y = bbx[1]+int(bbx[3]/2) draw.point((c_x,c_y), fill=(255,255,0)) print('center:({0},{1})'.format(c_x, c_y)) #get the 3D position of center cam_intrinsic = np.zeros((3,3)) cam_intrinsic.itemset(0, test_dataset.cam_fx) cam_intrinsic.itemset(4, test_dataset.cam_fy) cam_intrinsic.itemset(2, test_dataset.cam_cx) cam_intrinsic.itemset(5, test_dataset.cam_cy) cam_intrinsic.itemset(8, 1) cam_extrinsic = my_mat_final[0:3, :] cam2d_3d = np.matmul(cam_intrinsic, cam_extrinsic) cen_3d = np.matmul(np.linalg.pinv(cam2d_3d), [[c_x],[c_y],[1]]) # replace img.show() with plt.imshow(img) #transpose three 3D axis point into 2D x_3d = cen_3d + [[axis_range],[0],[0],[0]] y_3d = cen_3d + [[0],[axis_range],[0],[0]] z_3d = cen_3d + [[0],[0],[axis_range],[0]] x_2d = np.matmul(cam2d_3d, x_3d) y_2d = np.matmul(cam2d_3d, y_3d) z_2d = np.matmul(cam2d_3d, z_3d) #draw the axis on 2D draw.line((c_x, c_y, x_2d[0], x_2d[1]), fill=(255,255,0), width=5) draw.line((c_x, c_y, y_2d[0], y_2d[1]), fill=(0,255,0), width=5) draw.line((c_x, c_y, z_2d[0], z_2d[1]), fill=(0,0,255), width=5) #g13: draw the estimate pred obj for pti in pred: pti.transpose() pti_2d = np.matmul(cam_intrinsic, pti) #print('({0},{1})\n'.format(int(pti_2d[0]),int(pti_2d[1]))) draw.point([int(pti_2d[0]),int(pti_2d[1])], fill=(255,255,0)) #g13: show image #img.show() #save file under file img_file_name = '{0}/batch{1}_pred_obj{2}_pic{3}.png'.format(vimg_dir, j, test_dataset.list_obj[j], which_item) img.save( img_file_name, "PNG" ) img.close() # plot ground true ---------------------------- img = Image.open(test_dataset.list_rgb[j]) draw = ImageDraw.Draw(img) draw.line((bbx[0],bbx[1], bbx[0], bbx[1]+bbx[3]), fill=(255,0,0), width=5) draw.line((bbx[0],bbx[1], bbx[0]+bbx[2], bbx[1]), fill=(255,0,0), width=5) draw.line((bbx[0],bbx[1]+bbx[3], bbx[0]+bbx[2], bbx[1]+bbx[3]), fill=(255,0,0), width=5) draw.line((bbx[0]+bbx[2],bbx[1], bbx[0]+bbx[2], bbx[1]+bbx[3]), fill=(255,0,0), width=5) target_r = np.resize(np.array(meta[which_item][k_idx]['cam_R_m2c']), (3, 3)) target_t = np.array(meta[which_item][k_idx]['cam_t_m2c']) target_t = target_t[np.newaxis, :] cam_extrinsic_GT = np.concatenate((target_r, target_t.T), axis=1) #get center 3D cam2d_3d_GT = np.matmul(cam_intrinsic, cam_extrinsic_GT) cen_3d_GT = np.matmul(np.linalg.pinv(cam2d_3d_GT), [[c_x],[c_y],[1]]) #transpose three 3D axis point into 2D x_3d = cen_3d_GT + [[axis_range],[0],[0],[0]] y_3d = cen_3d_GT + [[0],[axis_range],[0],[0]] z_3d = cen_3d_GT + [[0],[0],[axis_range],[0]] x_2d = np.matmul(cam2d_3d_GT, x_3d) y_2d = np.matmul(cam2d_3d_GT, y_3d) z_2d = np.matmul(cam2d_3d_GT, z_3d) #draw the axis on 2D draw.line((c_x, c_y, x_2d[0], x_2d[1]), fill=(255,255,0), width=5) draw.line((c_x, c_y, y_2d[0], y_2d[1]), fill=(0,255,0), width=5) draw.line((c_x, c_y, z_2d[0], z_2d[1]), fill=(0,0,255), width=5) print('pred:\n{0}\nGT:\n{1}\n'.format(cam_extrinsic,cam_extrinsic_GT)) print('pred 3D:{0}\nGT 3D:{1}\n'.format(cen_3d, cen_3d_GT)) img_file_name = '{0}/batch{1}_pred_obj{2}_pic{3}_gt.png'.format(vimg_dir, j, test_dataset.list_obj[j], which_item) img.save( img_file_name, "PNG" ) img.close() meta_file.close() print('\nplot_result_img.py completed the task\n')
refiner.cuda() estimator.load_state_dict(torch.load(opt.model)) refiner.load_state_dict(torch.load(opt.refine_model)) estimator.eval() refiner.eval() testdataset = PoseDataset_linemod('test', num_points, False, opt.dataset_root, 0.0, True) testdataloader = torch.utils.data.DataLoader(testdataset, batch_size=1, shuffle=False, num_workers=8) sym_list = testdataset.get_sym_list() num_points_mesh = testdataset.get_num_points_mesh() criterion = Loss(num_points_mesh, sym_list) criterion_refine = Loss_refine(num_points_mesh, sym_list) diameter = [] meta_file = open('{0}/models_info.yml'.format(dataset_config_dir), 'r') meta = yaml.load(meta_file) for obj in objlist: diameter.append(meta[obj]['diameter'] / 1000.0 * 0.1) print(diameter) success_count = [0 for i in range(num_objects)] success_count_cpy = [0 for i in range(num_objects)] num_count = [0 for i in range(num_objects)] fw = open('{0}/eval_result_logs_ICP_DEL2.txt'.format(output_result_dir), 'w') import time
def main(): opt.manualSeed = random.randint(1, 10000) random.seed(opt.manualSeed) torch.manual_seed(opt.manualSeed) opt.num_objects = 3 opt.num_points = 500 opt.outf = 'trained_models' opt.log_dir = 'experiments/logs' opt.repeat_epoch = 20 estimator = PoseNet(num_points=opt.num_points, num_obj=opt.num_objects) estimator.cuda() refiner = PoseRefineNet(num_points=opt.num_points, num_obj=opt.num_objects) refiner.cuda() if opt.resume_posenet != '': estimator.load_state_dict( torch.load('{0}/{1}'.format(opt.outf, opt.resume_posenet))) if opt.resume_refinenet != '': refiner.load_state_dict( torch.load('{0}/{1}'.format(opt.outf, opt.resume_refinenet))) opt.refine_start = True opt.decay_start = True opt.lr *= opt.lr_rate opt.w *= opt.w_rate opt.batch_size = int(opt.batch_size / opt.iteration) optimizer = optim.Adam(refiner.parameters(), lr=opt.lr) else: opt.refine_start = False opt.decay_start = False optimizer = optim.Adam(estimator.parameters(), lr=opt.lr) dataset = PoseDataset('train', opt.num_points, True, opt.dataset_root, opt.noise_trans, opt.refine_start) dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True, num_workers=opt.workers) test_dataset = PoseDataset('test', opt.num_points, False, opt.dataset_root, 0.0, opt.refine_start) testdataloader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=opt.workers) opt.sym_list = dataset.get_sym_list() opt.num_points_mesh = 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}' .format(len(dataset), len(test_dataset), opt.num_points_mesh)) criterion = Loss(opt.num_points_mesh, opt.sym_list) criterion_refine = Loss_refine(opt.num_points_mesh, opt.sym_list) best_test = np.Inf if opt.start_epoch == 1: for log in os.listdir(opt.log_dir): os.remove(os.path.join(opt.log_dir, log)) st_time = time.time() for epoch in range(opt.start_epoch, opt.nepoch): logger = setup_logger( 'epoch%d' % epoch, os.path.join(opt.log_dir, '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_avg = 0.0 if opt.refine_start: estimator.eval() # affects dropout and batch normalization refiner.train() else: estimator.train() optimizer.zero_grad() for rep in range(opt.repeat_epoch): for i, data in enumerate(dataloader, 0): points, choose, img, target, model_points, idx = data #points ->torch.Size([500, 3]) ->在crop出来的像素区域随机选取500个点,利用相机内参结合深度值算出来的点云cloud #choose ->torch.Size([1, 500]) #img ->torch.Size([3, 80, 80]) #target ->torch.Size([500, 3]) ->真实模型上随机选取的mesh点进行ground truth pose变换后得到的点 #model_points ->torch.Size([500, 3]) ->真实模型上随机选取的mesh点在进行pose变换前的点 #idx ->torch.Size([1]) #tensor([4], device='cuda:0') #img和points对应rgb和点云信息,需要在网络内部fusion points, choose, img, target, model_points, idx = Variable(points).cuda(), \ Variable(choose).cuda(), \ Variable(img).cuda(), \ Variable(target).cuda(), \ Variable(model_points).cuda(), \ Variable(idx).cuda() pred_r, pred_t, pred_c, emb = estimator( img, points, choose, idx) loss, dis, new_points, new_target = criterion( pred_r, pred_t, pred_c, target, model_points, idx, points, opt.w, opt.refine_start) if opt.refine_start: for ite in range(0, opt.iteration): pred_r, pred_t = refiner(new_points, emb, idx) dis, new_points, new_target = criterion_refine( pred_r, pred_t, new_target, model_points, idx, new_points) dis.backward() else: loss.backward() train_dis_avg += dis.item() train_count += 1 if train_count % opt.batch_size == 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.batch_size), train_count, train_dis_avg / opt.batch_size)) optimizer.step() optimizer.zero_grad() train_dis_avg = 0 if train_count != 0 and train_count % 1000 == 0: if opt.refine_start: torch.save( refiner.state_dict(), '{0}/pose_refine_model_current.pth'.format( opt.outf)) else: torch.save( estimator.state_dict(), '{0}/pose_model_current.pth'.format(opt.outf)) print( '>>>>>>>>----------epoch {0} train finish---------<<<<<<<<'.format( epoch)) logger = setup_logger( 'epoch%d_test' % epoch, os.path.join(opt.log_dir, '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() refiner.eval() for j, data in enumerate(testdataloader, 0): points, choose, img, target, model_points, idx = data points, choose, img, target, model_points, idx = Variable(points).cuda(), \ Variable(choose).cuda(), \ Variable(img).cuda(), \ Variable(target).cuda(), \ Variable(model_points).cuda(), \ Variable(idx).cuda() pred_r, pred_t, pred_c, emb = estimator(img, points, choose, idx) _, dis, new_points, new_target = criterion(pred_r, pred_t, pred_c, target, model_points, idx, points, opt.w, opt.refine_start) if opt.refine_start: for ite in range(0, opt.iteration): pred_r, pred_t = refiner(new_points, emb, idx) dis, new_points, new_target = criterion_refine( pred_r, pred_t, new_target, model_points, idx, new_points) test_dis += dis.item() 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_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_test: best_test = test_dis if opt.refine_start: torch.save( refiner.state_dict(), '{0}/pose_refine_model_{1}_{2}.pth'.format( opt.outf, epoch, test_dis)) else: torch.save( estimator.state_dict(), '{0}/pose_model_{1}_{2}.pth'.format( opt.outf, epoch, test_dis)) print(epoch, '>>>>>>>>----------BEST TEST MODEL SAVED---------<<<<<<<<') if best_test < opt.decay_margin and not opt.decay_start: opt.decay_start = True opt.lr *= opt.lr_rate opt.w *= opt.w_rate optimizer = optim.Adam(estimator.parameters(), lr=opt.lr) if best_test < opt.refine_margin and not opt.refine_start: opt.refine_start = True opt.batch_size = int(opt.batch_size / opt.iteration) optimizer = optim.Adam(refiner.parameters(), lr=opt.lr) dataset = PoseDataset('train', opt.num_points, True, opt.dataset_root, opt.noise_trans, opt.refine_start) dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True, num_workers=opt.workers) test_dataset = PoseDataset('test', opt.num_points, False, opt.dataset_root, 0.0, opt.refine_start) testdataloader = torch.utils.data.DataLoader( test_dataset, batch_size=1, shuffle=False, num_workers=opt.workers) opt.sym_list = dataset.get_sym_list() opt.num_points_mesh = 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}' .format(len(dataset), len(test_dataset), opt.num_points_mesh)) criterion = Loss(opt.num_points_mesh, opt.sym_list) criterion_refine = Loss_refine(opt.num_points_mesh, opt.sym_list)
estimator.load_state_dict(torch.load(opt.model)) estimator.eval() # data testdataset = PoseDataset_rbo('test', num_points, False, opt.dataset_root, 0.0, True) sym_list = testdataset.get_sym_list() num_points_mesh = testdataset.get_num_points_mesh() #>>>>>>>>>>>>>>>>> how to get the diameter for each objects?? <<<<<<<<<<<<<<<<<< diameter = [1, 1, 1] success_count = [0 for i in range(num_objects)] num_count = [0 for i in range(num_objects)] fw = open('{0}/eval_result_logs.txt'.format(output_result_dir), 'w') opt.num_points = 500 criterion = Loss(500, [7, 8]) testdataloader = torch.utils.data.DataLoader(testdataset, batch_size=1, shuffle=False, num_workers=1) # for i, data in enumerate(testdataloader, 0): index = 5 data = testdataset.__getitem__(5) img, points, cloud_canon, model_points, choose, mask, num_parts, idx = data img = img.unsqueeze(0) points = points.unsqueeze(0) cloud_canon = cloud_canon.unsqueeze(0) points = Variable(points).cuda(0) choose = choose.view(-1, 1, opt.num_points) choose = Variable(choose).cuda(0) img = Variable(img).cuda(0)
def main(): opt.manualSeed = random.randint(1, 10000) random.seed(opt.manualSeed) torch.manual_seed(opt.manualSeed) if opt.dataset == 'ycb': 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 = '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 = 'experiments/logs/linemod' opt.repeat_epoch = 20 else: print('Unknown dataset') return estimator = PoseNet(num_points=opt.num_points, num_obj=opt.num_objects) estimator.cuda() refiner = PoseRefineNet(num_points=opt.num_points, num_obj=opt.num_objects) refiner.cuda() if opt.resume_posenet != '': estimator.load_state_dict( torch.load('{0}/{1}'.format(opt.outf, opt.resume_posenet))) if opt.resume_refinenet != '': refiner.load_state_dict( torch.load('{0}/{1}'.format(opt.outf, opt.resume_refinenet))) opt.refine_start = True opt.decay_start = True opt.lr *= opt.lr_rate opt.w *= opt.w_rate opt.batch_size = int(opt.batch_size / opt.iteration) optimizer = optim.Adam(refiner.parameters(), lr=opt.lr) else: opt.refine_start = False opt.decay_start = False optimizer = optim.Adam(estimator.parameters(), lr=opt.lr) if opt.dataset == 'ycb': dataset = PoseDataset_ycb('train', opt.num_points, True, opt.dataset_root, opt.noise_trans, opt.refine_start) elif opt.dataset == 'linemod': dataset = PoseDataset_linemod('train', opt.num_points, True, opt.dataset_root, opt.noise_trans, opt.refine_start) dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True, num_workers=opt.workers) if opt.dataset == 'ycb': test_dataset = PoseDataset_ycb('test', opt.num_points, False, opt.dataset_root, 0.0, opt.refine_start) elif opt.dataset == 'linemod': test_dataset = PoseDataset_linemod('test', opt.num_points, False, opt.dataset_root, 0.0, opt.refine_start) testdataloader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=opt.workers) opt.sym_list = dataset.get_sym_list() opt.num_points_mesh = 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.num_points_mesh, opt.sym_list)) criterion = Loss(opt.num_points_mesh, opt.sym_list) criterion_refine = Loss_refine(opt.num_points_mesh, opt.sym_list) best_test = np.Inf if opt.start_epoch == 1: for log in os.listdir(opt.log_dir): os.remove(os.path.join(opt.log_dir, log)) st_time = time.time() for epoch in range(opt.start_epoch, opt.nepoch): logger = setup_logger( 'epoch%d' % epoch, os.path.join(opt.log_dir, '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_avg = 0.0 if opt.refine_start: estimator.eval() refiner.train() else: estimator.train() optimizer.zero_grad() for rep in range(opt.repeat_epoch): for i, data in enumerate(dataloader, 0): points, choose, img, target, model_points, idx = data points, choose, img, target, model_points, idx = Variable(points).cuda(), \ Variable(choose).cuda(), \ Variable(img).cuda(), \ Variable(target).cuda(), \ Variable(model_points).cuda(), \ Variable(idx).cuda() pred_r, pred_t, pred_c, emb = estimator( img, points, choose, idx) loss, dis, new_points, new_target = criterion( pred_r, pred_t, pred_c, target, model_points, idx, points, opt.w, opt.refine_start) if opt.refine_start: for ite in range(0, opt.iteration): pred_r, pred_t = refiner(new_points, emb, idx) dis, new_points, new_target = criterion_refine( pred_r, pred_t, new_target, model_points, idx, new_points) dis.backward() else: loss.backward() train_dis_avg += dis.item() train_count += 1 if train_count % opt.batch_size == 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.batch_size), train_count, train_dis_avg / opt.batch_size)) optimizer.step() optimizer.zero_grad() train_dis_avg = 0 if train_count != 0 and train_count % 1000 == 0: if opt.refine_start: torch.save( refiner.state_dict(), '{0}/pose_refine_model_current.pth'.format( opt.outf)) else: torch.save( estimator.state_dict(), '{0}/pose_model_current.pth'.format(opt.outf)) print( '>>>>>>>>----------epoch {0} train finish---------<<<<<<<<'.format( epoch)) logger = setup_logger( 'epoch%d_test' % epoch, os.path.join(opt.log_dir, '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() refiner.eval() for j, data in enumerate(testdataloader, 0): points, choose, img, target, model_points, idx = data points, choose, img, target, model_points, idx = Variable(points).cuda(), \ Variable(choose).cuda(), \ Variable(img).cuda(), \ Variable(target).cuda(), \ Variable(model_points).cuda(), \ Variable(idx).cuda() pred_r, pred_t, pred_c, emb = estimator(img, points, choose, idx) _, dis, new_points, new_target = criterion(pred_r, pred_t, pred_c, target, model_points, idx, points, opt.w, opt.refine_start) if opt.refine_start: for ite in range(0, opt.iteration): pred_r, pred_t = refiner(new_points, emb, idx) dis, new_points, new_target = criterion_refine( pred_r, pred_t, new_target, model_points, idx, new_points) test_dis += dis.item() 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_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_test: best_test = test_dis if opt.refine_start: torch.save( refiner.state_dict(), '{0}/pose_refine_model_{1}_{2}.pth'.format( opt.outf, epoch, test_dis)) else: torch.save( estimator.state_dict(), '{0}/pose_model_{1}_{2}.pth'.format( opt.outf, epoch, test_dis)) print(epoch, '>>>>>>>>----------BEST TEST MODEL SAVED---------<<<<<<<<') if best_test < opt.decay_margin and not opt.decay_start: opt.decay_start = True opt.lr *= opt.lr_rate opt.w *= opt.w_rate optimizer = optim.Adam(estimator.parameters(), lr=opt.lr) if best_test < opt.refine_margin and not opt.refine_start: opt.refine_start = True opt.batch_size = int(opt.batch_size / opt.iteration) optimizer = optim.Adam(refiner.parameters(), lr=opt.lr) if opt.dataset == 'ycb': dataset = PoseDataset_ycb('train', opt.num_points, True, opt.dataset_root, opt.noise_trans, opt.refine_start) elif opt.dataset == 'linemod': dataset = PoseDataset_linemod('train', opt.num_points, True, opt.dataset_root, opt.noise_trans, opt.refine_start) dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True, num_workers=opt.workers) if opt.dataset == 'ycb': test_dataset = PoseDataset_ycb('test', opt.num_points, False, opt.dataset_root, 0.0, opt.refine_start) elif opt.dataset == 'linemod': test_dataset = PoseDataset_linemod('test', opt.num_points, False, opt.dataset_root, 0.0, opt.refine_start) testdataloader = torch.utils.data.DataLoader( test_dataset, batch_size=1, shuffle=False, num_workers=opt.workers) opt.sym_list = dataset.get_sym_list() opt.num_points_mesh = 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.num_points_mesh, opt.sym_list)) criterion = Loss(opt.num_points_mesh, opt.sym_list) criterion_refine = Loss_refine(opt.num_points_mesh, opt.sym_list)
def main(): if opt.dataset == 'ycb': opt.num_obj = 21 opt.sym_list = [12, 15, 18, 19, 20] opt.num_points = 1000 writer = SummaryWriter('experiments/runs/ycb/{0}'.format(opt.experiment_name)) opt.outf = 'trained_models/ycb/{0}'.format(opt.experiment_name) opt.log_dir = 'experiments/logs/ycb/{0}'.format(opt.experiment_name) opt.repeat_num = 1 if not os.path.exists(opt.outf): os.mkdir(opt.outf) if not os.path.exists(opt.log_dir): os.mkdir(opt.log_dir) else: print('Unknown dataset') return estimator = PoseNet(num_points = opt.num_points, num_vote = 9, num_obj = opt.num_obj) estimator.cuda() refiner = PoseRefineNet(num_points = opt.num_points, num_obj = opt.num_obj) refiner.cuda() if opt.resume_posenet != '': estimator.load_state_dict(torch.load('{0}/{1}'.format(opt.outf, opt.resume_posenet))) if opt.resume_refinenet != '': refiner.load_state_dict(torch.load('{0}/{1}'.format(opt.outf, opt.resume_refinenet))) opt.refine_start = True opt.lr = opt.lr_refine opt.batch_size = int(opt.batch_size / opt.iteration) optimizer = optim.Adam(refiner.parameters(), lr=opt.lr) else: opt.refine_start = False optimizer = optim.Adam(estimator.parameters(), lr=opt.lr) dataset = PoseDataset_ycb('train', opt.num_points, True, opt.dataset_root) dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True, num_workers=opt.workers) test_dataset = PoseDataset_ycb('test', opt.num_points, False, opt.dataset_root) testdataloader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=opt.workers) print('>>>>>>>>----------Dataset loaded!---------<<<<<<<<\nlength of the training set: {0}\nlength of the testing set: {1}\nnumber of sample points on mesh: {2}'.format(len(dataset), len(test_dataset), opt.num_points)) criterion = Loss(opt.num_points, opt.sym_list) criterion_refine = Loss_refine(opt.num_points, opt.sym_list) best_test = np.Inf if opt.start_epoch == 1: for log in os.listdir(opt.log_dir): os.remove(os.path.join(opt.log_dir, log)) st_time = time.time() train_scalar = 0 for epoch in range(opt.start_epoch, opt.nepoch): logger = setup_logger('epoch%d' % epoch, os.path.join(opt.log_dir, '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_loss_avg = 0.0 train_loss = 0.0 train_dis_avg = 0.0 train_dis = 0.0 if opt.refine_start: estimator.eval() refiner.train() else: estimator.train() optimizer.zero_grad() for rep in range(opt.repeat_num): for i, data in enumerate(dataloader, 0): points, choose, img, target, model_points, model_kp, vertex_gt, idx, target_r, target_t = data points, choose, img, target, model_points, model_kp, vertex_gt, idx, target_r, target_t = points.cuda(), choose.cuda(), img.cuda(), target.cuda(), model_points.cuda(), model_kp.cuda(), vertex_gt.cuda(), idx.cuda(), target_r.cuda(), target_t.cuda() vertex_pred, c_pred, emb = estimator(img, points, choose, idx) vertex_loss, pose_loss, dis, new_points, new_target = criterion(vertex_pred, vertex_gt, c_pred, points, target, model_points, model_kp, idx, target_r, target_t) loss = 10 * vertex_loss + pose_loss if opt.refine_start: for ite in range(0, opt.iteration): pred_r, pred_t = refiner(new_points, emb, idx) dis, new_points, new_target = criterion_refine(pred_r, pred_t, new_points, new_target, model_points, idx) dis.backward() else: loss.backward() train_loss_avg += loss.item() train_loss += loss.item() train_dis_avg += dis.item() train_dis += dis.item() train_count += 1 train_scalar += 1 if train_count % opt.batch_size == 0: logger.info('Train time {0} Epoch {1} Batch {2} Frame {3} Avg_loss:{4} Avg_diss:{5}'.format(time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)), epoch, int(train_count / opt.batch_size), train_count, train_loss_avg / opt.batch_size, train_dis_avg / opt.batch_size)) writer.add_scalar('ycb training loss', train_loss_avg / opt.batch_size, train_scalar) writer.add_scalar('ycb training dis', train_dis_avg / opt.batch_size, train_scalar) optimizer.step() optimizer.zero_grad() train_loss_avg = 0 train_dis_avg = 0 if train_count != 0 and train_count % 1000 == 0: if opt.refine_start: torch.save(refiner.state_dict(), '{0}/pose_refine_model_current.pth'.format(opt.outf)) else: torch.save(estimator.state_dict(), '{0}/pose_model_current.pth'.format(opt.outf)) print('>>>>>>>>----------epoch {0} train finish---------<<<<<<<<'.format(epoch)) train_loss = train_loss / train_count train_dis = train_dis / train_count logger.info('Train time {0} Epoch {1} TRAIN FINISH Avg loss: {2} Avg dis: {3}'.format(time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)), epoch, train_loss, train_dis)) logger = setup_logger('epoch%d_test' % epoch, os.path.join(opt.log_dir, '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_loss = 0.0 test_vertex_loss = 0.0 test_pose_loss = 0.0 test_dis = 0.0 test_count = 0 success_count = 0 estimator.eval() refiner.eval() for j, data in enumerate(testdataloader, 0): points, choose, img, target, model_points, model_kp, vertex_gt, idx, target_r, target_t = data points, choose, img, target, model_points, model_kp, vertex_gt, idx, target_r, target_t = points.cuda(), choose.cuda(), img.cuda(), target.cuda(), model_points.cuda(), model_kp.cuda(), vertex_gt.cuda(), idx.cuda(), target_r.cuda(), target_t.cuda() vertex_pred, c_pred, emb = estimator(img, points, choose, idx) vertex_loss, pose_loss, dis, new_points, new_target = criterion(vertex_pred, vertex_gt, c_pred, points, target, model_points, model_kp, idx, target_r, target_t) loss = 10 * vertex_loss + pose_loss if opt.refine_start: for ite in range(0, opt.iteration): pred_r, pred_t = refiner(new_points, emb, idx) dis, new_points, new_target = criterion_refine(pred_r, pred_t, new_points, new_target, model_points, idx) test_loss += loss.item() test_vertex_loss += vertex_loss.item() test_pose_loss += pose_loss.item() test_dis += dis.item() logger.info('Test time {0} Test Frame No.{1} loss:{2} dis:{3}'.format(time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)), test_count, loss, dis)) test_count += 1 if dis.item() < 0.02: success_count += 1 test_loss = test_loss / test_count test_vertex_loss = test_vertex_loss / test_count test_pose_loss = test_pose_loss / test_count test_dis = test_dis / test_count logger.info('Test time {0} Epoch {1} TEST FINISH Avg loss: {2} Avg dis: {3}'.format(time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)), epoch, test_loss, test_dis)) logger.info('Success rate: {}'.format(float(success_count) / test_count)) writer.add_scalar('ycb test loss', test_loss, epoch) writer.add_scalar('ycb test vertex loss', test_vertex_loss, epoch) writer.add_scalar('ycb test pose loss', test_pose_loss, epoch) writer.add_scalar('ycb test dis', test_dis, epoch) writer.add_scalar('ycb success rate', float(success_count) / test_count, epoch) writer.add_scalar('lr', optimizer.param_groups[0]['lr'], epoch) if test_dis <= best_test: best_test = test_dis if opt.refine_start: torch.save(refiner.state_dict(), '{0}/pose_refine_model_{1}_{2}.pth'.format(opt.outf, epoch, test_dis)) else: torch.save(estimator.state_dict(), '{0}/pose_model_{1}_{2}.pth'.format(opt.outf, epoch, test_dis)) print(epoch, '>>>>>>>>----------MODEL SAVED---------<<<<<<<<') if best_test < opt.refine_margin and not opt.refine_start: opt.refine_start = True opt.lr = opt.lr_refine opt.batch_size = int(opt.batch_size / opt.iteration) optimizer = optim.Adam(refiner.parameters(), lr=opt.lr) print('>>>>>>>>----------Refine started---------<<<<<<<<') writer.close()
def main(): opt.manualSeed = random.randint(1, 10000) random.seed(opt.manualSeed) torch.manual_seed(opt.manualSeed) 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_rot' #folder to save trained models opt.log_dir = 'experiments/logs/ycb_rot' #folder to save logs opt.repeat_epoch = 1 #number of repeat times for one epoch training estimator = PoseNet(num_points=opt.num_points, num_obj=opt.num_objects) estimator.cuda() refiner = PoseRefineNet(num_points=opt.num_points, num_obj=opt.num_objects) refiner.cuda() if opt.resume_posenet != '': estimator.load_state_dict( torch.load('{0}/{1}'.format(opt.outf, opt.resume_posenet))) if opt.resume_refinenet != '': refiner.load_state_dict( torch.load('{0}/{1}'.format(opt.outf, opt.resume_refinenet))) opt.refine_start = True opt.decay_start = True opt.lr *= opt.lr_rate opt.w *= opt.w_rate opt.batch_size = int(opt.batch_size / opt.iteration) optimizer = optim.Adam(refiner.parameters(), lr=opt.lr) else: opt.refine_start = False opt.decay_start = False optimizer = optim.Adam(estimator.parameters(), lr=opt.lr) object_list = list(range(1, 22)) output_format = [ otypes.DEPTH_POINTS_MASKED_AND_INDEXES, otypes.IMAGE_CROPPED, otypes.MODEL_POINTS_TRANSFORMED, otypes.MODEL_POINTS, otypes.OBJECT_LABEL, ] dataset = YCBDataset(opt.dataset_root, mode='train_syn_grid_valid', object_list=object_list, output_data=output_format, resample_on_error=True, preprocessors=[ YCBOcclusionAugmentor(opt.dataset_root), ColorJitter(), InplaneRotator() ], postprocessors=[ImageNormalizer(), PointShifter()], refine=opt.refine_start, image_size=[640, 480], num_points=1000) dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True, num_workers=opt.workers - 1) test_dataset = YCBDataset(opt.dataset_root, mode='valid', object_list=object_list, output_data=output_format, resample_on_error=True, preprocessors=[], postprocessors=[ImageNormalizer()], refine=opt.refine_start, image_size=[640, 480], num_points=1000) testdataloader = torch.utils.data.DataLoader(test_dataset, shuffle=True, batch_size=1, num_workers=1) opt.sym_list = [12, 15, 18, 19, 20] opt.num_points_mesh = dataset.num_pt_mesh_small 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.num_points_mesh, opt.sym_list)) criterion = Loss(opt.num_points_mesh, opt.sym_list) criterion_refine = Loss_refine(opt.num_points_mesh, opt.sym_list) best_test = np.Inf if opt.start_epoch == 1: for log in os.listdir(opt.log_dir): os.remove(os.path.join(opt.log_dir, log)) st_time = time.time() for epoch in range(opt.start_epoch, opt.nepoch): logger = setup_logger( 'epoch%d' % epoch, os.path.join(opt.log_dir, '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_avg = 0.0 if opt.refine_start: estimator.eval() refiner.train() else: estimator.train() optimizer.zero_grad() for rep in range(opt.repeat_epoch): for i, data in enumerate(dataloader, 0): points, choose, img, target, model_points, idx = data idx = idx - 1 points, choose, img, target, model_points, idx = Variable(points).cuda(), \ Variable(choose).cuda(), \ Variable(img).cuda(), \ Variable(target).cuda(), \ Variable(model_points).cuda(), \ Variable(idx).cuda() pred_r, pred_t, pred_c, emb = estimator( img, points, choose, idx) loss, dis, new_points, new_target = criterion( pred_r, pred_t, pred_c, target, model_points, idx, points, opt.w, opt.refine_start) if opt.refine_start: for ite in range(0, opt.iteration): pred_r, pred_t = refiner(new_points, emb, idx) dis, new_points, new_target = criterion_refine( pred_r, pred_t, new_target, model_points, idx, new_points) dis.backward() else: loss.backward() train_dis_avg += dis.item() train_count += 1 if train_count % opt.batch_size == 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.batch_size), train_count, train_dis_avg / opt.batch_size)) optimizer.step() optimizer.zero_grad() train_dis_avg = 0 if train_count != 0 and train_count % 1000 == 0: if opt.refine_start: torch.save( refiner.state_dict(), '{0}/pose_refine_model_current.pth'.format( opt.outf)) else: torch.save( estimator.state_dict(), '{0}/pose_model_current.pth'.format(opt.outf)) if (train_count >= 100000): break print( '>>>>>>>>----------epoch {0} train finish---------<<<<<<<<'.format( epoch)) logger = setup_logger( 'epoch%d_test' % epoch, os.path.join(opt.log_dir, '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() refiner.eval() for j, data in enumerate(testdataloader, 0): points, choose, img, target, model_points, idx = data idx = idx - 1 points, choose, img, target, model_points, idx = Variable(points).cuda(), \ Variable(choose).cuda(), \ Variable(img).cuda(), \ Variable(target).cuda(), \ Variable(model_points).cuda(), \ Variable(idx).cuda() pred_r, pred_t, pred_c, emb = estimator(img, points, choose, idx) _, dis, new_points, new_target = criterion(pred_r, pred_t, pred_c, target, model_points, idx, points, opt.w, opt.refine_start) if opt.refine_start: for ite in range(0, opt.iteration): pred_r, pred_t = refiner(new_points, emb, idx) dis, new_points, new_target = criterion_refine( pred_r, pred_t, new_target, model_points, idx, new_points) test_dis += dis.item() 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_count += 1 if (test_count >= 3000): break 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_test: best_test = test_dis if opt.refine_start: torch.save( refiner.state_dict(), '{0}/pose_refine_model_{1}_{2}.pth'.format( opt.outf, epoch, test_dis)) else: torch.save( estimator.state_dict(), '{0}/pose_model_{1}_{2}.pth'.format( opt.outf, epoch, test_dis)) print(epoch, '>>>>>>>>----------BEST TEST MODEL SAVED---------<<<<<<<<') if best_test < opt.decay_margin and not opt.decay_start: opt.decay_start = True opt.lr *= opt.lr_rate opt.w *= opt.w_rate optimizer = optim.Adam(estimator.parameters(), lr=opt.lr) if best_test < opt.refine_margin and not opt.refine_start: opt.refine_start = True opt.batch_size = int(opt.batch_size / opt.iteration) optimizer = optim.Adam(refiner.parameters(), lr=opt.lr) dataset = YCBDataset( opt.dataset_root, mode='train_syn_grid', object_list=object_list, output_data=output_format, resample_on_error=True, preprocessors=[ YCBOcclusionAugmentor(opt.dataset_root), ColorJitter(), InplaneRotator() ], postprocessors=[ImageNormalizer(), PointShifter()], refine=opt.refine_start, image_size=[640, 480], num_points=1000) dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True, num_workers=opt.workers) test_dataset = YCBDataset(opt.dataset_root, mode='valid', object_list=object_list, output_data=output_format, resample_on_error=True, preprocessors=[], postprocessors=[ImageNormalizer()], refine=opt.refine_start, image_size=[640, 480], num_points=1000) testdataloader = torch.utils.data.DataLoader( test_dataset, batch_size=1, shuffle=False, num_workers=opt.workers) opt.num_points_mesh = dataset.num_pt_mesh_large 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.num_points_mesh, opt.sym_list)) criterion = Loss(opt.num_points_mesh, opt.sym_list) criterion_refine = Loss_refine(opt.num_points_mesh, opt.sym_list)
def main(): # opt.manualSeed = random.randint(1, 10000) # # opt.manualSeed = 1 # random.seed(opt.manualSeed) # torch.manual_seed(opt.manualSeed) torch.set_printoptions(threshold=5000) # device_ids = [0,1] cudnn.benchmark = True if opt.dataset == 'ycb': 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 = 'experiments/logs/ycb' #folder to save logs opt.repeat_epoch = 3 #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 = 'experiments/logs/linemod' opt.repeat_epoch = 20 else: print('Unknown dataset') return estimator = PoseNet(num_points=opt.num_points, num_obj=opt.num_objects) estimator.cuda() refiner = PoseRefineNet(num_points=opt.num_points, num_obj=opt.num_objects) refiner.cuda() # estimator = nn.DataParallel(estimator, device_ids=device_ids) if opt.resume_posenet != '': estimator.load_state_dict( torch.load('{0}/{1}'.format(opt.outf, opt.resume_posenet))) if opt.resume_refinenet != '': refiner.load_state_dict( torch.load('{0}/{1}'.format(opt.outf, opt.resume_refinenet))) opt.refine_start = True opt.decay_start = True opt.lr *= opt.lr_rate opt.w *= opt.w_rate opt.batch_size = int(opt.batch_size / opt.iteration) optimizer = optim.Adam(refiner.parameters(), lr=opt.lr) else: print('no refinement') opt.refine_start = False opt.decay_start = False optimizer = optim.Adam(estimator.parameters(), lr=opt.lr) # optimizer = nn.DataParallel(optimizer, device_ids=device_ids) if opt.dataset == 'ycb': dataset = PoseDataset_ycb('train', opt.num_points, False, opt.dataset_root, opt.noise_trans, opt.refine_start) # print(dataset.list) elif opt.dataset == 'linemod': dataset = PoseDataset_linemod('train', opt.num_points, True, opt.dataset_root, opt.noise_trans, opt.refine_start) dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True, num_workers=opt.workers) if opt.dataset == 'ycb': test_dataset = PoseDataset_ycb('test', opt.num_points, False, opt.dataset_root, 0.0, opt.refine_start) elif opt.dataset == 'linemod': test_dataset = PoseDataset_linemod('test', opt.num_points, False, opt.dataset_root, 0.0, opt.refine_start) testdataloader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=opt.workers) opt.sym_list = dataset.get_sym_list() opt.num_points_mesh = 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.num_points_mesh, opt.sym_list)) criterion = Loss(opt.num_points_mesh, opt.sym_list) # criterion_refine = Loss_refine(opt.num_points_mesh, opt.sym_list) best_test = np.Inf best_epoch = 0 if opt.start_epoch == 1: for log in os.listdir(opt.log_dir): os.remove(os.path.join(opt.log_dir, log)) st_time = time.time() count_gen = 0 mode = 1 if mode == 1: for epoch in range(opt.start_epoch, opt.nepoch): logger = setup_logger( 'epoch%d' % epoch, os.path.join(opt.log_dir, '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_avg = 0.0 if opt.refine_start: estimator.eval() refiner.train() else: estimator.train() optimizer.zero_grad() for rep in range(opt.repeat_epoch): for i, data in enumerate(dataloader, 0): points, choose, img, target_sym, target_cen, idx, file_list_idx = data if idx is 9 or idx is 16: continue # points, choose, img, target_sym, target_cen, target, idx, file_list_idx = data # generate_obj_file(target_sym, target_cen, target, idx.squeeze()) # import pdb;pdb.set_trace() points, choose, img, target_sym, target_cen, idx = Variable(points).cuda(), \ Variable(choose).cuda(), \ Variable(img).cuda(), \ Variable(target_sym).cuda(), \ Variable(target_cen).cuda(), \ Variable(idx).cuda() # points, choose, img, target_sym, target_cen, idx = Variable(points), \ # Variable(choose), \ # Variable(img), \ # Variable(target_sym), \ # Variable(target_cen), \ # Variable(idx) pred_norm, pred_on_plane, emb = estimator( img, points, choose, idx) # pred_norm_new = torch.cat((pred_norm, torch.zeros(1,pred_norm.size(1),1)),2) # for i in range(pred_norm.size(1)): # pred_norm_new[0,i,2] = torch.sqrt(1 - pred_norm[0,i,0] * pred_norm[0,i,0] - pred_norm[0,i,1] * pred_norm[0,i,1]) # if epoch % 10 == 0: # generate_obj_file_pred(pred_norm, pred_on_plane, points, count_gen, idx) # count_gen += 1 # print(pred_norm[0,0,:]) loss = criterion(pred_norm, pred_on_plane, target_sym, target_cen, idx, points, opt.w, opt.refine_start) # scene_idx = dataset.list[file_list_idx] loss.backward() # train_dis_avg += dis.item() train_count += 1 if train_count % opt.batch_size == 0: logger.info( 'Train time {0} Epoch {1} Batch {2} Frame {3}'. format( time.strftime( "%Hh %Mm %Ss", time.gmtime(time.time() - st_time)), epoch, int(train_count / opt.batch_size), train_count)) optimizer.step() # for param_lr in optimizer.module.param_groups: # param_lr['lr'] /= 2 optimizer.zero_grad() train_dis_avg = 0 if train_count % 5000 == 0: print(pred_on_plane.max()) print(pred_on_plane.mean()) if train_count != 0 and train_count % 1000 == 0: if opt.refine_start: torch.save( refiner.state_dict(), '{0}/pose_refine_model_current.pth'.format( opt.outf)) else: torch.save( estimator.state_dict(), '{0}/pose_model_current.pth'.format(opt.outf)) print('>>>>>>>>----------epoch {0} train finish---------<<<<<<<<'. format(epoch)) logger = setup_logger( 'epoch%d_test' % epoch, os.path.join(opt.log_dir, '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_loss = 0.0 test_count = 0 estimator.eval() # refiner.eval() # for rep in range(opt.repeat_epoch): # for j, data in enumerate(testdataloader, 0): # points, choose, img, target_sym, target_cen, idx, img_idx = data # # points, choose, img, target, model_points, idx = Variable(points).cuda(), \ # # Variable(choose).cuda(), \ # # Variable(img).cuda(), \ # # Variable(target).cuda(), \ # # Variable(model_points).cuda(), \ # # Variable(idx).cuda() # points, choose, img, target_sym, target_cen, idx = Variable(points), \ # Variable(choose), \ # Variable(img), \ # Variable(target_sym), \ # Variable(target_cen), \ # Variable(idx) # pred_norm, pred_on_plane, emb = estimator(img, points, choose, idx) # loss = criterion(pred_norm, pred_on_plane, target_sym, target_cen, idx, points, opt.w, opt.refine_start) # test_loss += loss # logger.info('Test time {0} Test Frame No.{1}'.format(time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)), test_count)) # test_count += 1 # test_loss = test_loss / 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_loss)) print(pred_on_plane.max()) print(pred_on_plane.mean()) bs, num_p, _ = pred_on_plane.size() # if epoch % 40 == 0: # import pdb;pdb.set_trace() best_test = test_loss best_epoch = epoch if opt.refine_start: torch.save( refiner.state_dict(), '{0}/pose_refine_model_{1}_{2}.pth'.format( opt.outf, epoch, test_loss)) else: torch.save( estimator.state_dict(), '{0}/pose_model_{1}_{2}.pth'.format( opt.outf, epoch, test_loss)) print(epoch, '>>>>>>>>----------BEST TEST MODEL SAVED---------<<<<<<<<') if best_test < opt.decay_margin and not opt.decay_start: opt.decay_start = True opt.lr *= opt.lr_rate # opt.w *= opt.w_rate optimizer = optim.Adam(estimator.parameters(), lr=opt.lr) estimator.load_state_dict( torch.load('{0}/pose_model_{1}_{2}.pth'.format( opt.outf, best_epoch, best_test))) else: estimator.load_state_dict( torch.load('{0}/pose_model_11_0.0.pth'.format(opt.outf))) product_list = [] dist_list = [] true_positives = 0 false_positives = 0 false_negatives = 0 for index in range(len(test_dataset.list)): img = Image.open('{0}/data_v1/{1}-color.png'.format( test_dataset.root, test_dataset.list[index])) depth = np.array( Image.open('{0}/data_v1/{1}-depth.png'.format( test_dataset.root, test_dataset.list[index]))) label = np.array( Image.open('{0}/data_v1/{1}-label.png'.format( test_dataset.root, test_dataset.list[index]))) meta = scio.loadmat('{0}/data_v1/{1}-meta.mat'.format( test_dataset.root, test_dataset.list[index])) cam_cx = test_dataset.cam_cx_1 cam_cy = test_dataset.cam_cy_1 cam_fx = test_dataset.cam_fx_1 cam_fy = test_dataset.cam_fy_1 mask_back = ma.getmaskarray(ma.masked_equal(label, 0)) obj = meta['cls_indexes'].flatten().astype(np.int32) for idx in range(0, len(obj)): print('object index: ', obj[idx]) mask_depth = ma.getmaskarray(ma.masked_not_equal(depth, 0)) mask_label = ma.getmaskarray(ma.masked_equal(label, obj[idx])) mask = mask_label * mask_depth if not (len(mask.nonzero()[0]) > test_dataset.minimum_num_pt and len(test_dataset.symmetry[obj[idx]]['mirror']) > 0): continue rmin, rmax, cmin, cmax = get_bbox(mask_label) img_temp = np.transpose(np.array(img)[:, :, :3], (2, 0, 1))[:, rmin:rmax, cmin:cmax] img_masked = img_temp target_r = meta['poses'][:, :, idx][:, 0:3] target_t = np.array(meta['poses'][:, :, idx][:, 3:4].flatten()) add_t = np.array([ random.uniform(-test_dataset.noise_trans, test_dataset.noise_trans) for i in range(3) ]) choose = mask[rmin:rmax, cmin:cmax].flatten().nonzero()[0] if len(choose) > test_dataset.num_pt: c_mask = np.zeros(len(choose), dtype=int) c_mask[:test_dataset.num_pt] = 1 np.random.shuffle(c_mask) choose = choose[c_mask.nonzero()] else: choose = np.pad(choose, (0, test_dataset.num_pt - len(choose)), 'wrap') depth_masked = depth[ rmin:rmax, cmin:cmax].flatten()[choose][:, np.newaxis].astype(np.float32) xmap_masked = test_dataset.xmap[ rmin:rmax, cmin:cmax].flatten()[choose][:, np.newaxis].astype(np.float32) ymap_masked = test_dataset.ymap[ rmin:rmax, cmin:cmax].flatten()[choose][:, np.newaxis].astype(np.float32) choose = np.array([choose]) cam_scale = meta['factor_depth'][0][0] pt2 = depth_masked / cam_scale pt0 = (ymap_masked - cam_cx) * pt2 / cam_fx pt1 = (xmap_masked - cam_cy) * pt2 / cam_fy cloud = np.concatenate((pt0, pt1, pt2), axis=1) dellist = [j for j in range(0, len(test_dataset.cld[obj[idx]]))] # dellist = random.sample(dellist, len(test_dataset.cld[obj[idx]]) - test_dataset.num_pt_mesh_small) # model_points = np.delete(test_dataset.cld[obj[idx]], dellist, axis=0) model_points = test_dataset.cld[obj[idx]] target_sym = [] for sym in test_dataset.symmetry[obj[idx]]['mirror']: target_sym.append(np.dot(sym, target_r.T)) target_sym = np.array(target_sym) target_cen = np.add(test_dataset.symmetry[obj[idx]]['center'], target_t) target = np.dot(model_points, target_r.T) target = np.add(target, target_t) print('ground truth norm: ', target_sym) print('ground truth center: ', target_cen) points_ten, choose_ten, img_ten, target_sym_ten, target_cen_ten, target_ten, idx_ten = \ torch.from_numpy(cloud.astype(np.float32)).unsqueeze(0), \ torch.LongTensor(choose.astype(np.int32)).unsqueeze(0), \ test_dataset.norm(torch.from_numpy(img_masked.astype(np.float32))).unsqueeze(0), \ torch.from_numpy(target_sym.astype(np.float32)).unsqueeze(0), \ torch.from_numpy(target_cen.astype(np.float32)).unsqueeze(0), \ torch.from_numpy(target.astype(np.float32)).unsqueeze(0), \ torch.LongTensor([obj[idx]-1]).unsqueeze(0) # print(img_ten.size()) # print(points_ten.size()) # print(choose_ten.size()) # print(idx_ten.size()) points_ten, choose_ten, img_ten, target_sym_ten, target_cen_ten, idx_ten = Variable(points_ten).cuda(), \ Variable(choose_ten).cuda(), \ Variable(img_ten).cuda(), \ Variable(target_sym_ten).cuda(), \ Variable(target_cen_ten).cuda(), \ Variable(idx_ten).cuda() pred_norm, pred_on_plane, emb = estimator(img_ten, points_ten, choose_ten, idx_ten) # import pdb;pdb.set_trace() bs, num_p, _ = pred_on_plane.size() # pred_norm = torch.cat((pred_norm, torch.zeros(1,pred_norm.size(1),1)),2) # for i in range(pred_norm.size(1)): # pred_norm[0,i,2] = torch.sqrt(1 - pred_norm[0,i,0] * pred_norm[0,i,0] - pred_norm[0,i,1] * pred_norm[0,i,1]) # pred_norm = pred_norm / (torch.norm(pred_norm, dim=2).view(bs, num_p, 1)) generate_obj_file_norm_pred( pred_norm / (torch.norm(pred_norm, dim=2).view(bs, num_p, 1)), pred_on_plane, points_ten, test_dataset.list[index].split('/')[0], test_dataset.list[index].split('/')[1], obj[idx]) loss = criterion(pred_norm, pred_on_plane, target_sym_ten, target_cen_ten, idx, points_ten, opt.w, opt.refine_start) # print('test loss: ', loss) # bs, num_p, _ = pred_on_plane.size() pred_norm = pred_norm / (torch.norm(pred_norm, dim=2).view( bs, num_p, 1)) pred_norm = pred_norm.cpu().detach().numpy() pred_on_plane = pred_on_plane.cpu().detach().numpy() points = points_ten.cpu().detach().numpy() clustering_points_idx = np.where( pred_on_plane > pred_on_plane.max() * PRED_ON_PLANE_FACTOR + pred_on_plane.mean() * (1 - PRED_ON_PLANE_FACTOR))[1] clustering_norm = pred_norm[0, clustering_points_idx, :] clustering_points = points[0, clustering_points_idx, :] num_points = len(clustering_points_idx) # import pdb;pdb.set_trace() close_thresh = 5e-3 broad_thresh = 7e-3 sym_flag = [0 for i in range(target_sym.shape[0])] sym_max_product = [0.0 for i in range(target_sym.shape[0])] sym_dist = [0.0 for i in range(target_sym.shape[0])] count_pred = 0 while True: if num_points == 0: break count_pred += 1 if count_pred > target_sym.shape[0]: break best_fit_num = 0 count_try = 0 while True: if count_try > 3 or num_points <= 1: break pick_idx = np.random.randint(0, num_points - 1) pick_point = clustering_points[pick_idx] # proposal_norm = np.array(Plane(Point3D(pick_points[0]),Point3D(pick_points[1]),Point3D(pick_points[2])).normal_vector).astype(np.float32) proposal_norm = clustering_norm[pick_idx] proposal_norm = proposal_norm[:, np.newaxis] # import pdb;pdb.set_trace() proposal_point = pick_point # highest_pred_idx = np.argmax(pred_on_plane[0,clustering_points_idx,:]) # highest_pred_loc = clustering_points[highest_pred_idx] # proposal_norm = clustering_norm[highest_pred_idx][:,np.newaxis] clustering_diff = clustering_points - proposal_point clustering_dist = np.abs( np.matmul(clustering_diff, proposal_norm)) broad_inliers = np.where(clustering_dist < broad_thresh)[0] broad_inlier_num = len(broad_inliers) close_inliers = np.where(clustering_dist < close_thresh)[0] close_inlier_num = len(close_inliers) if broad_inlier_num > num_points / (5 - count_pred): best_fit_num = close_inlier_num best_fit_norm = proposal_norm best_fit_cen = clustering_points[close_inliers].mean(0) best_fit_idx = clustering_points_idx[close_inliers] scrub_idx = clustering_points_idx[broad_inliers] break else: count_try += 1 # else: # np.delete(clustering_points_idx, highest_pred_idx) # num_points -= 1 if count_try > 3 or num_points <= 1: break for i in range(2): def f(x): dist = 0 x = x / LA.norm(x) for point in clustering_points[broad_inliers]: dist += np.abs(point[0] * x[0] + point[1] * x[1] + point[2] * np.sqrt(1 - x[0] * x[0] - x[1] * x[1]) + x[2]) return dist start_point = np.copy(proposal_norm) start_point[2] = (-proposal_point * proposal_norm[:, 0]).sum() min_point = fmin(f, start_point) new_pred_loc = np.array([ 0, 0, -min_point[2] / np.sqrt(1 - min_point[0] * min_point[0] - min_point[1] * min_point[1]) ]) min_point[2] = np.sqrt(1 - min_point[0] * min_point[0] - min_point[1] * min_point[1]) new_proposal_norm = min_point clustering_diff = clustering_points - new_pred_loc clustering_dist = np.abs( np.matmul(clustering_diff, new_proposal_norm)) close_inliers = np.where(clustering_dist < close_thresh)[0] new_close_inlier_num = len(close_inliers) broad_inliers = np.where(clustering_dist < broad_thresh)[0] new_broad_inlier_num = len(broad_inliers) # import pdb;pdb.set_trace() if new_close_inlier_num > close_inlier_num: best_fit_num = new_close_inlier_num # proposal_point = clustering_points_idx[clustering_dist.argmin()] proposal_point = new_pred_loc best_fit_norm = new_proposal_norm[:, np.newaxis] best_fit_idx = clustering_points_idx[close_inliers] scrub_idx = clustering_points_idx[broad_inliers] best_fit_cen = new_pred_loc inlier_num = new_inlier_num proposal_norm = best_fit_norm # other_idx_pick = other_idx[other_idx_pick] # if len(other_idx_pick) > num_points//6: # pick_idx = np.concatenate((pick_idx, other_idx_pick), 0) # norm_proposal_new = clustering_norm[pick_idx,:].mean(0) # norm_proposal_new = norm_proposal_new / LA.norm(norm_proposal_new) # inlier_num_new = len(np.where(np.abs(clustering_norm-norm_proposal_new).sum(1) < thresh)[0]) # if inlier_num_new > inlier_num: # best_fit_num = inlier_num_new # best_fit_idx = np.where(np.abs(clustering_norm-norm_proposal_new).sum(1) < thresh_scrap) # best_fit_norm = norm_proposal_new # best_fit_cen = clustering_points[best_fit_idx].mean(0) if best_fit_num == 0: break else: print('predicted norm:{}, predicted point:{}'.format( best_fit_norm, best_fit_cen)) max_idx = np.argmax(np.matmul(target_sym, best_fit_norm)) sym_flag[max_idx] += 1 sym_product = np.abs((target_sym[max_idx] * (best_fit_cen - target_cen)).sum()) if sym_max_product[max_idx] < sym_product: sym_max_product[max_idx] = sym_product sym_dist[max_idx] = np.matmul(target_sym, best_fit_norm)[max_idx] # generate_obj_file_sym_pred(best_fit_norm, best_fit_cen, target_ten, test_dataset.list[index].split('/')[0], test_dataset.list[index].split('/')[1], obj[idx], count_pred) # import pdb;pdb.set_trace() clustering_points_idx = np.setdiff1d( clustering_points_idx, scrub_idx) clustering_norm = pred_norm[0, clustering_points_idx, :] clustering_points = points[0, clustering_points_idx, :] num_points = len(clustering_points_idx) for i in range(target_sym.shape[0]): if sym_flag[i] >= 1: dist_list.append(sym_dist[i]) product_list.append(sym_max_product[i]) false_positives += sym_flag[i] - 1 else: false_negatives += 1 product_list = np.array(product_list) dist_list = np.array(dist_list) # import pdb;pdb.set_trace() total_num = len(product_list) prec = [] recall = [] for t in range(1000): good_ones = len( np.logical_and(dist_list < 0.5 * t / 1000, product_list > math.cos(math.pi * 0.25 * t / 1000))) prec.append(good_ones * 1.0 / (false_positives + total_num)) recall.append(good_ones * 1.0 / (good_ones + false_negatives)) print(prec) print(recall) plt.plot(recall, prec, 'r') plt.axis([0, 1, 0, 1]) plt.savefig('prec-recall.png')
def main(): opt.manualSeed = random.randint(1, 10000) random.seed(opt.manualSeed) torch.manual_seed(opt.manualSeed) if opt.dataset == 'linemod': opt.num_objects = 13 opt.num_points = 500 opt.outf = 'trained_models/linemod' opt.log_dir = 'experiments/logs/linemod' output_results = 'check_linemod.txt' opt.repeat_epoch = 20 elif opt.dataset == 'ycb': 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 = 'experiments/logs/ycb' #folder to save logs opt.repeat_epoch = 1 #number of repeat times for one epoch training elif opt.dataset == 'ycb-syn': opt.num_objects = 31 # number of object classes in the dataset opt.num_points = 1000 # number of points on the input pointcloud opt.dataset_root = '/data/Akeaveny/Datasets/ycb_syn' opt.outf = 'trained_models/ycb_syn/ycb_syn2' # folder to save trained models opt.log_dir = 'experiments/logs/ycb_syn/ycb_syn2' # folder to save logs output_results = 'check_ycb_syn.txt' opt.w = 0.05 opt.refine_margin = 0.01 elif opt.dataset == 'arl': opt.num_objects = 10 # number of object classes in the dataset opt.num_points = 1000 # number of points on the input pointcloud opt.dataset_root = '/data/Akeaveny/Datasets/arl_dataset' opt.outf = 'trained_models/arl/clutter/arl_finetune_syn_2' # folder to save trained models opt.log_dir = '/home/akeaveny/catkin_ws/src/object-rpe-ak/DenseFusion/experiments/logs/arl/clutter/arl_finetune_syn_2' # folder to save logs output_results = 'check_arl_syn.txt' opt.nepoch = 750 opt.w = 0.05 opt.refine_margin = 0.0045 # TODO opt.repeat_epoch = 20 opt.start_epoch = 0 opt.resume_posenet = 'pose_model_1_0.012397416144377301.pth' opt.resume_refinenet = 'pose_refine_model_153_0.004032851301599294.pth' elif opt.dataset == 'arl1': opt.num_objects = 5 # number of object classes in the dataset opt.num_points = 1000 # number of points on the input pointcloud opt.dataset_root = '/data/Akeaveny/Datasets/arl_dataset' opt.outf = 'trained_models/arl1/clutter/arl_real_2' # folder to save trained models opt.log_dir = '/home/akeaveny/catkin_ws/src/object-rpe-ak/DenseFusion/experiments/logs/arl1/clutter/arl_real_2' # folder to save logs output_results = 'check_arl_syn.txt' opt.nepoch = 750 opt.w = 0.05 opt.refine_margin = 0.015 # opt.start_epoch = 120 # opt.resume_posenet = 'pose_model_current.pth' # opt.resume_refinenet = 'pose_refine_model_115_0.008727498716640046.pth' elif opt.dataset == 'elevator': opt.num_objects = 1 # number of object classes in the dataset opt.num_points = 1000 # number of points on the input pointcloud opt.dataset_root = '/data/Akeaveny/Datasets/elevator_dataset' opt.outf = 'trained_models/elevator/elevator_2' # folder to save trained models opt.log_dir = '/home/akeaveny/catkin_ws/src/object-rpe-ak/DenseFusion/experiments/logs/elevator/elevator_2' # folder to save logs output_results = 'check_arl_syn.txt' opt.nepoch = 750 opt.w = 0.05 opt.refine_margin = 0.015 opt.nepoch = 750 opt.w = 0.05 opt.refine_margin = 0.015 # TODO opt.repeat_epoch = 40 # opt.start_epoch = 47 # opt.resume_posenet = 'pose_model_current.pth' # opt.resume_refinenet = 'pose_refine_model_46_0.007581770288279472.pth' else: print('Unknown dataset') return estimator = PoseNet(num_points=opt.num_points, num_obj=opt.num_objects) estimator.cuda() refiner = PoseRefineNet(num_points=opt.num_points, num_obj=opt.num_objects) refiner.cuda() if opt.resume_posenet != '': estimator.load_state_dict( torch.load('{0}/{1}'.format(opt.outf, opt.resume_posenet))) if opt.resume_refinenet != '': refiner.load_state_dict( torch.load('{0}/{1}'.format(opt.outf, opt.resume_refinenet))) opt.refine_start = False opt.decay_start = False opt.lr *= opt.lr_rate opt.w *= opt.w_rate opt.batch_size = int(opt.batch_size / opt.iteration) optimizer = optim.Adam(refiner.parameters(), lr=opt.lr) else: opt.refine_start = False opt.decay_start = False optimizer = optim.Adam(estimator.parameters(), lr=opt.lr) if opt.dataset == 'ycb': dataset = PoseDataset_ycb('train', opt.num_points, True, opt.dataset_root, opt.noise_trans, opt.refine_start) elif opt.dataset == 'linemod': dataset = PoseDataset_linemod('train', opt.num_points, True, opt.dataset_root, opt.noise_trans, opt.refine_start) elif opt.dataset == 'ycb-syn': dataset = PoseDataset_ycb_syn('train', opt.num_points, True, opt.dataset_root, opt.noise_trans, opt.refine_start) elif opt.dataset == 'arl': dataset = PoseDataset_arl('train', opt.num_points, True, opt.dataset_root, opt.noise_trans, opt.refine_start) elif opt.dataset == 'arl1': dataset = PoseDataset_arl1('train', opt.num_points, True, opt.dataset_root, opt.noise_trans, opt.refine_start) elif opt.dataset == 'elevator': dataset = PoseDataset_elevator('train', opt.num_points, True, opt.dataset_root, opt.noise_trans, opt.refine_start) dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True, num_workers=opt.workers) if opt.dataset == 'ycb': test_dataset = PoseDataset_ycb('test', opt.num_points, False, opt.dataset_root, 0.0, opt.refine_start) elif opt.dataset == 'linemod': test_dataset = PoseDataset_linemod('test', opt.num_points, False, opt.dataset_root, 0.0, opt.refine_start) elif opt.dataset == 'ycb-syn': test_dataset = PoseDataset_ycb_syn('test', opt.num_points, True, opt.dataset_root, 0.0, opt.refine_start) elif opt.dataset == 'arl': test_dataset = PoseDataset_arl('test', opt.num_points, True, opt.dataset_root, 0.0, opt.refine_start) elif opt.dataset == 'arl1': test_dataset = PoseDataset_arl1('test', opt.num_points, True, opt.dataset_root, 0.0, opt.refine_start) elif opt.dataset == 'elevator': test_dataset = PoseDataset_elevator('test', opt.num_points, True, opt.dataset_root, 0.0, opt.refine_start) testdataloader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=opt.workers) opt.sym_list = dataset.get_sym_list() opt.num_points_mesh = 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.num_points_mesh, opt.sym_list)) criterion = Loss(opt.num_points_mesh, opt.sym_list) criterion_refine = Loss_refine(opt.num_points_mesh, opt.sym_list) best_test = np.Inf if opt.start_epoch == 1: for log in os.listdir(opt.log_dir): os.remove(os.path.join(opt.log_dir, log)) st_time = time.time() ###################### ###################### # TODO (ak): set up tensor board # if not os.path.exists(opt.log_dir): # os.makedirs(opt.log_dir) # # writer = SummaryWriter(opt.log_dir) ###################### ###################### for epoch in range(opt.start_epoch, opt.nepoch): logger = setup_logger( 'epoch%d' % epoch, os.path.join(opt.log_dir, '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_avg = 0.0 if opt.refine_start: estimator.eval() refiner.train() else: estimator.train() optimizer.zero_grad() for rep in range(opt.repeat_epoch): ################## # train ################## for i, data in enumerate(dataloader, 0): points, choose, img, target, model_points, idx = data # TODO: txt file # fw = open(test_folder + output_results, 'w') # fw.write('Points\n{0}\n\nchoose\n{1}\n\nimg\n{2}\n\ntarget\n{3}\n\nmodel_points\n{4}'.format(points, choose, img, target, model_points)) # fw.close() points, choose, img, target, model_points, idx = Variable(points).cuda(), \ Variable(choose).cuda(), \ Variable(img).cuda(), \ Variable(target).cuda(), \ Variable(model_points).cuda(), \ Variable(idx).cuda() pred_r, pred_t, pred_c, emb = estimator( img, points, choose, idx) loss, dis, new_points, new_target = criterion( pred_r, pred_t, pred_c, target, model_points, idx, points, opt.w, opt.refine_start) if opt.refine_start: for ite in range(0, opt.iteration): pred_r, pred_t = refiner(new_points, emb, idx) dis, new_points, new_target = criterion_refine( pred_r, pred_t, new_target, model_points, idx, new_points) dis.backward() else: loss.backward() train_dis_avg += dis.item() train_count += 1 if train_count % opt.batch_size == 0: logger.info( 'Train time {} Epoch {} Batch {} Frame {}/{} Avg_dis: {:.2f} [cm]' .format( time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)), epoch, int(train_count / opt.batch_size), train_count, len(dataset.list), train_dis_avg / opt.batch_size * 100)) optimizer.step() optimizer.zero_grad() # TODO: tensorboard # if train_count != 0 and train_count % 250 == 0: # scalar_info = {'loss': loss.item(), # 'dis': train_dis_avg / opt.batch_size} # for key, val in scalar_info.items(): # writer.add_scalar(key, val, train_count) train_dis_avg = 0 if train_count != 0 and train_count % 1000 == 0: if opt.refine_start: torch.save( refiner.state_dict(), '{0}/pose_refine_model_current.pth'.format( opt.outf)) else: torch.save( estimator.state_dict(), '{0}/pose_model_current.pth'.format(opt.outf)) # TODO: tensorboard # scalar_info = {'loss': loss.item(), # 'dis': dis.item()} # for key, val in scalar_info.items(): # writer.add_scalar(key, val, train_count) print( '>>>>>>>>----------epoch {0} train finish---------<<<<<<<<'.format( epoch)) logger = setup_logger( 'epoch%d_test' % epoch, os.path.join(opt.log_dir, '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() refiner.eval() for j, data in enumerate(testdataloader, 0): points, choose, img, target, model_points, idx = data points, choose, img, target, model_points, idx = Variable(points).cuda(), \ Variable(choose).cuda(), \ Variable(img).cuda(), \ Variable(target).cuda(), \ Variable(model_points).cuda(), \ Variable(idx).cuda() pred_r, pred_t, pred_c, emb = estimator(img, points, choose, idx) _, dis, new_points, new_target = criterion(pred_r, pred_t, pred_c, target, model_points, idx, points, opt.w, opt.refine_start) if opt.refine_start: for ite in range(0, opt.iteration): pred_r, pred_t = refiner(new_points, emb, idx) dis, new_points, new_target = criterion_refine( pred_r, pred_t, new_target, model_points, idx, new_points) test_dis += dis.item() logger.info('Test time {} Test Frame No.{} dis: {} [cm]'.format( time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)), test_count, dis * 100)) test_count += 1 test_dis = test_dis / test_count logger.info( 'Test time {} Epoch {} TEST FINISH Avg dis: {} [cm]'.format( time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)), epoch, test_dis * 100)) # TODO: tensorboard # scalar_info = {'test dis': test_dis} # for key, val in scalar_info.items(): # writer.add_scalar(key, val, train_count) if test_dis <= best_test: best_test = test_dis if opt.refine_start: torch.save( refiner.state_dict(), '{0}/pose_refine_model_{1}_{2}.pth'.format( opt.outf, epoch, test_dis)) else: torch.save( estimator.state_dict(), '{0}/pose_model_{1}_{2}.pth'.format( opt.outf, epoch, test_dis)) print(epoch, '>>>>>>>>----------BEST TEST MODEL SAVED---------<<<<<<<<') if best_test < opt.decay_margin and not opt.decay_start: opt.decay_start = True opt.lr *= opt.lr_rate opt.w *= opt.w_rate optimizer = optim.Adam(estimator.parameters(), lr=opt.lr) if best_test < opt.refine_margin and not opt.refine_start: opt.refine_start = True opt.batch_size = int(opt.batch_size / opt.iteration) optimizer = optim.Adam(refiner.parameters(), lr=opt.lr) if opt.dataset == 'ycb': dataset = PoseDataset_ycb('train', opt.num_points, True, opt.dataset_root, opt.noise_trans, opt.refine_start) elif opt.dataset == 'linemod': dataset = PoseDataset_linemod('train', opt.num_points, True, opt.dataset_root, opt.noise_trans, opt.refine_start) elif opt.dataset == 'ycb-syn': dataset = PoseDataset_ycb_syn('train', opt.num_points, True, opt.dataset_root, opt.noise_trans, opt.refine_start) elif opt.dataset == 'arl': dataset = PoseDataset_arl('train', opt.num_points, True, opt.dataset_root, opt.noise_trans, opt.refine_start) elif opt.dataset == 'arl1': dataset = PoseDataset_arl1('train', opt.num_points, True, opt.dataset_root, opt.noise_trans, opt.refine_start) elif opt.dataset == 'elevator': dataset = PoseDataset_elevator('train', opt.num_points, True, opt.dataset_root, opt.noise_trans, opt.refine_start) dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True, num_workers=opt.workers) if opt.dataset == 'ycb': test_dataset = PoseDataset_ycb('test', opt.num_points, False, opt.dataset_root, 0.0, opt.refine_start) elif opt.dataset == 'linemod': test_dataset = PoseDataset_linemod('test', opt.num_points, False, opt.dataset_root, 0.0, opt.refine_start) elif opt.dataset == 'ycb-syn': test_dataset = PoseDataset_ycb_syn('test', opt.num_points, True, opt.dataset_root, 0.0, opt.refine_start) elif opt.dataset == 'arl': test_dataset = PoseDataset_arl('test', opt.num_points, True, opt.dataset_root, 0.0, opt.refine_start) elif opt.dataset == 'arl1': test_dataset = PoseDataset_arl1('test', opt.num_points, True, opt.dataset_root, 0.0, opt.refine_start) elif opt.dataset == 'elevator': test_dataset = PoseDataset_elevator('test', opt.num_points, True, opt.dataset_root, 0.0, opt.refine_start) testdataloader = torch.utils.data.DataLoader( test_dataset, batch_size=1, shuffle=False, num_workers=opt.workers) opt.sym_list = dataset.get_sym_list() opt.num_points_mesh = 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.num_points_mesh, opt.sym_list)) criterion = Loss(opt.num_points_mesh, opt.sym_list) criterion_refine = Loss_refine(opt.num_points_mesh, opt.sym_list)
def train_net(): # set result directory if not os.path.exists(opt.result_dir): os.makedirs(opt.result_dir) tb_writer = tf.summary.FileWriter(opt.result_dir) logger = setup_logger('train_log', os.path.join(opt.result_dir, 'log.txt')) for key, value in vars(opt).items(): logger.info(key + ': ' + str(value)) os.environ['CUDA_VISIBLE_DEVICES'] = opt.gpu # model & loss estimator = DeformNet(opt.n_cat, opt.nv_prior) estimator.cuda() # pdb.set_trace() criterion = Loss(opt.corr_wt, opt.cd_wt, opt.entropy_wt, opt.deform_wt) chamferD = ChamferLoss() if opt.resume_model != '': estimator.load_state_dict(torch.load(opt.resume_model)) # dataset # 253445 images found. = [249127, 4318] # 1101 models loaded. train_dataset = PoseDataset(opt.dataset, 'train', opt.data_dir, opt.n_pts, opt.img_size) # 2754 images found. # 18 models loaded. val_dataset = PoseDataset(opt.dataset, 'test', opt.data_dir, opt.n_pts, opt.img_size) # start training st_time = time.time() train_steps = 1500 global_step = train_steps * (opt.start_epoch - 1) n_decays = len(opt.decay_epoch) assert len(opt.decay_rate) == n_decays for i in range(n_decays): if opt.start_epoch > opt.decay_epoch[i]: decay_count = i # pdb.set_trace() train_size = train_steps * opt.batch_size indices = [] page_start = -train_size for epoch in range(opt.start_epoch, opt.max_epoch + 1): # # train one epoch # logger.info('Time {0}'.format(time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)) + \ # ', ' + 'Epoch %02d' % epoch + ', ' + 'Training started')) # # create optimizer and adjust learning rate if needed # if decay_count < len(opt.decay_rate): # if epoch > opt.decay_epoch[decay_count]: # current_lr = opt.lr * opt.decay_rate[decay_count] # optimizer = torch.optim.Adam(estimator.parameters(), lr=current_lr) # decay_count += 1 # # sample train subset # page_start += train_size # len_last = len(indices) - page_start # if len_last < train_size: # indices = indices[page_start:] # if opt.dataset == 'CAMERA+Real': # # CAMERA : Real = 3 : 1 # camera_len = train_dataset.subset_len[0] # real_len = train_dataset.subset_len[1] # real_indices = list(range(camera_len, camera_len+real_len)) # camera_indices = list(range(camera_len)) # n_repeat = (train_size - len_last) // (4 * real_len) + 1 # data_list = random.sample(camera_indices, 3*n_repeat*real_len) + real_indices*n_repeat # random.shuffle(data_list) # indices += data_list # else: # data_list = list(range(train_dataset.length)) # for i in range((train_size - len_last) // train_dataset.length + 1): # random.shuffle(data_list) # indices += data_list # page_start = 0 # train_idx = indices[page_start:(page_start+train_size)] # train_sampler = torch.utils.data.sampler.SubsetRandomSampler(train_idx) # train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=opt.batch_size, sampler=train_sampler, # num_workers=opt.num_workers, pin_memory=True) # estimator.train() # for i, data in enumerate(train_dataloader, 1): # points, rgb, choose, cat_id, model, prior, sRT, nocs = data # points = points.cuda() # rgb = rgb.cuda() # choose = choose.cuda() # cat_id = cat_id.cuda() # model = model.cuda() # prior = prior.cuda() # sRT = sRT.cuda() # nocs = nocs.cuda() # assign_mat, deltas = estimator(points, rgb, choose, cat_id, prior) # loss, corr_loss, cd_loss, entropy_loss, deform_loss = criterion(assign_mat, deltas, prior, nocs, model) # optimizer.zero_grad() # loss.backward() # optimizer.step() # global_step += 1 # # write results to tensorboard # summary = tf.Summary(value=[tf.Summary.Value(tag='learning_rate', simple_value=current_lr), # tf.Summary.Value(tag='train_loss', simple_value=loss), # tf.Summary.Value(tag='corr_loss', simple_value=corr_loss), # tf.Summary.Value(tag='cd_loss', simple_value=cd_loss), # tf.Summary.Value(tag='entropy_loss', simple_value=entropy_loss), # tf.Summary.Value(tag='deform_loss', simple_value=deform_loss)]) # tb_writer.add_summary(summary, global_step) # if i % 10 == 0: # logger.info('Batch {0} Loss:{1:f}, corr_loss:{2:f}, cd_loss:{3:f}, entropy_loss:{4:f}, deform_loss:{5:f}'.format( # i, loss.item(), corr_loss.item(), cd_loss.item(), entropy_loss.item(), deform_loss.item())) # # logger.info('>>>>>>>>----------Epoch {:02d} train finish---------<<<<<<<<'.format(epoch)) # evaluate one epoch logger.info('Time {0}'.format( time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)) + ', ' + 'Epoch %02d' % epoch + ', ' + 'Testing started')) val_loss = 0.0 total_count = np.zeros((opt.n_cat, ), dtype=int) strict_success = np.zeros((opt.n_cat, ), dtype=int) # 5 degree and 5 cm easy_success = np.zeros((opt.n_cat, ), dtype=int) # 10 degree and 5 cm iou_success = np.zeros((opt.n_cat, ), dtype=int) # relative scale error < 0.1 # sample validation subset # val_size = 1500 # val_idx = random.sample(list(range(val_dataset.length)), val_size) # val_sampler = torch.utils.data.sampler.SubsetRandomSampler(val_idx) # val_dataloader = torch.utils.data.DataLoader(val_dataset, batch_size=1, sampler=val_sampler, # num_workers=opt.num_workers, pin_memory=True) val_dataloader = torch.utils.data.DataLoader( val_dataset, batch_size=1, num_workers=opt.num_workers, pin_memory=True) estimator.eval() cd_num = torch.zeros(6) prior_cd = torch.zeros(6) deform_cd = torch.zeros(6) # pdb.set_trace() for i, data in enumerate(val_dataloader, 1): points, rgb, choose, cat_id, model, prior, sRT, nocs = data points = points.cuda() rgb = rgb.cuda() choose = choose.cuda() cat_id = cat_id.cuda() model = model.cuda() prior = prior.cuda() sRT = sRT.cuda() nocs = nocs.cuda() assign_mat, deltas = estimator(points, rgb, choose, cat_id, prior) loss, _, _, _, _ = criterion(assign_mat, deltas, prior, nocs, model) # pdb.set_trace() prior_loss, _, _ = chamferD(prior, model) deform_loss, _, _ = chamferD(prior + deltas, model) idx = cat_id.item() cd_num[idx] += 1 prior_cd[idx] += prior_loss.item() deform_cd[idx] += deform_loss.item() # estimate pose and scale inst_shape = prior + deltas assign_mat = F.softmax(assign_mat, dim=2) nocs_coords = torch.bmm(assign_mat, inst_shape) nocs_coords = nocs_coords.detach().cpu().numpy()[0] points = points.cpu().numpy()[0] # use choose to remove repeated points choose = choose.cpu().numpy()[0] _, choose = np.unique(choose, return_index=True) nocs_coords = nocs_coords[choose, :] points = points[choose, :] _, _, _, pred_sRT = estimateSimilarityTransform( nocs_coords, points) # evaluate pose cat_id = cat_id.item() if pred_sRT is not None: sRT = sRT.detach().cpu().numpy()[0] R_error, T_error, IoU = compute_sRT_errors(pred_sRT, sRT) if R_error < 5 and T_error < 0.05: strict_success[cat_id] += 1 if R_error < 10 and T_error < 0.05: easy_success[cat_id] += 1 if IoU < 0.1: iou_success[cat_id] += 1 total_count[cat_id] += 1 val_loss += loss.item() if i % 100 == 0: logger.info('Batch {0} Loss:{1:f}'.format(i, loss.item())) # pdb.set_trace() deform_cd_metric = (deform_cd / cd_num) * 1000 print( "recon: {:.2f} : {:.2f} : {:.2f} : {:.2f} : {:.2f} : {:.2f} : {:.2f}" .format(deform_cd_metric[0], deform_cd_metric[1], deform_cd_metric[2], deform_cd_metric[3], deform_cd_metric[4], deform_cd_metric[5], torch.mean(deform_cd_metric))) prior_cd_metric = (prior_cd / cd_num) * 1000 print( "prior: {:.2f} : {:.2f} : {:.2f} : {:.2f} : {:.2f} : {:.2f} : {:.2f}" .format(prior_cd_metric[0], prior_cd_metric[1], prior_cd_metric[2], prior_cd_metric[3], prior_cd_metric[4], prior_cd_metric[5], torch.mean(prior_cd_metric)))
def main(): # opt.manualSeed = random.randint(1, 10000) # # opt.manualSeed = 1 # random.seed(opt.manualSeed) # torch.manual_seed(opt.manualSeed) torch.set_printoptions(threshold=5000) if opt.dataset == 'ycb': 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 = 'experiments/logs/ycb' #folder to save logs opt.repeat_epoch = 3 #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 = 'experiments/logs/linemod' opt.repeat_epoch = 20 else: print('Unknown dataset') return opt.estimator = PoseNet(num_points=opt.num_points, num_obj=opt.num_objects) # estimator.cuda() refiner = PoseRefineNet(num_points=opt.num_points, num_obj=opt.num_objects) # refiner.cuda() class_id = 1 opt.symmetry = {} with open('symmetries_ordered.txt', 'r') as f: while 1: line = f.readline() line = line[:-1] if not line: break opt.symmetry[class_id] = {} opt.symmetry[class_id]['center'] = list( map(lambda x: float(x), line.split(' '))) opt.symmetry[class_id]['mirror'] = [] for i in range(3): line = f.readline() line = line[:-1] x, y, z = list(map(lambda x: float(x), line.split(' '))) if not (x == 0 and y == 0 and z == 0): opt.symmetry[class_id]['mirror'].append((x, y, z)) f.readline() f.readline() class_id += 1 opt.norm = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) criterion = Loss(opt.num_points, opt.symmetry) opt.estimator.load_state_dict( torch.load('{0}/pose_model_64_0.0.pth'.format(opt.outf), map_location='cpu')) # import pdb;pdb.set_trace() print('start parallelization') pool = Pool(4) results = [ pool.apply_async(printCurve, [take_idx, criterion]) for take_idx in range(3) ] for take_idx in range(3): prec, recall = results[take_idx].get() # prec,recall = printCurve(take_idx, criterion) for dist_idx in range(5): plt.plot(recall[dist_idx], prec[dist_idx], label='dis={:.2f}'.format((dist_idx + 1) * 0.01)) plt.axis([0, 1, 0, 1]) plt.xlabel('Recall') plt.ylabel('Precision') plt.title(title_list[take_idx]) plt.legend() plt.savefig('prec-recall-{}.png'.format(title_list[take_idx])) plt.clf()
def main(): opt.manualSeed = random.randint(1, 10000) random.seed(opt.manualSeed) torch.manual_seed(opt.manualSeed) 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_plus_bing' #folder to save trained models opt.log_dir = 'experiments/logs/ycb_plus_bing' #folder to save logs opt.repeat_epoch = 1 #number of repeat times for one epoch training estimator = PoseNetPlusDuelBing(num_points=opt.num_points, num_obj=opt.num_objects) estimator.cuda() train_writer = SummaryWriter(comment='duel_binham_train') valid_writer = SummaryWriter(comment='duel_binham_valid') if opt.resume_posenet != '': estimator.load_state_dict( torch.load('{0}/{1}'.format(opt.outf, opt.resume_posenet))) elif opt.finetune_posenet != '': pretrained_dict = torch.load(opt.finetune_posenet) model_dict = estimator.state_dict() pretrained_dict = { k: v for k, v in pretrained_dict.items() if k in model_dict } model_dict.update(pretrained_dict) estimator.load_state_dict(model_dict) for k, v in estimator.named_parameters(): if (k in pretrained_dict): v.requires_grad = False opt.log_dir += '_cont' opt.outf += '_cont' opt.refine_start = False opt.decay_start = False optimizer = optim.Adam(estimator.parameters(), lr=opt.lr) object_list = list(range(1, 22)) output_format = [ otypes.DEPTH_POINTS_MASKED_AND_INDEXES, otypes.IMAGE_CROPPED, otypes.QUATERNION, otypes.MODEL_POINTS_TRANSFORMED, otypes.MODEL_POINTS, otypes.OBJECT_LABEL, ] dataset = YCBDataset( opt.dataset_root, mode='train_syn_grid_valid', object_list=object_list, output_data=output_format, resample_on_error=True, preprocessors=[ YCBOcclusionAugmentor(opt.dataset_root), ColorJitter(), #InplaneRotator(), ], postprocessors=[ImageNormalizer(), PointShifter()], image_size=[640, 480], num_points=1000) dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True, num_workers=opt.workers - 1) test_dataset = YCBDataset(opt.dataset_root, mode='valid', object_list=object_list, output_data=output_format, resample_on_error=True, preprocessors=[], postprocessors=[ImageNormalizer()], image_size=[640, 480], num_points=1000) testdataloader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=1) opt.sym_list = [12, 15, 18, 19, 20] opt.num_points_mesh = dataset.num_pt_mesh_small 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.num_points_mesh, opt.sym_list)) criterion_dist = Loss(opt.num_points_mesh, opt.sym_list) criterion_lik = DuelLoss(opt.num_points_mesh, opt.sym_list) best_dis = np.Inf best_lik = -np.Inf if opt.start_epoch == 1: for log in os.listdir(opt.log_dir): os.remove(os.path.join(opt.log_dir, log)) st_time = time.time() cum_batch_count = 0 mean_err = 0 for epoch in range(opt.start_epoch, opt.nepoch): logger = setup_logger( 'epoch%d' % epoch, os.path.join(opt.log_dir, '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_avg = 0.0 train_lik_avg = 0.0 estimator.train() optimizer.zero_grad() for rep in range(opt.repeat_epoch): for i, data in enumerate(dataloader, 0): points, choose, img, quat, target, model_points, idx = data idx = idx - 1 points, choose, img, quat, target, model_points, idx = Variable(points).cuda(), \ Variable(choose).cuda(), \ Variable(img).cuda(), \ Variable(quat).cuda(), \ Variable(target).cuda(), \ Variable(model_points).cuda(), \ Variable(idx).cuda() pred_r, pred_t, pred_c, pred_bq, pred_bz, emb = estimator( img, points, choose, idx) loss_dist, dis, new_points, new_target = criterion_dist( pred_r, pred_t, pred_c, target, model_points, idx, points, opt.w, opt.refine_start) how_max, which_max = torch.max(pred_c.detach(), 1) pred_q = pred_r[0, :, [1, 2, 3, 0]].detach() pred_q /= torch.norm(pred_q, dim=1).view(-1, 1) max_q = pred_q[which_max.item()] max_bq = pred_bq[0, which_max.item()] / torch.norm( pred_bq[0, which_max.item()]) max_bz = pred_bz[0, which_max.item()] loss_lik, lik = criterion_lik(max_q.view(-1), max_bq.view(-1), -torch.abs(max_bz.view(-1)), quat) loss = loss_dist + loss_lik loss.backward() train_dis_avg += dis.item() train_lik_avg += np.log(lik.item()) train_count += 1 if train_count % opt.batch_size == 0: logger.info( 'Train time {0} Epoch {1} Batch {2} Frame {3} Avg_dis:{4} Avg_lik:{5}' .format( time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)), epoch, int(train_count / opt.batch_size), train_count, train_dis_avg / opt.batch_size, train_lik_avg / opt.batch_size)) optimizer.step() optimizer.zero_grad() train_dis_avg = 0 train_lik_avg = 0 if train_count != 0 and train_count % 1000 == 0: torch.save(estimator.state_dict(), '{0}/pose_model_current.pth'.format(opt.outf)) print( '>>>>>>>>----------epoch {0} train finish---------<<<<<<<<'.format( epoch)) logger = setup_logger( 'epoch%d_test' % epoch, os.path.join(opt.log_dir, '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_lik = 0.0 test_count = 0 estimator.eval() for j, data in enumerate(testdataloader, 0): points, choose, img, quat, target, model_points, idx = data idx = idx - 1 points, choose, img, quat, target, model_points, idx = Variable(points).cuda(), \ Variable(choose).cuda(), \ Variable(img).cuda(), \ Variable(quat).cuda(), \ Variable(target).cuda(), \ Variable(model_points).cuda(), \ Variable(idx).cuda() pred_r, pred_t, pred_c, pred_bq, pred_bz, emb = estimator( img, points, choose, idx) _, dis, new_points, new_target = criterion_dist( pred_r, pred_t, pred_c, target, model_points, idx, points, opt.w, opt.refine_start) how_max, which_max = torch.max(pred_c.detach(), 1) pred_q = pred_r[0, :, [1, 2, 3, 0]].detach() pred_q /= torch.norm(pred_q, dim=1).view(-1, 1) max_q = pred_q[which_max.item()] max_bq = pred_bq[0, which_max.item()] / torch.norm( pred_bq[0, which_max.item()]) max_bz = pred_bz[0, which_max.item()] _, lik = criterion_lik(max_q.view(-1), max_bq.view(-1), -torch.abs(max_bz.view(-1)), quat) test_dis += dis.item() test_lik += np.log(lik.item()) logger.info( 'Test time {0} Test Frame No.{1} dis:{2} lik:{3}'.format( time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)), test_count, dis, lik)) test_count += 1 test_dis = test_dis / test_count test_lik = test_lik / test_count logger.info( 'Test time {0} Epoch {1} TEST FINISH Avg dis: {2} Avg lik: {3}'. format( time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)), epoch, test_dis, test_lik)) if test_dis <= best_dis or test_lik >= best_lik: best_dis = min(test_dis, best_dis) best_lik = max(test_lik, best_lik) torch.save( estimator.state_dict(), '{0}/pose_model_{1}_{2}_{3}.pth'.format( opt.outf, epoch, test_dis, test_lik)) print(epoch, '>>>>>>>>----------BEST TEST MODEL SAVED---------<<<<<<<<') if best_dis < opt.decay_margin and not opt.decay_start: opt.decay_start = True opt.lr *= opt.lr_rate opt.w *= opt.w_rate optimizer = optim.Adam(estimator.parameters(), lr=opt.lr)
def main(): opt.manualSeed = random.randint(1, 10000) random.seed(opt.manualSeed) torch.manual_seed(opt.manualSeed) opt.num_objects = 21 opt.num_points = 1000 opt.outf = 'trained_models/ycb/' + opt.output_dir opt.log_dir = 'experiments/logs/ycb/' + opt.output_dir opt.train_dir = 'experiments/tb/ycb/' + opt.output_dir + '/train' opt.test_dir = 'experiments/tb/ycb/' + opt.output_dir + '/test' opt.repeat_epoch = 1 if not os.path.exists(opt.outf): os.makedirs(opt.outf, exist_ok=True) if not os.path.exists(opt.log_dir): os.makedirs(opt.log_dir, exist_ok=True) if not os.path.exists(opt.train_dir): os.makedirs(opt.train_dir, exist_ok=True) if not os.path.exists(opt.test_dir): os.makedirs(opt.test_dir, exist_ok=True) opt.repeat_epoch = 1 estimator = PoseNet(num_points=opt.num_points, num_obj=opt.num_objects, object_max=opt.object_max) estimator.cuda() isFirstInitLastDatafolder = True if opt.resume_posenet != '': psp_estimator = torch.load( 'trained_models/ycb/pose_model_26_0.012863246640872631.pth') pretrained_estimator = torch.load('{0}/{1}'.format( opt.outf, opt.resume_posenet)) estimator_dict = estimator.state_dict() psp_dict = { k: v for k, v in psp_estimator.items() if k.find('cnn.model') == 0 } pretrained_dict = { k: v for k, v in pretrained_estimator.items() if k.find('cnn.model') != 0 } estimator_dict.update(psp_dict) estimator_dict.update(pretrained_dict) estimator.load_state_dict(estimator_dict) else: psp_estimator = torch.load( 'trained_models/ycb/pose_model_26_0.012863246640872631.pth') psp_dict = { k: v for k, v in psp_estimator.items() if k.find('cnn.model') == 0 } estimator_dict = estimator.state_dict() estimator_dict.update(psp_dict) estimator.load_state_dict(estimator_dict) opt.decay_start = False optimizer = optim.Adam(estimator.parameters(), lr=opt.lr) dataset = PoseDataset_ycb('train', opt.num_points, False, opt.dataset_root, opt.noise_trans, 'ori', False) dataloader = torch.utils.data.DataLoader(dataset, shuffle=False, num_workers=opt.workers) test_dataset = PoseDataset_ycb('test', opt.num_points, False, opt.dataset_root, 0.0, 'ori', False) testdataloader = torch.utils.data.DataLoader(test_dataset, shuffle=False, num_workers=opt.workers) opt.sym_list = dataset.get_sym_list() opt.num_points_mesh = 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.num_points_mesh, opt.sym_list)) criterion = Loss(opt.num_points_mesh, opt.sym_list) dis_vector_last_map = {key: [] for key in range(0, opt.num_objects)} for i in range(0, opt.num_objects): dis_vector_last_map[i] = None best_test = np.Inf if opt.start_epoch == 1: for log in os.listdir(opt.log_dir): os.remove(os.path.join(opt.log_dir, log)) st_time = time.time() for epoch in range(opt.start_epoch, opt.nepoch): logger = setup_logger( 'epoch%d' % epoch, os.path.join(opt.log_dir, '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_avg = 0.0 global_train_dis = 0.0 estimator.train() optimizer.zero_grad() for rep in range(opt.repeat_epoch): for i, data in enumerate(dataloader, 0): list_points, list_choose, list_img, list_target, list_model_points, list_idx, list_filename, \ list_full_img, list_focal_length, list_principal_point, list_motion = data for list_index in range(len(list_points)): if opt.dataset == 'ycb': points, choose, img, target, model_points, idx, filename, full_img, focal_length, principal_point \ , motion = list_points[list_index], list_choose[list_index], list_img[list_index], \ list_target[list_index], list_model_points[list_index], list_idx[list_index], \ list_filename[list_index], list_full_img[list_index], list_focal_length[ list_index], \ list_principal_point[list_index], list_motion[list_index] datafolder = filename[0].split('/')[1] if isFirstInitLastDatafolder: lastdatafolder = datafolder isFirstInitLastDatafolder = False if datafolder != lastdatafolder: for i in range(0, opt.num_objects): dis_vector_last_map[i] = None optimizer.step() optimizer.zero_grad() train_dis_avg = 0 estimator.temporalClear(opt.object_max, opt.mem_length) lastdatafolder = datafolder elif opt.dataset == 'linemod': list_points, list_choose, list_img, list_target, list_model_points, list_idx, list_filename = data points, choose, img, target, model_points, idx, filename = list_points[ 0] points, choose, img, target, model_points, idx = points.cuda(), \ choose.cuda(), \ img.cuda(), \ target.cuda(), \ model_points.cuda(), \ idx.cuda() pred_r, pred_t, pred_c, x_return = estimator( img, points, choose, idx, focal_length, principal_point, motion, True) loss, dis, new_points, new_target, dis_vector = criterion( pred_r, pred_t, pred_c, dis_vector_last_map[idx.item()], target, model_points, idx, x_return, opt.w, False, float(opt.loss_stable_alpha)) dis_vector_last_map[idx.item()] = dis_vector loss.backward(retain_graph=True) logger.info( 'Train time {0} Frame {1} Object {2}, Loss = {3}'. format( time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)), filename, idx.item(), dis)) train_dis_avg += dis.item() global_train_dis += dis.item() train_count += 1 if train_count % (len(list_points) * opt.batch_size) == 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.batch_size), train_count, train_dis_avg / (len(list_points) * opt.batch_size))) optimizer.step() optimizer.zero_grad() train_dis_avg = 0 if train_count != 0 and train_count % 1000 == 0: torch.save( estimator.state_dict(), '{0}/pose_model_current.pth'.format(opt.outf)) print( '>>>>>>>>----------epoch {0} train finish---------<<<<<<<<'.format( epoch)) global_train_dis = 0.0 logger = setup_logger( 'epoch%d_test' % epoch, os.path.join(opt.log_dir, '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 i in range(0, opt.num_objects): dis_vector_last_map[i] = None with torch.no_grad(): isFirstInitLastDatafolder = True for j, data in enumerate(testdataloader, 0): if opt.dataset == 'ycb': list_points, list_choose, list_img, list_target, list_model_points, list_idx, list_filename, \ list_full_img, list_focal_length, list_principal_point, list_motion = data for list_index in range(len(list_points)): points, choose, img, target, model_points, idx, filename, full_img, focal_length, principal_point, motion \ = list_points[list_index], list_choose[list_index], list_img[list_index], \ list_target[list_index], list_model_points[list_index], list_idx[list_index], \ list_filename[list_index], list_full_img[list_index], list_focal_length[list_index], \ list_principal_point[list_index], list_motion[list_index] datafolder = filename[0].split('/')[1] filehead = filename[0].split('/')[2] if isFirstInitLastDatafolder: lastdatafolder = datafolder isFirstInitLastDatafolder = False if datafolder != lastdatafolder: train_dis_avg = 0 estimator.temporalClear(opt.object_max) lastdatafolder = datafolder points, choose, img, target, model_points, idx = points.cuda(), \ choose.cuda(), \ img.cuda(), \ target.cuda(), \ model_points.cuda(), \ idx.cuda() cloud_path = "experiments/clouds/ycb/{0}/{1}/{2}/{3}_{4}".format( opt.output_dir, epoch, datafolder, filehead, int(idx)) # folder to save logs if not os.path.exists( "experiments/clouds/ycb/{0}/{1}/{2}".format( opt.output_dir, epoch, datafolder)): os.makedirs( "experiments/clouds/ycb/{0}/{1}/{2}".format( opt.output_dir, epoch, datafolder), exist_ok=True) pred_r, pred_t, pred_c, x_return = estimator( img, points, choose, idx, focal_length, principal_point, motion, cloud_path) _, dis, new_points, new_target, dis_vector = criterion( pred_r, pred_t, pred_c, dis_vector_last_map[idx.item()], target, model_points, idx, x_return, opt.w, opt.refine_start, float(opt.loss_stable_alpha)) dis_vector_last_map[idx.item()] = dis_vector test_dis += dis.item() logger.info( 'Test time {0} Test Frame No.{1} {2} {3} dis:{4}'. format( time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)), test_count, filename, idx.item(), dis)) test_count += 1 test_dis = test_dis / test_count logger.info('Test time {0} Epoch {1} TEST FINISH Avg dis: {2}'.format( time.strftime("%d %Hh %Mm %Ss", time.gmtime(time.time() - st_time)), epoch, test_dis)) if test_dis <= best_test: best_test = test_dis torch.save( estimator.state_dict(), '{0}/pose_model_ori_{1}_{2}.pth'.format( opt.outf, epoch, test_dis)) print(epoch, '>>>>>>>>----------BEST TEST MODEL SAVED---------<<<<<<<<') if best_test < opt.decay_margin and not opt.decay_start: opt.decay_start = True opt.lr *= opt.lr_rate opt.w *= opt.w_rate optimizer = optim.Adam(estimator.parameters(), lr=opt.lr)
def main(): opt.manualSeed = random.randint(1, 10000) random.seed(opt.manualSeed) torch.manual_seed(opt.manualSeed) if opt.dataset == 'ycb': opt.dataset_root = 'datasets/ycb/YCB_Video_Dataset' opt.num_objects = 21 opt.num_points = 1000 opt.result_dir = 'results/ycb' opt.repeat_epoch = 1 elif opt.dataset == 'linemod': opt.dataset_root = 'datasets/linemod/Linemod_preprocessed' opt.num_objects = 13 opt.num_points = 500 opt.result_dir = 'results/linemod' opt.repeat_epoch = 1 else: print('unknown dataset') return if opt.dataset == 'ycb': dataset = PoseDataset_ycb('train', opt.num_points, True, opt.dataset_root, opt.noise_trans) test_dataset = PoseDataset_ycb('test', opt.num_points, False, opt.dataset_root, 0.0) elif opt.dataset == 'linemod': dataset = PoseDataset_linemod('train', opt.num_points, True, opt.dataset_root, opt.noise_trans) test_dataset = PoseDataset_linemod('test', opt.num_points, False, opt.dataset_root, 0.0) 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.sym_list = dataset.get_sym_list() opt.num_points_mesh = dataset.get_num_points_mesh() opt.diameters = dataset.get_diameter() print('>>>>>>>>----------Dataset loaded!---------<<<<<<<<') print('length of the training set: {0}'.format(len(dataset))) print('length of the testing set: {0}'.format(len(test_dataset))) print('number of sample points on mesh: {0}'.format(opt.num_points_mesh)) print('symmetrical object list: {0}'.format(opt.sym_list)) if not os.path.exists(opt.result_dir): os.makedirs(opt.result_dir) tb_writer = tf.summary.FileWriter(opt.result_dir) os.environ['CUDA_VISIBLE_DEVICES'] = opt.gpu_id # network estimator = PoseNet(num_points=opt.num_points, num_obj=opt.num_objects, num_rot=opt.num_rot) estimator.cuda() # loss criterion = Loss(opt.sym_list, estimator.rot_anchors) knn = KNearestNeighbor(1) # learning rate decay best_test = np.Inf opt.first_decay_start = False opt.second_decay_start = False # if resume training if opt.resume_posenet != '': estimator.load_state_dict(torch.load(opt.resume_posenet)) model_name_parsing = (opt.resume_posenet.split('.')[0]).split('_') best_test = float(model_name_parsing[-1]) opt.start_epoch = int(model_name_parsing[-2]) + 1 if best_test < 0.016 and not opt.first_decay_start: opt.first_decay_start = True opt.lr *= 0.6 if best_test < 0.013 and not opt.second_decay_start: opt.second_decay_start = True opt.lr *= 0.5 # optimizer optimizer = torch.optim.Adam(estimator.parameters(), lr=opt.lr) global_step = (len(dataset) // opt.batch_size) * opt.repeat_epoch * (opt.start_epoch - 1) # train st_time = time.time() for epoch in range(opt.start_epoch, opt.nepoch): logger = setup_logger( 'epoch%02d' % epoch, os.path.join(opt.result_dir, 'epoch_%02d_train_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_loss_avg = 0.0 train_loss_r_avg = 0.0 train_loss_t_avg = 0.0 train_loss_reg_avg = 0.0 estimator.train() optimizer.zero_grad() for rep in range(opt.repeat_epoch): for i, data in enumerate(dataloader, 0): points, choose, img, target_t, target_r, model_points, idx, gt_t = data obj_diameter = opt.diameters[idx] points, choose, img, target_t, target_r, model_points, idx = Variable(points).cuda(), \ Variable(choose).cuda(), \ Variable(img).cuda(), \ Variable(target_t).cuda(), \ Variable(target_r).cuda(), \ Variable(model_points).cuda(), \ Variable(idx).cuda() pred_r, pred_t, pred_c = estimator(img, points, choose, idx) loss, loss_r, loss_t, loss_reg = criterion( pred_r, pred_t, pred_c, target_r, target_t, model_points, idx, obj_diameter) loss.backward() train_loss_avg += loss.item() train_loss_r_avg += loss_r.item() train_loss_t_avg += loss_t.item() train_loss_reg_avg += loss_reg.item() train_count += 1 if train_count % opt.batch_size == 0: global_step += 1 lr = opt.lr optimizer.step() optimizer.zero_grad() # write results to tensorboard summary = tf.Summary(value=[ tf.Summary.Value(tag='learning_rate', simple_value=lr), tf.Summary.Value(tag='loss', simple_value=train_loss_avg / opt.batch_size), tf.Summary.Value(tag='loss_r', simple_value=train_loss_r_avg / opt.batch_size), tf.Summary.Value(tag='loss_t', simple_value=train_loss_t_avg / opt.batch_size), tf.Summary.Value(tag='loss_reg', simple_value=train_loss_reg_avg / opt.batch_size) ]) tb_writer.add_summary(summary, global_step) logger.info( 'Train time {0} Epoch {1} Batch {2} Frame {3} Avg_loss:{4:f}' .format( time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)), epoch, int(train_count / opt.batch_size), train_count, train_loss_avg / opt.batch_size)) train_loss_avg = 0.0 train_loss_r_avg = 0.0 train_loss_t_avg = 0.0 train_loss_reg_avg = 0.0 print( '>>>>>>>>----------epoch {0} train finish---------<<<<<<<<'.format( epoch)) logger = setup_logger( 'epoch%02d_test' % epoch, os.path.join(opt.result_dir, 'epoch_%02d_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 save_model = False estimator.eval() success_count = [0 for i in range(opt.num_objects)] num_count = [0 for i in range(opt.num_objects)] for j, data in enumerate(testdataloader, 0): points, choose, img, target_t, target_r, model_points, idx, gt_t = data obj_diameter = opt.diameters[idx] points, choose, img, target_t, target_r, model_points, idx = Variable(points).cuda(), \ Variable(choose).cuda(), \ Variable(img).cuda(), \ Variable(target_t).cuda(), \ Variable(target_r).cuda(), \ Variable(model_points).cuda(), \ Variable(idx).cuda() pred_r, pred_t, pred_c = estimator(img, points, choose, idx) loss, _, _, _ = criterion(pred_r, pred_t, pred_c, target_r, target_t, model_points, idx, obj_diameter) test_count += 1 # evalaution how_min, which_min = torch.min(pred_c, 1) pred_r = pred_r[0][which_min[0]].view(-1).cpu().data.numpy() pred_r = quaternion_matrix(pred_r)[:3, :3] pred_t, pred_mask = ransac_voting_layer(points, pred_t) pred_t = pred_t.cpu().data.numpy() model_points = model_points[0].cpu().detach().numpy() pred = np.dot(model_points, pred_r.T) + pred_t target = target_r[0].cpu().detach().numpy() + gt_t[0].cpu( ).data.numpy() if idx[0].item() in opt.sym_list: pred = torch.from_numpy(pred.astype( np.float32)).cuda().transpose(1, 0).contiguous() target = torch.from_numpy(target.astype( np.float32)).cuda().transpose(1, 0).contiguous() inds = knn(target.unsqueeze(0), pred.unsqueeze(0)) target = torch.index_select(target, 1, inds.view(-1) - 1) dis = torch.mean(torch.norm( (pred.transpose(1, 0) - target.transpose(1, 0)), dim=1), dim=0).item() else: dis = np.mean(np.linalg.norm(pred - target, axis=1)) logger.info( 'Test time {0} Test Frame No.{1} loss:{2:f} confidence:{3:f} distance:{4:f}' .format( time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)), test_count, loss, how_min[0].item(), dis)) if dis < 0.1 * opt.diameters[idx[0].item()]: success_count[idx[0].item()] += 1 num_count[idx[0].item()] += 1 test_dis += dis # compute accuracy accuracy = 0.0 for i in range(opt.num_objects): accuracy += float(success_count[i]) / num_count[i] logger.info('Object {0} success rate: {1}'.format( test_dataset.objlist[i], float(success_count[i]) / num_count[i])) accuracy = accuracy / opt.num_objects test_dis = test_dis / test_count # log results logger.info( 'Test time {0} Epoch {1} TEST FINISH Avg dis: {2:f}, Accuracy: {3:f}' .format( time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)), epoch, test_dis, accuracy)) # tensorboard summary = tf.Summary(value=[ tf.Summary.Value(tag='accuracy', simple_value=accuracy), tf.Summary.Value(tag='test_dis', simple_value=test_dis) ]) tb_writer.add_summary(summary, global_step) # save model if test_dis < best_test: best_test = test_dis torch.save( estimator.state_dict(), '{0}/pose_model_{1:02d}_{2:06f}.pth'.format( opt.result_dir, epoch, best_test)) # adjust learning rate if necessary if best_test < 0.016 and not opt.first_decay_start: opt.first_decay_start = True opt.lr *= 0.6 optimizer = torch.optim.Adam(estimator.parameters(), lr=opt.lr) if best_test < 0.013 and not opt.second_decay_start: opt.second_decay_start = True opt.lr *= 0.5 optimizer = torch.optim.Adam(estimator.parameters(), lr=opt.lr) print( '>>>>>>>>----------epoch {0} test finish---------<<<<<<<<'.format( epoch))