Esempio n. 1
0
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)
Esempio n. 2
0
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()
Esempio n. 3
0
 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
Esempio n. 4
0
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()
Esempio n. 5
0
    corners_t = np.matmul(cam2optical, corners.transpose()).transpose()
    n = mesh_pts_t.shape[0]
    mesh_pts_t = np.concatenate(( mesh_pts_t[:,2].reshape(n,1) , -mesh_pts_t[:, 0].reshape(n,1), -mesh_pts_t[:, 1].reshape(n,1) ), axis=1)
    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):