ethnhe / FFB6D

[CVPR2021 Oral] FFB6D: A Full Flow Bidirectional Fusion Network for 6D Pose Estimation.
MIT License
290 stars 72 forks source link

Bad results were obtained on clamps #58

Open lyl-why opened 2 years ago

lyl-why commented 2 years ago

Hi, @ethnhe,this is really a nice work.When I tested with the code you provided, I achieved excellent results on the YCB dataset.But when I loaded the trained model and tested it on my another code, the AUC of the 051_large_clamp and 052_extra_large_clamp decreased significantly. I don't understand why this problem occurs,can you give me a hand? This is the result: ....................... 051_large_clamp add: 55.918468356315955 adds: 72.36640695549146 add(-s): 72.36640695549146 052_extra_large_clamp add: 46.069596855946074 adds: 74.85321034800057 add(-s): 74.85321034800057 061_foam_brick add: 74.87968248217449 adds: 94.66454086630006 add(-s): 94.66454086630006 mean kps errs: 0.0269345583775544 Average of all object: add: 80.76280976699599 adds: 92.01773639975028 add(-s): 86.96163893180336 All object (following PoseCNN): add: 82.18698476500266 adds: 92.31935751142115 ***add(-s): 86.41984918273243

This is the code for testing: ######## import argparse import numpy as np import torch import torch.nn.parallel from utils.basic_utils import Basic_Utils import torch.utils.data from datasets.ycb.dataset import Dataset as PoseDataset_ycb from models.ffb6d import FFB6D from tensorboardX import SummaryWriter import torch.cuda.amp from common import Config from utils.pvn3d_eval_utils_kpls import TorchEval from common import ConfigRandLA import torch.distributed

config = Config(ds_name='ycb') bs_utils = Basic_Utils(config) cls_lst = config.ycb_cls_lst writer = SummaryWriter()

parser = argparse.ArgumentParser() parser.add_argument("--local_rank", type=int, default=0) parser.add_argument('--resume_posenet', type=str, default = 'FFB6D_best.pth.tar', help='resume PoseNet model') opt = parser.parse_args() opt.test_gt = False

def main(): opt.outf = '/home/lyl/lyl/FFB6D-transformer/trained_models/ycb' # folder to save trained models opt.log_dir = 'experiments/logs/ycb' # folder to save logs

torch.cuda.set_device(opt.local_rank)
device = torch.device("cuda", opt.local_rank)
torch.distributed.init_process_group(backend='nccl', init_method='env://')

rndla_cfg = ConfigRandLA
estimator = FFB6D(n_classes=config.n_objects, n_pts=config.n_sample_points, rndla_cfg=rndla_cfg, n_kps=config.n_keypoints)
estimator.to(device)

checkpoint = torch.load('{0}/{1}'.format(opt.outf, opt.resume_posenet))
epoch = checkpoint["epoch"]
it = checkpoint.get("it", 0.0)
best_prec = checkpoint["best_prec"]
if estimator is not None and checkpoint["model_state"] is not None:
    ck_st = checkpoint['model_state']
    if 'module' in list(ck_st.keys())[0]:
        tmp_ck_st = {}
        for k, v in ck_st.items():
            tmp_ck_st[k.replace("module.", "")] = v
        ck_st = tmp_ck_st
    estimator.load_state_dict(ck_st)

test_dataset = PoseDataset_ycb('test')
testdataloader = torch.utils.data.DataLoader(test_dataset, batch_size=3, shuffle=False, num_workers=0)

teval = TorchEval()

with torch.no_grad():
    for i, inputs in enumerate(testdataloader, 0):
        inputs_torch = {}
        for key in inputs.keys():
            if inputs[key].dtype in [np.float32, np.uint8]:
                inputs_torch[key] = torch.from_numpy(inputs[key].astype(np.float32)).cuda()
            elif inputs[key].dtype in [np.int32, np.uint32]:
                inputs_torch[key] = torch.LongTensor(inputs[key].astype(np.int32)).cuda()
            elif inputs[key].dtype in [torch.uint8, torch.float32]:
                inputs_torch[key] = inputs[key].float().cuda()
            elif inputs[key].dtype in [torch.int32, torch.int16]:
                inputs_torch[key] = inputs[key].long().cuda()
            else:
                inputs_torch[key] = torch.LongTensor(inputs[key].astype(np.int32)).cuda()

        labels = inputs_torch['labels']
        end_points = estimator(inputs_torch)
        cld = inputs_torch['cld_rgb_nrm'][:, :3, :].permute(0, 2, 1).contiguous()
        _, cls_rgbd = torch.max(end_points['pred_rgbd_segs'], 1)

        if not opt.test_gt:
            # eval pose from point cloud prediction.
            teval.eval_pose_parallel(
                cld, inputs_torch['rgb'], cls_rgbd, end_points['pred_ctr_ofs'],
                inputs_torch['ctr_targ_ofst'], labels, 1, inputs_torch['cls_ids'],
                inputs_torch['RTs'], end_points['pred_kp_ofs'],
                inputs_torch['kp_3ds'], inputs_torch['ctr_3ds'], min_cnt=1,
                use_ctr_clus_flter=True, use_ctr=True, ds='ycb'
            )

        else:
            # test GT labels, keypoint and center point offset
            gt_ctr_ofs = inputs_torch['ctr_targ_ofst'].unsqueeze(2).permute(0, 2, 1, 3)
            gt_kp_ofs = inputs_torch['kp_targ_ofst'].permute(0, 2, 1, 3)
            teval.eval_pose_parallel(
                cld, inputs_torch['rgb'],inputs_torch['labels'], gt_ctr_ofs,
                inputs_torch['ctr_targ_ofst'], inputs_torch['labels'], 1, inputs_torch['cls_ids'],
                inputs_torch['RTs'], gt_kp_ofs, inputs_torch['kp_3ds'], inputs_torch['ctr_3ds'],
                min_cnt=1, use_ctr_clus_flter=True, use_ctr=True, ds='ycb'
            )
    teval.cal_auc()
HannahHaensen commented 2 years ago

@lyl-why did you had errors while training ycb?

sth like: https://github.com/ethnhe/FFB6D/issues/79