def main():
parser = argparse.ArgumentParser(description="ReID Baseline Training")
parser.add_argument("--config_file",
default="",
help="path to config file",
type=str)
parser.add_argument("opts",
help="Modify config options using the command-line",
default=None,
nargs=argparse.REMAINDER)
args = parser.parse_args()
if args.config_file != "":
cfg.merge_from_file(args.config_file)
cfg.merge_from_list(args.opts)
cfg.freeze()
output_dir = cfg.OUTPUT_DIR
if output_dir and not os.path.exists(output_dir):
os.makedirs(output_dir)
logger = setup_logger("reid_baseline", output_dir, 0)
logger.info(args)
if args.config_file != "":
logger.info("Loaded configuration file {}".format(args.config_file))
logger.info("Running with config:\n{}".format(cfg))
os.environ['CUDA_VISIBLE_DEVICES'] = cfg.MODEL.DEVICE_ID
cudnn.benchmark = True
feats, dataset = inference(cfg, logger)
distmat = compute_jaccard_distance(feats)
DBSCAN_cluster(distmat, dataset, output_dir)
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 = PointCloudAE(opt.emb_dim, opt.num_point)
# estimator.cuda()
criterion = ChamferLoss()
# if opt.resume_model != '':
# estimator.load_state_dict(torch.load(opt.resume_model))
# dataset
# train: 1102, val: 206
# train_dataset = ShapeDataset(opt.h5_file, mode='train', augment=True)
# train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=opt.batch_size,
# shuffle=True, num_workers=opt.num_workers)
val_dataset = ShapeDataset(opt.h5_file, mode='val', augment=False)
val_dataloader = torch.utils.data.DataLoader(val_dataset,
batch_size=1,
shuffle=False,
num_workers=opt.num_workers)
val_loss = 0.0
cd_num = torch.zeros(6)
cd = torch.zeros(6)
for i, data in enumerate(val_dataloader, 1):
batch_xyz, batch_label = data
idx = batch_label.item()
batch_xyz = batch_xyz[:, :, :3].cuda()
prior = torch.from_numpy(mean_shapes[idx]).cuda().float().unsqueeze(0)
loss, _, _ = criterion(prior, batch_xyz)
cd_num[idx] += 1
cd[idx] += loss.item()
# zero divider
cd_metric = (cd / cd_num) * 1000
print(
"{:.2f} : {:.2f} : {:.2f} : {:.2f} : {:.2f} : {:.2f} : {:.2f}".format(
cd_metric[0], cd_metric[1], cd_metric[2], cd_metric[3],
cd_metric[4], cd_metric[5], torch.mean(cd_metric)))
def main():
parser = argparse.ArgumentParser(description="ReID Baseline Training")
parser.add_argument("--local_rank", default=0, type=int)
parser.add_argument("--config_file",
default="",
help="path to config file",
type=str)
parser.add_argument("opts",
help="Modify config options using the command-line",
default=None,
nargs=argparse.REMAINDER)
args = parser.parse_args()
if args.config_file != "":
cfg.merge_from_file(args.config_file)
cfg.merge_from_list(args.opts)
cfg.freeze()
output_dir = cfg.OUTPUT_DIR
if output_dir and not os.path.exists(output_dir):
os.makedirs(output_dir)
logger = setup_logger("reid_baseline", output_dir, 0)
logger.info(args)
if args.config_file != "":
logger.info("Loaded configuration file {}".format(args.config_file))
logger.info("Running with config:\n{}".format(cfg))
cudnn.benchmark = True
distributed = int(
os.environ['WORLD_SIZE']) > 1 if 'WORLD_SIZE' in os.environ else False
if distributed:
gpu = args.local_rank
torch.cuda.set_device(gpu)
torch.distributed.init_process_group(backend='nccl',
init_method='env://')
args.world_size = torch.distributed.get_world_size()
else:
os.environ['CUDA_VISIBLE_DEVICES'] = cfg.MODEL.DEVICE_ID
naive_train(cfg, logger, distributed, args.local_rank)
def main():
parser = argparse.ArgumentParser(description="ReID Baseline Training")
parser.add_argument("--config_file",
default="",
help="path to config file",
type=str)
parser.add_argument("opts",
help="Modify config options using the command-line",
default=None,
nargs=argparse.REMAINDER)
args = parser.parse_args()
if args.config_file != "":
cfg.merge_from_file(args.config_file)
cfg.merge_from_list(args.opts)
cfg.freeze()
output_dir = cfg.OUTPUT_DIR
if output_dir and not os.path.exists(output_dir):
os.makedirs(output_dir)
logger = setup_logger("reid_baseline", output_dir, 0)
logger.info(args)
if args.config_file != "":
logger.info("Loaded configuration file {}".format(args.config_file))
logger.info("Running with config:\n{}".format(cfg))
os.environ['CUDA_VISIBLE_DEVICES'] = cfg.MODEL.DEVICE_ID
cudnn.benchmark = True
batch, num_query = inference(cfg, logger)
cmc, mAP, indices_np = post_processor(cfg, batch, num_query)
logger.info('Validation Results')
logger.info("mAP: {:.1%}".format(mAP))
for r in [1, 5, 10]:
logger.info("CMC curve, Rank-{:<3}:{:.1%}".format(r, cmc[r - 1]))
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 = PointCloudAE(opt.emb_dim, opt.num_point)
estimator.cuda()
criterion = ChamferLoss()
if opt.resume_model != '':
estimator.load_state_dict(torch.load(opt.resume_model))
# dataset
train_dataset = ShapeDataset(opt.h5_file, mode='train', augment=True)
train_dataloader = torch.utils.data.DataLoader(train_dataset,
batch_size=opt.batch_size,
shuffle=True,
num_workers=opt.num_workers)
val_dataset = ShapeDataset(opt.h5_file, mode='val', augment=False)
val_dataloader = torch.utils.data.DataLoader(val_dataset,
batch_size=opt.batch_size,
shuffle=False,
num_workers=opt.num_workers)
# train
st_time = time.time()
global_step = ((train_dataset.length + opt.batch_size - 1) //
opt.batch_size) * opt.repeat_epoch * (opt.start_epoch - 1)
decay_count = -1
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 global_step // opt.decay_step > decay_count:
decay_count += 1
if decay_count < len(opt.decay_rate):
current_lr = opt.lr * opt.decay_rate[decay_count]
optimizer = torch.optim.Adam(estimator.parameters(),
lr=current_lr)
batch_idx = 0
estimator.train()
for rep in range(opt.repeat_epoch):
for i, data in enumerate(train_dataloader):
# label must be zero_indexed
batch_xyz, batch_label = data
batch_xyz = batch_xyz[:, :, :3].cuda()
optimizer.zero_grad()
embedding, point_cloud = estimator(batch_xyz)
loss, _, _ = criterion(point_cloud, batch_xyz)
summary = tf.Summary(value=[
tf.Summary.Value(tag='learning_rate',
simple_value=current_lr),
tf.Summary.Value(tag='train_loss', simple_value=loss)
])
# backward
loss.backward()
optimizer.step()
global_step += 1
batch_idx += 1
# write results to tensorboard
tb_writer.add_summary(summary, global_step)
if batch_idx % 10 == 0:
logger.info('Batch {0} Loss:{1:f}'.format(batch_idx, loss))
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'))
estimator.eval()
val_loss = 0.0
for i, data in enumerate(val_dataloader, 1):
batch_xyz, batch_label = data
batch_xyz = batch_xyz[:, :, :3].cuda()
embedding, point_cloud = estimator(batch_xyz)
loss, _, _ = criterion(point_cloud, batch_xyz)
val_loss += loss.item()
logger.info('Batch {0} Loss:{1:f}'.format(i, loss))
val_loss = val_loss / i
summary = tf.Summary(
value=[tf.Summary.Value(tag='val_loss', simple_value=val_loss)])
tb_writer.add_summary(summary, global_step)
logger.info('Epoch {0:02d} test average loss: {1:06f}'.format(
epoch, val_loss))
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 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))
batch_time.update(time.time() - end)
end = time.time()
if i % 25 == 0:
logger.info('Validate: [{0}/{1}]\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'.format(
i,
len(val_loader),
batch_time=batch_time,
loss=losses))
logger.info('Finished validation')
return losses.avg
if __name__ == '__main__':
opt = parse_options()
torch.cuda.set_device(opt.local_rank)
torch.distributed.init_process_group(backend='nccl', init_method='env://')
cudnn.benchmark = True
logger = setup_logger(output=opt.save_model,
distributed_rank=dist.get_rank(),
name='vos')
main(opt)
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 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_global_duel_bing' #folder to save trained models
opt.log_dir = 'experiments/logs/ycb_global_duel_bing' #folder to save logs
opt.repeat_epoch = 1 #number of repeat times for one epoch training
estimator = PoseNetBinghamDuel(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)))
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.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 = DuelLoss(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()
cum_batch_count = 0
mean_z = 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
estimator.train()
optimizer.zero_grad()
for rep in range(opt.repeat_epoch):
for i, data in enumerate(dataloader, 0):
points, choose, img, target, idx = data
idx = idx - 1
points, choose, img, target, idx = Variable(points).cuda(), \
Variable(choose).cuda(), \
Variable(img).cuda(), \
Variable(target).cuda(), \
Variable(idx).cuda()
pred_q1, pred_q2, pred_z, _ = estimator(
img, points, choose, idx)
loss, dis = criterion(pred_q1.view(-1), pred_q2.view(-1),
torch.abs(pred_z.view(-1)), target)
loss.backward()
mean_z += torch.sum(torch.abs(pred_z.view(-1).detach()))
pred_q = makeBinghamM(
pred_q1.view(-1).detach(),
pred_q2.view(-1).detach())[:, 0]
pred_q = pred_q / pred_q.norm()
mean_err += tensorAngularDiff(pred_q, target) * 180 / np.pi
train_dis_avg += dis.item()
train_count += 1
if train_count % opt.batch_size == 0:
cum_batch_count += 1
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 / len(dataset),
train_dis_avg / opt.batch_size))
optimizer.step()
optimizer.zero_grad()
if (cum_batch_count % 100 == 0):
train_writer.add_scalar('loss', loss, cum_batch_count)
train_writer.add_scalar('lik', dis, cum_batch_count)
train_writer.add_scalar('mean_lik',
train_dis_avg / opt.batch_size,
cum_batch_count)
train_writer.add_scalar(
'mean_z', mean_z / (100 * opt.batch_size),
cum_batch_count)
train_writer.add_scalar(
'mean_err', mean_err / (100 * opt.batch_size),
cum_batch_count)
mean_sig = 0
mean_err = 0
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))
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()
mean_z = 0
mean_err = 0
for j, data in enumerate(testdataloader, 0):
points, choose, img, target, idx = data
idx = idx - 1
points, choose, img, target, idx = Variable(points).cuda(), \
Variable(choose).cuda(), \
Variable(img).cuda(), \
Variable(target).cuda(), \
Variable(idx).cuda()
pred_q1, pred_q2, pred_z, _ = estimator(img, points, choose, idx)
loss, dis = criterion(pred_q1.view(-1), pred_q2.view(-1),
torch.abs(pred_z.view(-1)), target)
mean_z += torch.sum(torch.abs(pred_z.view(-1).detach()))
pred_q = makeBinghamM(
pred_q1.view(-1).detach(),
pred_q2.view(-1).detach())[:, 0]
pred_q = pred_q / pred_q.norm()
mean_err += tensorAngularDiff(pred_q, target) * 180 / np.pi
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))
valid_writer.add_scalar('loss', loss, cum_batch_count)
valid_writer.add_scalar('lik', dis, cum_batch_count)
valid_writer.add_scalar('mean_lik', test_dis, cum_batch_count)
valid_writer.add_scalar('mean_z', mean_z / test_count, cum_batch_count)
valid_writer.add_scalar('mean_err', mean_err / test_count,
cum_batch_count)
mean_sig = 0
mean_err = 0
if test_dis <= best_test:
best_test = test_dis
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)
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_global_mnorm' #folder to save trained models
opt.log_dir = 'experiments/logs/ycb_global_mnorm' #folder to save logs
opt.repeat_epoch = 1 #number of repeat times for one epoch training
if not os.path.exists(opt.outf):
os.makedirs(opt.outf)
if not os.path.exists(opt.log_dir):
os.makedirs(opt.log_dir)
estimator = PoseNetGlobal(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',
object_list = object_list,
output_data = output_format,
resample_on_error = True,
preprocessors = [YCBOcclusionAugmentor(opt.dataset_root),
ColorJitter(),
InplaneRotator()],
postprocessors = [ImageNormalizer(), PointMeanNormalizer(0)], #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(), PointMeanNormalizer(0)],
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=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))
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
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)
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():
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)
if i%20==0:
logger.info('Epoch is [{}/{}], mini-batch is [{}/{}], time consumption is {:.8f}, batch_loss is {:.8f}'.format( \
epoch + 1, args.epoch_iter, i + 1, int(len(train_loader_source)), time.time() - start_time, loss.item()))
# if i>4000 and i%1000==0:
# f_score = test(epoch, model, args.t_eval_path, args.t_output_path, f_score, args.save_model)
logger.info('epoch_loss is {:.8f}, epoch_time is {:.8f}'.format(epoch_loss / int(7200 / args.batch_size),time.time() - epoch_time))
logger.info(time.asctime(time.localtime(time.time())))
if __name__ == '__main__':
args.workspace = os.path.join(args.workspace, args.exp_name)
os.makedirs(args.workspace, exist_ok=True)
logger = setup_logger(os.path.join(args.workspace, 'train_MSRA_log'))
criterion = Loss()
device = torch.device("cuda")
model = EAST()
# model = nn.DataParallel(model)
data_parallel = False
if torch.cuda.device_count() > 1:
model = nn.DataParallel(model)
data_parallel = True
model.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=[150,220], gamma=0.1)
# 先产生第一次的pseudo-label
logger.info("loading pretrained model from "+args.resume)
# model.load_state_dict(torch.load(args.resume))
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))
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)
# print("\n")
# print(" ---------------------------------------------------------")
# print(" best_f_score:", f_score)
# print(" ---------------------------------------------------------")
return f_score_new
if __name__ == '__main__':
train_img_path = os.path.abspath('/data/data_weijiawu/SynthText')
train_gt_path = os.path.abspath('/data/data_weijiawu/SynthText/gt.mat')
args.workspace = os.path.join(args.workspace, args.exp_name)
os.makedirs(args.workspace, exist_ok=True)
logger = setup_logger(os.path.join(args.workspace, 'train_synthtext_log'))
criterion = Loss()
device = torch.device("cuda")
model = EAST()
data_parallel = False
if torch.cuda.device_count() > 1:
model = nn.DataParallel(model)
data_parallel = True
model.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
scheduler = lr_scheduler.MultiStepLR(optimizer,
milestones=[args.epoch_iter // 2],
gamma=0.1)
# 先产生第一次的pseudo-label
# logger.info("loading pretrained model from ",args.resume)
logger.info('Epoch is [{}/{}], mini-batch is [{}/{}], time consumption is {:.8f}, batch_loss is {:.8f}'.format( \
epoch + 1, args.epoch_iter, i + 1, int(len(train_loader_source)), time.time() - start_time, loss.item()))
# if i>4000 and i%1000==0:
# f_score = test(epoch, model, args.t_eval_path, args.t_output_path, f_score, args.save_model)
logger.info('epoch_loss is {:.8f}, epoch_time is {:.8f}'.format(
epoch_loss / int(7200 / args.batch_size),
time.time() - epoch_time))
logger.info(time.asctime(time.localtime(time.time())))
if __name__ == '__main__':
args.workspace = os.path.join(args.workspace, args.exp_name)
os.makedirs(args.workspace, exist_ok=True)
logger = setup_logger(os.path.join(args.workspace, 'train_icdar15_log'))
criterion = Loss()
device = torch.device("cuda")
model = EAST()
data_parallel = False
if torch.cuda.device_count() > 1:
model = nn.DataParallel(model)
data_parallel = True
model.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
scheduler = lr_scheduler.MultiStepLR(optimizer,
milestones=[args.epoch_iter // 2],
gamma=0.1)
# 先产生第一次的pseudo-label
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)
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 main(args):
this_dir = osp.join(osp.dirname(__file__), '.')
save_dir = osp.join(this_dir, 'checkpoints')
if not osp.isdir(save_dir):
os.makedirs(save_dir)
command = 'python ' + ' '.join(sys.argv)
logger = utl.setup_logger(osp.join(this_dir, 'log.txt'), command=command)
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
utl.set_seed(int(args.seed))
model = build_model(args)
if osp.isfile(args.checkpoint):
checkpoint = torch.load(args.checkpoint,
map_location=torch.device('cpu'))
model.load_state_dict(checkpoint['model_state_dict'])
else:
model.apply(utl.weights_init)
if args.distributed:
model = nn.DataParallel(model)
model = model.to(device)
criterion = utl.MultiCrossEntropyLoss(ignore_index=21).to(device)
optimizer = optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
if osp.isfile(args.checkpoint):
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
for param_group in optimizer.param_groups:
param_group['lr'] = args.lr
args.start_epoch += checkpoint['epoch']
softmax = nn.Softmax(dim=1).to(device)
for epoch in range(args.start_epoch, args.start_epoch + args.epochs):
if epoch == 21:
args.lr = args.lr * 0.1
for param_group in optimizer.param_groups:
param_group['lr'] = args.lr
data_loaders = {
phase: utl.build_data_loader(args, phase)
for phase in args.phases
}
enc_losses = {phase: 0.0 for phase in args.phases}
enc_score_metrics = []
enc_target_metrics = []
enc_mAP = 0.0
dec_losses = {phase: 0.0 for phase in args.phases}
dec_score_metrics = []
dec_target_metrics = []
dec_mAP = 0.0
start = time.time()
for phase in args.phases:
training = phase == 'train'
if training:
model.train(True)
elif not training and args.debug:
model.train(False)
else:
continue
with torch.set_grad_enabled(training):
for batch_idx, (camera_inputs, motion_inputs, enc_target, dec_target) \
in enumerate(data_loaders[phase], start=1):
batch_size = camera_inputs.shape[0]
camera_inputs = camera_inputs.to(device)
motion_inputs = motion_inputs.to(device)
enc_target = enc_target.to(device).view(
-1, args.num_classes)
dec_target = dec_target.to(device).view(
-1, args.num_classes)
enc_score, dec_score = model(camera_inputs, motion_inputs)
enc_loss = criterion(enc_score, enc_target)
dec_loss = criterion(dec_score, dec_target)
enc_losses[phase] += enc_loss.item() * batch_size
dec_losses[phase] += dec_loss.item() * batch_size
if args.verbose:
print(
'Epoch: {:2} | iteration: {:3} | enc_loss: {:.5f} dec_loss: {:.5f}'
.format(epoch, batch_idx, enc_loss.item(),
dec_loss.item()))
if training:
optimizer.zero_grad()
loss = enc_loss + dec_loss
loss.backward()
optimizer.step()
else:
# Prepare metrics for encoder
enc_score = softmax(enc_score).cpu().numpy()
enc_target = enc_target.cpu().numpy()
enc_score_metrics.extend(enc_score)
enc_target_metrics.extend(enc_target)
# Prepare metrics for decoder
dec_score = softmax(dec_score).cpu().numpy()
dec_target = dec_target.cpu().numpy()
dec_score_metrics.extend(dec_score)
dec_target_metrics.extend(dec_target)
end = time.time()
if args.debug:
result_file = 'inputs-{}-epoch-{}.json'.format(args.inputs, epoch)
# Compute result for encoder
enc_mAP = utl.compute_result_multilabel(
args.class_index,
enc_score_metrics,
enc_target_metrics,
save_dir,
result_file,
ignore_class=[0, 21],
save=True,
)
# Compute result for decoder
dec_mAP = utl.compute_result_multilabel(
args.class_index,
dec_score_metrics,
dec_target_metrics,
save_dir,
result_file,
ignore_class=[0, 21],
save=False,
)
# Output result
logger.output(epoch,
enc_losses,
dec_losses,
len(data_loaders['train'].dataset),
len(data_loaders['test'].dataset),
enc_mAP,
dec_mAP,
end - start,
debug=args.debug)
# Save model
checkpoint_file = 'inputs-{}-epoch-{}.pth'.format(args.inputs, epoch)
torch.save(
{
'epoch':
epoch,
'model_state_dict':
model.module.state_dict()
if args.distributed else model.state_dict(),
'optimizer_state_dict':
optimizer.state_dict(),
}, osp.join(save_dir, checkpoint_file))
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(args, cfg):
# Set logger
logger = setup_logger(args.mode,
cfg.DIRS.LOGS,
0,
filename=f"{cfg.EXP}.txt")
# Declare variables
best_metric = 0.
start_cycle = 0
start_epoch = 0
# Define model
if args.mode == "distil":
student_model = build_sem_seg_model(cfg)
teacher_model = build_sem_seg_model(cfg)
optimizer = make_optimizer(cfg, student_model)
else:
model = build_sem_seg_model(cfg)
if args.mode == "swa":
swa_model = build_sem_seg_model(cfg)
if cfg.DATA.AUGMENT == "augmix":
from timm.models import convert_splitbn_model
model = convert_splitbn_model(model, 3)
swa_model = convert_splitbn_model(model, 3)
optimizer = make_optimizer(cfg, model)
# Define loss
loss_name = cfg.LOSS.NAME
if loss_name == "bce":
train_criterion = nn.BCEWithLogitsLoss()
elif loss_name == "focal":
train_criterion = SigmoidFocalLoss(1.25, 0.25)
elif loss_name == "dice":
train_criterion = BinaryDiceLoss()
elif loss_name == "iou":
train_criterion = BinaryIoULoss()
# CUDA & Mixed Precision
if cfg.SYSTEM.CUDA:
if args.mode == "distil":
student_model = student_model.cuda()
teacher_model = teacher_model.cuda()
elif args.mode == "swa":
model = model.cuda()
swa_model = swa_model.cuda()
else:
model = model.cuda()
train_criterion = train_criterion.cuda()
if cfg.SYSTEM.FP16:
bn_fp32 = True if cfg.SYSTEM.OPT_L == "O2" else None
if args.mode == "distil":
[student_model, teacher_model], optimizer = amp.initialize(
models=[student_model, teacher_model],
optimizers=optimizer,
opt_level=cfg.SYSTEM.OPT_L,
keep_batchnorm_fp32=bn_fp32)
if args.mode == "swa":
[model, swa_model
], optimizer = amp.initialize(models=[model, swa_model],
optimizers=optimizer,
opt_level=cfg.SYSTEM.OPT_L,
keep_batchnorm_fp32=bn_fp32)
else:
model, optimizer = amp.initialize(models=model,
optimizers=optimizer,
opt_level=cfg.SYSTEM.OPT_L,
keep_batchnorm_fp32=bn_fp32)
# Load checkpoint
if args.load != "":
if os.path.isfile(args.load):
logger.info(f"=> loading checkpoint {args.load}")
ckpt = torch.load(args.load, "cpu")
model.load_state_dict(ckpt.pop('state_dict'))
if args.swa:
swa_model.load_state_dict(model.state_dict())
if not args.finetune:
logger.info("resuming optimizer ...")
optimizer.load_state_dict(ckpt.pop('optimizer'))
if args.mode == "cycle":
start_cycle = ckpt["cycle"]
start_epoch, best_metric = ckpt['epoch'], ckpt['best_metric']
logger.info(
f"=> loaded checkpoint '{args.load}' (epoch {ckpt['epoch']}, best_metric: {ckpt['best_metric']})"
)
if args.mode == "swa":
ckpt = torch.load(args.load, "cpu")
swa_model.load_state_dict(ckpt.pop('state_dict'))
else:
logger.info(f"=> no checkpoint found at '{args.load}'")
if cfg.SYSTEM.MULTI_GPU:
model = nn.DataParallel(model)
# Load and split data
train_loader = make_dataloader(cfg, "train")
valid_loader = make_dataloader(cfg, "valid")
scheduler = WarmupCyclicalLR("cos",
cfg.OPT.BASE_LR,
cfg.TRAIN.EPOCHS,
iters_per_epoch=len(train_loader) //
cfg.OPT.GD_STEPS,
warmup_epochs=cfg.OPT.WARMUP_EPOCHS)
if args.mode == "train":
for epoch in range(start_epoch, cfg.TRAIN.EPOCHS):
train_loop(logger.info, cfg, model, train_loader, train_criterion,
optimizer, scheduler, epoch)
best_metric = valid_model(logger.info, cfg, model, valid_loader,
optimizer, epoch, None, best_metric,
True)
elif args.mode == "cycle":
for cycle in range(start_cycle, cfg.TRAIN.NUM_CYCLES):
for epoch in range(start_epoch, cfg.TRAIN.EPOCHS):
train_loop(logger.info, cfg, model, train_loader,
train_criterion, optimizer, scheduler, epoch)
best_metric = valid_model(logger.info, cfg, model,
valid_loader, optimizer, epoch,
cycle, best_metric, True)
# reset scheduler for new cycle
scheduler = WarmupCyclicalLR("cos",
cfg.OPT.BASE_LR,
cfg.TRAIN.EPOCHS,
iters_per_epoch=len(train_loader) //
cfg.OPT.GD_STEPS,
warmup_epochs=cfg.OPT.WARMUP_EPOCHS)
elif args.mode == "distil":
for epoch in range(start_epoch, cfg.TRAIN.EPOCHS):
distil_train_loop(logger.info, cfg, student_model, teacher_model,
train_loader, train_criterion, optimizer,
scheduler, epoch)
best_metric = valid_model(logger.info, cfg, student_model,
valid_loader, optimizer, epoch, None,
best_metric, True)
elif args.mode == "swa":
for epoch in range(start_epoch, cfg.TRAIN.EPOCHS):
train_loop(logger.info, cfg, model, train_loader, train_criterion,
optimizer, scheduler, epoch)
best_metric = valid_model(logger.info, cfg, model, valid_loader,
optimizer, epoch, None, best_metric,
True)
if (epoch + 1) == cfg.OPT.SWA.START:
copy_model(swa_model, model)
swa_n = 0
if ((epoch + 1) >= cfg.OPT.SWA.START) and ((epoch + 1) %
cfg.OPT.SWA.FREQ == 0):
moving_average(swa_model, model, 1.0 / (swa_n + 1))
swa_n += 1
bn_update(train_loader, swa_model)
best_metric = valid_model(logger.info, cfg, swa_model,
valid_loader, optimizer, epoch, None,
best_metric, True)
elif args.mode == "valid":
valid_model(logger.info, cfg, model, valid_loader, optimizer,
start_epoch)
else:
test_loader = make_dataloader(cfg, "test")
predictions = test_model(logger.info, cfg, model, test_loader)
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')
opt.resume_model))
model.load_state_dict(checkpoint)
for log in os.listdir(opt.log_dir):
os.remove(os.path.join(opt.log_dir, log))
optimizer = optim.Adam(model.parameters(), lr=opt.lr)
criterion = Loss()
best_val_cost = np.Inf
st_time = time.time()
for epoch in range(1, opt.n_epochs):
model.train()
train_all_cost = 0.0
train_time = 0
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'))
for i, data in enumerate(dataloader, 0):
rgb, target = data
rgb, target = Variable(rgb).cuda(), Variable(target).cuda()
semantic = model(rgb)
optimizer.zero_grad()
semantic_loss = criterion(semantic, target)
train_all_cost += semantic_loss.item()
semantic_loss.backward()
optimizer.step()
logger.info('Train time {0} Batch {1} CEloss {2}'.format(
def main(args):
this_dir = osp.join(osp.dirname(__file__), '.')
### make directory for each step size
# '/dataset/volume1/users/yumin/result'
#'/data/yumin/result'
save_dir = osp.join(
'/dataset/volume1/users/yumin/result',
'delta_{}_checkpoints_method{}_noenc_smoothbeta0.5'.format(
args.dataset, args.method))
if not osp.isdir(save_dir):
os.makedirs(save_dir)
command = 'python ' + ' '.join(sys.argv)
logger = utl.setup_logger(osp.join(this_dir, 'lstm_log.txt'),
command=command)
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
utl.set_seed(int(args.seed))
model = build_model(args)
if osp.isfile(args.checkpoint):
checkpoint = torch.load(args.checkpoint,
map_location=torch.device('cpu'))
model.load_state_dict(checkpoint['model_state_dict'])
else:
model.apply(utl.weights_init)
if args.distributed: ### !!!
model = nn.DataParallel(model)
model = model.to(device)
if args.dataset == 'THUMOS':
criterion1 = utl.MultiCrossEntropyLoss_Delta(
num_class=args.num_classes,
dirichlet=args.dirichlet,
ignore_index=21).to(device)
# criterion2 = nn.MSELoss()
# criterion2 = nn.L1Loss()
criterion2 = nn.SmoothL1Loss()
# criterion2 = nn.HuberLoss()
elif args.dataset == "TVSeries":
criterion = utl.MultiCrossEntropyLoss_Delta(
num_class=args.num_classes, dirichlet=args.dirichlet).to(device)
optimizer = optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
if osp.isfile(args.checkpoint):
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
for param_group in optimizer.param_groups:
param_group['lr'] = args.lr
args.start_epoch += checkpoint['epoch']
softmax = nn.Softmax(dim=1).to(device)
for epoch in range(args.start_epoch, args.start_epoch + args.epochs):
if epoch == 21:
args.lr = args.lr * 0.1
for param_group in optimizer.param_groups:
param_group['lr'] = args.lr
data_loaders = {
phase: utl.build_data_loader(args, phase)
for phase in args.phases
}
enc_losses = {phase: 0.0 for phase in args.phases}
enc_score_metrics = []
enc_target_metrics = []
delta_score_metrics = []
delta_target_metrics = []
enc_mAP = 0.0
delta_mAP = 0.0
start = time.time()
for phase in args.phases:
training = phase == 'train'
if training:
model.train(True)
elif not training and args.debug:
model.train(False)
else:
continue
with torch.set_grad_enabled(training):
for batch_idx, (camera_inputs, motion_inputs, enc_target, smooth_target) \
in enumerate(data_loaders[phase], start=1):
batch_size = camera_inputs.shape[0]
camera_inputs = camera_inputs.to(device)
motion_inputs = motion_inputs.to(device)
extend_target = enc_target.to(device)
enc_target = enc_target.to(device).view(
-1, args.num_classes)
smooth_target = smooth_target.to(device)
oad_score, delta_score = model(camera_inputs,
motion_inputs)
oad_before = oad_score.clone().detach()
oad_before = oad_before[:, 1::, :]
## have to make delta target and compute delta loss
new_target = smooth_target[:, 1::, :] - oad_before
# print('***** DELTA TARGET')
# print(new_target)
# print('***** DELTA SCORE')
# print(delta_score[:,1::,:])
oad_loss = criterion1(oad_score, extend_target)
delta_loss = criterion2(delta_score[:, 1::, :],
new_target) # ignore the first
# delta_loss = criterion2(delta_score[:,1::,:], extend_target[:,1::,:]) # without labelsmoothing
enc_losses[phase] += oad_loss.item() * batch_size
if args.verbose:
print(
'Epoch: {:2} | iteration: {:3} | enc_loss: {:.5f} | delta_loss: {:.5f}'
.format(epoch, batch_idx, oad_loss.item(),
delta_loss.item() * 10))
if training:
optimizer.zero_grad()
loss = oad_loss + delta_loss * 10
loss.backward()
optimizer.step()
else:
# Prepare metrics for encoder
enc_score = oad_score.cpu().numpy() ## softmax check
enc_target = extend_target.cpu().numpy()
enc_score_metrics.extend(enc_score)
enc_target_metrics.extend(enc_target)
delta_score_c = delta_score[:, 1::, :].reshape(
-1, args.num_classes)
delta = delta_score_c.cpu().numpy()
new_target_c = new_target.reshape(-1, args.num_classes)
delta_target = new_target_c.cpu().numpy()
delta_score_metrics.extend(delta)
delta_target_metrics.extend(delta_target)
end = time.time()
if args.debug:
if epoch % 1 == 0:
result_file = osp.join(
this_dir,
'delta-inputs-{}-epoch-{}.json'.format(args.inputs, epoch))
# Compute result for encoder
enc_mAP = utl.compute_result_multilabel(
args.dataset,
args.class_index,
enc_score_metrics,
enc_target_metrics,
save_dir,
result_file,
ignore_class=[0, 21],
save=True,
)
delta_mAP = utl.compute_result_multilabel(
args.dataset,
args.class_index,
delta_score_metrics,
delta_target_metrics,
save_dir,
result_file,
ignore_class=[0, 21],
save=True,
smooth=True,
)
# Output result
logger.delta_output(epoch,
enc_losses,
len(data_loaders['train'].dataset),
len(data_loaders['test'].dataset),
enc_mAP,
delta_mAP,
end - start,
debug=args.debug)
# Save model
checkpoint_file = 'delta-inputs-{}-epoch-{}.pth'.format(
args.inputs, epoch)
torch.save(
{
'epoch':
epoch,
'model_state_dict':
model.module.state_dict()
if args.distributed else model.state_dict(),
'optimizer_state_dict':
optimizer.state_dict(),
}, osp.join(save_dir, checkpoint_file))
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)
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)