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():
#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)
if __name__=="__main__":
import sys
sys.path.append('.')
from lib.utils.data_utils import LineModImageDB,LineModModelDB,Projector
from lib.datasets.linemod_dataset import LineModDatasetRealAug,ImageSizeBatchSampler,VotingType
from lib.ransac_voting_gpu_layer.ransac_voting_gpu import ransac_voting_layer
from torch.utils.data import RandomSampler,DataLoader
from lib.utils.draw_utils import pts_to_img_pts
from lib.utils.evaluation_utils import pnp
import random
image_db = LineModImageDB('duck', has_ro_set=False, has_ra_set=False, has_plane_set=False, has_render_set=False,
has_ms_set=False,has_fuse_set=False)
random.shuffle(image_db.real_set)
dataset = LineModDatasetRealAug(image_db.real_set[:5], data_prefix=image_db.linemod_dir,
vote_type=VotingType.Extreme, augment=False)
sampler = RandomSampler(dataset)
batch_sampler = ImageSizeBatchSampler(sampler, 5, False)
loader = DataLoader(dataset, batch_sampler=batch_sampler, num_workers=8)
modeldb=LineModModelDB()
camera_matrix=Projector().intrinsic_matrix['linemod'].astype(np.float32)
for i, data in enumerate(loader):
rgb, mask, vertex, vertex_weight, pose, gt_corners = data
pts2d=gt_corners[0].numpy()[:,:2].astype(np.float32)
pts3d=modeldb.get_extreme_3d('duck')
pts3d=np.concatenate([pts3d,modeldb.get_centers_3d('duck')[None,:]],0).astype(np.float32)
wgt2d=np.zeros([pts2d.shape[0],3]).astype(np.float32)
wgt2d[:,(0,2)]=1.0
for k in range(pts2d.shape[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)
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)
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=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=False, 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=4)
#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=4)
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=4)
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:
print("No evaluation today!")
#val(net, occ_val_loader, epoch, 'occ_val',use_motion=motion_model)
save_model(net.module.net, optimizer, epoch, model_dir)
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)
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)))