예제 #1
0
 def _initialize_model(self, model: BaseModel, weight_name):
     if not self._checkpoint.is_empty:
         state_dict = self._checkpoint.get_state_dict(weight_name)
         model.load_state_dict(state_dict)
         if self._resume:
             model.optimizer = self._checkpoint.get_optimizer(model)
             model.schedulers = self._checkpoint.get_schedulers(model)
예제 #2
0
def test_epoch(
    model: BaseModel,
    dataset,
    device,
    tracker: BaseTracker,
    checkpoint: ModelCheckpoint,
    voting_runs=1,
    tracker_options={},
):

    loaders = dataset.test_dataloaders

    for loader in loaders:
        stage_name = loader.dataset.name
        tracker.reset(stage_name)
        for i in range(voting_runs):
            with Ctq(loader) as tq_test_loader:
                for data in tq_test_loader:
                    with torch.no_grad():
                        model.set_input(data, device)
                        model.forward()

                    tracker.track(model, **tracker_options)
                    tq_test_loader.set_postfix(**tracker.get_metrics(),
                                               color=COLORS.TEST_COLOR)

        tracker.finalise(**tracker_options)
        tracker.print_summary()
    def track(self, model: BaseModel, **kwargs):
        """ Add current model predictions (usually the result of a batch) to the tracking
        """
        super().track(model)

        outputs = model.get_output()
        targets = model.get_labels()

        # Mask ignored label
        mask = targets != self._ignore_label
        outputs = outputs[mask]
        targets = targets[mask]

        outputs = SegmentationTracker.detach_tensor(outputs)
        targets = SegmentationTracker.detach_tensor(targets)
        if not torch.is_tensor(targets):
            targets = torch.from_numpy(targets)
        self._ap_meter.add(outputs,
                           F.one_hot(targets, self._num_classes).bool())

        outputs = self._convert(outputs)
        targets = self._convert(targets)

        if len(targets) == 0:
            return

        assert outputs.shape[0] == len(targets)
        self._confusion_matrix.count_predicted_batch(targets,
                                                     np.argmax(outputs, 1))

        self._acc = 100 * self._confusion_matrix.get_overall_accuracy()
        self._macc = 100 * self._confusion_matrix.get_mean_class_accuracy()
        self._miou = 100 * self._confusion_matrix.get_average_intersection_union(
        )
        self._map = 100 * self._ap_meter.value().mean().item()
예제 #4
0
    def track(self, model: BaseModel, **kwargs):
        """ Add current model predictions (usually the result of a batch) to the tracking
        """
        super().track(model)

        outputs = model.get_output()
        targets = model.get_labels()

        # Mask ignored label
        mask = targets != self._ignore_label
        outputs = outputs[mask]
        targets = targets[mask]

        outputs = self._convert(outputs)
        targets = self._convert(targets)

        if len(targets) == 0:
            return

        assert outputs.shape[0] == len(targets)
        self._confusion_matrix.count_predicted_batch(targets, np.argmax(outputs, 1))

        self._acc = 100 * self._confusion_matrix.get_overall_accuracy()
        self._macc = 100 * self._confusion_matrix.get_mean_class_accuracy()
        self._miou = 100 * self._confusion_matrix.get_average_intersection_union()
예제 #5
0
def run(cfg, model: BaseModel, dataset: BaseDataset, device,
        measurement_name: str):
    measurements = {}

    num_batches = getattr(cfg.debugging, "num_batches", np.inf)

    run_epoch(model, dataset.train_dataloader, device, num_batches)
    measurements["train"] = extract_histogram(model.get_spatial_ops(),
                                              normalize=False)

    if dataset.has_val_loader:
        run_epoch(model, dataset.val_dataloader, device, num_batches)
        measurements["val"] = extract_histogram(model.get_spatial_ops(),
                                                normalize=False)

    for loader in dataset.test_dataloaders:
        run_epoch(model, dataset.test_dataloaders, device, num_batches)
        measurements[loader.dataset.name] = extract_histogram(
            model.get_spatial_ops(), normalize=False)

    with open(
            os.path.join(DIR,
                         "measurements/{}.pickle".format(measurement_name)),
            "wb") as f:
        pickle.dump(measurements, f)
예제 #6
0
    def track(self, model: BaseModel, full_res=False, **kwargs):
        """ Add current model predictions (usually the result of a batch) to the tracking
        """
        super().track(model)

        # Train mode or low res, nothing special to do
        if self._stage == "train" or not full_res:
            return

        # Test mode, compute votes in order to get full res predictions
        if self._test_area is None:
            self._test_area = self._dataset.test_data.clone()
            if self._test_area.y is None:
                raise ValueError("It seems that the test area data does not have labels (attribute y).")
            self._test_area.prediction_count = torch.zeros(self._test_area.y.shape[0], dtype=torch.int)
            self._test_area.votes = torch.zeros((self._test_area.y.shape[0], self._num_classes), dtype=torch.float)
            self._test_area.to(model.device)

        # Gather input to the model and check that it fits with the test set
        inputs = model.get_input()
        if inputs[SaveOriginalPosId.KEY] is None or inputs[SaveOriginalPosId.KEY].max() >= self._test_area.pos.shape[0]:
            raise ValueError(
                "The inputs given to the model do not have a %s attribute or this attribute does\
                     not correspond to the number of points in the test area point cloud."
                % SaveOriginalPosId.KEY
            )

        # Set predictions
        outputs = model.get_output()
        self._test_area.votes[inputs[SaveOriginalPosId.KEY]] += outputs
        self._test_area.prediction_count[inputs[SaveOriginalPosId.KEY]] += 1
예제 #7
0
    def __init__(self, option, model_type, dataset, modules):
        BaseModel.__init__(self, option)
        self.mode = option.loss_mode
        self.normalize_feature = option.normalize_feature
        self.loss_names = ["loss_reg", "loss"]
        self.metric_loss_module, self.miner_module = BaseModel.get_metric_loss_and_miner(
            getattr(option, "metric_loss", None), getattr(option, "miner", None)
        )
        # Last Layer

        if option.mlp_cls is not None:
            last_mlp_opt = option.mlp_cls
            in_feat = last_mlp_opt.nn[0]
            self.FC_layer = Seq()
            for i in range(1, len(last_mlp_opt.nn)):
                self.FC_layer.append(
                    str(i),
                    Sequential(
                        *[
                            Linear(in_feat, last_mlp_opt.nn[i], bias=False),
                            FastBatchNorm1d(last_mlp_opt.nn[i], momentum=last_mlp_opt.bn_momentum),
                            LeakyReLU(0.2),
                        ]
                    ),
                )
                in_feat = last_mlp_opt.nn[i]

            if last_mlp_opt.dropout:
                self.FC_layer.append(Dropout(p=last_mlp_opt.dropout))

            self.FC_layer.append(Linear(in_feat, in_feat, bias=False))
        else:
            self.FC_layer = torch.nn.Identity()
예제 #8
0
def run_epoch(model: BaseModel, loader, device: str, num_batches: int):
    model.eval()
    with Ctq(loader) as tq_loader:
        for batch_idx, data in enumerate(tq_loader):
            if batch_idx < num_batches:
                process(model, data, device)
            else:
                break
    def track(self, model: BaseModel):
        super().track(model)
        if self._stage != "train":
            batch_idx, batch_idx_target = model.get_batch_idx()
            batch_xyz, batch_xyz_target = model.get_xyz()
            batch_ind, batch_ind_target = model.get_ind()
            batch_feat, batch_feat_target = model.get_outputs()

            nb_batches = batch_idx.max() + 1
            cum_sum = 0
            cum_sum_target = 0
            for b in range(nb_batches):
                xyz = batch_xyz[batch_idx == b]
                xyz_target = batch_xyz_target[batch_idx_target == b]
                feat = batch_feat[batch_idx == b]
                feat_target = batch_feat_target[batch_idx_target == b]
                # as we have concatenated ind,
                # we need to substract the cum_sum because we deal
                # with each batch independently
                ind = batch_ind[batch_idx == b] - cum_sum
                ind_target = batch_ind_target[batch_idx_target ==
                                              b] - cum_sum_target
                cum_sum += len(xyz)
                cum_sum_target += len(xyz_target)
                rand = torch.randperm(len(feat))[:self.num_points]
                rand_target = torch.randperm(
                    len(feat_target))[:self.num_points]

                matches_gt = torch.stack([ind, ind_target]).T
                T_gt = estimate_transfo(xyz[matches_gt[:, 0]],
                                        xyz_target[matches_gt[:, 1]])
                matches_pred = get_matches(feat[rand],
                                           feat_target[rand_target])
                T_pred = fast_global_registration(
                    xyz[rand][matches_pred[:, 0]],
                    xyz_target[rand_target][matches_pred[:, 1]])

                hit_ratio = compute_hit_ratio(
                    xyz[rand][matches_pred[:, 0]],
                    xyz_target[rand_target][matches_pred[:,
                                                         1]], T_gt, self.tau_1)
                trans_error, rot_error = compute_transfo_error(T_pred, T_gt)

                self._hit_ratio.add(hit_ratio.item())
                self._feat_match_ratio.add(
                    float(hit_ratio.item() > self.tau_2))
                self._trans_error.add(trans_error.item())
                self._rot_error.add(rot_error.item())
예제 #10
0
def run(model: BaseModel, dataset: BaseDataset, device, output_path, cfg):
    # Set dataloaders
    num_fragment = dataset.num_fragment
    if cfg.data.is_patch:
        for i in range(num_fragment):
            dataset.set_patches(i)
            dataset.create_dataloaders(
                model,
                cfg.batch_size,
                False,
                cfg.num_workers,
                False,
            )
            loader = dataset.test_dataloaders()[0]
            features = []
            scene_name, pc_name = dataset.get_name(i)

            with Ctq(loader) as tq_test_loader:
                for data in tq_test_loader:
                    # pcd = open3d.geometry.PointCloud()
                    # pcd.points = open3d.utility.Vector3dVector(data.pos[0].numpy())
                    # open3d.visualization.draw_geometries([pcd])
                    with torch.no_grad():
                        model.set_input(data, device)
                        model.forward()
                        features.append(model.get_output().cpu())
            features = torch.cat(features, 0).numpy()
            log.info("save {} from {} in  {}".format(pc_name, scene_name,
                                                     output_path))
            save(output_path, scene_name, pc_name,
                 dataset.base_dataset[i].to("cpu"), features)
    else:
        dataset.create_dataloaders(
            model,
            1,
            False,
            cfg.num_workers,
            False,
        )
        loader = dataset.test_dataloaders()[0]
        with Ctq(loader) as tq_test_loader:
            for i, data in enumerate(tq_test_loader):
                with torch.no_grad():
                    model.set_input(data, device)
                    model.forward()
                    features = model.get_output()[0]  # batch of 1
                    save(output_path, scene_name, pc_name, data.to("cpu"),
                         features)
예제 #11
0
    def save_best_models_under_current_metrics(
        self, model: BaseModel, metrics_holder: dict, metric_func: dict, **kwargs
    ):
        """[This function is responsible to save checkpoint under the current metrics and their associated DEFAULT_METRICS_FUNC]
        Arguments:
            model {[BaseModel]} -- [Model]
            metrics_holder {[Dict]} -- [Need to contain stage, epoch, current_metrics]
        """
        metrics = metrics_holder["current_metrics"]
        stage = metrics_holder["stage"]
        epoch = metrics_holder["epoch"]

        stats = self._checkpoint.stats
        state_dict = copy.deepcopy(model.state_dict())

        current_stat = {}
        current_stat["epoch"] = epoch

        models_to_save = self._checkpoint.models
        if stage not in stats:
            stats[stage] = []

        if stage == "train":
            models_to_save[Checkpoint._LATEST] = state_dict
        else:
            if len(stats[stage]) > 0:
                latest_stats = stats[stage][-1]

                msg = ""
                improved_metric = 0

                for metric_name, current_metric_value in metrics.items():
                    current_stat[metric_name] = current_metric_value

                    metric_func = self.find_func_from_metric_name(metric_name, metric_func)
                    best_metric_from_stats = latest_stats.get("best_{}".format(metric_name), current_metric_value)
                    best_value = metric_func(best_metric_from_stats, current_metric_value)
                    current_stat["best_{}".format(metric_name)] = best_value

                    # This new value seems to be better under metric_func
                    if (self._selection_stage == stage) and (
                        current_metric_value == best_value
                    ):  # Update the model weights
                        models_to_save["best_{}".format(metric_name)] = state_dict

                        msg += "{}: {} -> {}, ".format(metric_name, best_metric_from_stats, best_value)
                        improved_metric += 1

                if improved_metric > 0:
                    colored_print(COLORS.VAL_COLOR, msg[:-2])
            else:
                # stats[stage] is empty.
                for metric_name, metric_value in metrics.items():
                    current_stat[metric_name] = metric_value
                    current_stat["best_{}".format(metric_name)] = metric_value
                    models_to_save["best_{}".format(metric_name)] = state_dict

        self._checkpoint.stats[stage].append(current_stat)
        self._checkpoint.save_objects(models_to_save, stage, current_stat, model.optimizer, model.schedulers, **kwargs)
예제 #12
0
def run(model: BaseModel, dataset: BaseDataset, device, output_path):
    loaders = dataset.test_dataloaders
    predicted = {}
    for loader in loaders:
        loader.dataset.name
        with Ctq(loader) as tq_test_loader:
            for data in tq_test_loader:
                with torch.no_grad():
                    model.set_input(data, device)
                    model.forward()
                predicted = {
                    **predicted,
                    **dataset.predict_original_samples(data, model.conv_type,
                                                       model.get_output())
                }

    save(output_path, predicted)
    def track(self, model: BaseModel):
        """ Add model predictions (accuracy)
        """
        super().track(model)

        outputs = self._convert(model.get_output())
        N = len(outputs) // 2

        self._acc = compute_accuracy(outputs[:N], outputs[N:])
    def _init_from_compact_format(self, opt, model_type, dataset, modules_lib):
        """Create a unetbasedmodel from the compact options format - where the
        same convolution is given for each layer, and arguments are given
        in lists
        """

        self.down_modules = nn.ModuleList()
        self.inner_modules = nn.ModuleList()
        self.up_modules = nn.ModuleList()

        self.save_sampling_id = opt.down_conv.get('save_sampling_id')

        # Factory for creating up and down modules
        factory_module_cls = self._get_factory(model_type, modules_lib)
        down_conv_cls_name = opt.down_conv.module_name
        up_conv_cls_name = opt.up_conv.module_name if opt.get(
            'up_conv') is not None else None
        self._factory_module = factory_module_cls(
            down_conv_cls_name, up_conv_cls_name,
            modules_lib)  # Create the factory object

        # Loal module
        contains_global = hasattr(opt,
                                  "innermost") and opt.innermost is not None
        if contains_global:
            inners = self._create_inner_modules(opt.innermost, modules_lib)
            for inner in inners:
                self.inner_modules.append(inner)
        else:
            self.inner_modules.append(Identity())

        # Down modules
        for i in range(len(opt.down_conv.down_conv_nn)):
            args = self._fetch_arguments(opt.down_conv, i, "DOWN")
            conv_cls = self._get_from_kwargs(args, "conv_cls")
            down_module = conv_cls(**args)
            self._save_sampling_and_search(down_module)
            self.down_modules.append(down_module)

        # Up modules
        if up_conv_cls_name:
            for i in range(len(opt.up_conv.up_conv_nn)):
                args = self._fetch_arguments(opt.up_conv, i, "UP")
                conv_cls = self._get_from_kwargs(args, "conv_cls")
                up_module = conv_cls(**args)
                self._save_upsample(up_module)
                self.up_modules.append(up_module)

        self.metric_loss_module, self.miner_module = BaseModel.get_metric_loss_and_miner(
            getattr(opt, "metric_loss", None), getattr(opt, "miner", None))
예제 #15
0
def run(model: BaseModel, dataset: BaseDataset, device, cfg):
    dataset.create_dataloaders(
        model,
        1,
        False,
        cfg.training.num_workers,
        False,
    )
    loader = dataset.test_dataloaders[0]
    list_res = []
    with Ctq(loader) as tq_test_loader:
        for i, data in enumerate(tq_test_loader):
            with torch.no_grad():
                model.set_input(data, device)
                model.forward()

                name_scene, name_pair_source, name_pair_target = dataset.test_dataset[
                    0].get_name(i)
                input, input_target = model.get_input()
                xyz, xyz_target = input.pos, input_target.pos
                ind, ind_target = input.ind, input_target.ind
                matches_gt = torch.stack([ind, ind_target]).transpose(0, 1)
                feat, feat_target = model.get_output()
                rand = torch.randperm(len(feat))[:cfg.data.num_points]
                rand_target = torch.randperm(
                    len(feat_target))[:cfg.data.num_points]
                res = dict(name_scene=name_scene,
                           name_pair_source=name_pair_source,
                           name_pair_target=name_pair_target)
                T_gt = estimate_transfo(xyz[matches_gt[:, 0]],
                                        xyz_target[matches_gt[:, 1]])
                metric = compute_metrics(
                    xyz[rand],
                    xyz_target[rand_target],
                    feat[rand],
                    feat_target[rand_target],
                    T_gt,
                    sym=cfg.data.sym,
                    tau_1=cfg.data.tau_1,
                    tau_2=cfg.data.tau_2,
                    rot_thresh=cfg.data.rot_thresh,
                    trans_thresh=cfg.data.trans_thresh,
                    use_ransac=cfg.data.use_ransac,
                    ransac_thresh=cfg.data.first_subsampling,
                    use_teaser=cfg.data.use_teaser,
                    noise_bound_teaser=cfg.data.noise_bound_teaser,
                )
                res = dict(**res, **metric)
                list_res.append(res)

    df = pd.DataFrame(list_res)
    output_path = os.path.join(cfg.training.checkpoint_dir, cfg.data.name,
                               "matches")
    if not os.path.exists(output_path):
        os.makedirs(output_path, exist_ok=True)
    df.to_csv(osp.join(output_path, "final_res.csv"))
    print(df.groupby("name_scene").mean())
예제 #16
0
def train_epoch(
    epoch: int,
    model: BaseModel,
    dataset,
    device: str,
    tracker: BaseTracker,
    checkpoint: ModelCheckpoint,
    visualizer: Visualizer,
    debugging,
):

    early_break = getattr(debugging, "early_break", False)
    profiling = getattr(debugging, "profiling", False)

    model.train()
    tracker.reset("train")
    visualizer.reset(epoch, "train")
    train_loader = dataset.train_dataloader

    iter_data_time = time.time()
    with Ctq(train_loader) as tq_train_loader:
        for i, data in enumerate(tq_train_loader):
            model.set_input(data, device)
            t_data = time.time() - iter_data_time

            iter_start_time = time.time()
            model.optimize_parameters(epoch, dataset.batch_size)
            if i % 10 == 0:
                tracker.track(model)

            tq_train_loader.set_postfix(**tracker.get_metrics(),
                                        data_loading=float(t_data),
                                        iteration=float(time.time() -
                                                        iter_start_time),
                                        color=COLORS.TRAIN_COLOR)

            if visualizer.is_active:
                visualizer.save_visuals(model.get_current_visuals())

            iter_data_time = time.time()

            if early_break:
                break

            if profiling:
                if i > getattr(debugging, "num_batches", 50):
                    return 0

    metrics = tracker.publish(epoch)
    checkpoint.save_best_models_under_current_metrics(model, metrics,
                                                      tracker.metric_func)
    log.info("Learning rate = %f" % model.learning_rate)
예제 #17
0
    def __init__(self, option, model_type, dataset, modules):
        BaseModel.__init__(self, option)
        option_unet = option.option_unet
        self.normalize_feature = option.normalize_feature
        self.grid_size = option_unet.grid_size
        self.unet = UnetMSparseConv3d(
            option_unet.backbone,
            input_nc=option_unet.input_nc,
            pointnet_nn=option_unet.pointnet_nn,
            post_mlp_nn=option_unet.post_mlp_nn,
            pre_mlp_nn=option_unet.pre_mlp_nn,
            add_pos=option_unet.add_pos,
            add_pre_x=option_unet.add_pre_x,
            aggr=option_unet.aggr,
            backend=option.backend,
        )
        if option.mlp_cls is not None:
            last_mlp_opt = option.mlp_cls

            self.FC_layer = Seq()
            for i in range(1, len(last_mlp_opt.nn)):
                self.FC_layer.append(
                    nn.Sequential(*[
                        nn.Linear(last_mlp_opt.nn[i - 1],
                                  last_mlp_opt.nn[i],
                                  bias=False),
                        FastBatchNorm1d(last_mlp_opt.nn[i],
                                        momentum=last_mlp_opt.bn_momentum),
                        nn.LeakyReLU(0.2),
                    ]))
            if last_mlp_opt.dropout:
                self.FC_layer.append(nn.Dropout(p=last_mlp_opt.dropout))
        else:
            self.FC_layer = torch.nn.Identity()
        self.head = nn.Sequential(
            nn.Linear(option.output_nc, dataset.num_classes))
        self.loss_names = ["loss_seg"]
예제 #18
0
def eval_epoch(
    model: BaseModel,
    dataset,
    device,
    tracker: BaseTracker,
    checkpoint: ModelCheckpoint,
    voting_runs=1,
    tracker_options={},
):
    tracker.reset("val")
    loader = dataset.val_dataloader
    for i in range(voting_runs):
        with Ctq(loader) as tq_val_loader:
            for data in tq_val_loader:
                with torch.no_grad():
                    model.set_input(data, device)
                    model.forward()

                tracker.track(model, **tracker_options)
                tq_val_loader.set_postfix(**tracker.get_metrics(),
                                          color=COLORS.VAL_COLOR)

    tracker.finalise(**tracker_options)
    tracker.print_summary()
예제 #19
0
    def _init_from_compact_format(self, opt, model_type, dataset, modules_lib):
        """Create a backbonebasedmodel from the compact options format - where the
        same convolution is given for each layer, and arguments are given
        in lists
        """
        num_convs = len(opt.down_conv.down_conv_nn)
        self.down_modules = nn.ModuleList()
        factory_module_cls = self._get_factory(model_type, modules_lib)
        down_conv_cls_name = opt.down_conv.module_name
        self._factory_module = factory_module_cls(down_conv_cls_name, None,
                                                  modules_lib)
        # Down modules
        for i in range(num_convs):
            args = self._fetch_arguments(opt.down_conv, i, "DOWN")
            conv_cls = self._get_from_kwargs(args, "conv_cls")
            down_module = conv_cls(**args)
            self._save_sampling_and_search(down_module)
            self.down_modules.append(down_module)

        self.metric_loss_module, self.miner_module = BaseModel.get_metric_loss_and_miner(
            getattr(opt, "metric_loss", None), getattr(opt, "miner", None))
예제 #20
0
def test_epoch(
    epoch: int,
    model: BaseModel,
    dataset,
    device,
    tracker: BaseTracker,
    checkpoint: ModelCheckpoint,
    visualizer: Visualizer,
    debugging,
):
    early_break = getattr(debugging, "early_break", False)
    model.eval()

    loaders = dataset.test_dataloaders

    for loader in loaders:
        stage_name = loader.dataset.name
        tracker.reset(stage_name)
        visualizer.reset(epoch, stage_name)
        with Ctq(loader) as tq_test_loader:
            for data in tq_test_loader:
                with torch.no_grad():
                    model.set_input(data, device)
                    model.forward()

                tracker.track(model)
                tq_test_loader.set_postfix(**tracker.get_metrics(),
                                           color=COLORS.TEST_COLOR)

                if visualizer.is_active:
                    visualizer.save_visuals(model.get_current_visuals())

                if early_break:
                    break

        tracker.finalise()
        metrics = tracker.publish(epoch)
        tracker.print_summary()
        checkpoint.save_best_models_under_current_metrics(
            model, metrics, tracker.metric_func)
예제 #21
0
def eval_epoch(
    epoch: int,
    model: BaseModel,
    dataset,
    device,
    tracker: BaseTracker,
    checkpoint: ModelCheckpoint,
    visualizer: Visualizer,
    debugging,
):

    early_break = getattr(debugging, "early_break", False)

    model.eval()
    tracker.reset("val")
    visualizer.reset(epoch, "val")
    loader = dataset.val_dataloader
    with Ctq(loader) as tq_val_loader:
        for data in tq_val_loader:
            with torch.no_grad():
                model.set_input(data, device)
                model.forward()

            tracker.track(model)
            tq_val_loader.set_postfix(**tracker.get_metrics(),
                                      color=COLORS.VAL_COLOR)

            if visualizer.is_active:
                visualizer.save_visuals(model.get_current_visuals())

            if early_break:
                break

    metrics = tracker.publish(epoch)
    tracker.print_summary()
    checkpoint.save_best_models_under_current_metrics(model, metrics,
                                                      tracker.metric_func)
예제 #22
0
def run(model: BaseModel, dataset: BaseDataset, device, cfg):

    reg_thresh = cfg.data.registration_recall_thresh
    if reg_thresh is None:
        reg_thresh = 0.2
    print(time.strftime("%Y%m%d-%H%M%S"))
    dataset.create_dataloaders(
        model, 1, False, cfg.training.num_workers, False,
    )
    loader = dataset.test_dataloaders[0]
    list_res = []
    with Ctq(loader) as tq_test_loader:
        for i, data in enumerate(tq_test_loader):
            with torch.no_grad():
                t0 = time.time()
                model.set_input(data, device)
                model.forward()
                t1 = time.time()
                name_scene, name_pair_source, name_pair_target = dataset.test_dataset[0].get_name(i)
                input, input_target = model.get_input()
                xyz, xyz_target = input.pos, input_target.pos
                ind, ind_target = input.ind, input_target.ind
                matches_gt = torch.stack([ind, ind_target]).transpose(0, 1)
                feat, feat_target = model.get_output()
                # rand = voxel_selection(xyz, grid_size=0.06, min_points=cfg.data.min_points)
                # rand_target = voxel_selection(xyz_target, grid_size=0.06, min_points=cfg.data.min_points)

                rand = torch.randperm(len(feat))[: cfg.data.num_points]
                rand_target = torch.randperm(len(feat_target))[: cfg.data.num_points]
                res = dict(name_scene=name_scene, name_pair_source=name_pair_source, name_pair_target=name_pair_target)
                T_gt = estimate_transfo(xyz[matches_gt[:, 0]], xyz_target[matches_gt[:, 1]])
                t2 = time.time()
                metric = compute_metrics(
                    xyz[rand],
                    xyz_target[rand_target],
                    feat[rand],
                    feat_target[rand_target],
                    T_gt,
                    sym=cfg.data.sym,
                    tau_1=cfg.data.tau_1,
                    tau_2=cfg.data.tau_2,
                    rot_thresh=cfg.data.rot_thresh,
                    trans_thresh=cfg.data.trans_thresh,
                    use_ransac=cfg.data.use_ransac,
                    ransac_thresh=cfg.data.first_subsampling,
                    use_teaser=cfg.data.use_teaser,
                    noise_bound_teaser=cfg.data.noise_bound_teaser,
                    xyz_gt=xyz[matches_gt[:, 0]],
                    xyz_target_gt=xyz_target[matches_gt[:, 1]],
                    registration_recall_thresh=reg_thresh,
                )
                res = dict(**res, **metric)
                res["time_feature"] = t1 - t0
                res["time_feature_per_point"] = (t1 - t0) / (len(input.pos) + len(input_target.pos))
                res["time_prep"] = t2 - t1

                list_res.append(res)

    df = pd.DataFrame(list_res)
    output_path = os.path.join(cfg.training.checkpoint_dir, cfg.data.name, "matches")
    if not os.path.exists(output_path):
        os.makedirs(output_path, exist_ok=True)
    df.to_csv(osp.join(output_path, "final_res_{}.csv".format(time.strftime("%Y%m%d-%H%M%S"))))
    print(df.groupby("name_scene").mean())
예제 #23
0
 def __init__(self, option):
     BaseModel.__init__(self, option)
예제 #24
0
def run(model: BaseModel, dataset: BaseDataset, device, cfg):
    print(time.strftime("%Y%m%d-%H%M%S"))
    dataset.create_dataloaders(
        model,
        1,
        False,
        cfg.training.num_workers,
        False,
    )
    loader = dataset.test_dataset[0]

    ind = 0
    if cfg.ind is not None:
        ind = cfg.ind
    t = 5
    if cfg.t is not None:
        t = cfg.t
    r = 0.1
    if cfg.r is not None:
        r = cfg.r
    print(loader)
    print(ind)
    data = loader[ind]
    data.batch = torch.zeros(len(data.pos)).long()
    data.batch_target = torch.zeros(len(data.pos_target)).long()
    print(data)
    with torch.no_grad():
        model.set_input(data, device)
        model.forward()

        name_scene, name_pair_source, name_pair_target = dataset.test_dataset[
            0].get_name(ind)
        print(name_scene, name_pair_source, name_pair_target)
        input, input_target = model.get_input()
        xyz, xyz_target = input.pos, input_target.pos
        ind, ind_target = input.ind, input_target.ind
        matches_gt = torch.stack([ind, ind_target]).transpose(0, 1)
        feat, feat_target = model.get_output()
        # rand = voxel_selection(xyz, grid_size=0.06, min_points=cfg.data.min_points)
        # rand_target = voxel_selection(xyz_target, grid_size=0.06, min_points=cfg.data.min_points)

        rand = torch.randperm(len(feat))[:cfg.data.num_points]
        rand_target = torch.randperm(len(feat_target))[:cfg.data.num_points]
        T_gt = estimate_transfo(xyz[matches_gt[:, 0]].clone(),
                                xyz_target[matches_gt[:, 1]].clone())
        matches_pred = get_matches(feat[rand],
                                   feat_target[rand_target],
                                   sym=cfg.data.sym)
        # For color
        inliers = (torch.norm(
            xyz[rand][matches_pred[:, 0]] @ T_gt[:3, :3].T + T_gt[:3, 3] -
            xyz_target[rand_target][matches_pred[:, 1]],
            dim=1,
        ) < cfg.data.tau_1)
        # compute transformation
        T_teaser = teaser_pp_registration(
            xyz[rand][matches_pred[:, 0]],
            xyz_target[rand_target][matches_pred[:, 1]],
            noise_bound=cfg.data.tau_1)
        pcd_source = torch2o3d(input, [1, 0.7, 0.1])

        pcd_target = torch2o3d(input_target, [0, 0.15, 0.9])
        open3d.visualization.draw_geometries([pcd_source, pcd_target])
        pcd_source.transform(T_teaser.cpu().numpy())
        open3d.visualization.draw_geometries([pcd_source, pcd_target])
        pcd_source.transform(np.linalg.inv(T_teaser.cpu().numpy()))
        rand_ind = torch.randperm(len(rand[matches_pred[:, 0]]))[:250]
        pcd_source.transform(T_gt.cpu().numpy())
        kp_s = torch2o3d(input, ind=rand[matches_pred[:, 0]][rand_ind])
        kp_s.transform(T_gt.cpu().numpy())
        kp_t = torch2o3d(input_target,
                         ind=rand_target[matches_pred[:, 1]][rand_ind])
        match_visualizer(pcd_source,
                         kp_s,
                         pcd_target,
                         kp_t,
                         inliers[rand_ind].cpu().numpy(),
                         radius=r,
                         t=t)