コード例 #1
0
def export_detector_phototourism_gpu(config, output_dir, args):
    """
    input 1 images, output pseudo ground truth by homography adaptation.
    Save labels:
        pred:
            'prob' (keypoints): np (N1, 3)
    """
    from utils.utils import pltImshow
    from utils.utils import saveImg
    from utils.draw import draw_keypoints

    proj_path = "/data/projects/pytorch-superpoint"
    splits = ["train", "val"]

    # basic setting
    task = config["data"]["dataset"]
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    logging.info("train on device: %s", device)
    with open(osp.join(proj_path, output_dir, "config.yml"), "w") as f:
        yaml.dump(config, f, default_flow_style=False)

    ## parameters
    nms_dist = config["model"]["nms"]  # 4
    top_k = config["model"]["top_k"]
    homoAdapt_iter = config["data"]["homography_adaptation"]["num"]
    conf_thresh = config["model"]["detection_threshold"]
    nn_thresh = 0.7
    count = 0
    output_images = args.outputImg
    check_exist = True

    ## save data
    '''
    save_path = Path(output_dir)
    save_output = save_path
    save_output = save_output / "predictions" / export_task
    os.makedirs(save_output, exist_ok=True)
    '''
    def _create_loader(dataset, n_workers=8):
        return torch.utils.data.DataLoader(
            dataset,
            shuffle=False,
            pin_memory=True,
            num_workers=n_workers,
        )

    data_loaders = {}
    # create the dataset and dataloader classes
    for split in splits:
        dataset = Phototourism(split=split, **config["data"])
        data_loaders[split] = _create_loader(dataset)

    # model loading
    ## load pretrained
    try:
        path = config["pretrained"]
        print("==> Loading pre-trained network.")
        print("path: ", path)
        # This class runs the SuperPoint network and processes its outputs.

        fe = SuperPointFrontend_torch(
            config=config,
            weights_path=path,
            nms_dist=nms_dist,
            conf_thresh=conf_thresh,
            nn_thresh=nn_thresh,
            cuda=False,
            device=device,
        )
        print("==> Successfully loaded pre-trained network.")

        fe.net_parallel()
        print(path)
        # save to files
        '''
        save_file = save_output / "export.txt"
        with open(save_file, "a") as myfile:
            myfile.write("load model: " + path + "\n")
        '''
    except Exception:
        print(f"load model: {path} failed! ")
        raise

    ## loop through all images
    for split in splits:
        save_path = osp.join(proj_path, output_dir, "predictions", split)
        det_path = osp.join(save_path, "detection")
        if not osp.isdir(det_path):
            os.makedirs(det_path)

        if output_images:
            quality_res_path = osp.join(save_path, "quality_res")
            if not osp.isdir(quality_res_path):
                os.makedirs(quality_res_path)

        print(len(data_loaders[split]))
        for i, sample in tqdm(enumerate(data_loaders[split])):
            img, mask_2D = sample["image"], sample["valid_mask"]
            img = img.transpose(0, 1)
            img_2D = sample["image_2D"].numpy().squeeze()
            mask_2D = mask_2D.transpose(0, 1)

            inv_homographies, homographies = (
                sample["homographies"],
                sample["inv_homographies"],
            )
            img, mask_2D, homographies, inv_homographies = (
                img.to(device),
                mask_2D.to(device),
                homographies.to(device),
                inv_homographies.to(device),
            )
            # sample = test_set[i]
            name = sample["name"][0]
            fname_out = osp.join(det_path,
                                 "{}.npz".format(str(name).replace('/', '_')))
            if osp.isfile(fname_out):
                continue

            # pass through network
            heatmap = fe.run(img, onlyHeatmap=True, train=False)
            outputs = combine_heatmap(heatmap,
                                      inv_homographies,
                                      mask_2D,
                                      device=device)
            pts = fe.getPtsFromHeatmap(
                outputs.detach().cpu().squeeze())  # (x,y, prob)

            # subpixel prediction
            if config["model"]["subpixel"]["enable"]:
                fe.heatmap = outputs  # tensor [batch, 1, H, W]
                pts = fe.soft_argmax_points([pts])
                pts = pts[0]

            ## top K points
            pts = pts.transpose()
            if top_k:
                if pts.shape[0] > top_k:
                    pts = pts[:top_k, :]
                    print("topK filter: ", pts.shape)

            ## save keypoints
            pred = {}
            pred.update({"pts": pts})

            ## - make directories
            np.savez_compressed(fname_out, **pred)

            ## output images for visualization labels
            if output_images:
                img_pts = draw_keypoints(img_2D * 255, pts.transpose())
                fname_out_det = osp.join(quality_res_path,
                                         str(name).replace('/', '_') + ".png")
                saveImg(img_pts, fname_out_det)
            count += 1
            print(
                str(i + 1) + " out of " + str(len(data_loaders[split])) +
                " done.")

        print("output pseudo ground truth, ", split.capitalize(), ": ", count)

    print("Done")
コード例 #2
0
def export_detector_homoAdapt_gpu(config, output_dir, args):
    """
    input 1 images, output pseudo ground truth by homography adaptation.
    Save labels:
        pred:
            'prob' (keypoints): np (N1, 3)
    """
    from utils.utils import pltImshow
    from utils.utils import saveImg
    from utils.draw import draw_keypoints

    # basic setting
    task = config["data"]["dataset"]
    export_task = config["data"]["export_folder"]
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    logging.info("train on device: %s", device)
    with open(os.path.join(output_dir, "config.yml"), "w") as f:
        yaml.dump(config, f, default_flow_style=False)
    writer = SummaryWriter(
        getWriterPath(task=args.command, exper_name=args.exper_name, date=True)
    )

    ## parameters
    nms_dist = config["model"]["nms"]  # 4
    top_k = config["model"]["top_k"]
    homoAdapt_iter = config["data"]["homography_adaptation"]["num"]
    conf_thresh = config["model"]["detection_threshold"]
    nn_thresh = 0.7
    outputMatches = True
    count = 0
    max_length = 5
    output_images = args.outputImg
    check_exist = True

    ## save data
    save_path = Path(output_dir)
    save_output = save_path
    save_output = save_output / "predictions" / export_task
    save_path = save_path / "checkpoints"
    logging.info("=> will save everything to {}".format(save_path))
    os.makedirs(save_path, exist_ok=True)
    os.makedirs(save_output, exist_ok=True)

    # data loading
    from utils.loader import dataLoader_test as dataLoader

    data = dataLoader(config, dataset=task, export_task=export_task)
    print("Data is: ",data)
    test_set, test_loader = data["test_set"], data["test_loader"]
    print("Size test: ",len(test_set))
    print("Size loader: ",len(test_loader))
    # model loading
    ## load pretrained
    try:
        path = config["pretrained"]
        print("==> Loading pre-trained network.")
        print("path: ", path)
        # This class runs the SuperPoint network and processes its outputs.

        fe = SuperPointFrontend_torch(
            config=config,
            weights_path=path,
            nms_dist=nms_dist,
            conf_thresh=conf_thresh,
            nn_thresh=nn_thresh,
            cuda=False,
            device=device,
        )
        print("==> Successfully loaded pre-trained network.")

        fe.net_parallel()
        print(path)
        # save to files
        save_file = save_output / "export.txt"
        with open(save_file, "a") as myfile:
            myfile.write("load model: " + path + "\n")
    except Exception:
        print(f"load model: {path} failed! ")
        raise

    def load_as_float(path):
        return imread(path).astype(np.float32) / 255
    print("Tracker")
    tracker = PointTracker(max_length, nn_thresh=fe.nn_thresh)
    with open(save_file, "a") as myfile:
        myfile.write("homography adaptation: " + str(homoAdapt_iter) + "\n")
    print("Load save file")
    '''
    print(len(test_loader))
    for i,sample in enumerate(test_loader):
        print("Hello world")
        print("Img: ",sample["image"].size())
        print("Name: ",test_set[i]["name"])
        print("valid mask: ",test_set[i]["valid_mask"].size())
        print("valid img_2D: ",test_set[i]["image_2D"].size())
        print("valid mask: ",test_set[i]["valid_mask"].size())
        print("homograpgy: ",test_set[i]["homographies"].size())
        print("inverse: ",test_set[i]["inv_homographies"].size())
        print("scene name: ",test_set[i]["scene_name"])
        print()
    '''
    ## loop through all images
    for i, sample in tqdm(enumerate(test_loader)):
        img, mask_2D = sample["image"], sample["valid_mask"]
        img = img.transpose(0, 1)
        img_2D = sample["image_2D"].numpy().squeeze()
        mask_2D = mask_2D.transpose(0, 1)

        inv_homographies, homographies = (
            sample["homographies"],
            sample["inv_homographies"],
        )
        img, mask_2D, homographies, inv_homographies = (
            img.to(device),
            mask_2D.to(device),
            homographies.to(device),
            inv_homographies.to(device),
        )
        # sample = test_set[i]
        name = sample["name"][0]
        logging.info(f"name: {name}")
        if check_exist:
            p = Path(save_output, "{}.npz".format(name))
            if p.exists():
                logging.info("file %s exists. skip the sample.", name)
                continue
        print("Pass img to network")
        # pass through network
        heatmap = fe.run(img, onlyHeatmap=True, train=False)
        outputs = combine_heatmap(heatmap, inv_homographies, mask_2D, device=device)
        pts = fe.getPtsFromHeatmap(outputs.detach().cpu().squeeze())  # (x,y, prob)

        # subpixel prediction
        if config["model"]["subpixel"]["enable"]:
            fe.heatmap = outputs  # tensor [batch, 1, H, W]
            print("outputs: ", outputs.shape)
            print("pts: ", pts.shape)
            pts = fe.soft_argmax_points([pts])
            pts = pts[0]

        ## top K points
        pts = pts.transpose()
        print("total points: ", pts.shape)
        print("pts: ", pts[:5])
        if top_k:
            if pts.shape[0] > top_k:
                pts = pts[:top_k, :]
                print("topK filter: ", pts.shape)

        ## save keypoints
        pred = {}
        pred.update({"pts": pts})

        ## - make directories
        filename = str(name)
        if task == "Kitti" or "Kitti_inh":
            scene_name = sample["scene_name"][0]
            os.makedirs(Path(save_output, scene_name), exist_ok=True)

        path = Path(save_output, "{}.npz".format(filename))
        np.savez_compressed(path, **pred)

        ## output images for visualization labels
        if output_images:
            img_pts = draw_keypoints(img_2D * 255, pts.transpose())
            f = save_output / (str(count) + ".png")
            if task == "Coco" or "Kitti":
                f = save_output / (name + ".png")
            saveImg(img_pts, str(f))
        count += 1

    print("output pseudo ground truth: ", count)
    save_file = save_output / "export.txt"
    with open(save_file, "a") as myfile:
        myfile.write("Homography adaptation: " + str(homoAdapt_iter) + "\n")
        myfile.write("output pairs: " + str(count) + "\n")
    pass
コード例 #3
0
class KittiOdoLoader(object):
    def __init__(self,
                 dataset_dir,
                 img_height=375,
                 img_width=1242,
                 cam_ids=['02'],
                 get_X=False,
                 get_pose=False,
                 get_sift=False,
                 get_SP=False,
                 sift_num=2000,
                 if_BF_matcher=False,
                 save_npy=True):
                 # depth_size_ratio=1):
        dir_path = Path(__file__).realpath().dirname()

        self.dataset_dir = Path(dataset_dir)
        self.img_height = img_height
        self.img_width = img_width
        self.cam_ids = cam_ids
        # assert self.cam_ids == ['02'], 'Support left camera only!'
        self.cid_to_num = {'00': 0, '01': 1, '02': 2, '03': 3}
        self.train_seqs = [0, 1, 2, 3, 4, 5, 6, 7, 8]
        self.test_seqs = [9, 10]
        # self.train_seqs = [4]
        # self.test_seqs = []
        # self.train_seqs = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
        # self.test_seqs = []
        self.map_to_raw = {'00': '2011_10_03_drive_0027', '01': '2011_10_03_drive_0042', '02': '2011_10_03_drive_0034', '03': '2011_09_26_drive_0067', \
            '04': '2011_09_30_drive_0016', '05': '2011_09_30_drive_0018', '06': '2011_09_30_drive_0020', '07': '2011_09_30_drive_0027', \
            '08': '2011_09_30_drive_0028', '09': '2011_09_30_drive_0033', '10': '2011_09_30_drive_0034'}

        self.get_X = get_X
        self.get_pose = get_pose
        self.get_sift = get_sift
        self.get_SP = get_SP
        self.save_npy = save_npy
        if self.save_npy:
            logging.info('+++ Dumping as npy')
        else:
            logging.info('+++ Dumping as h5')
        if self.get_sift:
            self.sift_num = sift_num
            self.if_BF_matcher = if_BF_matcher
            self.sift = cv2.xfeatures2d.SIFT_create(nfeatures=self.sift_num, contrastThreshold=1e-5)
            # self.bf = cv2.BFMatcher(normType=cv2.NORM_L2)
            # FLANN_INDEX_KDTREE = 0
            # index_params = dict(algorithm = FLANN_INDEX_KDTREE, trees = 5)
            # search_params = dict(checks = 50)
            # self.flann = cv2.FlannBasedMatcher(index_params, search_params)
            # self.sift_matcher = self.bf if BF_matcher else self.flann

        self.scenes = {'train': [], 'test': []}
        if self.get_SP:
            self.prapare_SP()
        self.collect_train_folders()
        self.collect_test_folders()

    def prapare_SP(self):
        logging.info('Preparing SP inference.')
        with open(DEEPSFM_PATH + '/configs/superpoint_coco_train.yaml', 'r') as f:
            self.config_SP = yaml.load(f, Loader=yaml.FullLoader)
            nms_dist = self.config_SP['model']['nms']
            conf_thresh = self.config_SP['model']['detection_threshold']
            # nn_thresh = config_SP['model']['nn_thresh']
            nn_thresh = 1.0
            path = DEEPSFM_PATH + '/' + self.config_SP['pretrained']
            device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

            self.fe = SuperPointFrontend_torch(weights_path=path,
                                        nms_dist=nms_dist,
                                        conf_thresh=conf_thresh,
                                        nn_thresh=nn_thresh,
                                        cuda=False,
                                        device=device)


    def collect_train_folders(self):
        for seq in self.train_seqs:
            seq_dir = os.path.join(self.dataset_dir, 'sequences', '%.2d' % seq)
            self.scenes['train'].append(seq_dir)

    def collect_test_folders(self):
        for seq in self.test_seqs:
            seq_dir = os.path.join(self.dataset_dir, 'sequences', '%.2d' % seq)
            self.scenes['test'].append(seq_dir)

    def collect_scene_from_drive(self, drive_path):
        train_scenes = []
        logging.info('Gathering info for %s...'%drive_path)
        for c in self.cam_ids:
            scene_data = {'cid': c, 'cid_num': self.cid_to_num[c], 'dir': Path(drive_path), 'rel_path': Path(drive_path).name + '_' + c}
            img_dir = os.path.join(drive_path, 'image_%d'%scene_data['cid_num'])
            scene_data['img_files'] = sorted(glob(img_dir + '/*.png'))
            scene_data['N_frames'] = len(scene_data['img_files'])
            assert scene_data['N_frames'] != 0, 'No file found for %s!'%drive_path
            scene_data['frame_ids'] = ['{:06d}'.format(i) for i in range(scene_data['N_frames'])]

            img_shape = None
            zoom_xy = None
            show_zoom_info = True
            for idx in tqdm(range(scene_data['N_frames'])):
                img, zoom_xy, _ = self.load_image(scene_data, idx, show_zoom_info)
                show_zoom_info = False
                if img is None and idx==0:
                    logging.warning('0 images in %s. Skipped.'%drive_path)
                    return []
                else:
                    if img_shape is not None:
                        assert img_shape == img.shape, 'Inconsistent image shape in seq %s!'%drive_path
                    else:
                        img_shape = img.shape
            # print(img_shape)
            scene_data['calibs'] = {'im_shape': [img_shape[0], img_shape[1]], 'zoom_xy': zoom_xy, 'rescale': True if zoom_xy != (1., 1.) else False}
            # Get geo params from the RAW dataset calibs
            P_rect_ori_dict = self.get_P_rect(scene_data, scene_data['calibs'])
            intrinsics = P_rect_ori_dict[c][:,:3]
            calibs_rects = self.get_rect_cams(intrinsics, P_rect_ori_dict['02'])

            drive_in_raw = self.map_to_raw[drive_path[-2:]]
            date = drive_in_raw[:10]
            seq = drive_in_raw[-4:]
            calib_path_in_raw = Path(self.dataset_dir)/'raw'/date
            imu2velo_dict = read_calib_file(calib_path_in_raw/'calib_imu_to_velo.txt')
            velo2cam_dict = read_calib_file(calib_path_in_raw/'calib_velo_to_cam.txt')
            cam2cam_dict = read_calib_file(calib_path_in_raw/'calib_cam_to_cam.txt')
            velo2cam_mat = transform_from_rot_trans(velo2cam_dict['R'], velo2cam_dict['T'])
            imu2velo_mat = transform_from_rot_trans(imu2velo_dict['R'], imu2velo_dict['T'])
            cam_2rect_mat = transform_from_rot_trans(cam2cam_dict['R_rect_00'], np.zeros(3))
            scene_data['calibs'].update({'K': intrinsics, 'P_rect_ori_dict': P_rect_ori_dict, 'cam_2rect': cam_2rect_mat, 'velo2cam': velo2cam_mat})
            scene_data['calibs'].update(calibs_rects)

            # Get pose
            poses = np.genfromtxt(self.dataset_dir/'poses'/'{}.txt'.format(drive_path[-2:])).astype(np.float32).reshape(-1, 3, 4)
            assert scene_data['N_frames']==poses.shape[0], 'scene_data[N_frames]!=poses.shape[0], %d!=%d'%(scene_data['N_frames'], poses.shape[0])
            scene_data['poses'] = poses

            scene_data['Rt_cam2_gt'] = scene_data['calibs']['Rtl_gt']

            train_scenes.append(scene_data)
        return train_scenes

    def construct_sample(self, scene_data, idx, frame_id, show_zoom_info):
        img, zoom_xy, img_ori = self.load_image(scene_data, idx, show_zoom_info)
        # print(img.shape, img_ori.shape)
        sample = {"img":img, "id":frame_id}
        if self.get_X:
            velo = load_velo(scene_data, idx)
            if velo is None:
                logging.error('0 velo in %s. Skipped.'%scene_data['dir'])
            velo_homo = utils_misc.homo_np(velo)
            val_idxes, X_rect, X_cam0 = rectify(velo_homo, scene_data['calibs']) # list, [N, 3]
            sample['X_cam2_vis'] = X_rect[val_idxes].astype(np.float32)
            sample['X_cam0_vis'] = X_cam0[val_idxes].astype(np.float32)
        if self.get_pose:
            sample['pose'] = scene_data['poses'][idx].astype(np.float32)
        if self.get_sift:
            # logging.info('Getting sift for frame %d/%d.'%(idx, scene_data['N_frames']))
            kp, des = self.sift.detectAndCompute(img_ori, None) ## IMPORTANT: normalize these points
            x_all = np.array([p.pt for p in kp])
            # print(zoom_xy)
            x_all = (x_all * np.array([[zoom_xy[0], zoom_xy[1]]])).astype(np.float32)
            # print(x_all.shape, np.amax(x_all, axis=0), np.amin(x_all, axis=0))
            if x_all.shape[0] != self.sift_num:
                choice = crop_or_pad_choice(x_all.shape[0], self.sift_num, shuffle=True)
                x_all = x_all[choice]
                des = des[choice]
            sample['sift_kp'] = x_all
            sample['sift_des'] = des
        if self.get_SP:
            img_ori_gray = cv2.cvtColor(img_ori, cv2.COLOR_RGB2GRAY)
            img = torch.from_numpy(img_ori_gray).float().unsqueeze(0).unsqueeze(0).float() / 255.
            pts, desc, _, heatmap = self.fe.run(img)
            pts = pts[0].T # [N, 3]
            pts[:, :2] = (pts[:, :2] * np.array([[zoom_xy[0], zoom_xy[1]]])).astype(np.float32)
            desc = desc[0].T # [N, 256]
            sample['SP_kp'] = pts
            sample['SP_des'] = desc
        return sample

    def dump_drive(self, args, drive_path, split, scene_data=None):
        assert split in ['train', 'test']
        if scene_data is None:
            train_scenes = self.collect_scene_from_drive(drive_path)
            if not train_scenes:
                logging.warning('Empty scene data for %s. Skipped.'%drive_path)
                return
            assert len(train_scenes)==1, 'More than one camera not supported! %d'%len(train_scenes)
            scene_data = train_scenes[0]

        dump_dir = Path(args.dump_root)/scene_data['rel_path']
        dump_dir.mkdir_p()
        intrinsics = scene_data['calibs']['K']
        dump_cam_file = dump_dir/'cam'
        np.save(dump_cam_file+'.npy', intrinsics.astype(np.float32))
        dump_Rt_cam2_gt_file = dump_dir/'Rt_cam2_gt'
        np.save(dump_Rt_cam2_gt_file, scene_data['Rt_cam2_gt'].astype(np.float32))
        poses_file = dump_dir/'poses'
        poses = []

        logging.info('Dumping %d samples to %s...'%(scene_data['N_frames'], dump_dir))
        sample_name_list = []
        # sift_des_list = []
        for idx in tqdm(range(scene_data['N_frames'])):
            frame_id = scene_data['frame_ids'][idx]
            assert int(frame_id)==idx
            sample = self.construct_sample(scene_data, idx, frame_id, show_zoom_info=False)

            img, frame_nb = sample["img"], sample["id"]
            dump_img_file = dump_dir/'{}.jpg'.format(frame_nb)
            scipy.misc.imsave(dump_img_file, img)
            if "pose" in sample.keys():
                poses.append(sample["pose"].astype(np.float32))
            if "X_cam0_vis" in sample.keys():
                dump_X_cam0_file = dump_dir/'X_cam0_{}'.format(frame_nb)
                dump_X_cam2_file = dump_dir/'X_cam2_{}'.format(frame_nb)
                if self.save_npy:
                    np.save(dump_X_cam0_file+'.npy', sample["X_cam0_vis"])
                    np.save(dump_X_cam2_file+'.npy', sample["X_cam2_vis"])
                else:
                    saveh5({"X_cam0_vis": sample["X_cam0_vis"], "X_cam2_vis": sample["X_cam2_vis"]}, dump_X_file+'.h5')
            if "sift_kp" in sample.keys():
                dump_sift_file = dump_dir/'sift_{}'.format(frame_nb)
                if self.save_npy:
                    np.save(dump_sift_file+'.npy', np.hstack((sample['sift_kp'], sample['sift_des'])))
                else:
                    saveh5({'sift_kp': sample['sift_kp'], 'sift_des': sample['sift_des']}, dump_sift_file+'.h5')
                # sift_des_list.append(sample['sift_des'])
            if "SP_kp" in sample.keys():
                dump_sift_file = dump_dir/'SP_{}'.format(frame_nb)
                if self.save_npy:
                    np.save(dump_sift_file+'.npy', np.hstack((sample['SP_kp'], sample['SP_des'])))
                    # print(sample['SP_kp'].shape, sample['SP_des'].shape)
                else:
                    pass

            sample_name_list.append('%s %s'%(dump_dir[-5:], frame_nb))

        # Get all poses    
        if "pose" in sample.keys():      
            if len(poses) != 0:
                # np.savetxt(poses_file, np.array(poses).reshape(-1, 16), fmt='%.20e')a
                if self.save_npy:
                    np.save(poses_file+'.npy', np.stack(poses).reshape(-1, 3, 4))
                else:
                    saveh5({"poses": np.array(poses).reshape(-1, 3, 4)}, poses_file+'.h5')

        # Get SIFT matches
        if self.get_sift:
            delta_ijs = [1, 2, 3, 5, 8, 10]
            # delta_ijs = [1]
            num_tasks = len(delta_ijs)
            num_workers = min(len(delta_ijs), default_number_of_process)
            # num_workers = 1
            logging.info('Getting SIFT matches on %d workers for delta_ijs = %s'%(num_workers, ' '.join(str(e) for e in delta_ijs)))

            with ProcessPool(max_workers=num_workers) as pool:
                tasks = pool.map(dump_sift_match_idx, delta_ijs, [scene_data['N_frames']]*num_tasks, \
                    [dump_dir]*num_tasks, [self.save_npy]*num_tasks, [self.if_BF_matcher]*num_tasks)
                try:
                    for _ in tqdm(tasks.result(), total=num_tasks):
                        pass
                except KeyboardInterrupt as e:
                    tasks.cancel()
                    raise e

        # Get SP matches
        if self.get_SP:
            delta_ijs = [1, 2, 3, 5, 8, 10]
            nn_threshes = [0.7, 1.0]
            # delta_ijs = [1]
            num_tasks = len(delta_ijs)
            num_workers = min(len(delta_ijs), default_number_of_process)
            # num_workers = 1
            logging.info('Getting SP matches on %d workers for delta_ijs = %s'%(num_workers, ' '.join(str(e) for e in delta_ijs)))

            with ProcessPool(max_workers=num_workers) as pool:
                tasks = pool.map(dump_SP_match_idx, delta_ijs, [scene_data['N_frames']]*num_tasks, \
                    [dump_dir]*num_tasks, [self.save_npy]*num_tasks, [nn_threshes]*num_tasks)
                try:
                    for _ in tqdm(tasks.result(), total=num_tasks):
                        pass
                except KeyboardInterrupt as e:
                    tasks.cancel()
                    raise e

            # for delta_ij in delta_ijs:
            #     dump_match_idx(delta_ij, scene_data['N_frames'], sift_des_list, dump_dir, self.save_npy, self.if_BF_matcher)

        if len(dump_dir.files('*.jpg')) < 2:
            dump_dir.rmtree()

        return sample_name_list

    def load_image(self, scene_data, tgt_idx, show_zoom_info=True):
        img_file = scene_data['dir']/'image_{}'.format(scene_data['cid_num'])/scene_data['frame_ids'][tgt_idx]+'.png'
        if not img_file.isfile():
            logging.warning('Image %s not found!'%img_file)
            return None, None, None
        img_ori = scipy.misc.imread(img_file)
        if [self.img_height, self.img_width] == [img_ori.shape[0], img_ori.shape[1]]:
            return img_ori, (1., 1.), img_ori
        else:
            zoom_y = self.img_height/img_ori.shape[0]
            zoom_x = self.img_width/img_ori.shape[1]
            if show_zoom_info:
                logging.warning('[%s] Zooming the image (H%d, W%d) with zoom_yH=%f, zoom_xW=%f to (H%d, W%d).'%\
                    (img_file, img_ori.shape[0], img_ori.shape[1], zoom_y, zoom_x, self.img_height, self.img_width))
            img = scipy.misc.imresize(img_ori, (self.img_height, self.img_width))
            return img, (zoom_x, zoom_y), img_ori

    def get_P_rect(self, scene_data, calibs, get_2cam_dict=True):
        # calib_file = scene_data['dir'].parent/'calib_cam_to_cam.txt'
        calib_file = scene_data['dir']/'calib.txt'
        if get_2cam_dict:
            P_rect = {}
            for cid in ['00', '01', '02', '03']:
                P_rect[cid], _ = read_odo_calib_file(calib_file, cid=self.cid_to_num[cid])
                if calibs['rescale']:
                    P_rect[cid] = scale_P(P_rect[cid], calibs['zoom_xy'][0], calibs['zoom_xy'][1])
            return P_rect
        else:
            P_rect, _ = read_odo_calib_file(calib_file, cid=self.cid_to_num[cid])
            if calibs['rescale']:
                P_rect = scale_P(P_rect, calibs['zoom_xy'][0], calibs['zoom_xy'][1])
        return P_rect

    def get_rect_cams(self, K, P_rect_20):
        Ml_gt = np.matmul(np.linalg.inv(K), P_rect_20)
        tl_gt = Ml_gt[:, 3:4]
        Rl_gt = Ml_gt[:, :3]
        Rtl_gt = np.vstack((np.hstack((Rl_gt, tl_gt)), np.array([0., 0., 0., 1.], dtype=np.float32)))
        calibs_rects = {'Rtl_gt': Rtl_gt}
        return calibs_rects