def demo():
net = Resnet18_8s(ver_dim=vote_num * 2, seg_dim=2)
net = NetWrapper(net).cuda()
net = DataParallel(net)
optimizer = optim.Adam(net.parameters(), lr=train_cfg['lr'])
model_dir = os.path.join(cfg.MODEL_DIR, "cat_demo")
load_model(net.module.net, optimizer, model_dir, args.load_epoch)
data, points_3d, bb8_3d = read_data()
image, mask, vertex, vertex_weights, pose, corner_target = [
d.unsqueeze(0).cuda() for d in data
]
seg_pred, vertex_pred, loss_seg, loss_vertex, precision, recall = net(
image, mask, vertex, vertex_weights)
eval_net = DataParallel(EvalWrapper().cuda())
corner_pred = eval_net(seg_pred, vertex_pred).cpu().detach().numpy()[0]
camera_matrix = np.array([[572.4114, 0., 325.2611],
[0., 573.57043, 242.04899], [0., 0., 1.]])
pose_pred = pnp(points_3d, corner_pred, camera_matrix)
projector = Projector()
bb8_2d_pred = projector.project(bb8_3d, pose_pred, 'linemod')
bb8_2d_gt = projector.project(bb8_3d, pose[0].detach().cpu().numpy(),
'linemod')
image = imagenet_to_uint8(image.detach().cpu().numpy())[0]
visualize_bounding_box(image[None, ...], bb8_2d_pred[None, None, ...],
bb8_2d_gt[None, None, ...])
def demo():
net = Resnet18_8s(ver_dim=vote_num * 2, seg_dim=2)
net = NetWrapper(net).cuda()
net = DataParallel(net)
optimizer = optim.Adam(net.parameters(), lr=train_cfg['lr'])
model_dir = os.path.join(cfg.MODEL_DIR, "switch_linemod_train")
load_model(net.module.net, optimizer, model_dir, -1)
image, points_3d, bb8_3d = read_data()
image = image[None, ...]
seg_pred, vertex_pred = net(image)
# visualize_mask(mask)
# visualize_vertex(vertex, vertex_weights)
# visualize_hypothesis(image, seg_pred, vertex_pred, corner_target)
# visualize_voting_ellipse(image, seg_pred, vertex_pred, corner_target)
eval_net = DataParallel(EvalWrapper().cuda())
corner_pred = eval_net(seg_pred, vertex_pred).cpu().detach().numpy()[0]
camera_matrix = np.array([[572.4114, 0., 325.2611],
[0., 573.57043, 242.04899], [0., 0., 1.]])
pose_pred = pnp(points_3d, corner_pred, camera_matrix)
projector = Projector()
bb8_2d_pred = projector.project(bb8_3d, pose_pred, 'linemod')
print(bb8_2d_pred)
image = imagenet_to_uint8(image.detach().cpu().numpy())[0]
visualize_bounding_box(image[None, ...], bb8_2d_pred[None, None, ...])
def demo():
net = Resnet18_8s(ver_dim=vote_num * 2, seg_dim=2)
net = NetWrapper(net).cuda()
net = DataParallel(net)
optimizer = optim.Adam(net.parameters(), lr=train_cfg['lr'])
model_dir = os.path.join(cfg.MODEL_DIR, "cat_linemod_train") #cat_demo
load_model(net.module.net, optimizer, model_dir, args.load_epoch)
data, points_3d, bb8_3d = read_data()
image, mask, vertex, vertex_weights, pose, corner_target = [
d.unsqueeze(0).cuda() for d in data
]
seg_pred, vertex_pred, loss_seg, loss_vertex, precision, recall = net(
image, mask, vertex, vertex_weights)
eval_net = DataParallel(EvalWrapper().cuda())
#向量方形图，语义分割图，然后 ransac 计算 kp，，向量方向图一旦准了，kp也就准了
corner_pred = eval_net(seg_pred, vertex_pred).cpu().detach().numpy()[0]
camera_matrix = np.array([[572.4114, 0., 325.2611],
[0., 573.57043, 242.04899], [0., 0., 1.]])
pose_pred = pnp(points_3d, corner_pred, camera_matrix)
projector = Projector()
#
bb8_2d_pred = projector.project(bb8_3d, pose_pred, 'linemod')
bb8_2d_gt = projector.project(bb8_3d, pose[0].detach().cpu().numpy(),
'linemod')
image = imagenet_to_uint8(image.detach().cpu().numpy())[0]
print("loss_seg:{} , loss_vertex:{} , precision:{},recall:{}, ".format(
loss_seg, loss_vertex, precision, recall))
#399.pth
#loss_seg:tensor([0.0015], device='cuda:0', grad_fn=<MeanBackward0>) , loss_vertex:tensor([0.0016], device='cuda:0', grad_fn=<DivBackward1>) ,
#precision:tensor([0.9434], device='cuda:0'),recall:tensor([0.9677], device='cuda:0'),
#199.pth
# loss_seg:tensor([0.0015], device='cuda:0', grad_fn=<MeanBackward0>) , loss_vertex:tensor([0.0016], device='cuda:0', grad_fn=<DivBackward1>) ,
# precision:tensor([0.9583], device='cuda:0'),recall:tensor([0.9524], device='cuda:0'),
erro = bb8_2d_pred - bb8_2d_gt
erro = np.abs(erro)
err = np.reshape(erro, (erro.size, ))
#abserr = map(abs,err)
print("reproject sum_error:{} ".format(np.sum(err)))
## 199 reproject sum_error:13.385891544820552
## 399 reproject sum_error:12.718721049803733
##看了是有提高 准召定义 准确下降，召回上升
visualize_bounding_box(image[None, ...], bb8_2d_pred[None, None, ...],
bb8_2d_gt[None, None, ...])
def train(cfg, network):
trainer = make_trainer(cfg, network)
optimizer = make_optimizer(cfg, network)
scheduler = make_lr_scheduler(cfg, optimizer)
recorder = make_recorder(cfg)
evaluator = make_evaluator(cfg)
begin_epoch = load_model(network, optimizer, scheduler, recorder, cfg.model_dir, resume=cfg.resume)
# set_lr_scheduler(cfg, scheduler)
train_loader = make_data_loader(cfg, is_train=True)
val_loader = make_data_loader(cfg, is_train=False)
for epoch in range(begin_epoch, cfg.train.epoch):
recorder.epoch = epoch
trainer.train(epoch, train_loader, optimizer, recorder)
scheduler.step()
if (epoch + 1) % cfg.save_ep == 0:
save_model(network, optimizer, scheduler, recorder, epoch, cfg.model_dir)
if (epoch + 1) % cfg.eval_ep == 0:
trainer.val(epoch, val_loader, evaluator, recorder)
return network
def train(cfg, network):
# cfg.train.num_workers = 0
if cfg.train.dataset[:4] != 'City':
torch.multiprocessing.set_sharing_strategy('file_system')
train_loader = make_data_loader(cfg, is_train=True, max_iter=cfg.ep_iter)
val_loader = make_data_loader(cfg, is_train=False)
# train_loader = make_data_loader(cfg, is_train=True, max_iter=100)
trainer = make_trainer(cfg, network)
optimizer = make_optimizer(cfg, network)
scheduler = make_lr_scheduler(cfg, optimizer)
recorder = make_recorder(cfg)
evaluator = make_evaluator(cfg)
begin_epoch = load_model(network, optimizer, scheduler, recorder, cfg.model_dir, resume=cfg.resume)
# set_lr_scheduler(cfg, scheduler)
for epoch in range(begin_epoch, cfg.train.epoch):
recorder.epoch = epoch
trainer.train(epoch, train_loader, optimizer, recorder)
scheduler.step()
if (epoch + 1) % cfg.save_ep == 0:
save_model(network, optimizer, scheduler, recorder, epoch, cfg.model_dir)
if (epoch + 1) % cfg.eval_ep == 0:
trainer.val(epoch, val_loader, evaluator, recorder)
return network
def demo():
net = Resnet18_8s(ver_dim=vote_num * 2, seg_dim=2)
net = NetWrapper(net).cuda()
net = DataParallel(net)
optimizer = optim.Adam(net.parameters(), lr=train_cfg['lr'])
model_dir = os.path.join(cfg.MODEL_DIR, "cat_demo")
load_model(net.module.net, optimizer, model_dir, -1)
data, points_3d, bb8_3d = read_data()
#print("BB8_3D: ",bb8_3d)
image, mask, vertex, vertex_weights, pose, corner_target = [
d.unsqueeze(0).cuda() for d in data
]
seg_pred, vertex_pred, loss_seg, loss_vertex, precision, recall = net(
image, mask, vertex, vertex_weights)
seg_mask = torch.argmax(seg_pred, 1)
print("seg_mask", seg_mask, type(seg_mask), seg_mask.shape, seg_mask[0])
visualize_mask(seg_mask)
visualize_mask(mask)
#visualize_vertex(vertex, vertex_weights)
#visualize_hypothesis(image, seg_pred, vertex_pred, corner_target)
visualize_voting_ellipse(image, seg_pred, vertex_pred, corner_target)
eval_net = DataParallel(EvalWrapper().cuda())
corner_pred = eval_net(seg_pred, vertex_pred).cpu().detach().numpy()[0]
print("Corner Predictions: ", corner_pred)
camera_matrix = np.array([[572.4114, 0., 325.2611],
[0., 573.57043, 242.04899], [0., 0., 1.]])
pose_pred = pnp(points_3d, corner_pred, camera_matrix)
projector = Projector()
bb8_2d_pred = projector.project(bb8_3d, pose_pred, 'linemod')
print("Pose prediction :\n", pose_pred)
print("GT pose: \n", pose[0].detach().cpu().numpy())
bb8_2d_gt = projector.project(bb8_3d, pose[0].detach().cpu().numpy(),
'linemod')
print(bb8_2d_gt)
image = imagenet_to_uint8(image.detach().cpu().numpy())[0]
visualize_bounding_box(image[None, ...], bb8_2d_pred[None, None, ...],
bb8_2d_gt[None, None, ...])
def train(cfg, network):
if cfg.train.dataset[:4] != 'City':
torch.multiprocessing.set_sharing_strategy('file_system')
trainer = make_trainer(cfg, network)
optimizer = make_optimizer(cfg, network)
scheduler = make_lr_scheduler(cfg, optimizer)
recorder = make_recorder(cfg)
if 'Coco' not in cfg.train.dataset:
evaluator = make_evaluator(cfg)
begin_epoch = load_model(network,
optimizer,
scheduler,
recorder,
cfg.model_dir,
resume=cfg.resume)
# set_lr_scheduler(cfg, scheduler)
train_loader = make_data_loader(cfg, is_train=True)
val_loader = make_data_loader(cfg, is_train=False)
# train_loader = make_data_loader(cfg, is_train=True, max_iter=100)
global_steps = None
if cfg.neptune:
global_steps = {
'train_global_steps': 0,
'valid_global_steps': 0,
}
neptune.init('hccccccccc/clean-pvnet')
neptune.create_experiment(cfg.model_dir.split('/')[-1])
neptune.append_tag('pose')
for epoch in range(begin_epoch, cfg.train.epoch):
recorder.epoch = epoch
trainer.train(epoch, train_loader, optimizer, recorder, global_steps)
scheduler.step()
if (epoch + 1) % cfg.save_ep == 0:
save_model(network, optimizer, scheduler, recorder, epoch,
cfg.model_dir)
if (epoch + 1) % cfg.eval_ep == 0:
if 'Coco' in cfg.train.dataset:
trainer.val_coco(val_loader, global_steps)
else:
trainer.val(epoch, val_loader, evaluator, recorder)
if cfg.neptune:
neptune.stop()
return network
def train_vgg16():
device = torch.device('cuda:0')
dataset = LineModDatasetReal()
dataloader = DataLoader(dataset, batch_size=4, shuffle=True, num_workers=4)
net = VGG16Convs()
net = net.to(device)
optimizer = optim.Adam(net.parameters(), lr=2e-5, weight_decay=5e-4)
model_dir = os.path.join(cfg.MODEL_DIR, 'vgg16')
epoch = load_model(net, optimizer, model_dir)
for epoch in range(epoch, epoch + 100):
train(net, optimizer, dataloader, device, epoch)
save_model(net, optimizer, epoch, model_dir)
def train(cfg, network):
trainer = make_trainer(cfg, network)
optimizer = make_optimizer(cfg, network)
scheduler = make_lr_scheduler(cfg, optimizer)
recorder = make_recorder(cfg)
evaluator = make_evaluator(cfg)
begin_epoch = load_model(network,
optimizer,
scheduler,
recorder,
cfg.model_dir,
resume=cfg.resume)
# begin_epoch = 0 #如果要继续训练那么请注释这一行
# set_lr_scheduler(cfg, scheduler)
# print("before train loader")
train_loader = make_data_loader(cfg, is_train=True) #到这里才读取的数据
# print("under train loader")
val_loader = make_data_loader(cfg, is_train=False)
# #这里是查看train_loader的相关参数个结构
# tmp_file = open('/home/tianhao.lu/code/Deep_snake/snake/Result/Contour/contour.log', 'w')
# tmp_file.writelines("train_loader type:" + str(type(train_loader)) + "\n")
# tmp_file.writelines("train_loader len:" + str(len(train_loader)) + "\n")
# tmp_file.writelines("train_loader data:" + str(train_loader) + "\n")
# for tmp_data in train_loader:
# tmp_file.writelines("one train_loader data type:" + str(type(tmp_data)) + "\n")
# for key in tmp_data:
# tmp_file.writelines("one train_loader data key:" + str(key) + "\n")
# tmp_file.writelines("one train_loader data len:" + str(len(tmp_data[key])) + "\n")
# # tmp_file.writelines("one train_loader data:" + str(tmp_data) + "\n")
# break
# tmp_file.writelines(str("*************************************************************** \n"))
# tmp_file.close()
for epoch in range(begin_epoch, cfg.train.epoch):
recorder.epoch = epoch
trainer.train(epoch, train_loader, optimizer, recorder)
scheduler.step() #optimizer.step()模型才会更新，scheduler.step()是用来调整lr的
if (epoch + 1) % cfg.save_ep == 0:
save_model(network, optimizer, scheduler, recorder, epoch,
cfg.model_dir)
if (epoch + 1) % cfg.eval_ep == 0:
trainer.val(epoch, val_loader, evaluator, recorder)
return network
def train(cfg, network):
trainer = make_trainer(cfg, network)
optimizer = make_optimizer(cfg, network)
scheduler = make_lr_scheduler(cfg, optimizer)
recorder = make_recorder(cfg)
evaluator = make_evaluator(cfg)
begin_epoch = load_model(network,
optimizer,
scheduler,
recorder,
cfg.trained_model_dir,
resume=cfg.resume)
set_lr_scheduler(cfg, scheduler)
train_loader = make_data_loader(cfg,
is_train=True,
is_distributed=cfg.distributed,
max_iter=cfg.ep_iter)
val_loader = make_data_loader(cfg, is_train=False)
for epoch in range(begin_epoch, cfg.train.epoch):
recorder.epoch = epoch
if cfg.distributed:
train_loader.batch_sampler.sampler.set_epoch(epoch)
trainer.train(epoch, train_loader, optimizer, recorder)
scheduler.step()
if (epoch + 1) % cfg.save_ep == 0 and cfg.local_rank == 0:
save_model(network, optimizer, scheduler, recorder,
cfg.trained_model_dir, epoch)
if (epoch + 1) % cfg.save_latest_ep == 0 and cfg.local_rank == 0:
save_model(network,
optimizer,
scheduler,
recorder,
cfg.trained_model_dir,
epoch,
last=True)
if (epoch + 1) % cfg.eval_ep == 0:
trainer.val(epoch, val_loader, evaluator, recorder)
return network
def train(cfg, network):
trainer = make_trainer(cfg, network)
optimizer = make_optimizer(cfg, network)
scheduler = make_lr_scheduler(cfg, optimizer)
recorder = make_recorder(cfg)
#evaluator = make_evaluator(cfg)
if cfg.network_full_init:
begin_epoch = load_network(network,
cfg.model_dir,
resume=cfg.resume,
epoch=cfg.test.epoch)
begin_epoch = 0
else:
begin_epoch = load_model(network,
optimizer,
scheduler,
recorder,
cfg.model_dir,
resume=cfg.resume)
# set_lr_scheduler(cfg, scheduler)
if DEBUG:
print('------------------Loading training set-------------------')
train_loader = make_data_loader(cfg, is_train=True)
if DEBUG:
print('Loading training set done...')
print('---------------------------------------------------------')
#val_loader = make_data_loader(cfg, is_train=False)
for epoch in range(begin_epoch, cfg.train.epoch):
recorder.epoch = epoch
trainer.train(epoch, train_loader, optimizer, recorder)
scheduler.step()
if (epoch + 1) % cfg.save_ep == 0:
save_model(network, optimizer, scheduler, recorder, epoch,
cfg.model_dir)
#if (epoch + 1) % cfg.eval_ep == 0:
# trainer.val(epoch, val_loader, evaluator, recorder)
return network
def test_net():
net=Resnet18_8s(ver_dim=vote_num*2, seg_dim=2)
net=NetWrapper(net)
net=DataParallel(net).cuda()
optimizer = optim.Adam(net.parameters(), lr=train_cfg['lr'])
model_dir=os.path.join(cfg.MODEL_DIR,train_cfg['model_name'])
motion_model=train_cfg['motion_model']
print('motion state {}'.format(motion_model))
# if args.test_model:
begin_epoch=load_model(net.module.net, optimizer, model_dir, args.load_epoch)
print('testing normal linemod ...')
image_db = LineModImageDB(args.linemod_cls,has_render_set=False,
has_fuse_set=False)
test_db = image_db.test_real_set+image_db.val_real_set
test_set = LineModDatasetRealAug(test_db, cfg.LINEMOD, vote_type, augment=False, use_motion=motion_model)
test_sampler = SequentialSampler(test_set)
test_batch_sampler = ImageSizeBatchSampler(test_sampler, train_cfg['test_batch_size'], False)
test_loader = DataLoader(test_set, batch_sampler=test_batch_sampler, num_workers=0)
args.use_test_set = True
prefix='test'
val(net, test_loader, begin_epoch, prefix, use_motion=motion_model)
if args.linemod_cls in cfg.occ_linemod_cls_names:
print('testing occluded linemod ...')
occ_image_db = OcclusionLineModImageDB(args.linemod_cls)
occ_test_db = occ_image_db.test_real_set
occ_test_set = LineModDatasetRealAug(occ_test_db, cfg.OCCLUSION_LINEMOD, vote_type,
augment=False, use_motion=motion_model)
occ_test_sampler = SequentialSampler(occ_test_set)
occ_test_batch_sampler = ImageSizeBatchSampler(occ_test_sampler, train_cfg['test_batch_size'], False)
occ_test_loader = DataLoader(occ_test_set, batch_sampler=occ_test_batch_sampler, num_workers=0)
prefix='occ_test' if args.use_test_set else 'occ_val'
val(net, occ_test_loader, begin_epoch, prefix, use_motion=motion_model)
def train_net():
net = Resnet18_8s_modified(ver_dim=vote_num * 2, seg_dim=2)
#net=NetWrapper(net)
net = MappingNetWrapper(net)
net = DataParallel(net).cuda()
optimizer = optim.Adam(net.parameters(), lr=train_cfg['lr'])
model_dir = os.path.join(cfg.MODEL_DIR, train_cfg['model_name'])
motion_model = train_cfg['motion_model']
print('motion state {}'.format(motion_model))
if args.test_model:
begin_epoch = load_model(net.module.net, optimizer, model_dir,
args.load_epoch)
if args.normal:
print('testing normal linemod ...')
image_db = LineModImageDB(args.linemod_cls,
has_render_set=False,
has_fuse_set=False)
test_db = image_db.test_real_set + image_db.val_real_set
test_set = LineModDatasetAug(test_db,
cfg.LINEMOD,
vote_type,
augment=False,
use_motion=motion_model)
test_sampler = SequentialSampler(test_set)
test_batch_sampler = ImageSizeBatchSampler(
test_sampler, train_cfg['test_batch_size'], False)
test_loader = DataLoader(test_set,
batch_sampler=test_batch_sampler,
num_workers=0)
prefix = 'test' if args.use_test_set else 'val'
val(net, test_loader, begin_epoch, prefix, use_motion=motion_model)
# if args.occluded and args.linemod_cls in cfg.occ_linemod_cls_names:
# print('testing occluded linemod ...')
# occ_image_db = OcclusionLineModImageDB(args.linemod_cls)
# occ_test_db = occ_image_db.test_real_set
# occ_test_set = LineModDatasetAug(occ_test_db,
# cfg.OCCLUSION_LINEMOD,
# vote_type,
# augment=False,
# use_motion=motion_model)
# occ_test_sampler = SequentialSampler(occ_test_set)
# occ_test_batch_sampler = ImageSizeBatchSampler(
# occ_test_sampler, train_cfg['test_batch_size'], False)
# occ_test_loader = DataLoader(occ_test_set,
# batch_sampler=occ_test_batch_sampler,
# num_workers=0)
# prefix = 'occ_test' if args.use_test_set else 'occ_val'
# val(net,
# occ_test_loader,
# begin_epoch,
# prefix,
# use_motion=motion_model)
# if args.truncated:
# print('testing truncated linemod ...')
# trun_image_db = TruncatedLineModImageDB(args.linemod_cls)
# print(len(trun_image_db.set))
# trun_image_set = LineModDatasetRealAug(trun_image_db.set,
# cfg.LINEMOD,
# vote_type,
# augment=False,
# use_intrinsic=True,
# use_motion=motion_model)
# trun_test_sampler = SequentialSampler(trun_image_set)
# trun_test_batch_sampler = ImageSizeBatchSampler(
# trun_test_sampler, train_cfg['test_batch_size'], False)
# trun_test_loader = DataLoader(
# trun_image_set,
# batch_sampler=trun_test_batch_sampler,
# num_workers=0)
# prefix = 'trun_test'
# val(net,
# trun_test_loader,
# begin_epoch,
# prefix,
# True,
# use_motion=motion_model)
else:
begin_epoch = 0
if train_cfg['resume']:
begin_epoch = load_model(net.module.net, optimizer, model_dir)
image_db = LineModImageDB(
args.linemod_cls,
has_fuse_set=False, #train_cfg['use_fuse'],
has_render_set=False) #True)
# train_db=[]
# train_db+=image_db.render_set
# if train_cfg['use_real_train']:
# train_db+=image_db.train_real_set
# if train_cfg['use_fuse']:
# train_db+=image_db.fuse_set
train_db = image_db.train_real_set
train_set = LineModDatasetAug(train_db,
cfg.LINEMOD,
vote_type,
augment=True,
cfg=train_cfg['aug_cfg'],
use_motion=motion_model)
train_sampler = RandomSampler(train_set)
train_batch_sampler = ImageSizeBatchSampler(
train_sampler,
train_cfg['train_batch_size'],
False,
cfg=train_cfg['aug_cfg'])
train_loader = DataLoader(train_set,
batch_sampler=train_batch_sampler,
num_workers=12)
val_db = image_db.val_real_set
val_set = LineModDatasetAug(val_db,
cfg.LINEMOD,
vote_type,
augment=False,
cfg=train_cfg['aug_cfg'],
use_motion=motion_model)
val_sampler = SequentialSampler(val_set)
val_batch_sampler = ImageSizeBatchSampler(val_sampler,
train_cfg['test_batch_size'],
False,
cfg=train_cfg['aug_cfg'])
val_loader = DataLoader(val_set,
batch_sampler=val_batch_sampler,
num_workers=12)
# if args.linemod_cls in cfg.occ_linemod_cls_names:
# occ_image_db=OcclusionLineModImageDB(args.linemod_cls)
# occ_val_db=occ_image_db.test_real_set[:len(occ_image_db.test_real_set)//2]
# occ_val_set = LineModDatasetRealAug(occ_val_db, cfg.OCCLUSION_LINEMOD, vote_type, augment=False, cfg=train_cfg['aug_cfg'], use_motion=motion_model)
# occ_val_sampler = SequentialSampler(occ_val_set)
# occ_val_batch_sampler = ImageSizeBatchSampler(occ_val_sampler, train_cfg['test_batch_size'], False, cfg=train_cfg['aug_cfg'])
# occ_val_loader = DataLoader(occ_val_set, batch_sampler=occ_val_batch_sampler, num_workers=12)
for epoch in range(begin_epoch, train_cfg['epoch_num']):
adjust_learning_rate(optimizer, epoch, train_cfg['lr_decay_rate'],
train_cfg['lr_decay_epoch'])
train(net, optimizer, train_loader, epoch)
val(net, val_loader, epoch, use_motion=motion_model)
# if args.linemod_cls in cfg.occ_linemod_cls_names:
# val(net, occ_val_loader, epoch, 'occ_val',use_motion=motion_model)
save_model(net.module.net, optimizer, epoch, model_dir)
def train_net():
#seg_dim 为啥两张，kp 个数，×2，？ 一张一个x坐标，一张一个y坐标，不是一张图片估计 一个关键点吗
net = Resnet18_8s(ver_dim=vote_num * 2, seg_dim=2)
#前传 误差定义 评价
net = NetWrapper(net)
#多卡
#net=DataParallel(net,device_ids=[0,1]).cuda()
net = DataParallel(net, device_ids=[0]).cuda()
#优化方法
optimizer = optim.Adam(net.parameters(), lr=train_cfg['lr'])
model_dir = os.path.join(cfg.MODEL_DIR, train_cfg['model_name'])
motion_model = train_cfg['motion_model']
print('motion state {}'.format(motion_model))
if args.test_model:
begin_epoch = load_model(net.module.net, optimizer, model_dir,
args.load_epoch)
if args.normal:
print('testing normal linemod ...')
image_db = LineModImageDB(args.linemod_cls,
has_render_set=False,
has_fuse_set=False)
test_db = image_db.test_real_set + image_db.val_real_set
test_set = LineModDatasetRealAug(test_db,
cfg.LINEMOD,
vote_type,
augment=False,
use_motion=motion_model)
test_sampler = SequentialSampler(test_set)
test_batch_sampler = ImageSizeBatchSampler(
test_sampler, train_cfg['test_batch_size'], False)
test_loader = DataLoader(test_set,
batch_sampler=test_batch_sampler,
num_workers=0)
prefix = 'test' if args.use_test_set else 'val'
val(net, test_loader, begin_epoch, prefix, use_motion=motion_model)
if args.occluded and args.linemod_cls in cfg.occ_linemod_cls_names:
print('testing occluded linemod ...')
occ_image_db = OcclusionLineModImageDB(args.linemod_cls)
occ_test_db = occ_image_db.test_real_set
occ_test_set = LineModDatasetRealAug(occ_test_db,
cfg.OCCLUSION_LINEMOD,
vote_type,
augment=False,
use_motion=motion_model)
occ_test_sampler = SequentialSampler(occ_test_set)
occ_test_batch_sampler = ImageSizeBatchSampler(
occ_test_sampler, train_cfg['test_batch_size'], False)
occ_test_loader = DataLoader(occ_test_set,
batch_sampler=occ_test_batch_sampler,
num_workers=0)
prefix = 'occ_test' if args.use_test_set else 'occ_val'
val(net,
occ_test_loader,
begin_epoch,
prefix,
use_motion=motion_model)
if args.truncated:
print('testing truncated linemod ...')
trun_image_db = TruncatedLineModImageDB(args.linemod_cls)
print(len(trun_image_db.set))
trun_image_set = LineModDatasetRealAug(trun_image_db.set,
cfg.LINEMOD,
vote_type,
augment=False,
use_intrinsic=True,
use_motion=motion_model)
trun_test_sampler = SequentialSampler(trun_image_set)
trun_test_batch_sampler = ImageSizeBatchSampler(
trun_test_sampler, train_cfg['test_batch_size'], False)
trun_test_loader = DataLoader(
trun_image_set,
batch_sampler=trun_test_batch_sampler,
num_workers=0)
prefix = 'trun_test'
val(net,
trun_test_loader,
begin_epoch,
prefix,
True,
use_motion=motion_model)
else:
# //train from finetune,or train from 0
# "resume":true,
# //this para no use
# "finetune":false,
begin_epoch = 0
# train from finetune
if train_cfg['resume']:
begin_epoch = load_model(net.module.net, optimizer, model_dir)
#cat already render
# image_db = LineModImageDB(args.linemod_cls,
# has_fuse_set=train_cfg['use_fuse'],
# has_render_set=True)
image_db = LineModImageDB(args.linemod_cls,
has_fuse_set=train_cfg['use_fuse'],
has_render_set=False)
train_db = []
#三类训练数据吗，render，real，fuse？？ 阶段数据呢，，遮挡数据呢
train_db += image_db.render_set
if train_cfg['use_real_train']:
train_db += image_db.train_real_set
if train_cfg['use_fuse']:
train_db += image_db.fuse_set
train_set = LineModDatasetRealAug(train_db,
cfg.LINEMOD,
vote_type,
augment=True,
cfg=train_cfg['aug_cfg'],
use_motion=motion_model)
train_sampler = RandomSampler(train_set)
train_batch_sampler = ImageSizeBatchSampler(
train_sampler,
train_cfg['train_batch_size'],
False,
cfg=train_cfg['aug_cfg'])
train_loader = DataLoader(train_set,
batch_sampler=train_batch_sampler,
num_workers=12)
val_db = image_db.val_real_set
val_set = LineModDatasetRealAug(val_db,
cfg.LINEMOD,
vote_type,
augment=False,
cfg=train_cfg['aug_cfg'],
use_motion=motion_model)
val_sampler = SequentialSampler(val_set)
val_batch_sampler = ImageSizeBatchSampler(val_sampler,
train_cfg['test_batch_size'],
False,
cfg=train_cfg['aug_cfg'])
val_loader = DataLoader(val_set,
batch_sampler=val_batch_sampler,
num_workers=12)
if args.linemod_cls in cfg.occ_linemod_cls_names:
occ_image_db = OcclusionLineModImageDB(args.linemod_cls)
occ_val_db = occ_image_db.test_real_set[:len(occ_image_db.
test_real_set) // 2]
occ_val_set = LineModDatasetRealAug(occ_val_db,
cfg.OCCLUSION_LINEMOD,
vote_type,
augment=False,
cfg=train_cfg['aug_cfg'],
use_motion=motion_model)
occ_val_sampler = SequentialSampler(occ_val_set)
occ_val_batch_sampler = ImageSizeBatchSampler(
occ_val_sampler,
train_cfg['test_batch_size'],
False,
cfg=train_cfg['aug_cfg'])
occ_val_loader = DataLoader(occ_val_set,
batch_sampler=occ_val_batch_sampler,
num_workers=12)
for epoch in range(begin_epoch, train_cfg['epoch_num']):
adjust_learning_rate(optimizer, epoch, train_cfg['lr_decay_rate'],
train_cfg['lr_decay_epoch'])
print("xx")
train(net, optimizer, train_loader, epoch)
val(net, val_loader, epoch, use_motion=motion_model)
if args.linemod_cls in cfg.occ_linemod_cls_names:
val(net,
occ_val_loader,
epoch,
'occ_val',
use_motion=motion_model)
save_model(net.module.net, optimizer, epoch, model_dir)
def train_net():
# mp.spawn(demo_basic,
# args=(4,),
# nprocs=4,
# join=True)
tf_dir = './runs/' + train_cfg['exp_name']
writer = SummaryWriter(log_dir=tf_dir)
Path("/home/gerard/myPvnet/pvnet/{}".format(train_cfg["exp_name"])).mkdir(
parents=True, exist_ok=True)
model_dir = os.path.join(cfg.MODEL_DIR, train_cfg['model_name'])
imNet = ImageUNet(ver_dim=(vote_num * 2), seg_dim=2)
estNet = EstimateUNet(ver_dim=(vote_num * 2), seg_dim=2)
net = NetWrapper(imNet, estNet)
net = DataParallel(net).cuda()
# if train_cfg['exp_name'] == 'AE_ape':
# model_Dir=os.path.join(cfg.MODEL_DIR,'ape_linemod_train_GE')
# imNet=load_pretrained_imNet(ImageUNet(ver_dim=(vote_num*2), seg_dim=2), model_Dir, epoch=0)
# elif train_cfg['exp_name'] == 'GE_ape':
# model_Dir=os.path.join(cfg.MODEL_DIR,'ape_linemod_train_AE')
# # estNet=load_pretrained_estNet(EstimateUNet(ver_dim=(vote_num*2), seg_dim=2), model_Dir, epoch=25)
# elif train_cfg['exp_name'] == 'AE_GE_ape':
# model_Dir=os.path.join(cfg.PVModelDir)['net'])
# PVNet = PVNet.cuda(rank)
# PVNet=DistributedDataParallel(PVNet, device_ids=[rank])
PVModelDir = '/home/gerard/baseline_models/{}_baseline/199.pth'.format(
train_cfg['object'])
PVNet = PVnet(ver_dim=vote_num * 2, seg_dim=2)
PVNet.load_state_dict(torch.load(PVModelDir)['net'])
PVNet = PVNet.half()
PVNet = DataParallel(PVNet).cuda()
randomCropping = RandomScaleCrop()
optimizer = optim.Adam(net.parameters(), lr=train_cfg['lr'])
motion_model = train_cfg['motion_model']
print('motion state {}'.format(motion_model))
for param_group in optimizer.param_groups:
lr = param_group['lr']
image_db = LineModImageDB(args.linemod_cls,
has_render_set=False,
has_fuse_set=False)
test_db = image_db.test_real_set + image_db.val_real_set
test_set = LineModDatasetRealAug(test_db,
randomCropping,
cfg.LINEMOD,
vote_type,
augment=False,
use_motion=motion_model)
test_sampler = SequentialSampler(test_set)
test_batch_sampler = ImageSizeBatchSampler(test_sampler,
train_cfg['test_batch_size'],
False)
test_loader = DataLoader(test_set,
batch_sampler=test_batch_sampler,
num_workers=0)
prefix = 'test' #if args.use_test_set else 'val'
if args.test_model:
begin_epoch = load_model(net.module.imNet, net.module.estNet,
optimizer, model_dir, args.load_epoch)
if args.normal:
print('testing normal linemod ...')
# image_db = LineModImageDB(args.linemod_cls,has_render_set=False,
# has_fuse_set=False)
# test_db = image_db.test_real_set+image_db.val_real_set
# test_set = LineModDatasetRealAug(test_db, randomCropping, cfg.LINEMOD, vote_type, augment=False, use_motion=motion_model)
# test_sampler = SequentialSampler(test_set)
# test_batch_sampler = ImageSizeBatchSampler(test_sampler, train_cfg['test_batch_size'], False)
# test_loader = DataLoader(test_set, batch_sampler=test_batch_sampler, num_workers=0)
val_db = image_db.test_real_set + image_db.val_real_set
val_set = LineModDatasetRealAug(val_db,
randomCropping,
cfg.LINEMOD,
vote_type,
augment=False,
cfg=train_cfg['aug_cfg'],
use_motion=motion_model)
val_sampler = SequentialSampler(val_set)
val_batch_sampler = ImageSizeBatchSampler(
val_sampler,
train_cfg['test_batch_size'],
False,
cfg=train_cfg['aug_cfg'])
val_loader = DataLoader(val_set,
batch_sampler=val_batch_sampler,
num_workers=16)
print('test with val loader')
_, _, _, _, _, _, _, _, _ = val(net,
PVNet,
val_loader,
begin_epoch,
lr,
writer,
use_motion=motion_model)
prefix = 'test' if args.use_test_set else 'val'
print('test with test_loader')
_, _, _, _, _, _, _, _, _ = val(net,
PVNet,
test_loader,
begin_epoch,
lr,
writer,
use_motion=motion_model)
if args.occluded and args.linemod_cls in cfg.occ_linemod_cls_names:
print('testing occluded linemod ...')
occ_image_db = OcclusionLineModImageDB(args.linemod_cls)
occ_test_db = occ_image_db.test_real_set
occ_test_set = LineModDatasetRealAug(occ_test_db,
randomCropping,
cfg.OCCLUSION_LINEMOD,
vote_type,
augment=False,
use_motion=motion_model)
occ_test_sampler = SequentialSampler(occ_test_set)
occ_test_batch_sampler = ImageSizeBatchSampler(
occ_test_sampler, train_cfg['test_batch_size'], False)
occ_test_loader = DataLoader(occ_test_set,
batch_sampler=occ_test_batch_sampler,
num_workers=0)
prefix = 'occ_test' if args.use_test_set else 'occ_val'
_, _, _, _, _, _, _, _, _ = val(net,
PVNet,
occ_test_loader,
begin_epoch,
lr,
writer,
prefix,
use_motion=motion_model)
if args.truncated:
print('testing truncated linemod ...')
trun_image_db = TruncatedLineModImageDB(args.linemod_cls)
print(len(trun_image_db.set))
trun_image_set = LineModDatasetRealAug(trun_image_db.set,
cfg.LINEMOD,
vote_type,
randomCropping,
augment=False,
use_intrinsic=True,
use_motion=motion_model)
trun_test_sampler = SequentialSampler(trun_image_set)
trun_test_batch_sampler = ImageSizeBatchSampler(
trun_test_sampler, train_cfg['test_batch_size'], False)
trun_test_loader = DataLoader(
trun_image_set,
batch_sampler=trun_test_batch_sampler,
num_workers=16)
prefix = 'trun_test'
_, _, _, _, _, _, _, _, _ = val(net,
PVNet,
trun_test_loader,
begin_epoch,
lr,
writer,
prefix,
True,
use_motion=motion_model)
else:
begin_epoch = 0
if train_cfg['resume']:
begin_epoch = load_model(net.module.imNet, net.module.estNet,
optimizer, model_dir)
# reset learning rate
for param_group in optimizer.param_groups:
param_group['lr'] = train_cfg['lr']
lr = param_group['lr']
image_db = LineModImageDB(args.linemod_cls,
has_fuse_set=train_cfg['use_fuse'],
has_render_set=True)
train_db = []
train_db += image_db.render_set
if train_cfg['use_real_train']:
train_db += image_db.train_real_set
if train_cfg['use_fuse']:
train_db += image_db.fuse_set
train_set = LineModDatasetRealAug(train_db,
randomCropping,
cfg.LINEMOD,
vote_type,
augment=True,
cfg=train_cfg['aug_cfg'],
use_motion=motion_model)
train_sampler = RandomSampler(train_set)
# train_sampler = torch.utils.data.distributed.DistributedSampler(
# train_set,
# num_replicas=world_size,
# rank=rank
# )
# train_batch_sampler = ImageSizeBatchSampler(train_sampler, int(train_cfg['train_batch_size']/world_size), False, cfg=train_cfg['aug_cfg'])
train_batch_sampler = ImageSizeBatchSampler(
train_sampler,
train_cfg['train_batch_size'],
False,
cfg=train_cfg['aug_cfg'])
# train_loader = DataLoader(train_set, batch_sampler=train_batch_sampler, shuffle=False, num_workers=16, pin_memory=True)
train_loader = DataLoader(train_set,
batch_sampler=train_batch_sampler,
num_workers=16)
val_db = image_db.test_real_set + image_db.val_real_set
val_set = LineModDatasetRealAug(val_db,
randomCropping,
cfg.LINEMOD,
vote_type,
augment=False,
cfg=train_cfg['aug_cfg'],
use_motion=motion_model)
val_sampler = SequentialSampler(val_set)
val_batch_sampler = ImageSizeBatchSampler(val_sampler,
train_cfg['test_batch_size'],
False,
cfg=train_cfg['aug_cfg'])
val_loader = DataLoader(val_set,
batch_sampler=val_batch_sampler,
num_workers=16)
if args.linemod_cls in cfg.occ_linemod_cls_names:
occ_image_db = OcclusionLineModImageDB(args.linemod_cls)
occ_val_db = occ_image_db.test_real_set[:len(occ_image_db.
test_real_set) // 2]
occ_val_set = LineModDatasetRealAug(occ_val_db,
randomCropping,
cfg.OCCLUSION_LINEMOD,
vote_type,
augment=False,
cfg=train_cfg['aug_cfg'],
use_motion=motion_model)
occ_val_sampler = SequentialSampler(occ_val_set)
occ_val_batch_sampler = ImageSizeBatchSampler(
occ_val_sampler,
train_cfg['test_batch_size'],
False,
cfg=train_cfg['aug_cfg'])
occ_val_loader = DataLoader(occ_val_set,
batch_sampler=occ_val_batch_sampler,
num_workers=16)
add_list_list = []
p_inc_list = []
p_dec_v_list = []
p_dec_q_list = []
largest_a_list = []
smallest_v_list = []
smallest_q_list = []
first_a_list = []
first_v_list = []
epoch_count = 0
for epoch in range(begin_epoch, train_cfg['epoch_num']):
adjust_learning_rate(optimizer, epoch, train_cfg['lr_decay_rate'],
train_cfg['lr_decay_epoch'])
for param_group in optimizer.param_groups:
lr = param_group['lr']
train(net, PVNet, optimizer, train_loader, epoch)
# print('evaluate with test_loader')
# _,_,_,_,_,_,_,_,_ = val(net, PVNet, test_loader, epoch, lr, writer, use_motion=motion_model)
# print('evaluate with train_loader')
# _,_,_,_,_,_,_,_,_ = val(net, PVNet, train_loader, epoch, lr, writer, use_motion=motion_model)
print('evaluate with val_loader')
add_list, first_a, first_v, largest_a, smallest_v, smallest_q, p_inc_add, p_dec_v, p_dec_q = val(
net,
PVNet,
val_loader,
epoch,
lr,
writer,
use_motion=motion_model)
if (train_cfg['eval_epoch']
and epoch % train_cfg['eval_inter'] == 0 and
epoch >= train_cfg['eval_epoch_begin']) or args.test_model:
if epoch >= 30:
# add_list_list.append(add_list)
# first_a_list.append(first_a)
first_v_list.append(first_v)
# p_inc_list.append(p_inc_add)
p_dec_v_list.append(p_dec_v)
p_dec_q_list.append(p_dec_q)
# largest_a_list.append(largest_a)
smallest_v_list.append(smallest_v)
smallest_q_list.append(smallest_q)
# if args.linemod_cls in cfg.occ_linemod_cls_names:
# val(net, PVNet, occ_val_loader, epoch, lr, writer, 'occ_val',use_motion=motion_model)
save_model(net.module.imNet, net.module.estNet, optimizer, epoch,
model_dir)
epoch_count += 1
print(train_cfg['exp_name'])
# print('PVNet ADD. mean: {} +/- {}, max: {}'.format(np.mean(first_a_list),np.std(first_a_list),np.max(first_a_list)))
# print('PVNet X-X^. mean: {} +/- {}, max: {}'.format(np.mean(first_v_list),np.std(first_v_list),np.max(first_v_list)))
# print('ADD. mean: {} +/- {}, max: {}: '.format(np.mean(largest_a_list),np.std(largest_a_list),np.max(largest_a_list)))
# print('ADD perc increase. mean: {} +/- {}, max: {}'.format(np.mean(p_inc_list),np.std(p_inc_list),np.max(p_inc_list)))
print('X-X^ perc decrease. mean: {} +/- {}, max: {}'.format(
np.mean(p_dec_v_list), np.std(p_dec_v_list), np.max(p_dec_v_list)))
print('q-q^ perc decrease. mean: {} +/- {}, max: {}'.format(
np.mean(p_dec_q_list), np.std(p_dec_q_list), np.max(p_dec_q_list)))
def train_net():
#print("Out-location", out_location)
net = Resnet18_8s(ver_dim=vote_num * 2, seg_dim=2)
net = NetWrapper(net)
net = DataParallel(net).cuda()
optimizer = optim.Adam(net.parameters(), lr=train_cfg['lr'])
model_dir = os.path.join(cfg.MODEL_DIR, train_cfg['model_name'])
print(model_dir)
motion_model = train_cfg['motion_model']
print('motion state {}'.format(motion_model))
if args.test_model:
begin_epoch = load_model(net.module.net, optimizer, model_dir,
args.load_epoch)
print('testing ...')
val_db = ValidationDatasetIntake.getval_dataset(
) #image_db.val_real_set
print("val_db - ", len(val_db))
#print(val_db)
val_set = HomemadeDataset(val_db,
cfg.HOMEMADE,
vote_type,
augment=False,
cfg=train_cfg['homemade_cfg'],
use_motion=motion_model)
print("val set!", val_set)
val_sampler = SequentialSampler(val_set)
val_batch_sampler = ImageSizeBatchSampler(
val_sampler,
train_cfg['test_batch_size'],
False,
cfg=train_cfg['homemade_cfg'])
val_loader = DataLoader(val_set,
batch_sampler=val_batch_sampler,
num_workers=0)
val(net, val_loader, begin_epoch, use_motion=motion_model)
#val(net, val_loader, epoch,use_motion=motion_model)
if args.occluded and args.homemade_cls in cfg.occ_homemade_cls_names:
print('testing occluded dataset ...')
occ_image_db = OcclusionHomemadeImageDB(args.homemade_cls)
occ_test_db = occ_image_db.test_real_set
occ_test_set = HomemadeDataset(occ_test_db,
cfg.OCCLUSION_HOMEMADE,
vote_type,
augment=False,
use_motion=motion_model)
occ_test_sampler = SequentialSampler(occ_test_set)
occ_test_batch_sampler = ImageSizeBatchSampler(
occ_test_sampler, train_cfg['test_batch_size'], False)
occ_test_loader = DataLoader(occ_test_set,
batch_sampler=occ_test_batch_sampler,
num_workers=0)
prefix = 'occ_test' if args.use_test_set else 'occ_val'
val(net,
occ_test_loader,
begin_epoch,
prefix,
use_motion=motion_model)
if args.truncated:
print('testing truncated dataset ...')
trun_image_db = TruncatedHomemadeImageDB(args.homemade_cls)
print(len(trun_image_db.set))
trun_image_set = HomemadeDataset(trun_image_db.set,
cfg.HOMEMADE,
vote_type,
augment=False,
use_intrinsic=True,
use_motion=motion_model)
trun_test_sampler = SequentialSampler(trun_image_set)
trun_test_batch_sampler = ImageSizeBatchSampler(
trun_test_sampler, train_cfg['test_batch_size'], False)
trun_test_loader = DataLoader(
trun_image_set,
batch_sampler=trun_test_batch_sampler,
num_workers=0)
prefix = 'trun_test'
val(net,
trun_test_loader,
begin_epoch,
prefix,
True,
use_motion=motion_model)
else:
begin_epoch = 0
print("Train_cfg[resume] - ", train_cfg['resume'])
if train_cfg['resume']:
begin_epoch = load_model(net.module.net, optimizer, model_dir)
print("class", args.homemade_cls, "use fuse: ", train_cfg['use_fuse'])
image_db = HomemadeImageDB(args.homemade_cls,
has_fuse_set=train_cfg['use_fuse'],
has_render_set=True)
train_db = []
train_db += image_db.render_set
print("train_db after render set: ", len(train_db))
if train_cfg['use_real_train']:
train_db += image_db.train_real_set
print("DB After real train set: ",
len(train_db)) #we have no real images
if train_cfg['use_fuse']:
train_db += image_db.fuse_set
print("DB After fuse set: ", len(train_db))
train_set = HomemadeDataset(train_db,
cfg.HOMEMADE,
vote_type,
augment=True,
cfg=train_cfg['homemade_cfg'],
use_motion=motion_model)
train_sampler = RandomSampler(train_set)
train_batch_sampler = ImageSizeBatchSampler(
train_sampler,
train_cfg['train_batch_size'],
False,
cfg=train_cfg['homemade_cfg'])
train_loader = DataLoader(train_set,
batch_sampler=train_batch_sampler,
num_workers=12)
val_db = ValidationDatasetIntake.getval_dataset(
) #image_db.val_real_set
print("val_db - ", len(val_db))
#print(val_db)
val_set = HomemadeDataset(val_db,
cfg.HOMEMADE,
vote_type,
augment=False,
cfg=train_cfg['homemade_cfg'],
use_motion=motion_model)
print("val set!", val_set)
val_sampler = SequentialSampler(val_set)
val_batch_sampler = ImageSizeBatchSampler(
val_sampler,
train_cfg['test_batch_size'],
False,
cfg=train_cfg['homemade_cfg'])
val_loader = DataLoader(val_set,
batch_sampler=val_batch_sampler,
num_workers=12)
if args.homemade_cls in cfg.occ_homemade_cls_names:
occ_image_db = OcclusionHomemadeImageDB(args.homemade_cls)
occ_val_db = occ_image_db.test_real_set[:len(occ_image_db.
test_real_set) // 2]
occ_val_set = HomemadeDataset(occ_val_db,
cfg.OCCLUSION_LINEMOD,
vote_type,
augment=False,
cfg=train_cfg['homemade_cfg'],
use_motion=motion_model)
occ_val_sampler = SequentialSampler(occ_val_set)
occ_val_batch_sampler = ImageSizeBatchSampler(
occ_val_sampler,
train_cfg['test_batch_size'],
False,
cfg=train_cfg['homemade_cfg'])
occ_val_loader = DataLoader(occ_val_set,
batch_sampler=occ_val_batch_sampler,
num_workers=12)
for epoch in range(begin_epoch, train_cfg['epoch_num']):
adjust_learning_rate(optimizer, epoch, train_cfg['lr_decay_rate'],
train_cfg['lr_decay_epoch'])
train(net, optimizer, train_loader, epoch)
#val(net, val_loader, epoch,use_motion=motion_model)
if args.homemade_cls in cfg.occ_homemade_cls_names:
val(net,
occ_val_loader,
epoch,
'occ_val',
use_motion=motion_model)
save_model(net.module.net, optimizer, epoch, model_dir)
def demo(idx):
net = Resnet18_8s(ver_dim=vote_num * 2, seg_dim=2)
net = NetWrapper(net).cuda()
net = DataParallel(net)
optimizer = optim.Adam(net.parameters(), lr=train_cfg['lr'])
model_dir = os.path.join(cfg.MODEL_DIR, "intake_demo")
load_model(net.module.net, optimizer, model_dir, -1)
data, points_3d, bb8_3d = read_data(idx)
#print("BB8_3D: ",bb8_3d)
image, mask, vertex, vertex_weights, pose, corner_target = [
d.unsqueeze(0).cuda() for d in data
]
seg_pred, vertex_pred, loss_seg, loss_vertex, precision, recall = net(
image, mask, vertex, vertex_weights)
seg_mask = torch.argmax(seg_pred, 1)
visualize_mask(seg_mask)
visualize_mask(mask)
#visualize_vertex(vertex, vertex_weights)
#visualize_hypothesis(image, seg_pred, vertex_pred, corner_target)
#visualize_voting_ellipse(image, seg_pred, vertex_pred, corner_target)
#############
eval_net = DataParallel(EvalWrapper().cuda())
uncertain_eval_net = DataParallel(UncertaintyEvalWrapper().cuda())
corner_pred = eval_net(seg_pred, vertex_pred).cpu().detach().numpy()[0]
net.eval()
loss_seg, loss_vertex, precision, recall = [
torch.mean(val) for val in (loss_seg, loss_vertex, precision, recall)
]
print("LOSS SEG :", loss_seg, "\nLOSS VERTEX : ", loss_vertex,
"\nPRECISION :", precision, '\nRECALL :', recall)
###############
#print("Corner Predictions: ",corner_pred)
camera_matrix = np.array([[700, 0., 320.], [0., 700, 240.], [0., 0., 1.]])
pose_pred = pnp(points_3d, corner_pred, camera_matrix)
projector = Projector()
print("Pose prediction :\n", pose_pred)
pose_gt = pose[0].detach().cpu().numpy()
print("GT Pose :\n", pose[0].detach().cpu().numpy())
s = 0
import math as m
for i in range(3):
if pose_pred[2][3] < 0:
print('NB!')
s += (pose_pred[i][3] - pose_gt[i][3])**2
s = m.sqrt(s)
print("--->",
loss_seg.detach().cpu().numpy(),
loss_vertex.detach().cpu().numpy(),
precision.detach().cpu().numpy(),
recall.detach().cpu().numpy(), s)
bb8_2d_pred = projector.project(bb8_3d, pose_pred, 'blender')
bb8_2d_gt = projector.project(bb8_3d, pose[0].detach().cpu().numpy(),
'blender')
#print(bb8_2d_gt)
image = imagenet_to_uint8(image.detach().cpu().numpy())[0]
visualize_bounding_box(image[None, ...], bb8_2d_pred[None, None, ...],
bb8_2d_gt[None, None, ...])
def demo():
net = Resnet18_8s(ver_dim=vote_num * 2, seg_dim=2)
net = NetWrapper(net).cuda()
net = DataParallel(net)
optimizer = optim.Adam(net.parameters(), lr=train_cfg['lr'])
model_dir = os.path.join(cfg.MODEL_DIR, 'cat_demo')
load_model(net.module.net, optimizer, model_dir, args.load_epoch)
data, points_3d, bb8_3d = read_data()
image, mask, vertex, vertex_weights, pose, corner_target = [
d.unsqueeze(0).cuda() for d in data
]
# Run the net
seg_pred, vertex_pred, loss_seg, loss_vertex, precision, recall = net(
image, mask, vertex, vertex_weights)
print('vertex_pred.shape')
print(vertex_pred.shape)
print(' ')
print('vertex_pred[0]')
print(vertex_pred)
print(' ')
# Various visualizations
#visualize_vertex_field(vertex_pred,vertex_weights, keypointIdx=3)
print('asdasdsadas')
print(seg_pred.shape, mask.shape)
visualize_mask(np.squeeze(seg_pred.cpu().detach().numpy()),
mask.cpu().detach().numpy())
rgb = Image.open('data/demo/cat.jpg')
img = np.array(rgb)
#visualize_overlap_mask(img, np.squeeze(seg_pred.cpu().detach().numpy()), None)
# Run the ransac voting
eval_net = DataParallel(EvalWrapper2().cuda())
#corner_pred = eval_net(seg_pred, vertex_pred).cpu().detach().numpy()[0]
corner_pred, covar = [
x.cpu().detach().numpy()[0] for x in eval_net(seg_pred, vertex_pred)
]
print('Keypoint predictions:')
print(corner_pred)
print(' ')
print('covar: ', covar)
print(' ')
camera_matrix = np.array([[572.4114, 0., 325.2611],
[0., 573.57043, 242.04899], [0., 0., 1.]])
# Fit pose to points
#pose_pred = pnp(points_3d, corner_pred, camera_matrix)
#evaluator = Evaluator()
#pose_pred = evaluator.evaluate_uncertainty(corner_pred,covar,pose,'cat',intri_matrix=camera_matrix)
def getWeights(covar):
cov_invs = []
for vi in range(covar.shape[0]): # For every keypoint
if covar[vi, 0, 0] < 1e-6 or np.sum(np.isnan(covar)[vi]) > 0:
cov_invs.append(np.zeros([2, 2]).astype(np.float32))
continue
cov_inv = np.linalg.inv(scipy.linalg.sqrtm(covar[vi]))
cov_invs.append(cov_inv)
cov_invs = np.asarray(cov_invs) # pn,2,2
weights = cov_invs.reshape([-1, 4])
weights = weights[:, (0, 1, 3)]
return weights
weights = getWeights(covar)
pose_pred = uncertainty_pnp(corner_pred, weights, points_3d, camera_matrix)
print('Predicted pose: \n', pose_pred)
print('Ground truth pose: \n', pose[0].detach().cpu().numpy())
print(' ')
projector = Projector()
bb8_2d_pred = projector.project(bb8_3d, pose_pred, 'linemod')
bb8_2d_gt = projector.project(bb8_3d, pose[0].detach().cpu().numpy(),
'linemod')
image = imagenet_to_uint8(image.detach().cpu().numpy())[0]
visualize_points(image[None, ...],
corner_target.detach().cpu().numpy(),
pts_pred=corner_pred[None, :, :])
visualize_bounding_box(image[None, ...], bb8_2d_pred[None, None, ...],
bb8_2d_gt[None, None, ...])
def train_net():
net = Resnet18_8s(ver_dim=vote_num * 2, seg_dim=2)
net = NetWrapper(net)
net = DataParallel(net).cuda()
optimizer = optim.Adam(net.parameters(), lr=train_cfg["lr"])
model_dir = os.path.join(cfg.MODEL_DIR, train_cfg["model_name"])
motion_model = train_cfg["motion_model"]
print("motion state {}".format(motion_model))
if args.test_model:
begin_epoch = load_model(net.module.net, optimizer, model_dir,
args.load_epoch)
if args.normal:
print("testing normal linemod ...")
image_db = LineModImageDB(args.linemod_cls,
has_render_set=False,
has_fuse_set=False)
test_db = image_db.test_real_set + image_db.val_real_set
test_set = LineModDatasetRealAug(test_db,
cfg.LINEMOD,
vote_type,
augment=False,
use_motion=motion_model)
test_sampler = SequentialSampler(test_set)
test_batch_sampler = ImageSizeBatchSampler(
test_sampler, train_cfg["test_batch_size"], False)
test_loader = DataLoader(test_set,
batch_sampler=test_batch_sampler,
num_workers=0)
prefix = "test" if args.use_test_set else "val"
val(net, test_loader, begin_epoch, prefix, use_motion=motion_model)
if args.occluded and args.linemod_cls in cfg.occ_linemod_cls_names:
print("testing occluded linemod ...")
occ_image_db = OcclusionLineModImageDB(args.linemod_cls)
occ_test_db = occ_image_db.test_real_set
occ_test_set = LineModDatasetRealAug(
occ_test_db,
cfg.OCCLUSION_LINEMOD,
vote_type,
augment=False,
use_motion=motion_model,
)
occ_test_sampler = SequentialSampler(occ_test_set)
occ_test_batch_sampler = ImageSizeBatchSampler(
occ_test_sampler, train_cfg["test_batch_size"], False)
occ_test_loader = DataLoader(occ_test_set,
batch_sampler=occ_test_batch_sampler,
num_workers=0)
prefix = "occ_test" if args.use_test_set else "occ_val"
val(net,
occ_test_loader,
begin_epoch,
prefix,
use_motion=motion_model)
if args.truncated:
print("testing truncated linemod ...")
trun_image_db = TruncatedLineModImageDB(args.linemod_cls)
print(len(trun_image_db.set))
trun_image_set = LineModDatasetRealAug(
trun_image_db.set,
cfg.LINEMOD,
vote_type,
augment=False,
use_intrinsic=True,
use_motion=motion_model,
)
trun_test_sampler = SequentialSampler(trun_image_set)
trun_test_batch_sampler = ImageSizeBatchSampler(
trun_test_sampler, train_cfg["test_batch_size"], False)
trun_test_loader = DataLoader(
trun_image_set,
batch_sampler=trun_test_batch_sampler,
num_workers=0)
prefix = "trun_test"
val(
net,
trun_test_loader,
begin_epoch,
prefix,
True,
use_motion=motion_model,
)
else:
begin_epoch = 0
if train_cfg["resume"]:
begin_epoch = load_model(net.module.net, optimizer, model_dir)
image_db = LineModImageDB(args.linemod_cls,
has_fuse_set=train_cfg["use_fuse"],
has_render_set=False)
train_db = []
train_db += image_db.render_set
if train_cfg["use_real_train"]:
train_db += image_db.train_real_set
if train_cfg["use_fuse"]:
train_db += image_db.fuse_set
train_set = LineModDatasetRealAug(
train_db,
cfg.LINEMOD,
vote_type,
augment=True,
cfg=train_cfg["aug_cfg"],
use_motion=motion_model,
)
train_sampler = RandomSampler(train_set)
train_batch_sampler = ImageSizeBatchSampler(
train_sampler,
train_cfg["train_batch_size"],
False,
cfg=train_cfg["aug_cfg"],
)
train_loader = DataLoader(train_set,
batch_sampler=train_batch_sampler,
num_workers=12)
val_db = image_db.val_real_set
val_set = LineModDatasetRealAug(
val_db,
cfg.LINEMOD,
vote_type,
augment=False,
cfg=train_cfg["aug_cfg"],
use_motion=motion_model,
)
val_sampler = SequentialSampler(val_set)
val_batch_sampler = ImageSizeBatchSampler(val_sampler,
train_cfg["test_batch_size"],
False,
cfg=train_cfg["aug_cfg"])
val_loader = DataLoader(val_set,
batch_sampler=val_batch_sampler,
num_workers=12)
if args.use_occlussion and args.linemod_cls in cfg.occ_linemod_cls_names:
occ_image_db = OcclusionLineModImageDB(args.linemod_cls)
occ_val_db = occ_image_db.test_real_set[:len(occ_image_db.
test_real_set) // 2]
occ_val_set = LineModDatasetRealAug(
occ_val_db,
cfg.OCCLUSION_LINEMOD,
vote_type,
augment=False,
cfg=train_cfg["aug_cfg"],
use_motion=motion_model,
)
occ_val_sampler = SequentialSampler(occ_val_set)
occ_val_batch_sampler = ImageSizeBatchSampler(
occ_val_sampler,
train_cfg["test_batch_size"],
False,
cfg=train_cfg["aug_cfg"],
)
occ_val_loader = DataLoader(occ_val_set,
batch_sampler=occ_val_batch_sampler,
num_workers=12)
for epoch in range(begin_epoch, train_cfg["epoch_num"]):
adjust_learning_rate(
optimizer,
epoch,
train_cfg["lr_decay_rate"],
train_cfg["lr_decay_epoch"],
)
train(net, optimizer, train_loader, epoch)
val(net, val_loader, epoch, use_motion=motion_model)
if args.use_occlussion and args.linemod_cls in cfg.occ_linemod_cls_names:
val(net,
occ_val_loader,
epoch,
"occ_val",
use_motion=motion_model)
save_model(net.module.net, optimizer, epoch, model_dir)
bb8_2d_pred[None, None, ...],
save=True,
count=count) #,bb8_2d_gt[None, None, ...])
if __name__ == "__main__":
import torchvision.transforms as transforms
net = Resnet18_8s(ver_dim=vote_num * 2, seg_dim=2)
net = NetWrapper(net).cuda()
net = DataParallel(net)
optimizer = optim.Adam(net.parameters(), lr=train_cfg['lr'])
#load model
model_dir = os.path.join(cfg.MODEL_DIR,
"{}_demo".format(args.homemade_cls))
load_model(net.module.net, optimizer, model_dir, args.load_epoch)
demo_dir = os.path.join(cfg.DATA_DIR, 'demo',
'{}'.format(args.homemade_cls))
#load bb8
bb8_3d = np.loadtxt(
os.path.join(cfg.HOMEMADE, '{}'.format(args.homemade_cls),
'corners.txt')) #, delimiter=" ", usecols=range(3))
points_3d = np.loadtxt(
os.path.join(cfg.HOMEMADE, '{}'.format(args.homemade_cls),
'{}_points_3d.txt'.format(args.homemade_cls)))
camera_matrix = np.array([[542.9039, 0., 323.9369],
[0., 529.055241, 239.57233], [0., 0., 1.]])
"""
inference(os.path.join(demo_dir,'source_real','intake_logi_3.jpg'))
"""
#def train_net():
net=Resnet18_8s(ver_dim=vote_num*2, seg_dim=2)
net=NetWrapper(net)
net=DataParallel(net).cuda()
optimizer = optim.Adam(net.parameters(), lr=train_cfg['lr'])
model_dir=os.path.join(cfg.MODEL_DIR,train_cfg['model_name'])
motion_model=train_cfg['motion_model']
print('motion state {}'.format(motion_model))
#if args.test_model:
#else:
begin_epoch=0
if train_cfg['resume']:
begin_epoch=load_model(net.module.net, optimizer, model_dir)
image_db = LineModImageDB(args.linemod_cls,
has_fuse_set=train_cfg['use_fuse'],
has_render_set=True)
train_db=[]
train_db+=image_db.render_set
if train_cfg['use_real_train']:
train_db+=image_db.train_real_set
if train_cfg['use_fuse']:
train_db+=image_db.fuse_set
train_set = LineModDatasetRealAug(train_db, cfg.LINEMOD, vote_type, augment=True, cfg=train_cfg['aug_cfg'], use_motion=motion_model)
train_sampler = RandomSampler(train_set)
