def finalise(self, full_res=False, vote_miou=True, ply_output="", **kwargs): per_class_iou = self._confusion_matrix.get_intersection_union_per_class( )[0] self._iou_per_class = { self._dataset.INV_OBJECT_LABEL[k]: v for k, v in enumerate(per_class_iou) } if vote_miou and self._test_area: # Complete for points that have a prediction self._test_area = self._test_area.to("cpu") c = ConfusionMatrix(self._num_classes) has_prediction = self._test_area.prediction_count > 0 gt = self._test_area.y[has_prediction].numpy() pred = torch.argmax(self._test_area.votes[has_prediction], 1).numpy() c.count_predicted_batch(gt, pred) self._vote_miou = c.get_average_intersection_union() * 100 if full_res: self._compute_full_miou() if ply_output: has_prediction = self._test_area.prediction_count > 0 self._dataset.to_ply( self._test_area.pos[has_prediction].cpu(), torch.argmax(self._test_area.votes[has_prediction], 1).cpu().numpy(), ply_output, )
def _compute_full_miou(self): if self._full_vote_miou is not None: return has_prediction = self._test_area.prediction_count > 0 log.info( "Computing full res mIoU, we have predictions for %.2f%% of the points." % (torch.sum(has_prediction) / (1.0 * has_prediction.shape[0]) * 100)) self._test_area = self._test_area.to("cpu") # Full res interpolation full_pred = knn_interpolate( self._test_area.votes[has_prediction], self._test_area.pos[has_prediction], self._test_area.pos, k=1, ) # Full res pred c = ConfusionMatrix(self._num_classes) c.count_predicted_batch(self._test_area.y.numpy(), torch.argmax(full_pred, 1).numpy()) self._full_vote_miou = c.get_average_intersection_union() * 100
class ScannetSegmentationTracker(SegmentationTracker): def reset(self, stage="train"): super().reset(stage=stage) self._full_confusion_matrix = ConfusionMatrix(self._num_classes) self._raw_datas = {} self._votes = {} self._vote_counts = {} self._full_preds = {} self._full_acc = None def track(self, model: model_interface.TrackerInterface, full_res=False, **kwargs): """ Add current model predictions (usually the result of a batch) to the tracking """ super().track(model) # Set conv type self._conv_type = model.conv_type # Train mode or low res, nothing special to do if not full_res or self._stage == "train" or kwargs.get("data") is None: return self._vote(kwargs.get("data"), model.get_output()) def get_metrics(self, verbose=False) -> Dict[str, Any]: """ Returns a dictionnary of all metrics and losses being tracked """ metrics = super().get_metrics(verbose) if self._full_acc: metrics["{}_full_acc".format(self._stage)] = self._full_acc metrics["{}_full_macc".format(self._stage)] = self._full_macc metrics["{}_full_miou".format(self._stage)] = self._full_miou return metrics def finalise(self, full_res=False, make_submission=False, **kwargs): if not full_res and not make_submission: return self._predict_full_res() # Compute full res metrics if self._dataset.has_labels(self._stage): for scan_id in self._full_preds: full_labels = self._raw_datas[scan_id].y # Mask ignored labels mask = full_labels != self._ignore_label full_labels = full_labels[mask] full_preds = self._full_preds[scan_id].cpu()[mask].numpy() self._full_confusion_matrix.count_predicted_batch(full_labels, full_preds) self._full_acc = 100 * self._full_confusion_matrix.get_overall_accuracy() self._full_macc = 100 * self._full_confusion_matrix.get_mean_class_accuracy() self._full_miou = 100 * self._full_confusion_matrix.get_average_intersection_union() # Save files to disk if make_submission and self._stage == "test": self._make_submission() def _make_submission(self): orginal_class_ids = np.asarray(self._dataset.train_dataset.valid_class_idx) path_to_submission = self._dataset.path_to_submission for scan_id in self._full_preds: full_pred = self._full_preds[scan_id].cpu().numpy().astype(np.int8) full_pred = orginal_class_ids[full_pred] # remap labels to original labels between 0 and 40 scan_name = self._raw_datas[scan_id].scan_name path_file = osp.join(path_to_submission, "{}.txt".format(scan_name)) np.savetxt(path_file, full_pred, delimiter="/n", fmt="%d") def _vote(self, data, output): """ Populates scores for the points in data Parameters ---------- data : Data should contain `pos` and `SaveOriginalPosId.KEY` keys output : torch.Tensor probablities out of the model, shape: [N,nb_classes] """ id_scans = data.id_scan.squeeze() if self._conv_type == "DENSE": batch_size = len(id_scans) output = output.view(batch_size, -1, output.shape[-1]) for idx_batch, id_scan in enumerate(id_scans): # First time we see this scan if id_scan not in self._raw_datas: raw_data = self._dataset.get_raw_data(self._stage, id_scan, remap_labels=True) self._raw_datas[id_scan] = raw_data self._vote_counts[id_scan] = torch.zeros(raw_data.pos.shape[0], dtype=torch.int) self._votes[id_scan] = torch.zeros((raw_data.pos.shape[0], self._num_classes), dtype=torch.float) else: raw_data = self._raw_datas[id_scan] batch_mask = idx_batch if self._conv_type != "DENSE": batch_mask = data.batch == idx_batch idx = data[SaveOriginalPosId.KEY][batch_mask] self._votes[id_scan][idx] += output[batch_mask].cpu() self._vote_counts[id_scan][idx] += 1 def _predict_full_res(self): """ Predict full resolution results based on votes """ for id_scan in self._votes: has_prediction = self._vote_counts[id_scan] > 0 self._votes[id_scan][has_prediction] /= self._vote_counts[id_scan][has_prediction].unsqueeze(-1) # Upsample and predict full_pred = knn_interpolate( self._votes[id_scan][has_prediction], self._raw_datas[id_scan].pos[has_prediction], self._raw_datas[id_scan].pos, k=1, ) self._full_preds[id_scan] = full_pred.argmax(-1)
class SegmentationTracker(BaseTracker): def __init__( self, dataset, stage="train", wandb_log=False, use_tensorboard: bool = False, ignore_label: int = IGNORE_LABEL ): """ This is a generic tracker for segmentation tasks. It uses a confusion matrix in the back-end to track results. Use the tracker to track an epoch. You can use the reset function before you start a new epoch Arguments: dataset -- dataset to track (used for the number of classes) Keyword Arguments: stage {str} -- current stage. (train, validation, test, etc...) (default: {"train"}) wandb_log {str} -- Log using weight and biases """ super(SegmentationTracker, self).__init__(stage, wandb_log, use_tensorboard) self._num_classes = dataset.num_classes self._ignore_label = ignore_label self._dataset = dataset self.reset(stage) def reset(self, stage="train"): super().reset(stage=stage) self._confusion_matrix = ConfusionMatrix(self._num_classes) self._acc = 0 self._macc = 0 self._miou = 0 @staticmethod def detach_tensor(tensor): if torch.torch.is_tensor(tensor): tensor = tensor.detach() return tensor @property def confusion_matrix(self): return self._confusion_matrix.confusion_matrix def track(self, model: model_interface.TrackerInterface, **kwargs): """ Add current model predictions (usually the result of a batch) to the tracking """ if not self._dataset.has_labels(self._stage): return 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() def get_metrics(self, verbose=False) -> Dict[str, Any]: """ Returns a dictionnary of all metrics and losses being tracked """ metrics = super().get_metrics(verbose) metrics["{}_acc".format(self._stage)] = self._acc metrics["{}_macc".format(self._stage)] = self._macc metrics["{}_miou".format(self._stage)] = self._miou return metrics @property def metric_func(self): self._metric_func = { "miou": max, "macc": max, "acc": max, "loss": min, "map": max, } # Those map subsentences to their optimization functions return self._metric_func
class S3DISTracker(SegmentationTracker): def reset(self, *args, **kwargs): super().reset(*args, **kwargs) self._test_area = None self._full_vote_miou = None self._vote_miou = None self._full_confusion = None self._iou_per_class = {} def track(self, model: model_interface.TrackerInterface, full_res=False, data=None, **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 origin ids and check that it fits with the test set inputs = data if data is not None else model.get_input() if inputs[SaveOriginalPosId.KEY] is None: raise ValueError("The inputs given to the model do not have a %s attribute." % SaveOriginalPosId.KEY) originids = inputs[SaveOriginalPosId.KEY] if originids.dim() == 2: originids = originids.flatten() if originids.max() >= self._test_area.pos.shape[0]: raise ValueError("Origin ids are larger than the number of points in the original point cloud.") # Set predictions outputs = model.get_output() self._test_area.votes[originids] += outputs self._test_area.prediction_count[originids] += 1 def finalise(self, full_res=False, vote_miou=True, ply_output="", **kwargs): per_class_iou = self._confusion_matrix.get_intersection_union_per_class()[0] self._iou_per_class = {self._dataset.INV_OBJECT_LABEL[k]: v for k, v in enumerate(per_class_iou)} if vote_miou and self._test_area: # Complete for points that have a prediction self._test_area = self._test_area.to("cpu") c = ConfusionMatrix(self._num_classes) has_prediction = self._test_area.prediction_count > 0 gt = self._test_area.y[has_prediction].numpy() pred = torch.argmax(self._test_area.votes[has_prediction], 1).numpy() c.count_predicted_batch(gt, pred) self._vote_miou = c.get_average_intersection_union() * 100 if full_res: self._compute_full_miou() if ply_output: has_prediction = self._test_area.prediction_count > 0 self._dataset.to_ply( self._test_area.pos[has_prediction].cpu(), torch.argmax(self._test_area.votes[has_prediction], 1).cpu().numpy(), ply_output, ) def _compute_full_miou(self): if self._full_vote_miou is not None: return has_prediction = self._test_area.prediction_count > 0 log.info( "Computing full res mIoU, we have predictions for %.2f%% of the points." % (torch.sum(has_prediction) / (1.0 * has_prediction.shape[0]) * 100) ) self._test_area = self._test_area.to("cpu") # Full res interpolation full_pred = knn_interpolate( self._test_area.votes[has_prediction], self._test_area.pos[has_prediction], self._test_area.pos, k=1, ) # Full res pred self._full_confusion = ConfusionMatrix(self._num_classes) self._full_confusion.count_predicted_batch(self._test_area.y.numpy(), torch.argmax(full_pred, 1).numpy()) self._full_vote_miou = self._full_confusion.get_average_intersection_union() * 100 @property def full_confusion_matrix(self): return self._full_confusion def get_metrics(self, verbose=False) -> Dict[str, Any]: """ Returns a dictionnary of all metrics and losses being tracked """ metrics = super().get_metrics(verbose) if verbose: metrics["{}_iou_per_class".format(self._stage)] = self._iou_per_class if self._vote_miou: metrics["{}_full_vote_miou".format(self._stage)] = self._full_vote_miou metrics["{}_vote_miou".format(self._stage)] = self._vote_miou return metrics