def main():
opt.manualSeed = random.randint(1, 10000)
random.seed(opt.manualSeed)
torch.manual_seed(opt.manualSeed)
if not os.path.exists(opt.output_folder):
os.makedirs(opt.output_folder)
num_points = 1000 #number of points on the input pointcloud
num_objects = 21
estimator = PoseNetGlobal(num_points=num_points, num_obj=num_objects)
estimator.cuda()
estimator.load_state_dict(torch.load(opt.weights))
output_format = [
otypes.OBJECT_LABEL, otypes.QUATERNION, otypes.IMAGE_CROPPED,
otypes.DEPTH_POINTS_MASKED_AND_INDEXES
]
estimator.eval()
pbar = trange(2, num_objects + 1)
for cls in pbar:
dataset = YCBDataset(
opt.dataset_root,
mode='grid',
object_list=[cls],
output_data=output_format,
resample_on_error=True,
preprocessors=[
YCBOcclusionAugmentor(opt.dataset_root),
ColorJitter(),
],
postprocessors=[ImageNormalizer(),
PointShifter()],
image_size=[640, 480],
num_points=1000)
classes = dataset.classes
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=1,
shuffle=False,
num_workers=opt.workers)
pbar.set_description('Featurizing {}'.format(classes[cls]))
for aug_idx in trange(opt.num_augmentations):
for i, data in tqdm(enumerate(dataloader), total=len(dataloader)):
if (len(data) == 0 or len(data[0]) == 0):
continue
idx, quat, img, points, choose = data
data_path = dataset.image_list[i]
idx = idx - 1
img = Variable(img).cuda()
points = Variable(points).cuda()
choose = Variable(choose).cuda()
idx = Variable(idx).cuda()
assert cls == data_path[1]
assert cls - 1 == int(idx[0])
feat, _ = estimator.globalFeature(img, points, choose, idx)
output_filename = '{0}/{1}_{2}_{3}_feat.npz'.format(
opt.output_folder, data_path[0], classes[cls], aug_idx)
os.makedirs(os.path.dirname(output_filename), exist_ok=True)
np.savez(output_filename,
quat=to_np(quat)[0],
feat=to_np(feat)[0])
def main():
opt.manualSeed = random.randint(1, 10000)
random.seed(opt.manualSeed)
torch.manual_seed(opt.manualSeed)
if not os.path.exists(opt.output_folder):
os.makedirs(opt.output_folder)
num_points = 1000 #number of points on the input pointcloud
num_objects = 21
estimator = PoseNetGlobal(num_points=num_points, num_obj=num_objects)
estimator.cuda()
estimator.load_state_dict(torch.load(opt.weights))
output_format = [
otypes.QUATERNION, otypes.IMAGE_CROPPED,
otypes.DEPTH_POINTS_MASKED_AND_INDEXES
]
estimator.eval()
for cls in trange(1, num_objects + 1):
dataset = YCBDataset(opt.dataset_root,
mode=opt.mode,
object_list=[cls],
output_data=output_format,
postprocessor=ImageNormalizer,
image_size=[640, 480],
num_points=1000)
classes = dataset.classes
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=1,
shuffle=False,
num_workers=opt.workers)
for i, data in tqdm(enumerate(dataloader), total=len(dataloader)):
if (len(data) == 0 or len(data[0]) == 0):
continue
quat, img, points, choose = data
data_path = dataset.image_list[i]
img = Variable(img).cuda()
points = Variable(points).cuda()
choose = Variable(choose).cuda()
idx = Variable(torch.LongTensor(cls - 1)).cuda()
assert cls == data_path[1]
feat, _ = estimator.globalFeature(img, points, choose, idx)
output_filename = '{0}/{1}_{2}_feat.npz'.format(
opt.output_folder, data_path[0], classes[cls])
os.makedirs(os.path.dirname(output_filename), exist_ok=True)
np.savez(output_filename, quat=to_np(quat)[0], feat=to_np(feat)[0])
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)
opt.num_objects = 21 #number of object classes in the dataset
opt.num_points = 1000 #number of points on the input pointcloud
opt.outf = 'trained_models/ycb_plus_bing' #folder to save trained models
opt.log_dir = 'experiments/logs/ycb_plus_bing' #folder to save logs
opt.repeat_epoch = 1 #number of repeat times for one epoch training
estimator = PoseNetPlusDuelBing(num_points=opt.num_points,
num_obj=opt.num_objects)
estimator.cuda()
train_writer = SummaryWriter(comment='duel_binham_train')
valid_writer = SummaryWriter(comment='duel_binham_valid')
if opt.resume_posenet != '':
estimator.load_state_dict(
torch.load('{0}/{1}'.format(opt.outf, opt.resume_posenet)))
elif opt.finetune_posenet != '':
pretrained_dict = torch.load(opt.finetune_posenet)
model_dict = estimator.state_dict()
pretrained_dict = {
k: v
for k, v in pretrained_dict.items() if k in model_dict
}
model_dict.update(pretrained_dict)
estimator.load_state_dict(model_dict)
for k, v in estimator.named_parameters():
if (k in pretrained_dict):
v.requires_grad = False
opt.log_dir += '_cont'
opt.outf += '_cont'
opt.refine_start = False
opt.decay_start = False
optimizer = optim.Adam(estimator.parameters(), lr=opt.lr)
object_list = list(range(1, 22))
output_format = [
otypes.DEPTH_POINTS_MASKED_AND_INDEXES,
otypes.IMAGE_CROPPED,
otypes.QUATERNION,
otypes.MODEL_POINTS_TRANSFORMED,
otypes.MODEL_POINTS,
otypes.OBJECT_LABEL,
]
dataset = YCBDataset(
opt.dataset_root,
mode='train_syn_grid_valid',
object_list=object_list,
output_data=output_format,
resample_on_error=True,
preprocessors=[
YCBOcclusionAugmentor(opt.dataset_root),
ColorJitter(),
#InplaneRotator(),
],
postprocessors=[ImageNormalizer(), PointShifter()],
image_size=[640, 480],
num_points=1000)
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=1,
shuffle=True,
num_workers=opt.workers - 1)
test_dataset = YCBDataset(opt.dataset_root,
mode='valid',
object_list=object_list,
output_data=output_format,
resample_on_error=True,
preprocessors=[],
postprocessors=[ImageNormalizer()],
image_size=[640, 480],
num_points=1000)
testdataloader = torch.utils.data.DataLoader(test_dataset,
batch_size=1,
shuffle=False,
num_workers=1)
opt.sym_list = [12, 15, 18, 19, 20]
opt.num_points_mesh = dataset.num_pt_mesh_small
print(
'>>>>>>>>----------Dataset loaded!---------<<<<<<<<\nlength of the training set: {0}\nlength of the testing set: {1}\nnumber of sample points on mesh: {2}\nsymmetry object list: {3}'
.format(len(dataset), len(test_dataset), opt.num_points_mesh,
opt.sym_list))
criterion_dist = Loss(opt.num_points_mesh, opt.sym_list)
criterion_lik = DuelLoss(opt.num_points_mesh, opt.sym_list)
best_dis = np.Inf
best_lik = -np.Inf
if opt.start_epoch == 1:
for log in os.listdir(opt.log_dir):
os.remove(os.path.join(opt.log_dir, log))
st_time = time.time()
cum_batch_count = 0
mean_err = 0
for epoch in range(opt.start_epoch, opt.nepoch):
logger = setup_logger(
'epoch%d' % epoch,
os.path.join(opt.log_dir, 'epoch_%d_log.txt' % epoch))
logger.info('Train time {0}'.format(
time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)) +
', ' + 'Training started'))
train_count = 0
train_dis_avg = 0.0
train_lik_avg = 0.0
estimator.train()
optimizer.zero_grad()
for rep in range(opt.repeat_epoch):
for i, data in enumerate(dataloader, 0):
points, choose, img, quat, target, model_points, idx = data
idx = idx - 1
points, choose, img, quat, target, model_points, idx = Variable(points).cuda(), \
Variable(choose).cuda(), \
Variable(img).cuda(), \
Variable(quat).cuda(), \
Variable(target).cuda(), \
Variable(model_points).cuda(), \
Variable(idx).cuda()
pred_r, pred_t, pred_c, pred_bq, pred_bz, emb = estimator(
img, points, choose, idx)
loss_dist, dis, new_points, new_target = criterion_dist(
pred_r, pred_t, pred_c, target, model_points, idx, points,
opt.w, opt.refine_start)
how_max, which_max = torch.max(pred_c.detach(), 1)
pred_q = pred_r[0, :, [1, 2, 3, 0]].detach()
pred_q /= torch.norm(pred_q, dim=1).view(-1, 1)
max_q = pred_q[which_max.item()]
max_bq = pred_bq[0, which_max.item()] / torch.norm(
pred_bq[0, which_max.item()])
max_bz = pred_bz[0, which_max.item()]
loss_lik, lik = criterion_lik(max_q.view(-1), max_bq.view(-1),
-torch.abs(max_bz.view(-1)),
quat)
loss = loss_dist + loss_lik
loss.backward()
train_dis_avg += dis.item()
train_lik_avg += np.log(lik.item())
train_count += 1
if train_count % opt.batch_size == 0:
logger.info(
'Train time {0} Epoch {1} Batch {2} Frame {3} Avg_dis:{4} Avg_lik:{5}'
.format(
time.strftime("%Hh %Mm %Ss",
time.gmtime(time.time() - st_time)),
epoch, int(train_count / opt.batch_size),
train_count, train_dis_avg / opt.batch_size,
train_lik_avg / opt.batch_size))
optimizer.step()
optimizer.zero_grad()
train_dis_avg = 0
train_lik_avg = 0
if train_count != 0 and train_count % 1000 == 0:
torch.save(estimator.state_dict(),
'{0}/pose_model_current.pth'.format(opt.outf))
print(
'>>>>>>>>----------epoch {0} train finish---------<<<<<<<<'.format(
epoch))
logger = setup_logger(
'epoch%d_test' % epoch,
os.path.join(opt.log_dir, 'epoch_%d_test_log.txt' % epoch))
logger.info('Test time {0}'.format(
time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)) +
', ' + 'Testing started'))
test_dis = 0.0
test_lik = 0.0
test_count = 0
estimator.eval()
for j, data in enumerate(testdataloader, 0):
points, choose, img, quat, target, model_points, idx = data
idx = idx - 1
points, choose, img, quat, target, model_points, idx = Variable(points).cuda(), \
Variable(choose).cuda(), \
Variable(img).cuda(), \
Variable(quat).cuda(), \
Variable(target).cuda(), \
Variable(model_points).cuda(), \
Variable(idx).cuda()
pred_r, pred_t, pred_c, pred_bq, pred_bz, emb = estimator(
img, points, choose, idx)
_, dis, new_points, new_target = criterion_dist(
pred_r, pred_t, pred_c, target, model_points, idx, points,
opt.w, opt.refine_start)
how_max, which_max = torch.max(pred_c.detach(), 1)
pred_q = pred_r[0, :, [1, 2, 3, 0]].detach()
pred_q /= torch.norm(pred_q, dim=1).view(-1, 1)
max_q = pred_q[which_max.item()]
max_bq = pred_bq[0, which_max.item()] / torch.norm(
pred_bq[0, which_max.item()])
max_bz = pred_bz[0, which_max.item()]
_, lik = criterion_lik(max_q.view(-1), max_bq.view(-1),
-torch.abs(max_bz.view(-1)), quat)
test_dis += dis.item()
test_lik += np.log(lik.item())
logger.info(
'Test time {0} Test Frame No.{1} dis:{2} lik:{3}'.format(
time.strftime("%Hh %Mm %Ss",
time.gmtime(time.time() - st_time)),
test_count, dis, lik))
test_count += 1
test_dis = test_dis / test_count
test_lik = test_lik / test_count
logger.info(
'Test time {0} Epoch {1} TEST FINISH Avg dis: {2} Avg lik: {3}'.
format(
time.strftime("%Hh %Mm %Ss",
time.gmtime(time.time() - st_time)), epoch,
test_dis, test_lik))
if test_dis <= best_dis or test_lik >= best_lik:
best_dis = min(test_dis, best_dis)
best_lik = max(test_lik, best_lik)
torch.save(
estimator.state_dict(),
'{0}/pose_model_{1}_{2}_{3}.pth'.format(
opt.outf, epoch, test_dis, test_lik))
print(epoch,
'>>>>>>>>----------BEST TEST MODEL SAVED---------<<<<<<<<')
if best_dis < opt.decay_margin and not opt.decay_start:
opt.decay_start = True
opt.lr *= opt.lr_rate
opt.w *= opt.w_rate
optimizer = optim.Adam(estimator.parameters(), lr=opt.lr)
def main():
opt.manualSeed = random.randint(1, 10000)
random.seed(opt.manualSeed)
torch.manual_seed(opt.manualSeed)
if not os.path.exists(opt.output_folder):
os.makedirs(opt.output_folder)
num_points = 1000 #number of points on the input pointcloud
num_objects = 21
if (opt.object_indices is None):
opt.object_indices = list(range(1, num_objects + 1))
estimator = PoseNet(num_points=num_points, num_obj=num_objects)
estimator.cuda()
estimator.load_state_dict(torch.load(opt.weights))
output_format = [
otypes.OBJECT_LABEL, otypes.QUATERNION, otypes.IMAGE_CROPPED,
otypes.DEPTH_POINTS_MASKED_AND_INDEXES
]
estimator.eval()
with torch.no_grad():
with std_out_err_redirect_tqdm() as orig_stdout:
pbar = tqdm(opt.object_indices,
file=orig_stdout,
dynamic_ncols=True)
for cls in pbar:
preprocessors = [InplaneRotator(theta=np.pi / 2)]
postprocessors = [ImageNormalizer()]
if (opt.num_augmentations > 0):
preprocessors.extend([
YCBOcclusionAugmentor(opt.dataset_root),
ColorJitter(),
])
postprocessors.append(PointShifter())
dataset = YCBDataset(opt.dataset_root,
mode=opt.dataset_mode,
object_list=[cls],
output_data=output_format,
resample_on_error=False,
add_syn_background=opt.add_syn_background,
add_syn_noise=opt.add_syn_background,
preprocessors=preprocessors,
postprocessors=postprocessors,
image_size=[640, 480],
num_points=1000)
classes = dataset.classes
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=1,
shuffle=False,
num_workers=opt.workers)
pbar.set_description('Featurizing {}'.format(classes[cls]))
if (opt.num_augmentations > 0):
pbar_aug = trange(opt.start_index,
opt.num_augmentations,
file=orig_stdout,
dynamic_ncols=True)
else:
pbar_aug = [None]
for aug_idx in pbar_aug:
pbar_save = tqdm(enumerate(dataloader),
total=len(dataloader),
file=orig_stdout,
dynamic_ncols=True)
for i, data in pbar_save:
if (len(data) == 0 or len(data[0]) == 0):
continue
idx, quat, img, points, choose = data
data_path = dataset.image_list[i]
idx = idx - 1
img = Variable(img).cuda()
points = Variable(points).cuda()
choose = Variable(choose).cuda()
idx = Variable(idx).cuda()
assert cls == data_path[1]
assert cls - 1 == int(idx[0])
pred_r, pred_t, pred_c, emb, feat, feat_global = estimator.allFeatures(
img, points, choose, idx)
if (opt.num_augmentations > 0):
output_filename = '{0}/data/{1}_{2}_{3}_feat.npz'.format(
opt.output_folder, data_path[0], classes[cls],
aug_idx)
else:
output_filename = '{0}/data/{1}_{2}_feat.npz'.format(
opt.output_folder, data_path[0], classes[cls])
#pbar_save.set_description(output_filename)
os.makedirs(os.path.dirname(output_filename),
exist_ok=True)
how_max, which_max = torch.max(pred_c, 1)
max_feat = feat[0, :, which_max[0]].view(-1)
pred_t = pred_t[0, :]
pred_q = pred_r[0, :, [1, 2, 3, 0]]
pred_q /= torch.norm(pred_q, dim=1).view(-1, 1)
np.savez(
output_filename,
max_c=to_np(how_max),
max_q=to_np(pred_q[which_max.item()]),
max_t=to_np(pred_t[which_max.item()]),
pred_c=to_np(pred_c),
pred_q=to_np(pred_q),
pred_t=to_np(pred_t),
quat=to_np(quat)[0],
feat=to_np(max_feat),
#feat_all = to_np(feat)[0].T, i
feat_global=to_np(feat_global)[0],
#max_confidence = to_np(how_max),
#confidence = to_np(pred_c)[0],
)
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)
if not os.path.exists(opt.output_folder):
os.makedirs(opt.output_folder)
num_points = 1000 #number of points on the input pointcloud
num_objects = 21
if (opt.object_indices is None):
opt.object_indices = list(range(1, num_objects + 1))
estimator = PoseCNNFeaturizer()
output_format = [otypes.IMAGE, otypes.DEPTH_IMAGE]
with std_out_err_redirect_tqdm() as orig_stdout:
preprocessors = []
postprocessors = []
if (opt.num_augmentations > 0):
preprocessors.extend([
YCBOcclusionAugmentor(opt.dataset_root),
ColorJitter(),
])
postprocessors.append(PointShifter())
dataset = YCBDataset(opt.dataset_root,
mode=opt.dataset_mode,
object_list=opt.object_indices,
output_data=output_format,
resample_on_error=False,
preprocessors=preprocessors,
postprocessors=postprocessors,
image_size=[640, 480],
num_points=1000)
_, u_idxs = np.unique(zip(*dataset.image_list)[0], return_index=True)
dataset.image_list = np.array(dataset.image_list)[u_idxs].tolist()
dataset.list_obj = np.array(dataset.list_obj)[u_idxs].tolist()
classes = dataset.classes
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=1,
shuffle=False,
num_workers=opt.workers)
#pbar.set_description('Featurizing {}'.format(classes[cls]))
if (opt.num_augmentations > 0):
pbar_aug = trange(opt.start_index,
opt.num_augmentations,
file=orig_stdout,
dynamic_ncols=True)
else:
pbar_aug = [None]
for aug_idx in pbar_aug:
pbar_save = tqdm(enumerate(dataloader),
total=len(dataloader),
file=orig_stdout,
dynamic_ncols=True)
for i, data in pbar_save:
if (len(data) == 0 or len(data[0]) == 0):
continue
img, depth = data
img = toPoseCNNImage(img[0])
depth = to_np(depth[0])
data_path = dataset.image_list[i]
path = '{}/data/{}-meta.mat'.format(dataset.dataset_root,
dataset.getPath(i))
meta_data = scio.loadmat(path)
try:
seg = estimator(img, depth, meta_data)
except Exception as e:
print(e)
continue
for pose_idx, cls in enumerate(seg['rois'][:, 1]):
cls = int(cls)
quat = getObjectGTQuaternion(meta_data, cls)
feat = seg['feats'][pose_idx]
fc6 = seg['fc6'][pose_idx]
if (opt.num_augmentations > 0):
output_filename = '{0}/data/{1}_{2}_{3}_feat.npz'.format(
opt.output_folder, data_path[0], classes[cls],
aug_idx)
else:
output_filename = '{0}/data/{1}_{2}_feat.npz'.format(
opt.output_folder, data_path[0], classes[cls])
#pbar_save.set_description(output_filename)
if not os.path.exists(os.path.dirname(output_filename)):
os.makedirs(os.path.dirname(output_filename))
np.savez(output_filename, quat=quat, feat=feat, fc6=fc6)