repo_description: Official PyTorch implementation of Superpoint Transformer [ICCV'23], SuperCluster [3DV'24 Oral], and EZ-SP [ICRA'26]
repo_topics: 3d 3dv2024 deep-learning efficient fast graph-clustering hierarchical iccv2023 lightweight panoptic-segmentation partition partitioning point-cloud pytorch semantic-segmentation superpoint transformer
code: import torch
  import logging
  from typing import Any, List, Tuple, Dict
  from torchmetrics import MaxMetric, MeanMetric
  from torchmetrics.classification import BinaryAccuracy, BinaryF1Score
  from torch_geometric.nn.pool.consecutive import consecutive_cluster
  
  from src.utils import init_weights, PanopticSegmentationOutput, \
      PartitionParameterSearchStorage
  from src.metrics import MeanAveragePrecision3D, PanopticQuality3D, \
      ConfusionMatrix
  from src.models.semantic import SemanticSegmentationModule
  from src.loss import BCEWithLogitsLoss
  from src.data import NAG
  
  
  log = logging.getLogger(__name__)
  
  
  __all__ = ['PanopticSegmentationModule']
  
  
  class PanopticSegmentationModule(SemanticSegmentationModule):
      """A LightningModule for panoptic segmentation of point clouds.
  
      :param net: torch.nn.Module
          Backbone model. This can typically be an `SPT` object
      :param edge_affinity_head: torch.nn.Module
          Edge affinity prediction head for instance/panoptic graph
          clustering. This is typically an MLP
      :param partitioner: src.nn.instance.InstancePartitioner
          Instance partition head, expects a fully-fledged
          `InstancePartitioner` module as input. This module is only
          called when the actual instance/panoptic segmentation is
          required. At train time, it is not essential, since we do not
          propagate gradient to its parameters. However, we may still tune
          its parameters to maximize instance/panoptic metrics on the
          train set. This tuning involves a simple grid-search on a small
          range of parameters and needs to be called at least once at the
          very end of training
      :param criterion: torch.nn._Loss
          Loss
      :param optimizer: torch.optim.Optimizer
          Optimizer
      :param scheduler: torch.optim.lr_scheduler.LRScheduler
          Learning rate scheduler
      :param num_classes: int
          Number of classes in the dataset
      :param stuff_classes: List[int]
          Indices of the classes to be treated as 'stuff', as opposed to
          'thing'
      :param class_names: List[str]
          Name for each class
      :param sampling_loss:  bool
          If True, the target labels will be obtained from labels of
          the points sampled in the batch at hand. This affects
          training supervision where sampling augmentations may be
          used for dropping some points or superpoints. If False, the
          target labels will be based on exact superpoint-wise
          histograms of labels computed at preprocessing time,
          disregarding potential level-0 point down-sampling
      :param loss_type: str
          Type of loss applied.
          'ce': cross-entropy (if `multi_stage_loss_lambdas` is used,
          all 1+ levels will be supervised with cross-entropy).
          'kl': Kullback-Leibler divergence (if `multi_stage_loss_lambdas`
          is used, all 1+ levels will be supervised with cross-entropy).
          'ce_kl': cross-entropy on level 1 and Kullback-Leibler for
          all levels above
          'wce': not documented for now
          'wce_kl': not documented for now
      :param weighted_loss: bool
          If True, the loss will be weighted based on the class
          frequencies computed on the train dataset. See
          `BaseDataset.get_class_weight()` for more
      :param init_linear: str
          Initialization method for all linear layers. Supports
          'xavier_uniform', 'xavier_normal', 'kaiming_uniform',
          'kaiming_normal', 'trunc_normal'
      :param init_rpe: str
          Initialization method for all linear layers producing
          relative positional encodings. Supports 'xavier_uniform',
          'xavier_normal', 'kaiming_uniform', 'kaiming_normal',
          'trunc_normal'
      :param transformer_lr_scale: float
          Scaling parameter applied to the learning rate for the
          `TransformerBlock` in each `Stage` and for the pooling block
          in `DownNFuseStage` modules. Setting this to a value lower
          than 1 mitigates exploding gradients in attentive blocks
          during training
      :param multi_stage_loss_lambdas: List[float]
          List of weights for combining losses computed on the output
          of each partition level. If not specified, the loss will
          be computed on the level 1 outputs only
      :param edge_affinity_criterion: torch.nn._Loss
          Loss on the edges of the superpoint level 1 for affinity
          prediction
      :param edge_affinity_loss_weights: List[float]
          Weights for insisting on certain cases in the edge affinity
          loss:
           - 0: same-class same-object edges
           - 1: same-class different-object edges
           - 2: different-class same-object edges
           - 3: different-class different-object edges
      :param edge_affinity_loss_lambda: float
          Weight for combining the semantic segmentation loss with the
          node offset and edge affinity losses. The final loss will be:
          `L_node_classif + edge_affinity_loss_lambda * L_edge_affinity
          + node_offset_loss_lambda * L_node_offset`
      :param node_offset_criterion: torch.nn._Loss
          Loss on the nodes of the superpoint level 1 for node offset
          prediction
      :param node_offset_loss_lambda: float
          Weight for combining the semantic segmentation loss with the
          node offset and edge affinity losses. The final loss will be:
          `L_node_classif + edge_affinity_loss_lambda * L_edge_affinity
          + node_offset_loss_lambda * L_node_offset`
      :param gc_every_n_steps: int
          Explicitly call the garbage collector after a certain number
          of steps. May involve a computation overhead. Mostly hear
          for debugging purposes when observing suspicious GPU memory
          increase during training
      :param track_val_every_n_epoch: int
          If specified, the output for a validation batch of interest
          specified with `track_val_idx` will be stored to disk every
          `track_val_every_n_epoch` epochs. Must be a multiple of
          `check_val_every_n_epoch`. See `track_batch()` for more
      :param track_val_idx: int
          If specified, the output for the `track_val_idx`th
          validation batch will be saved to disk periodically based on
          `track_val_every_n_epoch`. Importantly, this index is expected
          to match the `Dataloader`'s index wrt the current epoch
          and NOT an index wrt the `Dataset`. Said otherwise, if the
          `Dataloader(shuffle=True)` then, the stored batch will not be
          the same at each epoch. For this reason, if tracking the same
          object across training is needed, the `Dataloader` and the
          transforms should be free from any stochasticity
      :param track_test_idx:
          If specified, the output for the `track_test_idx`th
          test batch will be saved to disk. If `track_test_idx=-1`,
          predictions for the entire test set will be saved to disk
      :param min_instance_size: int
          Minimum target instance size to consider when computing the
          metrics. If a target is smaller, it will be ignored, as well
          as its matched prediction, if any. See `MeanAveragePrecision3D`
      :param partition_every_n_epoch: int
          Since we do not need to compute the actual panoptic/instance
          segmentation to train the model, we can simply do so once in a
          while to track the training and validation metrics. This
          parameter rules the frequency at which the panoptic/instance
          partition and metrics are computed during training
      :param no_instance_metrics: bool
          Whether instance segmentation metrics should be computed. These
          may incur an overhead. Besides, the SuperCluster formulation is
          mainly targeted for panoptic segmentation, as the model is not
          specifically trained to maximize instance metrics, which, among
          other things, involve predicting an instance confidence score
      :param no_instance_metrics_on_train_set: bool
          Same as `no_instance_metrics` but specifically for the train
          set. This is in case we still want the instance metrics every
          partition_every_n_epoch` on the validation set, but want to
          avoid the compute overhead of computing the instance partition
          and metrics at every single training epoch
      :param kwargs: Dict
          Kwargs will be passed to `_load_from_checkpoint()`
      """
  
      _IGNORED_HYPERPARAMETERS = [
          'net',
          'edge_affinity_head',
          'partitioner',
          'criterion',
          'edge_affinity_criterion',
          'node_offset_criterion']
  
      def __init__(
              self,
              net: torch.nn.Module,
              edge_affinity_head: torch.nn.Module,
              partitioner: 'InstancePartitioner',
              criterion: 'torch.nn._Loss',
              optimizer: torch.optim.Optimizer,
              scheduler: torch.optim.lr_scheduler.LRScheduler,
              num_classes: int,
              stuff_classes: List[int],
              class_names: List[str] = None,
              sampling_loss: bool = False,
              loss_type: str = 'ce_kl',
              weighted_loss: bool = True,
              init_linear: str = None,
              init_rpe: str = None,
              transformer_lr_scale: float = 1,
              multi_stage_loss_lambdas: List[float] = None,
              edge_affinity_criterion: 'torch.nn._Loss' = None,
              edge_affinity_loss_weights: List[float] = None,
              edge_affinity_loss_lambda: float = 1,
              node_offset_criterion: 'torch.nn._Loss' = None,
              node_offset_loss_lambda: float = 1,
              gc_every_n_steps: int = 0,
              track_val_every_n_epoch: int = 1,
              track_val_idx: int = None,
              track_test_idx: int = None,
              min_instance_size: int = 100,
              partition_every_n_epoch: int = 50,
              no_instance_metrics: bool = True,
              no_instance_metrics_on_train_set: bool = True,
              **kwargs):
          super().__init__(
              net,
              criterion,
              optimizer,
              scheduler,
              num_classes,
              class_names=class_names,
              sampling_loss=sampling_loss,
              loss_type=loss_type,
              weighted_loss=weighted_loss,
              init_linear=init_linear,
              init_rpe=init_rpe,
              transformer_lr_scale=transformer_lr_scale,
              multi_stage_loss_lambdas=multi_stage_loss_lambdas,
              gc_every_n_steps=gc_every_n_steps,
              track_val_every_n_epoch=track_val_every_n_epoch,
              track_val_idx=track_val_idx,
              track_test_idx=track_test_idx,
              **kwargs)
  
          # Instance partition head, expects a fully-fledged
          # InstancePartitioner module as input.
          # This module is only called when the actual instance/panoptic
          # segmentation is required. At train time, it is not essential,
          # since we do not propagate gradient to its parameters. However,
          # we still tune its parameters to maximize instance/panoptic
          # metrics on the train set. This tuning involves a simple
          # grid-search on a small range of parameters and needs to be
          # called at least once at the very end of training
          self.partition_every_n_epoch = partition_every_n_epoch
          self.no_instance_metrics = no_instance_metrics
          self.no_instance_metrics_on_train_set = no_instance_metrics_on_train_set
          self.partitioner = partitioner
  
          # Store the stuff class indices
          self.stuff_classes = stuff_classes
  
          # Loss functions for edge affinity and node offset predictions.
          # NB: the semantic loss is already accounted for in the
          # SemanticSegmentationModule constructor
          self.edge_affinity_criterion = BCEWithLogitsLoss() \
              if edge_affinity_criterion is None else edge_affinity_criterion
          # self.node_offset_criterion = WeightedL2Loss() \
          #     if node_offset_criterion is None else node_offset_criterion
  
          # Model heads for edge affinity and node offset predictions
          # Initialize the model segmentation head (or heads)
          # out_dim = self.net.out_dim[0] if self.multi_stage_loss \
          #     else self.net.out_dim
          # self.edge_affinity_head = FFN(out_dim * 2, hidden_dim=32, out_dim=1)
          self.edge_affinity_head = edge_affinity_head
          # self.node_offset_head = FFN(out_dim, hidden_dim=32, out_dim=3)
  
          # Custom weight initialization. In particular, this applies
          # Xavier / Glorot initialization on Linear and RPE layers by
          # default, but can be tuned
          init = lambda m: init_weights(m, linear=init_linear, rpe=init_rpe)
          self.edge_affinity_head.apply(init)
          # self.node_offset_head.apply(init)
  
          # Metric objects for calculating panoptic segmentation scores on
          # each dataset split
          self.train_panoptic = PanopticQuality3D(
              self.num_classes,
              ignore_unseen_classes=True,
              stuff_classes=self.stuff_classes,
              compute_on_cpu=True,
              **kwargs)
          self.val_panoptic = PanopticQuality3D(
              self.num_classes,
              ignore_unseen_classes=True,
              stuff_classes=self.stuff_classes,
              compute_on_cpu=True,
              **kwargs)
          self.test_panoptic = PanopticQuality3D(
              self.num_classes,
              ignore_unseen_classes=True,
              stuff_classes=self.stuff_classes,
              compute_on_cpu=True,
              **kwargs)
  
          # Metric objects for calculating semantic segmentation scores on
          # predicted instances on each dataset split
          self.train_semantic = ConfusionMatrix(self.num_classes)
          self.val_semantic = ConfusionMatrix(self.num_classes)
          self.test_semantic = ConfusionMatrix(self.num_classes)
  
          # Metric objects for calculating instance segmentation scores on
          # each dataset split
          self.train_instance = MeanAveragePrecision3D(
              self.num_classes,
              stuff_classes=self.stuff_classes,
              min_size=min_instance_size,
              compute_on_cpu=True,
              remove_void=True,
              **kwargs)
          self.val_instance = MeanAveragePrecision3D(
              self.num_classes,
              stuff_classes=self.stuff_classes,
              min_size=min_instance_size,
              compute_on_cpu=True,
              remove_void=True,
              **kwargs)
          self.test_instance = MeanAveragePrecision3D(
              self.num_classes,
              stuff_classes=self.stuff_classes,
              min_size=min_instance_size,
              compute_on_cpu=True,
              remove_void=True,
              **kwargs)
  
          # Storage to accumulate multiple batch partition predictions, to
          # be used when searching for the best partition setting
          self.train_multi_partition_storage = []
  
          # Metric objects for calculating node offset prediction scores
          # on each dataset split
          # self.train_offset_wl2 = WeightedL2Error()
          # self.train_offset_wl1 = WeightedL1Error()
          # self.train_offset_l2 = L2Error()
          # self.train_offset_l1 = L1Error()
          # self.val_offset_wl2 = WeightedL2Error()
          # self.val_offset_wl1 = WeightedL1Error()
          # self.val_offset_l2 = L2Error()
          # self.val_offset_l1 = L1Error()
          # self.test_offset_wl2 = WeightedL2Error()
          # self.test_offset_wl1 = WeightedL1Error()
          # self.test_offset_l2 = L2Error()
          # self.test_offset_l1 = L1Error()
  
          # Metric objects for calculating edge affinity prediction scores
          # on each dataset split
          self.train_affinity_oa = BinaryAccuracy()
          self.train_affinity_f1 = BinaryF1Score()
          self.val_affinity_oa = BinaryAccuracy()
          self.val_affinity_f1 = BinaryF1Score()
          self.test_affinity_oa = BinaryAccuracy()
          self.test_affinity_f1 = BinaryF1Score()
  
          # For averaging losses across batches
          self.train_semantic_loss = MeanMetric()
          self.train_edge_affinity_loss = MeanMetric()
          # self.train_node_offset_loss = MeanMetric()
          self.val_semantic_loss = MeanMetric()
          self.val_edge_affinity_loss = MeanMetric()
          # self.val_node_offset_loss = MeanMetric()
          self.test_semantic_loss = MeanMetric()
          self.test_edge_affinity_loss = MeanMetric()
          # self.test_node_offset_loss = MeanMetric()
  
          # For tracking best-so-far validation metrics
          self.val_map_best = MaxMetric()
          self.val_pq_best = MaxMetric()
          self.val_pqmod_best = MaxMetric()
          self.val_mprec_best = MaxMetric()
          self.val_mrec_best = MaxMetric()
          self.val_instance_miou_best = MaxMetric()
          self.val_instance_oa_best = MaxMetric()
          self.val_instance_macc_best = MaxMetric()
          # self.val_offset_wl2_best = MinMetric()
          # self.val_offset_wl1_best = MinMetric()
          # self.val_offset_l2_best = MinMetric()
          # self.val_offset_l1_best = MinMetric()
          self.val_affinity_oa_best = MaxMetric()
          self.val_affinity_f1_best = MaxMetric()
  
      @property
      def needs_partition(self) -> bool:
          """Whether the `self.partitioner` should be called to compute
          the actual panoptic segmentation. During training, the actual
          partition is not really needed, as we do not learn to partition,
          but learn to predict inputs for the partition step instead. For
          this reason, we save compute and time during training by only
          computing the partition once in a while with
          `self.partition_every_n_epoch`.
          """
          # Get the current epoch. For the validation set, we alter the
          # epoch number so that `partition_every_n_epoch` can align
          # with `check_val_every_n_epoch`. Indeed, it seems the epoch
          # number during the validation step is always one increment
          # ahead
          epoch = self.current_epoch + 1
  
          # If no Trainer attached to the model, run the partition
          if self._trainer is None:
              return True
  
          # Come useful checks to decide whether the partition should be
          # triggered
          k = self.partition_every_n_epoch
          last_epoch = epoch == self.trainer.max_epochs
          first_epoch = epoch == 1
          kth_epoch = epoch % k == 0 if k > 0 else False
  
          # For training, the partition is computed based on
          # `partition_every_n_epoch`, or if we reached the last epoch.
          # The first epoch will be skipped, because trained weights are
          # unlikely to produce interesting inputs for the partition
          if self.trainer.training:
              return (kth_epoch and not first_epoch) or last_epoch
  
          # For validation, we have the same behavior as training, with
          # the difference that if `check_val_every_n_epoch` is larger
          # than `partition_every_n_epoch`, we automatically trigger the
          # partition
          if self.trainer.validating:
              k_val = self.trainer.check_val_every_n_epoch
              nearest_multiple = epoch % k < k_val if k > 0 else False
              if 0 < k <= k_val:
                  return not first_epoch or last_epoch
              else:
                  return (nearest_multiple and not first_epoch) or last_epoch
  
          # For all other Trainer stages, we run the partition by default
          return True
  
      @property
      def needs_instance(self) -> bool:
          """Returns True if the instance segmentation metrics need to be
          computed. In particular, since computing instance metrics can be
          computationally costly, we may want to skip it during training
          by setting `no_instance_metrics_on_train_set=True`, or all the
          time by setting `no_instance_metrics=True`.
          """
          if self.no_instance_metrics:
              return False
  
          if self._trainer is None:
              return self.needs_partition
  
          if self.trainer.training and self.no_instance_metrics_on_train_set:
              return False
  
          return self.needs_partition
  
      def forward(
              self,
              nag: NAG,
              grid: Any = None
      ) -> PanopticSegmentationOutput:
          # Extract features
          x = self.net(nag)
  
          # Compute level-1 or multi-level semantic predictions
          semantic_pred = [head(x_) for head, x_ in zip(self.head, x)] \
              if self.multi_stage_loss else self.head(x)
  
          # Recover level-1 features only
          x = x[0] if self.multi_stage_loss else x
  
          # Compute node offset predictions
          # node_offset_pred = self.node_offset_head(x)
  
          # # Forcefully set 0-offset for nodes with stuff predictions
          # node_logits = semantic_pred[0] if self.multi_stage_loss \
          #     else semantic_pred
          # is_stuff = get_stuff_mask(node_logits, self.stuff_classes)
          # node_offset_pred[is_stuff] = 0
  
          # TODO: OPTIONALLY REMOVE OFFSET
          # node_offset_pred = node_offset_pred * 0
  
          # TODO: offset soft-assigned to 0 based on the predicted
          #  stuff/thing probas. A stuff/thing classification loss could
          #  provide additional supervision
  
          # Compute edge affinity predictions
          # NB: we make edge features symmetric, since we want to compute
          # edge affinity, which is not directed
          x_edge = x[nag[1].obj_edge_index]
          x_edge = torch.cat(
              ((x_edge[0] - x_edge[1]).abs(), (x_edge[0] + x_edge[1]) / 2), dim=1)
          norm_index = torch.zeros(
              x_edge.shape[0], device=x_edge.device, dtype=torch.long)
          edge_affinity_logits = self.edge_affinity_head(
              x_edge, batch=norm_index).squeeze()
  
          # Gather results in an output object
          output = PanopticSegmentationOutput(
              semantic_pred,
              self.stuff_classes,
              edge_affinity_logits,
              # node_offset_pred,
              nag.get_sub_size(1,low = 0))
  
          # Compute the panoptic partition
          output = self._forward_partition(nag, output, grid=grid)
  
          return output
  
      def _forward_partition(
              self,
              nag: NAG,
              output: PanopticSegmentationOutput,
              grid: Any = None,
              force: bool = False
      ) -> PanopticSegmentationOutput:
          """Compute the panoptic partition based on the predicted node
          offsets, node semantic logits, and edge affinity logits.
  
          The partition will only be computed if required. In general,
          during training, the actual partition is not needed for the
          model to be supervised. We only run it once in a while to
          evaluate the panoptic/instance segmentation metrics or tune
          the partition hyperparameters on the train set.
  
          :param nag: NAG object
          :param output: PanopticSegmentationOutput
          :param grid: Dict
              A dictionary containing settings for grid-searching optimal
              partition parameters
          :param force: bool
              Whether to forcefully compute the partition, regardless of
              `self.needs_partition`. This mechanism is typically needed
              during training when we want to store or log predictions for
              a batch of interest at an epoch when `self.needs_partition`
              is False
  
          :return: output
          """
          if not self.needs_partition and not force:
              return output
  
          # Recover some useful information from the NAG and
          # PanopticSegmentationOutput objects
          batch = nag[1].batch
          # node_x = nag[1].pos + output.node_offset_pred
          node_x = nag[1].pos
          node_size = nag.get_sub_size(1, low = 0)
          node_logits = output.logits[0] if output.multi_stage else output.logits
          edge_index = nag[1].obj_edge_index
          edge_affinity_logits = output.edge_affinity_logits
  
          # Compute the instance partition
          # NB: we detach the tensors here: this operation runs on CPU and
          # is non-differentiable
          obj_index = self.partitioner(
              batch,
              node_x.detach(),
              node_logits.detach(),
              self.stuff_classes,
              node_size,
              edge_index,
              edge_affinity_logits.detach(),
              grid=grid)
  
          # Store the results in the output object
          output.obj_index_pred = obj_index
  
          return output
  
      def on_fit_start(self) -> None:
          super().on_fit_start()
  
          # Get the LightningDataModule stuff classes and make sure it
          # matches self.stuff_classes. We could also forcefully update
          # the LightningModule with this new information, but it could
          # easily become tedious to track all places where stuff_classes
          # affects the LightningModule object.
          stuff_classes = self.trainer.datamodule.test_dataset.stuff_classes
          assert sorted(stuff_classes) == sorted(self.stuff_classes), \
              f'LightningModule has the following stuff classes ' \
              f'{self.stuff_classes} while the LightningDataModule has ' \
              f'{stuff_classes}.'
  
      def on_train_start(self) -> None:
          # By default, lightning executes validation step sanity checks
          # before training starts, so we need to make sure `*_best`
          # metrics do not store anything from these checks
          super().on_train_start()
          self.val_panoptic.reset()
          self.val_semantic.reset()
          self.val_instance.reset()
          # self.val_offset_wl2.reset()
          # self.val_offset_wl1.reset()
          # self.val_offset_l2.reset()
          # self.val_offset_l1.reset()
          self.val_affinity_oa.reset()
          self.val_affinity_f1.reset()
          self.val_map_best.reset()
          self.val_pq_best.reset()
          self.val_pqmod_best.reset()
          self.val_mprec_best.reset()
          self.val_mrec_best.reset()
          self.val_instance_miou_best.reset()
          self.val_instance_oa_best.reset()
          self.val_instance_macc_best.reset()
          # self.val_offset_wl2_best.reset()
          # self.val_offset_wl1_best.reset()
          # self.val_offset_l2_best.reset()
          # self.val_offset_l1_best.reset()
          self.val_affinity_oa_best.reset()
          self.val_affinity_f1_best.reset()
          self.train_multi_partition_storage = []
  
      def _create_empty_output(self, nag: NAG) -> PanopticSegmentationOutput:
          """Local helper method to initialize an empty output for
          multi-run prediction.
          """
          # Prepare empty output for semantic segmentation
          output_semseg = super()._create_empty_output(nag)
  
          # Prepare empty edge affinity and node offset outputs
          num_edges = nag[1].obj_edge_index.shape[1]
          edge_affinity_logits = torch.zeros(num_edges, device=nag.device)
          # node_offset_pred = torch.zeros_like(nag[1].pos)
          node_size = nag.get_sub_size(1, low=0)
  
          return PanopticSegmentationOutput(
              output_semseg.logits,
              self.stuff_classes,
              edge_affinity_logits,
              # node_offset_pred,
              node_size)
  
      @staticmethod
      def _update_output_multi(
              output_multi: PanopticSegmentationOutput,
              nag: NAG, output: PanopticSegmentationOutput,
              nag_transformed: NAG,
              key: str
      ) -> PanopticSegmentationOutput:
          """Local helper method to accumulate multiple predictions on
          the same--or part of the same--point cloud.
          """
          raise NotImplementedError(
              "The current implementation does not properly support multi-run "
              "for instance/panoptic segmentation")
  
          # Update semantic segmentation logits only
          output_multi = super()._update_output_multi(
              output_multi, nag, output, nag_transformed, key)
  
          # Update node-wise predictions
          # TODO: this is INCORRECT accumulation of node offsets. Need to
          #  define the mean, not the mean of the successive predictions
          node_id = nag_transformed[1][key]
          output_multi.node_offset_pred[node_id] = \
              (output_multi.node_offset_pred[node_id]
               + output.node_offset_pred) / 2
  
          # Update edge-wise predictions
          edge_index_1 = nag[1].obj_edge_index
          edge_index_2 = node_id[nag_transformed[1].obj_edge_index]
          base = nag[1].num_points + 1
          edge_id_1 = edge_index_1[0] * base + edge_index_1[1]
          edge_id_2 = edge_index_2[0] * base + edge_index_2[1]
          edge_id_cat = consecutive_cluster(torch.cat((edge_id_1, edge_id_2)))[0]
          edge_id_1 = edge_id_cat[:edge_id_1.numel()]
          edge_id_2 = edge_id_cat[edge_id_1.numel():]
          pivot = torch.zeros(base ** 2, device=output.edge_affinity_logits)
          pivot[edge_id_1] = output_multi.edge_affinity_logits
          # TODO: this is INCORRECT accumulation of node offsets. Need to
          #  define the mean, not the mean of the successive predictions
          pivot[edge_id_2] = (pivot[edge_id_2] + output.edge_affinity_logits) / 2
          output_multi.edge_affinity_logits = pivot[edge_id_1]
  
          return output_multi
  
      @staticmethod
      def _propagate_output_to_unseen_neighbors(
              output: PanopticSegmentationOutput,
              nag: NAG, seen: torch.Tensor,
              neighbors: torch.Tensor
      ) -> PanopticSegmentationOutput:
          """Local helper method to propagate predictions to unseen
          neighbors.
          """
          # Propagate semantic segmentation to neighbors
          output = super()._propagate_output_to_unseen_neighbors(
              output, nag, seen, neighbors)
  
          # Heuristic for unseen node offsets: unseen nodes take the same
          # offset as their nearest neighbor
          seen_idx = torch.where(seen)[0]
          unseen_idx = torch.where(~seen)[0]
          output.node_offset_pred[unseen_idx] = \
              output.node_offset_pred[seen_idx][neighbors]
  
          # Heuristic for unseen edge affinity predictions: we set the
          # edge affinity to 0.5
          seen_edge = nag[1].obj_edge_index[seen]
          unseen_edge_idx = torch.where(~seen_edge)[0]
          output.edge_affinity_logits[unseen_edge_idx] = 0.5
  
          return output
  
      def get_target(
              self,
              nag: NAG,
              output: PanopticSegmentationOutput
      ) -> PanopticSegmentationOutput:
          """Recover the target data for semantic and panoptic
          segmentation and store it in the `output` object.
  
          More specifically:
            - label histogram(s) for semantic segmentation will be saved
              in `output.y_hist`
            - instance graph data `obj_edge_index` and `obj_edge_affinity`
              will be saved in `output.obj_edge_index` and
              `output.obj_edge_affinity`, respectively
            - node positions `pos` and `obj_pos` will be saved in
              `output.pos` and `output.obj_pos`, respectively. Besides,
              the `output.obj_offset` will carry the target offset,
              computed from those
          """
          # Recover targets for semantic segmentation
          output = super().get_target(nag, output)
  
          # Recover targets for instance/panoptic segmentation
          output.obj_edge_index = getattr(nag[1], 'obj_edge_index', None)
          output.obj_edge_affinity = getattr(nag[1], 'obj_edge_affinity', None)
          output.pos = nag[1].pos
          output.obj_pos = getattr(nag[1], 'obj_pos', None)
          output.obj = nag[1].obj
  
          return output
  
      def _edge_affinity_weights(
              self,
              is_same_class: torch.Tensor,
              is_same_obj: torch.Tensor
      ) -> torch.Tensor:
          """Helper function to compute edge weights to be used by the
          edge affinity loss. Each edge may have a different weight, based
          on whether its source and target nodes have the same class or
          belong to the same object. The weight given to each case
          (same-class and same-object, same-class and different object,
          etc..) is specified in `edge_affinity_loss_weights`.
  
          :param is_same_class: BoolTensor
              Mask indicating edges between nodes of the same semantic
              class
          :param is_same_obj: BoolTensor
              Mask indicating edges between nodes of the same object
          """
          # Recover the weights given to each case
          w = self.hparams.edge_affinity_loss_weights
  
          # If edge_affinity_loss_weights was not specified, no weighting
          # scheme will be applied to the edges
          if w is None or not len(w) == 4:
              return None
  
          # Compute the weight for each edge
          edge_weight = torch.ones_like(is_same_class).float()
          edge_weight[is_same_class * is_same_obj] = w[0]
          edge_weight[is_same_class * ~is_same_obj] = w[1]
          edge_weight[~is_same_class * is_same_obj] = w[2]
          edge_weight[~is_same_class * ~is_same_obj] = w[3]
          return edge_weight
  
      def model_step(
              self,
              batch: NAG
      ) -> Tuple[torch.Tensor, PanopticSegmentationOutput]:
          # Loss and predictions for semantic segmentation
          semantic_loss, output = super().model_step(batch)
  
          # Cannot compute losses if some target data are missing
          if not output.has_target:
              return None, output
  
          # Compute the node offset loss, weighted by the node size
          # node_offset_loss = self.node_offset_criterion(
          #     *output.sanitized_node_offsets)
  
          # Compute the edge affinity loss
          edge_affinity_pred, edge_affinity_target, is_same_class, is_same_obj = \
              output.sanitized_edge_affinities()
          edge_weight = self._edge_affinity_weights(is_same_class, is_same_obj)
          edge_affinity_loss = self.edge_affinity_criterion(
              edge_affinity_pred, edge_affinity_target, edge_weight)
  
          # Combine the losses together
          # TODO: remove node offset cleanly
          # loss = semantic_loss \
          #        + self.hparams.edge_affinity_loss_lambda * edge_affinity_loss \
          #        + self.hparams.node_offset_loss_lambda * node_offset_loss
          loss = semantic_loss \
                 + self.hparams.edge_affinity_loss_lambda * edge_affinity_loss
  
          # Save individual losses in the output object
          output.semantic_loss = semantic_loss
          # TODO: remove node offset cleanly
          # output.node_offset_loss = 0
          output.edge_affinity_loss = edge_affinity_loss
  
          return loss, output
  
      def train_step_update_metrics(
              self,
              loss: torch.Tensor,
              output: PanopticSegmentationOutput
      ) -> None:
          """Update train metrics with the content of the output object.
          """
          # Update semantic segmentation metrics
          super().train_step_update_metrics(loss, output)
  
          # Update instance and panoptic metrics
          if self.needs_partition and not output.has_multi_instance_pred:
              obj_score, obj_y, instance_data = output.panoptic_pred()
              obj_score = obj_score.detach().cpu()
              obj_y = obj_y.detach()
              obj_hist = instance_data.target_label_histogram(self.num_classes)
              self.train_panoptic.update(obj_y.cpu(), instance_data.cpu())
              self.train_semantic(obj_y, obj_hist)
              if self.needs_instance:
                  self.train_instance.update(obj_score, obj_y, instance_data.cpu())
          elif self.needs_partition:
              logits = output.logits[0] if output.multi_stage else output.logits
              storage = PartitionParameterSearchStorage(
                      logits.detach().cpu(),
                      self.stuff_classes,
                      output.node_size.detach().cpu(),
                      output.edge_affinity_logits.detach().cpu(),
                      output.obj.cpu(),
                      [(v[0], v[1].detach().cpu()) for v in output.obj_index_pred])
              self.train_multi_partition_storage.append(storage)
  
          # Update tracked losses
          self.train_semantic_loss(output.semantic_loss.detach())
          # self.train_node_offset_loss(output.node_offset_loss.detach())
          self.train_edge_affinity_loss(output.edge_affinity_loss.detach())
  
          # Update node offset metrics
          # node_offset_pred, node_offset, node_size = output.sanitized_node_offsets
          # node_offset_pred = node_offset_pred.detach()
          # node_offset = node_offset.detach()
          # node_size = node_size.detach()
          # self.train_offset_wl2(node_offset_pred, node_offset, node_size)
          # self.train_offset_wl1(node_offset_pred, node_offset, node_size)
          # self.train_offset_l2(node_offset_pred, node_offset)
          # self.train_offset_l1(node_offset_pred, node_offset)
  
          # Update edge affinity metrics
          ea_pred, ea_target, is_same_class, is_same_obj = \
              output.sanitized_edge_affinities()
          ea_pred = ea_pred.detach()
          ea_target_binary = (ea_target.detach() > 0.5).long()
          self.train_affinity_oa(ea_pred, ea_target_binary)
          self.train_affinity_f1(ea_pred, ea_target_binary)
  
      def train_step_log_metrics(self) -> None:
          """Log train metrics after a single step with the content of the
          output object.
          """
          super().train_step_log_metrics()
          self.log(
              "train/semantic_loss", self.train_semantic_loss, on_step=False,
              on_epoch=True, prog_bar=True)
          # self.log(
          #     "train/node_offset_loss", self.train_node_offset_loss, on_step=False,
          #     on_epoch=True, prog_bar=True)
          self.log(
              "train/edge_affinity_loss", self.train_edge_affinity_loss, on_step=False,
              on_epoch=True, prog_bar=True)
  
      def on_train_epoch_end(self) -> None:
          # Log semantic segmentation metrics and reset confusion matrix
          super().on_train_epoch_end()
  
          # TODO: support logging panoptic metrics for DDP
          if self.trainer.num_devices > 1:
              log.warning(
                  "Panoptic and instance segmentation metrics are not guaranteed "
                  "to be well-behaved on DDP yet.")
  
          if self.needs_partition:
              # If multiple partitions settings were tested during the
              # epoch, this will search for the best one, update the
              # internal states of train metrics with related predictions,
              # and update the partitioner's settings
              setting = self._compute_best_partition_settings()[0]
  
              # Compute the instance and panoptic metrics
              panoptic_results = self.train_panoptic.compute()
              if self.needs_instance:
                  instance_results = self.train_instance.compute()
  
              # Gather tracked metrics
              pq = panoptic_results.pq
              sq = panoptic_results.sq
              rq = panoptic_results.rq
              pq_thing = panoptic_results.pq_thing
              pq_stuff = panoptic_results.pq_stuff
              pqmod = panoptic_results.pq_modified
              mprec = panoptic_results.mean_precision
              mrec = panoptic_results.mean_recall
              pq_per_class = panoptic_results.pq_per_class
              if self.needs_instance:
                  map = instance_results.map
                  map_50 = instance_results.map_50
                  map_75 = instance_results.map_75
                  map_per_class = instance_results.map_per_class
  
              # Log metrics
              self.log("train/pq", 100 * pq, prog_bar=True)
              self.log("train/sq", 100 * sq, prog_bar=True)
              self.log("train/rq", 100 * rq, prog_bar=True)
              self.log("train/pq_thing", 100 * pq_thing, prog_bar=True)
              self.log("train/pq_stuff", 100 * pq_stuff, prog_bar=True)
              self.log("train/pqmod", 100 * pqmod, prog_bar=True)
              self.log("train/mprec", 100 * mprec, prog_bar=True)
              self.log("train/mrec", 100 * mrec, prog_bar=True)
              self.log("train/instance_miou", self.train_semantic.miou(), prog_bar=True)
              self.log("train/instance_oa", self.train_semantic.oa(), prog_bar=True)
              self.log("train/instance_macc", self.train_semantic.macc(), prog_bar=True)
              for iou, seen, name in zip(*self.train_semantic.iou(), self.class_names):
                  if seen:
                      self.log(f"train/instance_iou_{name}", iou, prog_bar=True)
              if self.needs_instance:
                  self.log("train/map", 100 * map, prog_bar=True)
                  self.log("train/map_50", 100 * map_50, prog_bar=True)
                  self.log("train/map_75", 100 * map_75, prog_bar=True)
              for pq_c, name in zip(pq_per_class, self.class_names):
                  self.log(f"train/pq_{name}", 100 * pq_c, prog_bar=True)
              if self.needs_instance:
                  for map_c, name in zip(map_per_class, self.class_names):
                      self.log(f"train/map_{name}", 100 * map_c, prog_bar=True)
              if setting is not None:
                  for k, v in setting.items():
                      self.log(f"partition_settings/{k}", v, prog_bar=True)
  
          # Log metrics
          # self.log("train/offset_wl2", self.train_offset_wl2.compute(), prog_bar=True)
          # self.log("train/offset_wl1", self.train_offset_wl1.compute(), prog_bar=True)
          # self.log("train/offset_l2", self.train_offset_l2.compute(), prog_bar=True)
          # self.log("train/offset_l1", self.train_offset_l1.compute(), prog_bar=True)
          self.log("train/affinity_oa", 100 * self.train_affinity_oa.compute(), prog_bar=True)
          self.log("train/affinity_f1", 100 * self.train_affinity_f1.compute(), prog_bar=True)
  
          # Reset metrics accumulated over the last epoch
          # self.train_offset_wl2.reset()
          # self.train_offset_wl1.reset()
          # self.train_offset_l2.reset()
          # self.train_offset_l1.reset()
          self.train_affinity_oa.reset()
          self.train_affinity_f1.reset()
          self.train_panoptic.reset()
          self.train_semantic.reset()
          self.train_instance.reset()
  
      def _compute_best_partition_settings(
              self,
              monitor: str = 'pq',
              maximize: bool = True
      ) -> Tuple[Dict, float]:
          """Compute the best partition settings from
          `self.train_multi_partition_storage`. This will have the
          following internal effects:
            - `self.partitioner` will be updated with the settings which
              produced the best metrics on the epoch
            - `self.train_panoptic` will be updated with the batch
              predictions with the best settings
            - `self.train_instance` will be updated with the batch
              predictions with the best settings, if required
  
          :param monitor: str
              The metric based on which we will select the best settings
          :param maximize: bool
              Whether the monitored metric should be maximized or
              minimized
          :return:
          """
          # Nothing happens if multi-partition was not activated during
          # the epoch
          if len(self.train_multi_partition_storage) == 0:
              return None, None
  
          # Reset the instance and panoptic metrics, these will be used to
          # compute metric performance
          self.train_panoptic.reset()
          self.train_instance.reset()
  
          # Check whether the metric to monitor is for the semantic or
          # panoptic segmentation task
          if monitor in self.train_panoptic.__slots__:
              task = 'panoptic'
              meter = self.train_panoptic
          elif monitor in self.train_instance.__slots__:
              task = 'instance'
              meter = self.train_instance
          else:
              raise ValueError(f"Unknown metric, cannot monitor '{monitor}'.")
          if task == 'instance' and not self.needs_instance:
              raise ValueError(
                  'Cannot compute the best partition settings on the train set '
                  'based on instance metrics if `self.needs_instance` is False')
  
          # Recover from the first PartitionParameterSearchStorage, which
          # settings were explored
          settings = self.train_multi_partition_storage[0].settings
  
          # Compute the metric for each partition setting while tracking
          # the best setting
          best_metric = -torch.inf if maximize else torch.inf
          best_setting = None
          for s in settings:
  
              # Accumulate batch predictions in the meter
              for storage in self.train_multi_partition_storage:
                  obj_score, obj_y, instance_data = \
                      storage.panoptic_pred(s)
                  if task == 'panoptic':
                      meter.update(obj_y, instance_data)
                  else:
                      meter.update(obj_score, obj_y, instance_data)
  
              # Compute the monitored metric on the whole epoch
              metric = getattr(meter.compute(), monitor)
  
              # Update the best metric and settings
              condition = (metric > best_metric) if maximize \
                  else (metric < best_setting)
              if condition:
                  best_metric = metric
                  best_setting = s
  
              # Reset the meter to avoid mixing predictions of different
              # settings
              meter.reset()
  
          # Update the partitioner with the best metrics
          for k, v in best_setting.items():
              setattr(self.partitioner, k, v)
  
          # Update the train meters with the data for computation of
          # logged metrics with the accumulated data from the best
          # setting, thus mimicking a normal epoch with a single partition
          # prediction per batch
          for storage in self.train_multi_partition_storage:
              obj_score, obj_y, instance_data = \
                  storage.panoptic_pred(best_setting)
              obj_hist = instance_data.target_label_histogram(self.num_classes)
              self.train_panoptic.update(obj_y, instance_data)
              self.train_semantic(
                  obj_y.to(self.train_semantic.device),
                  obj_hist.to(self.train_semantic.device))
              if self.needs_instance:
                  self.train_instance.update(obj_score, obj_y, instance_data)
  
          return best_setting, best_metric
  
      def validation_step_update_metrics(
              self,
              loss: torch.Tensor,
              output: PanopticSegmentationOutput
      ) -> None:
          """Update validation metrics with the content of the output
          object.
          """
          # Update semantic segmentation metrics
          super().validation_step_update_metrics(loss, output)
  
          # Update instance and panoptic metrics
          if self.needs_partition:
              obj_score, obj_y, instance_data = output.panoptic_pred()
              obj_score = obj_score.detach().cpu()
              obj_y = obj_y.detach()
              obj_hist = instance_data.target_label_histogram(self.num_classes)
              self.val_panoptic.update(obj_y.cpu(), instance_data.cpu())
              self.val_semantic(obj_y, obj_hist)
              if self.needs_instance:
                  self.val_instance.update(obj_score, obj_y, instance_data.cpu())
  
          # Update tracked losses
          self.val_semantic_loss(output.semantic_loss.detach())
          # self.val_node_offset_loss(output.node_offset_loss.detach())
          self.val_edge_affinity_loss(output.edge_affinity_loss.detach())
  
          # Update node offset metrics
          # node_offset_pred, node_offset, node_size = output.sanitized_node_offsets
          # node_offset_pred = node_offset_pred.detach()
          # node_offset = node_offset.detach()
          # node_size = node_size.detach()
          # self.val_offset_wl2(node_offset_pred, node_offset, node_size)
          # self.val_offset_wl1(node_offset_pred, node_offset, node_size)
          # self.val_offset_l2(node_offset_pred, node_offset)
          # self.val_offset_l1(node_offset_pred, node_offset)
  
          # Update edge affinity metrics
          ea_pred, ea_target, is_same_class, is_same_obj = \
              output.sanitized_edge_affinities()
          ea_pred = ea_pred.detach()
          ea_target_binary = (ea_target.detach() > 0.5).long()
          self.val_affinity_oa(ea_pred, ea_target_binary)
          self.val_affinity_f1(ea_pred, ea_target_binary)
  
      def validation_step_log_metrics(self) -> None:
          """Log validation metrics after a single step with the content
          of the output object.
          """
          super().validation_step_log_metrics()
          self.log(
              "val/semantic_loss", self.val_semantic_loss, on_step=False,
              on_epoch=True, prog_bar=True)
          # self.log(
          #     "val/node_offset_loss", self.val_node_offset_loss, on_step=False,
          #     on_epoch=True, prog_bar=True)
          self.log(
              "val/edge_affinity_loss", self.val_edge_affinity_loss, on_step=False,
              on_epoch=True, prog_bar=True)
  
      def on_validation_epoch_end(self) -> None:
          # Log semantic segmentation metrics and reset confusion matrix
          super().on_validation_epoch_end()
  
          # TODO: support logging panoptic metrics for DDP
          if self.trainer.num_devices > 1:
              log.warning(
                  "Panoptic and instance segmentation metrics are not guaranteed "
                  "to be well-behaved on DDP yet.")
  
          if self.needs_partition:
              # Compute the instance and panoptic metrics
              panoptic_results = self.val_panoptic.compute()
              if self.needs_instance:
                  instance_results = self.val_instance.compute()
  
              # Gather tracked metrics
              pq = panoptic_results.pq
              sq = panoptic_results.sq
              rq = panoptic_results.rq
              pq_thing = panoptic_results.pq_thing
              pq_stuff = panoptic_results.pq_stuff
              pqmod = panoptic_results.pq_modified
              mprec = panoptic_results.mean_precision
              mrec = panoptic_results.mean_recall
              pq_per_class = panoptic_results.pq_per_class
              if self.needs_instance:
                  map = instance_results.map
                  map_50 = instance_results.map_50
                  map_75 = instance_results.map_75
                  map_per_class = instance_results.map_per_class
  
              # Log metrics
              self.log("val/pq", 100 * pq, prog_bar=True)
              self.log("val/sq", 100 * sq, prog_bar=True)
              self.log("val/rq", 100 * rq, prog_bar=True)
              self.log("val/pq_thing", 100 * pq_thing, prog_bar=True)
              self.log("val/pq_stuff", 100 * pq_stuff, prog_bar=True)
              self.log("val/pqmod", 100 * pqmod, prog_bar=True)
              self.log("val/mprec", 100 * mprec, prog_bar=True)
              self.log("val/mrec", 100 * mrec, prog_bar=True)
              instance_miou = self.val_semantic.miou()
              instance_oa = self.val_semantic.oa()
              instance_macc = self.val_semantic.macc()
              self.log("val/instance_miou", instance_miou, prog_bar=True)
              self.log("val/instance_oa", instance_oa, prog_bar=True)
              self.log("val/instance_macc", instance_macc, prog_bar=True)
              for iou, seen, name in zip(*self.val_semantic.iou(), self.class_names):
                  if seen:
                      self.log(f"val/instance_iou_{name}", iou, prog_bar=True)
              if self.needs_instance:
                  self.log("val/map", 100 * map, prog_bar=True)
                  self.log("val/map_50", 100 * map_50, prog_bar=True)
                  self.log("val/map_75", 100 * map_75, prog_bar=True)
              for pq_c, name in zip(pq_per_class, self.class_names):
                  self.log(f"val/pq_{name}", 100 * pq_c, prog_bar=True)
              if self.needs_instance:
                  for map_c, name in zip(map_per_class, self.class_names):
                      self.log(f"val/map_{name}", 100 * map_c, prog_bar=True)
  
              # Update best-so-far metrics
              self.val_pq_best(pq)
              self.val_pqmod_best(pqmod)
              self.val_mprec_best(mprec)
              self.val_mrec_best(mrec)
              if self.needs_instance:
                  self.val_map_best(map)
              self.val_instance_miou_best(instance_miou)
              self.val_instance_oa_best(instance_oa)
              self.val_instance_macc_best(instance_macc)
  
              # Log best-so-far metrics, using `.compute()` instead of passing
              # the whole torchmetrics object, because otherwise metric would
              # be reset by lightning after each epoch
              self.log("val/pq_best", 100 * self.val_pq_best.compute(), prog_bar=True)
              self.log("val/pqmod_best", 100 * self.val_pqmod_best.compute(), prog_bar=True)
              self.log("val/mprec_best", 100 * self.val_mprec_best.compute(), prog_bar=True)
              self.log("val/mrec_best", 100 * self.val_mrec_best.compute(), prog_bar=True)
              if self.needs_instance:
                  self.log("val/map_best", 100 * self.val_map_best.compute(), prog_bar=True)
              self.log("val/instance_miou_best", self.val_instance_miou_best.compute(), prog_bar=True)
              self.log("val/instance_oa_best", self.val_instance_oa_best.compute(), prog_bar=True)
              self.log("val/instance_macc_best", self.val_instance_macc_best.compute(), prog_bar=True)
  
          # Compute the metrics tracked for model selection on validation
          # offset_wl2 = self.val_offset_wl2.compute()
          # offset_wl1 = self.val_offset_wl1.compute()
          # offset_l2 = self.val_offset_l2.compute()
          # offset_l1 = self.val_offset_l1.compute()
          affinity_oa = self.val_affinity_oa.compute()
          affinity_f1 = self.val_affinity_f1.compute()
  
          # Log metrics
          # self.log("val/offset_wl2", offset_wl2, prog_bar=True)
          # self.log("val/offset_wl1", offset_wl1, prog_bar=True)
          # self.log("val/offset_l2", offset_l2, prog_bar=True)
          # self.log("val/offset_l1", offset_l1, prog_bar=True)
          self.log("val/affinity_oa", 100 * affinity_oa, prog_bar=True)
          self.log("val/affinity_f1", 100 * affinity_f1, prog_bar=True)
  
          # Update best-so-far metrics
          # self.val_offset_wl2_best(offset_wl2)
          # self.val_offset_wl1_best(offset_wl1)
          # self.val_offset_l2_best(offset_l2)
          # self.val_offset_l1_best(offset_l1)
          self.val_affinity_oa_best(affinity_oa)
          self.val_affinity_f1_best(affinity_f1)
  
          # Log best-so-far metrics, using `.compute()` instead of passing
          # the whole torchmetrics object, because otherwise metric would
          # be reset by lightning after each epoch
          # self.log("val/offset_wl2_best", self.val_offset_wl2_best.compute(), prog_bar=True)
          # self.log("val/offset_wl1_best", self.val_offset_wl1_best.compute(), prog_bar=True)
          # self.log("val/offset_l2_best", self.val_offset_l2_best.compute(), prog_bar=True)
          # self.log("val/offset_l1_best", self.val_offset_l1_best.compute(), prog_bar=True)
          self.log("val/affinity_oa_best", 100 * self.val_affinity_oa_best.compute(), prog_bar=True)
          self.log("val/affinity_f1_best", 100 * self.val_affinity_f1_best.compute(), prog_bar=True)
  
          # Reset metrics accumulated over the last epoch
          # self.val_offset_wl2.reset()
          # self.val_offset_wl1.reset()
          # self.val_offset_l2.reset()
          # self.val_offset_l1.reset()
          self.val_affinity_oa.reset()
          self.val_affinity_f1.reset()
          self.val_panoptic.reset()
          self.val_semantic.reset()
          self.val_instance.reset()
  
      def test_step_update_metrics(
              self,
              loss: torch.Tensor,
              output: PanopticSegmentationOutput
      ) -> None:
          """Update test metrics with the content of the output object.
          """
          # Update semantic segmentation metrics
          super().test_step_update_metrics(loss, output)
  
          # If the test set misses targets, we keep track of it, to skip
          # metrics computation on the test set
          if not self.test_has_target:
              return
  
          # Update instance and panoptic metrics
          if self.needs_partition:
              obj_score, obj_y, instance_data = output.panoptic_pred()
              obj_score = obj_score.detach().cpu()
              obj_y = obj_y.detach()
              obj_hist = instance_data.target_label_histogram(self.num_classes)
              self.test_panoptic.update(obj_y.cpu(), instance_data.cpu())
              self.test_semantic(obj_y, obj_hist)
              if self.needs_instance:
                  self.test_instance.update(obj_score, obj_y, instance_data.cpu())
  
          # Update tracked losses
          self.test_semantic_loss(output.semantic_loss.detach())
          # self.test_node_offset_loss(output.node_offset_loss.detach())
          self.test_edge_affinity_loss(output.edge_affinity_loss.detach())
  
          # Update node offset metrics
          # node_offset_pred, node_offset, node_size = output.sanitized_node_offsets
          # node_offset_pred = node_offset_pred.detach()
          # node_offset = node_offset.detach()
          # node_size = node_size.detach()
          # self.test_offset_wl2(node_offset_pred, node_offset, node_size)
          # self.test_offset_wl1(node_offset_pred, node_offset, node_size)
          # self.test_offset_l2(node_offset_pred, node_offset)
          # self.test_offset_l1(node_offset_pred, node_offset)
  
          # Update edge affinity metrics
          ea_pred, ea_target, is_same_class, is_same_obj = \
              output.sanitized_edge_affinities()
          ea_pred = ea_pred.detach()
          ea_target_binary = (ea_target.detach() > 0.5).long()
          self.test_affinity_oa(ea_pred, ea_target_binary)
          self.test_affinity_f1(ea_pred, ea_target_binary)
  
      def test_step_log_metrics(self) -> None:
          """Log test metrics after a single step with the content of the
          output object.
          """
          super().test_step_log_metrics()
  
          # If the test set misses targets, we keep track of it, to skip
          # metrics computation on the test set
          if not self.test_has_target:
              return
  
          self.log(
              "test/semantic_loss", self.test_semantic_loss, on_step=False,
              on_epoch=True, prog_bar=True)
          # self.log(
          #     "test/node_offset_loss", self.test_node_offset_loss, on_step=False,
          #     on_epoch=True, prog_bar=True)
          self.log(
              "test/edge_affinity_loss", self.test_edge_affinity_loss, on_step=False,
              on_epoch=True, prog_bar=True)
  
      def on_test_epoch_end(self) -> None:
          # Log semantic segmentation metrics and reset confusion matrix
          super().on_test_epoch_end()
  
          # If test set misses target data, reset metrics and skip logging
          if not self.test_has_target:
              # self.test_offset_wl2.reset()
              # self.test_offset_wl1.reset()
              # self.test_offset_l2.reset()
              # self.test_offset_l1.reset()
              self.test_affinity_oa.reset()
              self.test_affinity_f1.reset()
              self.test_panoptic.reset()
              self.test_semantic.reset()
              self.test_instance.reset()
              return
  
          # TODO: support logging panoptic metrics for DDP
          if self.trainer.num_devices > 1:
              log.warning(
                  "Panoptic and instance segmentation metrics are not guaranteed "
                  "to be well-behaved on DDP yet.")
  
          if self.needs_partition:
              # Compute the instance and panoptic metrics
              panoptic_results = self.test_panoptic.compute()
              if self.needs_instance:
                  instance_results = self.test_instance.compute()
  
              # Gather tracked metrics
              pq = panoptic_results.pq
              sq = panoptic_results.sq
              rq = panoptic_results.rq
              pq_thing = panoptic_results.pq_thing
              pq_stuff = panoptic_results.pq_stuff
              pqmod = panoptic_results.pq_modified
              mprec = panoptic_results.mean_precision
              mrec = panoptic_results.mean_recall
              pq_per_class = panoptic_results.pq_per_class
              if self.needs_instance:
                  map = instance_results.map
                  map_50 = instance_results.map_50
                  map_75 = instance_results.map_75
                  map_per_class = instance_results.map_per_class
  
              # Log metrics
              self.log("test/pq", 100 * pq, prog_bar=True)
              self.log("test/sq", 100 * sq, prog_bar=True)
              self.log("test/rq", 100 * rq, prog_bar=True)
              self.log("test/pq_thing", 100 * pq_thing, prog_bar=True)
              self.log("test/pq_stuff", 100 * pq_stuff, prog_bar=True)
              self.log("test/pqmod", 100 * pqmod, prog_bar=True)
              self.log("test/mprec", 100 * mprec, prog_bar=True)
              self.log("test/mrec", 100 * mrec, prog_bar=True)
              self.log("test/instance_miou", self.test_semantic.miou(), prog_bar=True)
              self.log("test/instance_oa", self.test_semantic.oa(), prog_bar=True)
              self.log("test/instance_macc", self.test_semantic.macc(), prog_bar=True)
              for iou, seen, name in zip(*self.test_semantic.iou(), self.class_names):
                  if seen:
                      self.log(f"test/instance_iou_{name}", iou, prog_bar=True)
              if self.needs_instance:
                  self.log("test/map", 100 * map, prog_bar=True)
                  self.log("test/map_50", 100 * map_50, prog_bar=True)
                  self.log("test/map_75", 100 * map_75, prog_bar=True)
              for pq_c, name in zip(pq_per_class, self.class_names):
                  self.log(f"test/pq_{name}", 100 * pq_c, prog_bar=True)
              if self.needs_instance:
                  for map_c, name in zip(map_per_class, self.class_names):
                      self.log(f"test/map_{name}", 100 * map_c, prog_bar=True)
  
          # Log metrics
          # self.log("test/offset_wl2", self.test_offset_wl2.compute(), prog_bar=True)
          # self.log("test/offset_wl1", self.test_offset_wl1.compute(), prog_bar=True)
          # self.log("test/offset_l2", self.test_offset_l2.compute(), prog_bar=True)
          # self.log("test/offset_l1", self.test_offset_l1.compute(), prog_bar=True)
          self.log("test/affinity_oa", 100 * self.test_affinity_oa.compute(), prog_bar=True)
          self.log("test/affinity_f1", 100 * self.test_affinity_f1.compute(), prog_bar=True)
  
          # Reset metrics accumulated over the last epoch
          # self.test_offset_wl2.reset()
          # self.test_offset_wl1.reset()
          # self.test_offset_l2.reset()
          # self.test_offset_l1.reset()
          self.test_affinity_oa.reset()
          self.test_affinity_f1.reset()
          self.test_panoptic.reset()
          self.test_semantic.reset()
          self.test_instance.reset()
  
      def track_batch(
              self,
              batch: NAG,
              batch_idx: int,
              output: PanopticSegmentationOutput,
              folder: str = None
      ) -> None:
          """Store a batch prediction to disk. The corresponding `NAG`
          object will be populated with panoptic segmentation predictions
          for:
          - levels 1+ if `multi_stage` output (i.e. loss supervision on
            levels 1 and above)
          - only level 1 otherwise
  
          Besides, we also pre-compute the level-0 predictions as this is
          frequently required for downstream tasks. However, we choose not
          to compute the full-resolution predictions for the sake of disk
          memory.
  
          If a `folder` is provided, the NAG will be saved there under:
            <folder>/predictions/<stage>/<epoch>/batch_<batch_idx>.h5
          If not, the folder will be the logger's directory, if any.
          If not, the current working directory will be used.
  
          :param batch: NAG
              Object that will be stored to disk. Before that, the
              model predictions will be added to the attributes of each
              level, to facilitate downstream use of the stored `NAG`
          :param batch_idx: int
              Index of the batch to be stored
          :param output: PanopticSegmentationOutput
               Output of `self.model_step()`
          :param folder: str
              Path where to save the tracked batch. If not provided, the
              logger's saving directory will be used as fallback. If not
              logger is found, the current working directory will be used
          :return:
          """
          # Sanity check in case using multi-run inference
          if not isinstance(batch, NAG):
              raise NotImplementedError(
                  f"Expected as NAG, but received a {type(batch)}. Are you "
                  f"perhaps running multi-run inference ? If so, this is not "
                  f"compatible with batch_saving, please deactivate either one.")
  
          # Compute the panoptic partition if not already done
          if output.obj_index_pred is None:
              output = self._forward_partition(batch, output, force=True)
  
          # Store the output predictions in conveniently-accessible
          # attributes in the NAG, for easy downstream use of the saved
          # object
          sp_y_pred, sp_obj_index_pred, sp_obj_pred = (
              output.superpoint_panoptic_pred())
          
          batch[1].obj_y_pred = sp_y_pred
          batch[1].obj_index_pred = sp_obj_index_pred
          batch[1].obj_pred = sp_obj_pred
          batch[1].edge_affinity_logits = output.edge_affinity_logits
          
          if batch.has_atoms :
              vox_y_pred, vox_obj_index_pred, vox_obj_pred = (
                  output.voxel_panoptic_pred(super_index=batch[0].super_index))
              
              batch[0].obj_y_pred = vox_y_pred
              batch[0].obj_index_pred = vox_obj_index_pred
              batch[0].obj_pred = vox_obj_pred
  
          # Parent behavior for saving semantic segmentation prediction
          super().track_batch(batch, batch_idx, output, folder=folder)
  
      def load_state_dict(self, state_dict: Dict, strict: bool = True) -> None:
          """Basic `load_state_dict` from `torch.nn.Module` with a bit of
          acrobatics due to `criterion.weight`.
  
          This attribute, when present in the `state_dict`, causes
          `load_state_dict` to crash. More precisely, `criterion.weight`
          is holding the per-class weights for classification losses.
          """
          # Special treatment for BCEWithLogitsLoss
          if self.edge_affinity_criterion.pos_weight is not None:
              pos_weight_bckp = self.edge_affinity_criterion.pos_weight
              self.edge_affinity_criterion.pos_weight = None
  
          if 'edge_affinity_criterion.pos_weight' in state_dict.keys():
              pos_weight = state_dict.pop('edge_affinity_criterion.pos_weight')
          else:
              pos_weight = None
  
          # Load the state_dict
          super().load_state_dict(state_dict, strict=strict)
  
          # If need be, assign the class weights to the criterion
          if self.edge_affinity_criterion.pos_weight is not None:
              self.edge_affinity_criterion.pos_weight = pos_weight \
                  if pos_weight is not None else pos_weight_bckp
  
      def _load_from_checkpoint(
              self,
              checkpoint_path: str,
              **kwargs
      ) -> 'PanopticSegmentationModule':
          """Simpler version of `LightningModule.load_from_checkpoint()`
          for easier use: no need to explicitly pass `model.net`,
          `model.criterion`, etc.
          """
          return self.__class__.load_from_checkpoint(
              checkpoint_path,
              net=self.net,
              edge_affinity_head=self.edge_affinity_head,
              partitioner=self.partitioner,
              criterion=self.criterion,
              **kwargs)
  
  
  # TODO: gridsearch instance partition parameters
  
  if __name__ == "__main__":
      import hydra
      import omegaconf
      import pyrootutils
  
      root = str(pyrootutils.setup_root(__file__, pythonpath=True))
      cfg = omegaconf.OmegaConf.load(root + "/configs/model/panoptic/spt-2.yaml")
      _ = hydra.utils.instantiate(cfg)
  