parser.add_argument('--bin_size',
type=int,
default=15,
help='bin size for the euler angle classification')
opt = parser.parse_args()
print(opt)
# ========================================================== #
# ================CREATE NETWORK============================ #
azi_classes, ele_classes, inp_classes = int(360 / opt.bin_size), int(
180 / opt.bin_size), int(360 / opt.bin_size)
model = PoseEstimator(shape_feature_dim=opt.shape_feature_dim,
img_feature_dim=opt.img_feature_dim,
azi_classes=azi_classes,
ele_classes=ele_classes,
inp_classes=inp_classes)
model.cuda()
load_checkpoint(model, opt.model)
# ========================================================== #
# =============DEFINE stuff for logs ======================= #
# write basic information into the log file
if not os.path.isdir(opt.output):
os.mkdir(opt.output)
logname = os.path.join(opt.output, 'testing.txt')
f = open(logname, mode='w')
f.write('\n')
f.close()
drop_last=True)
# ========================================================== #
# ================CREATE NETWORK============================ #
if opt.shape is None:
model = BaselineEstimator(img_feature_dim=opt.img_feature_dim,
separate_branch=opt.separate_branch,
azi_classes=opt.azi_classes,
ele_classes=opt.ele_classes,
inp_classes=opt.inp_classes)
else:
model = PoseEstimator(shape=opt.shape,
shape_feature_dim=opt.shape_feature_dim,
img_feature_dim=opt.img_feature_dim,
azi_classes=opt.azi_classes,
ele_classes=opt.ele_classes,
inp_classes=opt.inp_classes,
render_number=opt.num_render,
separate_branch=opt.separate_branch,
channels=opt.channels)
model.cuda()
if opt.model is not None:
load_checkpoint(model, opt.model)
else:
model.apply(KaiMingInit)
# ========================================================== #
# ================CREATE OPTIMIZER AND LOSS================= #
optimizer = optim.Adam(model.parameters(), lr=opt.lr, weight_decay=0.0005)
lrScheduler = optim.lr_scheduler.MultiStepLR(optimizer, [opt.decrease],
gamma=0.1)
parser.add_argument('--render_path',
type=str,
default=None,
help='render images path')
parser.add_argument('--obj_path', type=str, default=None, help='obj path')
opt = parser.parse_args()
print(opt)
# ========================================================== #
# ================CREATE NETWORK============================ #
model = PoseEstimator(shape=opt.shape,
shape_feature_dim=opt.shape_feature_dim,
img_feature_dim=opt.img_feature_dim,
azi_classes=opt.azi_classes,
ele_classes=opt.ele_classes,
inp_classes=opt.inp_classes,
render_number=opt.num_render,
separate_branch=opt.separate_branch,
channels=opt.channels)
model.cuda()
if opt.model is not None:
checkpoint = torch.load(opt.model,
map_location=lambda storage, loc: storage.cuda())
pretrained_dict = checkpoint['state_dict']
model_dict = model.state_dict()
pretrained_dict = {
k: v
for k, v in pretrained_dict.items() if k in model_dict
}
model_dict.update(pretrained_dict)
parser.add_argument('--input_dim', type=int, default=224, help='input image dimension')
# path to load the necessary input for network inference
parser.add_argument('--image_path', type=str, default=None, help='real images path')
parser.add_argument('--render_path', type=str, default=None, help='render images path')
parser.add_argument('--obj_path', type=str, default=None, help='obj path')
opt = parser.parse_args()
print(opt)
# ========================================================== #
# ================CREATE NETWORK============================ #
azi_classes, ele_classes, inp_classes = int(360 / opt.bin_size), int(180 / opt.bin_size), int(360 / opt.bin_size)
model = PoseEstimator(shape=opt.shape, shape_feature_dim=opt.shape_feature_dim, img_feature_dim=opt.img_feature_dim,
azi_classes=azi_classes, ele_classes=ele_classes, inp_classes=inp_classes, view_num=opt.view_num)
model.cuda()
if opt.model is not None:
checkpoint = torch.load(opt.model, map_location=lambda storage, loc: storage.cuda())
pretrained_dict = checkpoint['state_dict']
model_dict = model.state_dict()
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
print('Previous weight loaded')
model.eval()
# ========================================================== #
# ==================INPUT IMAGE AND RENDERS================= #
# ========================================================== #
# ================CREATE NETWORK============================ #
azi_classes, ele_classes, inp_classes = int(360 / opt.bin_size), int(
180 / opt.bin_size), int(360 / opt.bin_size)
if opt.shape is None:
model = BaselineEstimator(img_feature_dim=opt.img_feature_dim,
azi_classes=azi_classes,
ele_classes=ele_classes,
inp_classes=inp_classes)
else:
model = PoseEstimator(shape=opt.shape,
shape_feature_dim=opt.shape_feature_dim,
img_feature_dim=opt.img_feature_dim,
azi_classes=azi_classes,
ele_classes=ele_classes,
inp_classes=inp_classes,
view_num=opt.view_num)
model.cuda()
if opt.model is not None:
load_checkpoint(model, opt.model)
else:
model.apply(KaiMingInit)
# ========================================================== #
# ================CREATE OPTIMIZER AND LOSS================= #
optimizer = optim.Adam(model.parameters(), lr=opt.lr, weight_decay=0.0005)
lrScheduler = optim.lr_scheduler.MultiStepLR(optimizer, [opt.decrease],
gamma=0.1)
dataset_train = Pascal3D(train=True, root_dir=opt.root_dir_train, annotation_file=opt.annot_train,
cls_choice=test_cats, input_dim=opt.input_dim, point_num=opt.point_num,
keypoint=opt.keypoint, novel=opt.novel, shot=opt.shot)
else:
sys.exit('Wrong setting!')
train_loader = DataLoader(dataset_train, batch_size=opt.batch_size, shuffle=True, num_workers=opt.workers, drop_last=True)
print('train data consist of {} samples'.format(len(dataset_train)))
# ========================================================== #
# ================CREATE NETWORK============================ #
azi_classes, ele_classes, inp_classes = int(360 / opt.bin_size), int(180 / opt.bin_size), int(360 / opt.bin_size)
model = PoseEstimator(shape_feature_dim=opt.shape_feature_dim, img_feature_dim=opt.img_feature_dim,
azi_classes=azi_classes, ele_classes=ele_classes, inp_classes=inp_classes)
model.cuda()
if opt.resume is not None:
load_checkpoint(model, opt.resume)
else:
model.apply(KaiMingInit)
# ========================================================== #
# ================CREATE OPTIMIZER AND LOSS================= #
optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=opt.lr, weight_decay=5e-4)
lrScheduler = optim.lr_scheduler.MultiStepLR(optimizer, [opt.lr_step], gamma=0.1)
criterion_azi = SmoothCELoss(360, 24, opt.smooth) if opt.smooth is not None else CELoss(360)
criterion_ele = SmoothCELoss(180, 12, opt.smooth) if opt.smooth is not None else CELoss(180)
criterion_inp = SmoothCELoss(360, 24, opt.smooth) if opt.smooth is not None else CELoss(360)
parser.add_argument('--num_render', type=int, default=12, help='number of render images used in each sample')
parser.add_argument('--tour', type=int, default=2, help='elevation tour for randomized references')
parser.add_argument('--random_model', action='store_true', help='whether use random model in testing')
opt = parser.parse_args()
print(opt)
# ========================================================== #
# ================CREATE NETWORK============================ #
if opt.shape is None:
model = BaselineEstimator(img_feature_dim=opt.img_feature_dim,
azi_classes=opt.azi_classes, ele_classes=opt.ele_classes, inp_classes=opt.inp_classes)
else:
model = PoseEstimator(shape=opt.shape, shape_feature_dim=opt.shape_feature_dim, img_feature_dim=opt.img_feature_dim,
azi_classes=opt.azi_classes, ele_classes=opt.ele_classes, inp_classes=opt.inp_classes,
render_number=opt.num_render)
model.cuda()
if not os.path.isfile(opt.model):
raise ValueError('Non existing file: {0}'.format(opt.model))
else:
load_checkpoint(model, opt.model)
bin_size = 360. / opt.azi_classes
# ========================================================== #
# =============DEFINE stuff for logs ======================= #
# write basic information into the log file
predictions_path = os.getcwd() if opt.output_dir is None else opt.output_dir
if not os.path.isdir(predictions_path):
os.mkdir(predictions_path)
