def _valid_epoch(self):
    # Change the network to evaluation mode
    self.model.eval()
    self.val_data_loader.dataset.reset_seed(0)
    num_data = 0
    hit_ratio_meter, feat_match_ratio, loss_meter, rte_meter, rre_meter = AverageMeter(
    ), AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter()
    data_timer, feat_timer, matching_timer = Timer(), Timer(), Timer()
    tot_num_data = len(self.val_data_loader.dataset)
    if self.val_max_iter > 0:
      tot_num_data = min(self.val_max_iter, tot_num_data)
    data_loader_iter = self.val_data_loader.__iter__()

    for batch_idx in range(tot_num_data):
      data_timer.tic()
      input_dict = data_loader_iter.next()
      data_timer.toc()

      # pairs consist of (xyz1 index, xyz0 index)
      feat_timer.tic()
      sinput0 = ME.SparseTensor(
          input_dict['sinput0_F'], coords=input_dict['sinput0_C']).to(self.device)
      F0 = self.model(sinput0).F

      sinput1 = ME.SparseTensor(
          input_dict['sinput1_F'], coords=input_dict['sinput1_C']).to(self.device)
      F1 = self.model(sinput1).F
      feat_timer.toc()

      matching_timer.tic()
      xyz0, xyz1, T_gt = input_dict['pcd0'], input_dict['pcd1'], input_dict['T_gt']
      xyz0_corr, xyz1_corr = self.find_corr(xyz0, xyz1, F0, F1, subsample_size=5000)
      T_est = te.est_quad_linear_robust(xyz0_corr, xyz1_corr)

      loss = corr_dist(T_est, T_gt, xyz0, xyz1, weight=None)
      loss_meter.update(loss)

      rte = np.linalg.norm(T_est[:3, 3] - T_gt[:3, 3])
      rte_meter.update(rte)
      rre = np.arccos((np.trace(T_est[:3, :3].t() @ T_gt[:3, :3]) - 1) / 2)
      if not np.isnan(rre):
        rre_meter.update(rre)

      hit_ratio = self.evaluate_hit_ratio(
          xyz0_corr, xyz1_corr, T_gt, thresh=self.config.hit_ratio_thresh)
      hit_ratio_meter.update(hit_ratio)
      feat_match_ratio.update(hit_ratio > 0.05)
      matching_timer.toc()

      num_data += 1
      torch.cuda.empty_cache()

      if batch_idx % 100 == 0 and batch_idx > 0:
        logging.info(' '.join([
            f"Validation iter {num_data} / {tot_num_data} : Data Loading Time: {data_timer.avg:.3f},",
            f"Feature Extraction Time: {feat_timer.avg:.3f}, Matching Time: {matching_timer.avg:.3f},",
            f"Loss: {loss_meter.avg:.3f}, RTE: {rte_meter.avg:.3f}, RRE: {rre_meter.avg:.3f},",
            f"Hit Ratio: {hit_ratio_meter.avg:.3f}, Feat Match Ratio: {feat_match_ratio.avg:.3f}"
        ]))
        data_timer.reset()

    logging.info(' '.join([
        f"Final Loss: {loss_meter.avg:.3f}, RTE: {rte_meter.avg:.3f}, RRE: {rre_meter.avg:.3f},",
        f"Hit Ratio: {hit_ratio_meter.avg:.3f}, Feat Match Ratio: {feat_match_ratio.avg:.3f}"
    ]))
    return {
        "loss": loss_meter.avg,
        "rre": rre_meter.avg,
        "rte": rte_meter.avg,
        'feat_match_ratio': feat_match_ratio.avg,
        'hit_ratio': hit_ratio_meter.avg
    }
Exemplo n.º 2
0
    def _valid_epoch(self):
        # Change the network to evaluation mode
        self.model.eval()
        self.val_data_loader.dataset.reset_seed(0)
        num_data = 0
        hit_ratio_meter, feat_match_ratio, loss_meter, rte_meter, rre_meter = AverageMeter(
        ), AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter()
        data_timer, feat_timer, matching_timer = Timer(), Timer(), Timer()
        tot_num_data = len(self.val_data_loader.dataset)
        if self.val_max_iter > 0:
            tot_num_data = min(self.val_max_iter, tot_num_data)
        data_loader_iter = self.val_data_loader.__iter__()

        for batch_idx in range(tot_num_data):
            data_timer.tic()
            input_dict = data_loader_iter.next()
            data_timer.toc()

            # pairs consist of (xyz1 index, xyz0 index)
            feat_timer.tic()

            coords = input_dict['sinput0_C'].to(self.device)
            sinput0 = ME.SparseTensor(
                input_dict['sinput0_F'].to(self.device),
                coordinates=input_dict['sinput0_C'].to(self.device).type(torch.float))

            F0 = self.model(sinput0).F

            sinput1 = ME.SparseTensor(
                input_dict['sinput1_F'].to(self.device),
                coordinates=input_dict['sinput1_C'].to(self.device).type(torch.float))
            F1 = self.model(sinput1).F
            feat_timer.toc()

            matching_timer.tic()
            xyz0, xyz1, T_gt = input_dict['pcd0'], input_dict['pcd1'], input_dict['T_gt']
            xyz0_corr, xyz1_corr = self.find_corr(xyz0, xyz1, F0, F1, subsample_size=5000)

            if False:

                from sklearn.decomposition import PCA
                import open3d as o3d

                pc0 = o3d.geometry.PointCloud()
                pc0.points = o3d.utility.Vector3dVector(xyz0.numpy())
                pca = PCA(n_components=3)

                colors = pca.fit_transform(torch.cat((F0, F1), axis=0).cpu().numpy())
                colors -= colors.min()
                colors /= colors.max()
                pc0.colors = o3d.utility.Vector3dVector(colors[0:F0.shape[0]])

                o3d.io.write_point_cloud("pc0.ply", pc0)
                pc0.transform(T_gt.numpy())
                o3d.io.write_point_cloud("pc0_trans.ply", pc0)

                pc1 = o3d.geometry.PointCloud()
                pc1.points = o3d.utility.Vector3dVector(xyz1.numpy())
                pc1.colors = o3d.utility.Vector3dVector(colors[F0.shape[0]:])
                o3d.io.write_point_cloud("pc1.ply", pc1)

                ind_0 = input_dict['correspondences'][:, 0].type(torch.long)
                ind_1 = input_dict['correspondences'][:, 1].type(torch.long)

                pc1.points = o3d.utility.Vector3dVector(xyz1[ind_1].numpy())
                pc1.colors = o3d.utility.Vector3dVector(
                    colors[F0.shape[0]:][ind_1])
                o3d.io.write_point_cloud("pc1_corr.ply", pc1)

                pc0.points = o3d.utility.Vector3dVector(xyz0[ind_0].numpy())
                pc0.colors = o3d.utility.Vector3dVector(colors[:F0.shape[0]][ind_0])
                pc0.transform(T_gt.numpy())
                o3d.io.write_point_cloud("pc0_trans_corr.ply", pc0)
                import pdb
                pdb.set_trace()

            #pc0.points = o3d.utility.Vector3dVector(xyz0_corr.numpy())
            # pc0.transform(T_gt.numpy())
            #o3d.io.write_point_cloud("xyz0_corr_trans.ply" , pc0)
#
            #pc0.points = o3d.utility.Vector3dVector(xyz1_corr.numpy())
            #o3d.io.write_point_cloud("xyz1_corr_trans.ply" , pc0)

            T_est = te.est_quad_linear_robust(xyz0_corr, xyz1_corr)

            loss = corr_dist(T_est, T_gt, xyz0, xyz1, weight=None)
            loss_meter.update(loss)

            rte = np.linalg.norm(T_est[:3, 3] - T_gt[:3, 3])
            rte_meter.update(rte)
            rre = np.arccos((np.trace(T_est[:3, :3].t() @ T_gt[:3, :3]) - 1) / 2)
            if not np.isnan(rre):
                rre_meter.update(rre)

            hit_ratio = self.evaluate_hit_ratio(xyz0_corr, xyz1_corr, T_gt, thresh=self.config.hit_ratio_thresh)
            hit_ratio_meter.update(hit_ratio)
            feat_match_ratio.update(hit_ratio > 0.05)
            matching_timer.toc()

            num_data += 1
            torch.cuda.empty_cache()

            if batch_idx % 100 == 0 and batch_idx > 0:
                logging.info(' '.join([
                    f"Validation iter {num_data} / {tot_num_data} : Data Loading Time: {data_timer.avg:.3f},",
                    f"Feature Extraction Time: {feat_timer.avg:.3f}, Matching Time: {matching_timer.avg:.3f},",
                    f"Loss: {loss_meter.avg:.3f}, RTE: {rte_meter.avg:.3f}, RRE: {rre_meter.avg:.3f},",
                    f"Hit Ratio: {hit_ratio_meter.avg:.3f}, Feat Match Ratio: {feat_match_ratio.avg:.3f}"
                ]))
                data_timer.reset()

        logging.info(' '.join([
            f"Final Loss: {loss_meter.avg:.3f}, RTE: {rte_meter.avg:.3f}, RRE: {rre_meter.avg:.3f},",
            f"Hit Ratio: {hit_ratio_meter.avg:.3f}, Feat Match Ratio: {feat_match_ratio.avg:.3f}"
        ]))
        return {
            "loss": loss_meter.avg,
            "rre": rre_meter.avg,
            "rte": rte_meter.avg,
            'feat_match_ratio': feat_match_ratio.avg,
            'hit_ratio': hit_ratio_meter.avg
        }