def main():
if args.dataset == "ycb":
test_ds = YCB_Dataset('test')
obj_id = -1
else:
test_ds = LM_Dataset('test', cls_type=args.cls)
obj_id = config.lm_obj_dict[args.cls]
test_loader = torch.utils.data.DataLoader(
test_ds, batch_size=config.test_mini_batch_size, shuffle=False,
num_workers=20
)
model = PVN3D(
num_classes=config.n_objects, pcld_input_channels=6, pcld_use_xyz=True,
num_points=config.n_sample_points
).cuda()
model = convert_model(model)
model.cuda()
# load status from checkpoint
if args.checkpoint is not None:
checkpoint_status = load_checkpoint(
model, None, filename=args.checkpoint[:-8]
)
model = nn.DataParallel(model)
for i, data in tqdm.tqdm(
enumerate(test_loader), leave=False, desc="val"
):
cal_view_pred_pose(model, data, epoch=i, obj_id=obj_id)
def main():
rospy.init_node('pvn3d_LM_pred')
'''
## Dataset and Cls_id args ##
if args.dataset == "ycb":
pass
else:
pass'''
## Model Initialization ##
model = PVN3D(
num_classes=config.n_objects, pcld_input_channels=6, pcld_use_xyz=True,
num_points=config.n_sample_points, num_kps=config.n_keypoints
).cuda()
model = convert_model(model)
model.cuda()
# load status from checkpoint
if args.checkpoint is not None:
checkpoint_status = load_checkpoint(
model, None, filename=args.checkpoint[:-8]
)
model = nn.DataParallel(model)
print('Stand-alone Demo started...')
while not rospy.is_shutdown():
demo = PVN3D_ROS(config, model, pub_all=True, auto_detect=True)
rospy.spin()
def define_network(self, weights_path):
# define model object on GPU
model = PVN3D(num_classes=self.config.n_objects,
pcld_input_channels=6,
pcld_use_xyz=True,
num_points=self.config.n_sample_points).cuda()
# convert batch norm into synchornized batch norm
model = convert_model(model)
# model to GPU
model.cuda()
# load weights
checkpoint_status = self.load_checkpoint(model,
None,
filename=weights_path[:-8])
# convert model to distributed mode
model = nn.DataParallel(model)
return model
def detection_server():
rospy.init_node('pvn3d_detection_server')
## Model Initialization ##
model = PVN3D(
num_classes=config.n_objects, pcld_input_channels=6, pcld_use_xyz=True,
num_points=config.n_sample_points, num_kps=config.n_keypoints
).cuda()
model = convert_model(model)
model.cuda()
# load status from checkpoint
if args.checkpoint is not None:
checkpoint_status = load_checkpoint(
model, None, filename=args.checkpoint[:-8]
)
model = nn.DataParallel(model)
demo = PVN3D_ROS(config, model, pub_all=True, auto_detect=False)
while not rospy.is_shutdown():
print('PVN3d Detection server initialized. Waiting for Requests...')
server = rospy.Service('get_poses_pvn3d', getPoses_srv, demo.getDetections )
server.spin()
corners_t = np.concatenate(( corners_t[:,2].reshape(8,1) , -corners_t[:, 0].reshape(8,1), -corners_t[:, 1].reshape(8,1) ), axis=1)
'''
test_ds = CrankSlider_Dataset('test')
test_loader = torch.utils.data.DataLoader(
test_ds,
batch_size=config.test_mini_batch_size,
shuffle=False,
num_workers=20)
model = PVN3D(num_classes=config.n_classes,
pcld_input_channels=6,
pcld_use_xyz=True,
num_points=config.n_sample_points,
num_kps=config.n_keypoints).cuda()
model = convert_model(model)
model.cuda()
optimizer = optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
# default value
it = -1 # for the initialize value of `LambdaLR` and `BNMomentumScheduler`
best_loss = 1e10
start_epoch = 1
cur_mdl_pth = os.path.join(config.log_model_dir, 'pvn3d.pth.tar')
if args.checkpoint is None and os.path.exists(cur_mdl_pth):
args.checkpoint = cur_mdl_pth
# load status from checkpoint
if args.checkpoint is not None:
