def track(self, model: model_interface.TrackerInterface, **kwargs):
        super().track(model)
        if self._stage != "train":
            batch_idx, batch_idx_target = model.get_batch()
            batch_xyz, batch_xyz_target = model.get_xyz()  # type: ignore
            batch_ind, batch_ind_target, batch_size_ind = model.get_ind(
            )  # type: ignore
            batch_feat, batch_feat_target = model.get_output()

            nb_batches = batch_idx.max() + 1
            cum_sum = 0
            cum_sum_target = 0
            begin = 0
            end = batch_size_ind[0].item()
            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[b * len(batch_ind) / nb_batches : (b + 1) * len(batch_ind) / nb_batches] - cum_sum
                # ind_target = (batch_ind_target[b * len(batch_ind_target) / nb_batches : (b + 1) * len(batch_ind_target) / nb_batches]- cum_sum_target)
                ind = batch_ind[begin:end] - cum_sum
                ind_target = batch_ind_target[begin:end] - cum_sum_target
                # print(begin, end)
                if b < nb_batches - 1:
                    begin = end
                    end = begin + batch_size_ind[b + 1].item()
                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]).transpose(0, 1)

                # print(matches_gt.max(0), len(xyz), len(xyz_target), len(matches_gt))
                # print(batch_ind.shape, nb_batches)
                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())
Beispiel #2
0
    def track(self, model: model_interface.TrackerInterface, **kwargs):
        """ Add current model predictions (usually the result of a batch) to the tracking
        """
        super().track(model)
        outputs = self._convert(model.get_output())
        targets = self._convert(model.get_labels())
        batch_idx = self._convert(model.get_batch())
        if batch_idx is None:
            raise ValueError(
                "Your model need to set the batch_idx variable in its set_input function."
            )

        nb_batches = batch_idx.max() + 1

        # pred to the groundtruth classes (selected by seg_classes[cat])
        for b in range(nb_batches):
            segl = targets[batch_idx == b]
            cat = self._seg_to_class[segl[0]]
            logits = outputs[batch_idx == b, :]  # (num_points, num_classes)
            segp = logits[:, self._class_seg_map[cat]].argmax(
                1) + self._class_seg_map[cat][0]
            part_ious = np.zeros(len(self._class_seg_map[cat]))
            for l in self._class_seg_map[cat]:
                if np.sum((segl == l) | (segp == l)) == 0:
                    # part is not present in this shape
                    part_ious[l - self._class_seg_map[cat][0]] = 1
                else:
                    part_ious[l - self._class_seg_map[cat][0]] = float(
                        np.sum((segl == l) & (segp == l))) / float(
                            np.sum((segl == l) | (segp == l)))
            self._shape_ious[cat].append(np.mean(part_ious))

        self._miou_per_class, self._Cmiou, self._Imiou = self._get_metrics_per_class(
        )
    def track(self,
              model: model_interface.TrackerInterface,
              full_res: bool = False,
              data: Data = None,
              **kwargs):
        """ Add current model predictions (usually the result of a batch) to the tracking
        """
        super().track(model)
        self._conv_type = model.conv_type
        outputs = self._convert(model.get_output())
        targets = self._convert(model.get_labels())
        batch_idx = self._convert(model.get_batch())
        if batch_idx is None:
            raise ValueError(
                "Your model need to set the batch_idx variable in its set_input function."
            )

        nb_batches = batch_idx.max() + 1

        if self._stage != "train" and full_res:
            self._add_votes(data, outputs, batch_idx)

        # pred to the groundtruth classes (selected by seg_classes[cat])
        for b in range(nb_batches):
            segl = targets[batch_idx == b]
            cat = self._seg_to_class[segl[0]]
            logits = outputs[batch_idx == b, :]  # (num_points, num_classes)
            segp = logits[:, self._class_seg_map[cat]].argmax(
                1) + self._class_seg_map[cat][0]
            part_ious = self._compute_part_ious(segl, segp, cat)
            self._shape_ious[cat].append(np.mean(part_ious))

        self._miou_per_class, self._Cmiou, self._Imiou = ShapenetPartTracker._get_metrics_per_class(
            self._shape_ious)