diff --git a/conf/model/segmentation/default.yaml b/conf/model/segmentation/default.yaml index 349caa4..ebf3207 100644 --- a/conf/model/segmentation/default.yaml +++ b/conf/model/segmentation/default.yaml @@ -1,6 +1,7 @@ # @package model defaults: - /model/default + - /tracker: segmentation/default model: _recursive_: false @@ -11,4 +12,4 @@ model: backbone: input_nc: ${dataset.cfg.feature_dimension} - architecture: unet \ No newline at end of file + architecture: unet diff --git a/conf/tracker/segmentation/default.yaml b/conf/tracker/segmentation/default.yaml new file mode 100644 index 0000000..81e6f05 --- /dev/null +++ b/conf/tracker/segmentation/default.yaml @@ -0,0 +1,2 @@ +_target_: torch_points3d.metrics.segmentation.segmentation_tracker.SegmentationTracker +num_classes: ${dataset.cfg.num_classes} diff --git a/test/__init__.py b/test/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/test/test_confusion_matrix.py b/test/test_confusion_matrix.py new file mode 100644 index 0000000..2a04757 --- /dev/null +++ b/test/test_confusion_matrix.py @@ -0,0 +1,59 @@ +import torch +import os +import sys +import unittest +import pytest +import numpy as np + + +DIR = os.path.dirname(os.path.realpath(__file__)) +ROOT = os.path.join(DIR, "..") +sys.path.insert(0, ROOT) +sys.path.append('.') + +from torch_points3d.metrics.segmentation.metrics import compute_intersection_union_per_class +from torch_points3d.metrics.segmentation.metrics import compute_average_intersection_union +from torch_points3d.metrics.segmentation.metrics import compute_overall_accuracy +from torch_points3d.metrics.segmentation.metrics import compute_mean_class_accuracy + + + +def test_compute_intersection_union_per_class(): + matrix = torch.tensor([[4, 1], [2, 10]]) + iou, _ = compute_intersection_union_per_class(matrix) + miou = compute_average_intersection_union(matrix) + np.testing.assert_allclose(iou[0].item(), 4 / (4.0 + 1.0 + 2.0)) + np.testing.assert_allclose(iou[1].item(), 10 / (10.0 + 1.0 + 2.0)) + np.testing.assert_allclose(iou.mean().item(), miou.item()) + +def test_compute_overall_accuracy(): + list_matrix = [ + torch.tensor([[4, 1], [2, 10]]).float(), + torch.tensor([[4, 1], [2, 10]]).int(), + torch.tensor([[0, 0], [0, 0]]).float() + ] + list_answer = [ + (4.0+10.0)/(4.0 + 10.0 + 1.0 +2.0), + (4.0+10.0)/(4.0 + 10.0 + 1.0 +2.0), + 0.0 + ] + for i in range(len(list_matrix)): + acc = compute_overall_accuracy(list_matrix[i]) + if(isinstance(acc, torch.Tensor)): + np.testing.assert_allclose(acc.item(), list_answer[i]) + else: + np.testing.assert_allclose(acc, list_answer[i]) + + +def test_compute_mean_class_accuracy(): + matrix = torch.tensor([[4, 1], [2, 10]]).float() + macc = compute_mean_class_accuracy(matrix) + np.testing.assert_allclose(macc.item(), (4/5 + 10/12)*0.5) + + + +@pytest.mark.parametrize("missing_as_one, answer", [pytest.param(False, (0.5 + 0.5) / 2), pytest.param(True, (0.5 + 1 + 0.5) / 3)]) +def test_test_getMeanIoUMissing(missing_as_one, answer): + matrix = torch.tensor([[1, 1, 0], [0, 1, 0], [0, 0, 0]]) + np.testing.assert_allclose(compute_average_intersection_union(matrix, missing_as_one=missing_as_one).item(), answer) + diff --git a/test/test_model.py b/test/test_model.py index 896940a..b97368b 100644 --- a/test/test_model.py +++ b/test/test_model.py @@ -1,4 +1,4 @@ -import unittest +import pytest import sys import os import torch @@ -9,27 +9,25 @@ DIR = os.path.dirname(os.path.realpath(__file__)) ROOT = os.path.join(DIR, "..") sys.path.insert(0, ROOT) +sys.path.append(".") -from torch_points3d.models.segmentation.sparseconv3d import APIModel +from torch_points3d.models.segmentation.base_model import SegmentationBaseModel +from torch_points3d.core.instantiator import HydraInstantiator -class TestAPIModel(unittest.TestCase): - def test_forward(self): - option_dataset = OmegaConf.create({"feature_dimension": 1, "num_classes": 10}) +@pytest.mark.skip("For now we skip the tests...") +def test_forward(self): + option_dataset = OmegaConf.create({"feature_dimension": 1, "num_classes": 10}) + option_criterion = OmegaConf.create({"_target_": "torch.nn.NLLLoss"}) + instantiator = HydraInstantiator() - option = OmegaConf.load(os.path.join(ROOT, "conf", "models", "segmentation", "sparseconv3d.yaml")) - name_model = list(option.keys())[0] - model = APIModel(option[name_model], option_dataset) + model = SegmentationBaseModel(instantiator, 10, option_backbone, option_criterion) - pos = torch.randn(1000, 3) - coords = torch.round(pos * 10000) - x = torch.ones(1000, 1) - batch = torch.zeros(1000).long() - y = torch.randint(0, 10, (1000,)) - data = Batch(pos=pos, x=x, batch=batch, y=y, coords=coords) - model.set_input(data) - model.forward() - - -if __name__ == "__main__": - unittest.main() + pos = torch.randn(1000, 3) + coords = torch.round(pos * 10000) + x = torch.ones(1000, 6) + batch = torch.zeros(1000).long() + y = torch.randint(0, 10, (1000,)) + data = Batch(pos=pos, x=x, batch=batch, y=y, coords=coords) + model.set_input(data) + model.forward() diff --git a/test/test_segmentation_tracker.py b/test/test_segmentation_tracker.py new file mode 100644 index 0000000..cb54eb6 --- /dev/null +++ b/test/test_segmentation_tracker.py @@ -0,0 +1,120 @@ +import numpy as np +import torch +import sys +import os + +import pytest + + +from torch_geometric.data import Data + +DIR = os.path.dirname(os.path.realpath(__file__)) +ROOT = os.path.join(DIR, "..") +sys.path.insert(0, ROOT) +sys.path.append(".") + +from torch_points3d.metrics.segmentation.segmentation_tracker import SegmentationTracker + + +class MockDataset: + INV_OBJECT_LABEL = {0: "first", 1: "wall", 2: "not", 3: "here", 4: "hoy"} + pos = torch.tensor([[1, 0, 0], [2, 0, 0], [3, 0, 0], [-1, 0, 0]]).float() + test_label = torch.tensor([1, 1, 0, 0]) + + def __init__(self): + self.num_classes = 2 + + @property + def test_data(self): + return Data(pos=self.pos, y=self.test_label) + + def has_labels(self, stage): + return True + + +class MockModel: + def __init__(self): + self.iter = 0 + self.losses = [ + {"loss_1": 1, "loss_2": 2}, + {"loss_1": 2, "loss_2": 2}, + {"loss_1": 1, "loss_2": 2}, + {"loss_1": 1, "loss_2": 2}, + ] + self.outputs = [ + torch.tensor([[0, 1], [0, 1]]), + torch.tensor([[1, 0], [1, 0]]), + torch.tensor([[1, 0], [1, 0]]), + torch.tensor([[1, 0], [1, 0], [1, 0]]), + ] + self.labels = [torch.tensor([1, 1]), torch.tensor([1, 1]), torch.tensor([1, 1]), torch.tensor([0, 0, -100])] + self.batch_idx = [torch.tensor([0, 1]), torch.tensor([0, 1]), torch.tensor([0, 1]), torch.tensor([0, 0, 1])] + + def get_input(self): + return Data(pos=MockDataset.pos[:2, :], origin_id=torch.tensor([0, 1])) + + def get_output(self): + return self.outputs[self.iter].float() + + def get_labels(self): + return self.labels[self.iter] + + def get_current_losses(self): + return self.losses[self.iter] + + def get_batch(self): + return self.batch_idx[self.iter] + + @property + def device(self): + return "cpu" + + +def test_forward(): + tracker = SegmentationTracker(num_classes=2, stage="train") + model = MockModel() + output = {"preds": model.get_output(), "labels": model.get_labels()} + losses = model.get_current_losses() + metrics = tracker(output, losses) + # metrics = tracker.get_metrics() + + for k in ["train_acc", "train_miou", "train_macc"]: + np.testing.assert_allclose(metrics[k], 100, rtol=1e-5) + model.iter += 1 + output = {"preds": model.get_output(), "labels": model.get_labels()} + losses = model.get_current_losses() + metrics = tracker(output, losses) + # metrics = tracker.get_metrics() + metrics = tracker.finalise() + for k in ["train_acc", "train_macc"]: + assert metrics[k] == 50 + np.testing.assert_allclose(metrics["train_miou"], 25, atol=1e-5) + assert metrics["train_loss_1"] == 1.5 + + tracker.reset("test") + model.iter += 1 + output = {"preds": model.get_output(), "labels": model.get_labels()} + losses = model.get_current_losses() + metrics = tracker(output, losses) + # metrics = tracker.get_metrics() + for name in ["test_acc", "test_miou", "test_macc"]: + np.testing.assert_allclose(metrics[name].item(), 0, atol=1e-5) + + +@pytest.mark.parametrize("finalise", [pytest.param(True), pytest.param(False)]) +def test_ignore_label(finalise): + tracker = SegmentationTracker(num_classes=2, ignore_label=-100) + tracker.reset("test") + model = MockModel() + model.iter = 3 + output = {"preds": model.get_output(), "labels": model.get_labels()} + losses = model.get_current_losses() + metrics = tracker(output, losses) + if not finalise: + # metrics = tracker.get_metrics() + for k in ["test_acc", "test_miou", "test_macc"]: + np.testing.assert_allclose(metrics[k], 100) + else: + tracker.finalise() + with pytest.raises(RuntimeError): + tracker(output) diff --git a/torch_points3d/core/instantiator.py b/torch_points3d/core/instantiator.py index e4b3d04..3be6d90 100644 --- a/torch_points3d/core/instantiator.py +++ b/torch_points3d/core/instantiator.py @@ -44,6 +44,9 @@ def litmodel(self, cfg: DictConfig) -> "PointCloudBaseModule": def model(self, cfg: DictConfig) -> "PointCloudBaseModel": return self.instantiate(cfg, self) + def tracker(self, cfg: DictConfig, stage: str = ""): + return self.instantiate(cfg, stage=stage) + def backbone(self, cfg: DictConfig): return self.instantiate(cfg) diff --git a/torch_points3d/metrics/__init__.py b/torch_points3d/metrics/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/torch_points3d/metrics/base_tracker.py b/torch_points3d/metrics/base_tracker.py new file mode 100644 index 0000000..dd65630 --- /dev/null +++ b/torch_points3d/metrics/base_tracker.py @@ -0,0 +1,59 @@ +from typing import Any, Dict, Optional +import torch +from torch import nn +from torchmetrics import AverageMeter + + +class BaseTracker(nn.Module): + """ + pytorch Module to manage the losses and the metrics + """ + + def __init__(self, stage: str = "train"): + super().__init__() + self.stage: str = stage + self._finalised: bool = False + self.loss_metrics: nn.ModuleDict = nn.ModuleDict() + + def track(self, output_model, *args, **kwargs) -> Dict[str, Any]: + raise NotImplementedError + + def track_loss(self, losses: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: + out_loss = dict() + for key, loss in losses.items(): + loss_key = f"{self.stage}_{key}" + if loss_key not in self.loss_metrics.keys(): + self.loss_metrics[loss_key] = AverageMeter().to(loss) + val = self.loss_metrics[loss_key](loss) + out_loss[loss_key] = val + return out_loss + + def forward( + self, output_model: Dict[str, Any], losses: Optional[Dict[str, torch.Tensor]] = None, *args, **kwargs + ) -> Dict[str, Any]: + if self._finalised: + raise RuntimeError("Cannot track new values with a finalised tracker, you need to reset it first") + tracked_metric = self.track(output_model, *args, **kwargs) + if losses is not None: + tracked_loss = self.track_loss(losses) + tracked_results = dict(**tracked_loss, **tracked_metric) + else: + tracked_results = tracked_metric + return tracked_results + + def _finalise(self) -> Dict[str, Any]: + raise NotImplementedError("method that aggregae metrics") + + def finalise(self) -> Dict[str, Any]: + metrics = self._finalise() + self._finalised = True + loss_metrics = self.get_final_loss_metrics() + final_metrics = {**loss_metrics, **metrics} + return final_metrics + + def get_final_loss_metrics(self): + metrics = dict() + for key, m in self.loss_metrics.items(): + metrics[key] = m.compute() + self.loss_metrics = nn.ModuleDict() + return metrics diff --git a/torch_points3d/metrics/segmentation/__init__.py b/torch_points3d/metrics/segmentation/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/torch_points3d/metrics/segmentation/metrics.py b/torch_points3d/metrics/segmentation/metrics.py new file mode 100644 index 0000000..a3aecc9 --- /dev/null +++ b/torch_points3d/metrics/segmentation/metrics.py @@ -0,0 +1,81 @@ +import torch +from typing import Optional, Tuple, Union + + +def compute_average_intersection_union(confusion_matrix: torch.Tensor, missing_as_one: bool = False) -> torch.Tensor: + """ + compute intersection over union on average from confusion matrix + Parameters + Parameters + ---------- + confusion_matrix: torch.Tensor + square matrix + missing_as_one: bool, default: False + """ + + values, existing_classes_mask = compute_intersection_union_per_class(confusion_matrix, return_existing_mask=True) + if torch.sum(existing_classes_mask) == 0: + return torch.sum(existing_classes_mask) + if missing_as_one: + values[~existing_classes_mask] = 1 + existing_classes_mask[:] = True + return torch.sum(values[existing_classes_mask]) / torch.sum(existing_classes_mask) + + +def compute_mean_class_accuracy(confusion_matrix: torch.Tensor) -> torch.Tensor: + """ + compute intersection over union on average from confusion matrix + + Parameters + ---------- + confusion_matrix: torch.Tensor + square matrix + """ + total_gts = confusion_matrix.sum(1) + labels_presents = torch.where(total_gts > 0)[0] + if len(labels_presents) == 0: + return total_gts[0] + ones = torch.ones_like(total_gts) + max_ones_total_gts = torch.cat([total_gts[None, :], ones[None, :]], 0).max(0)[0] + re = (torch.diagonal(confusion_matrix)[labels_presents].float() / max_ones_total_gts[labels_presents]).sum() + return re / float(len(labels_presents)) + + +def compute_overall_accuracy(confusion_matrix: torch.Tensor) -> Union[int, torch.Tensor]: + """ + compute overall accuracy from confusion matrix + + Parameters + ---------- + confusion_matrix: torch.Tensor + square matrix + """ + all_values = confusion_matrix.sum() + if all_values == 0: + return 0 + matrix_diagonal = torch.trace(confusion_matrix) + return matrix_diagonal.float() / all_values + + +def compute_intersection_union_per_class( + confusion_matrix: torch.Tensor, return_existing_mask: bool = False, eps: float = 1e-8 +) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """ + compute intersection over union per class from confusion matrix + + Parameters + ---------- + confusion_matrix: torch.Tensor + square matrix + """ + + TP_plus_FN = confusion_matrix.sum(0) + TP_plus_FP = confusion_matrix.sum(1) + TP = torch.diagonal(confusion_matrix) + union = TP_plus_FN + TP_plus_FP - TP + iou = eps + TP / (union + eps) + existing_class_mask = union > 1e-3 + if return_existing_mask: + return iou, existing_class_mask + else: + return iou, None diff --git a/torch_points3d/metrics/segmentation/segmentation_tracker.py b/torch_points3d/metrics/segmentation/segmentation_tracker.py new file mode 100644 index 0000000..d9c5a77 --- /dev/null +++ b/torch_points3d/metrics/segmentation/segmentation_tracker.py @@ -0,0 +1,51 @@ +import torch +from typing import Any, Dict +from torchmetrics import ConfusionMatrix +from torchmetrics import Metric + +from torch_points3d.metrics.base_tracker import BaseTracker +import torch_points3d.metrics.segmentation.metrics as mt + + +class SegmentationTracker(BaseTracker): + """ + track different registration metrics + """ + + def __init__( + self, + num_classes: int, + stage: str = "train", + ignore_label: int = -1, + eps: float = 1e-8, + ): + super().__init__(stage) + self._ignore_label = ignore_label + self._num_classes = num_classes + self.confusion_matrix_metric = ConfusionMatrix(num_classes=self._num_classes) + self.eps = eps + + def compute_metrics_from_cm(self, matrix: torch.Tensor) -> Dict[str, Any]: + acc = mt.compute_overall_accuracy(matrix) + macc = mt.compute_mean_class_accuracy(matrix) + miou = mt.compute_average_intersection_union(matrix) + iou_per_class, _ = mt.compute_intersection_union_per_class(matrix, eps=self.eps) + iou_per_class_dict = {f"{self.stage}_iou_class_{i}": (100 * v) for i, v in enumerate(iou_per_class)} + res = { + "{}_acc".format(self.stage): 100 * acc, + "{}_macc".format(self.stage): 100 * macc, + "{}_miou".format(self.stage): 100 * miou, + } + res = dict(**res, **iou_per_class_dict) + return res + + def track(self, output, **kwargs) -> Dict[str, Any]: + mask = output["labels"] != self._ignore_label + matrix = self.confusion_matrix_metric(output["preds"][mask], output["labels"][mask]) + segmentation_metrics = self.compute_metrics_from_cm(matrix) + return segmentation_metrics + + def _finalise(self): + matrix = self.confusion_matrix_metric.compute() + segmentation_metrics = self.compute_metrics_from_cm(matrix) + return segmentation_metrics diff --git a/torch_points3d/models/base_model.py b/torch_points3d/models/base_model.py index a800fac..2845996 100644 --- a/torch_points3d/models/base_model.py +++ b/torch_points3d/models/base_model.py @@ -1,4 +1,4 @@ -from typing import Union +from typing import Dict, Optional import torch import torch.nn as nn @@ -12,12 +12,22 @@ def __init__(self, instantiator: Instantiator): super().__init__() self.instantiator = instantiator + self._losses: Dict[str, float] = {} def set_input(self, data: Data) -> None: raise (NotImplementedError("set_input needs to be defined!")) - def forward(self) -> Union[torch.Tensor, None]: + def forward(self) -> Optional[torch.Tensor]: raise (NotImplementedError("forward needs to be defined!")) - def get_losses(self) -> Union[torch.Tensor, None]: + def compute_loss(self): raise (NotImplementedError("get_losses needs to be defined!")) + + def get_losses(self) -> Optional[Dict["str", torch.Tensor]]: + return self._losses + + def get_outputs(self) -> Dict[str, Optional[torch.Tensor]]: + """ + return the outputs to track for the metrics + """ + raise (NotImplementedError("outputs need to be defined")) diff --git a/torch_points3d/models/segmentation/base_model.py b/torch_points3d/models/segmentation/base_model.py index 8e11465..42ac312 100644 --- a/torch_points3d/models/segmentation/base_model.py +++ b/torch_points3d/models/segmentation/base_model.py @@ -1,5 +1,5 @@ from omegaconf import DictConfig -from typing import Union +from typing import Dict, Optional import torch import torch.nn as nn @@ -28,17 +28,25 @@ def set_input(self, data: Data) -> None: else: self.labels = None - def forward(self) -> Union[torch.Tensor, None]: + def forward(self) -> Optional[torch.Tensor]: features = self.backbone(self.input).x logits = self.head(features) - self.output = F.log_softmax(logits, dim=-1) - - return self.get_losses() - - def get_losses(self) -> Union[torch.Tensor, None]: - # only compute loss if loss is defined and the dset has labels - if self.labels is None or self.criterion is None: - return + self._output = F.log_softmax(logits, dim=-1) + loss = self.compute_losses() + return loss + + def compute_losses(self): + """ + compute every loss. store the total loss in an attribute _loss + """ + if self.labels is not None and self.criterion is not None: + self._losses["loss"] = self.criterion(self._output, self.labels) + return self._losses["loss"] + else: + return None - self.loss = self.criterion(self.output, self.labels) - return self.loss + def get_outputs(self) -> Dict[str, torch.Tensor]: + """ + return the outputs to track for the metrics + """ + return {"labels": self.labels, "preds": self._output} diff --git a/torch_points3d/tasks/base_model.py b/torch_points3d/tasks/base_model.py index 339678c..03fc2b3 100644 --- a/torch_points3d/tasks/base_model.py +++ b/torch_points3d/tasks/base_model.py @@ -16,12 +16,14 @@ def __init__( optimizer: OptimizerConfig, instantiator: Instantiator, scheduler: SchedulerConfig = None, # scheduler shouldn't be required + tracker: Optional[DictConfig] = None, ): super().__init__() # some optimizers/schedulers need parameters only known dynamically # allow users to override the getter to instantiate them lazily self.optimizer_cfg = optimizer self.scheduler_cfg = scheduler + self.tracker_cfg = tracker self.instantiator = instantiator self._init_model(model) @@ -91,15 +93,22 @@ def configure_metrics(self, stage: str) -> Optional[Any]: This is called on fit start to have access to the data module, and initialize any data specific metrics. """ + self.tracker = self.instantiator.tracker(self.tracker_cfg) - def training_step(self, batch, batch_idx): + def _step(self, batch, batch_idx, stage: str): self.model.set_input(batch) - return self.model.forward() + loss = self.model.forward() + losses = self.model.get_losses() + outputs = self.model.get_outputs() + metric_dict = self.tracker(outputs, losses) + self.log_dict(metric_dict, prog_bar=True, on_step=False, on_epoch=True) + return loss + + def training_step(self, batch, batch_idx): + return self._step(batch, batch_idx, "train") def validation_step(self, batch, batch_idx): - self.model.set_input(batch) - return self.model.forward() + return self._step(batch, batch_idx, "val") def testing_step(self, batch, batch_idx): - self.model.set_input(batch) - return self.model.forward() + return self._step(batch, batch_idx, "test")