src/utils/instance.py

import sys
import hydra
import torch
import numpy as np
import pandas as pd
import os.path as osp
from tqdm import tqdm
from copy import deepcopy
from itertools import product
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import matplotlib.cm as cm
from torch.nn.functional import one_hot
from torch_geometric.nn.pool.consecutive import consecutive_cluster
from src.utils.hydra import init_config
from src.utils.neighbors import knn_2
from src.utils.graph import to_trimmed
from src.utils.cpu import available_cpu_count
from src.utils.scatter import scatter_mean_weighted
from src.utils.semantic import _set_attribute_preserving_transforms


src_folder = osp.dirname(osp.dirname(osp.abspath(__file__)))
sys.path.append(src_folder)
sys.path.append(osp.join(src_folder, "dependencies/grid_graph/python/bin"))
sys.path.append(osp.join(src_folder, "dependencies/parallel_cut_pursuit/python/wrappers"))


from grid_graph import edge_list_to_forward_star
from cp_d0_dist import cp_d0_dist


__all__ = [
    'generate_random_bbox_data', 'generate_random_segment_data',
    'instance_cut_pursuit', 'oracle_superpoint_clustering', 'get_stuff_mask',
    'compute_panoptic_metrics', 'compute_panoptic_metrics_s3dis_6fold',
    'grid_search_panoptic_partition']


_MAX_NUM_EDGES = 4294967295


def generate_random_bbox_data(
        num_img=1,
        num_classes=1,
        height=128,
        width=128,
        h_split=1,
        w_split=2,
        det_gt_ratio=1):
    # Create some images with a ground truth partition
    instance_images = -torch.ones(num_img, height, width, dtype=torch.long)
    label_images = -torch.ones(num_img, height, width, dtype=torch.long)

    h_gt = height // h_split
    w_gt = width // w_split

    gt_boxes = torch.zeros(num_img * h_split * w_split, 4)
    gt_labels = torch.randint(0, num_classes, (num_img * h_split * w_split,))
    iterator = product(range(num_img), range(h_split), range(w_split))
    for idx, (i_img, i, j) in enumerate(iterator):
        h1 = i * h_gt
        h2 = (i + 1) * h_gt
        w1 = j * w_gt
        w2 = (j + 1) * w_gt
        instance_images[i_img, h1:h2, w1:w2] = idx
        label_images[i_img, h1:h2, w1:w2] = gt_labels[idx]
        gt_boxes[idx] = torch.tensor([h1, w1, h2, w2])

    # Create some random detection boxes
    num_gt = (instance_images.max() + 1).item()
    num_det = int(num_gt * det_gt_ratio)

    i_center_det = torch.randint(0, height, (num_det,))
    j_center_det = torch.randint(0, width, (num_det,))
    h_det = torch.randint(int(h_gt * 0.7), int(h_gt * 1.3), (num_det,))
    w_det = torch.randint(int(w_gt * 0.7), int(w_gt * 1.3), (num_det,))

    det_boxes = torch.vstack([
        (i_center_det - h_det / 2).clamp(min=0),
        (j_center_det - w_det / 2).clamp(min=0),
        (i_center_det + h_det / 2).clamp(max=height),
        (j_center_det + w_det / 2).clamp(max=width)]).T.round()
    det_img_idx = torch.randint(0, num_img, (num_det,))
    det_labels = torch.randint(0, num_classes, (num_det,))
    det_scores = torch.rand(num_det)

    # Display the images stacked along their height (first dim) and draw
    # the box for each detection
    fig, ax = plt.subplots()
    ax.imshow(instance_images.view(-1, width), cmap='jet')
    for idx_det in range(num_det):
        i = det_boxes[idx_det, 0] + det_img_idx[idx_det] * height
        j = det_boxes[idx_det, 1]
        h = det_boxes[idx_det, 2] - det_boxes[idx_det, 0]
        w = det_boxes[idx_det, 3] - det_boxes[idx_det, 1]
        rect = patches.Rectangle(
            (j, i),
            w,
            h,
            linewidth=3,
            edgecolor=cm.nipy_spectral(idx_det / num_det),
            facecolor='none')
        ax.add_patch(rect)
    plt.show()

    # Display the images stacked along their height (first dim) and draw the
    # box for each detection
    fig, ax = plt.subplots()
    ax.imshow(label_images.view(-1, width).float() / num_classes, cmap='jet')
    for idx_det in range(num_det):
        i = det_boxes[idx_det, 0] + det_img_idx[idx_det] * height
        j = det_boxes[idx_det, 1]
        h = det_boxes[idx_det, 2] - det_boxes[idx_det, 0]
        w = det_boxes[idx_det, 3] - det_boxes[idx_det, 1]
        c = cm.nipy_spectral(det_labels[idx_det].float().item() / num_classes)
        rect = patches.Rectangle(
            (j, i),
            w,
            h,
            linewidth=3,
            edgecolor=c,
            facecolor='none')
        ax.add_patch(rect)
    plt.show()

    # Compute the metrics using torchmetrics
    iterator = zip(gt_boxes.view(num_img, -1, 4), gt_labels.view(num_img, -1))
    targets = [
        dict(boxes=boxes, labels=labels)
        for boxes, labels in iterator]

    preds = [
        dict(
            boxes=det_boxes[det_img_idx == i_img],
            labels=det_labels[det_img_idx == i_img],
            scores=det_scores[det_img_idx == i_img])
        for i_img in range(num_img)]

    # For each predicted pixel, we compute the gt object idx, and the gt
    # label, to build an InstanceData.
    # NB: we cannot build this by creating a single pred_idx image,
    # because predictions may overlap in this toy setup, unlike our 3D
    # superpoint partition paradigm...
    pred_idx = []
    gt_idx = []
    gt_y = []
    for idx_det in range(num_det):
        i_img = det_img_idx[idx_det]
        x1, y1, x2, y2 = det_boxes[idx_det].long()
        num_points = (x2 - x1) * (y2 - y1)
        pred_idx.append(torch.full((num_points,), idx_det))
        gt_idx.append(instance_images[i_img, x1:x2, y1:y2].flatten())
        gt_y.append(label_images[i_img, x1:x2, y1:y2].flatten())
    pred_idx = torch.cat(pred_idx)
    gt_idx = torch.cat(gt_idx)
    gt_y = torch.cat(gt_y)
    count = torch.ones_like(pred_idx)

    from src.data.instance import InstanceData
    instance_data = InstanceData(pred_idx, gt_idx, count, gt_y, dense=True)

    return targets, preds, gt_idx, gt_y, count, instance_data


def generate_single_random_segment_image(
        num_gt=10,
        num_pred=12,
        num_classes=3,
        height=32,
        width=64,
        shift=5,
        random_pred_label=False,
        show=True,
        iterations=20):
    """Generate an image with random ground truth and predicted instance
    and semantic segmentation data. To make the images realisitc, and to
    ensure that the instances form a PARTITION of the image, we rely on
    voronoi cells. Besides, to encourage a realistic overalp between the
    predicted and target instances, the predcition cell centers are
    sampled near the target samples.
    """
    # Generate random pixel positions for the ground truth and the
    # prediction centers. To produce predictions with "controllable"
    # overlap with the targets, we use the gt's centers as seeds for the
    # prediction centers and randomly sample shift them
    x = torch.randint(0, height, (num_gt,))
    y = torch.randint(0, width, (num_gt,))
    gt_xy = torch.vstack((x, y)).T
    if num_pred <= num_gt:
        idx_ref_gt = torch.from_numpy(
            np.random.choice(num_gt, num_pred, replace=False))
    else:
        idx_ref_gt = torch.from_numpy(
            np.random.choice(num_gt, num_pred % num_gt, replace=False))
        idx_ref_gt = torch.cat((
            torch.arange(num_gt).repeat(num_pred // num_gt), idx_ref_gt))
    xy_shift = torch.randint(0, 2 * shift, (num_pred, 2)) - shift
    pred_xy = gt_xy[idx_ref_gt] + xy_shift
    clamp_min = torch.tensor([0, 0])
    clamp_max = torch.tensor([height, width])
    pred_xy = pred_xy.clamp(min=clamp_min, max=clamp_max)

    # The above prediction center generation process may produce
    # duplicates, which can in turn generate downstream errors. To avoid
    # this, we greedily search for duplicates and shift them
    already_used_xy_ids = []
    for i_pred, xy in enumerate(pred_xy):
        xy_id = xy[0] * width + xy[1]
        count = 0

        while xy_id in already_used_xy_ids and count < iterations:
            xy_shift = torch.randint(0, 2 * shift, (2,)) - shift
            xy = gt_xy[idx_ref_gt[i_pred]] + xy_shift
            xy = xy.clamp(min=clamp_min, max=clamp_max)
            xy_id = xy[0] * width + xy[1]
            count += 1

        if count == iterations:
            raise ValueError(
                f"Reached max iterations={iterations} while resampling "
                "duplicate prediction centers")

        already_used_xy_ids.append(xy_id)
        pred_xy[i_pred] = xy

    # Generate labels and scores
    gt_labels = torch.randint(0, num_classes, (num_gt,))
    if random_pred_label:
        pred_labels = torch.randint(0, num_classes, (num_pred,))
    else:
        pred_labels = gt_labels[idx_ref_gt]
    pred_scores = torch.rand(num_pred)

    # Generate a 3D point cloud representing the pixel coordinates of the
    # image. This will be used to compute the 1-NNs and, from there, a
    # partition into voronoi cells
    x, y = torch.meshgrid(
        torch.arange(height), torch.arange(width), indexing='ij')
    x = x.flatten()
    y = y.flatten()
    z = torch.zeros_like(x)
    xyz = torch.vstack((x, y, z)).T

    # Compute a gt segmentation image from the 1-NN of each pixel, wrt the
    # gt segment centers
    gt_xyz = torch.cat((gt_xy, torch.zeros_like(gt_xy[:, [0]])), dim=1).float()
    gt_nn = knn_2(gt_xyz, xyz.float(), 1, r_max=max(width, height))[0]
    gt_seg_image = gt_nn.view(height, width)
    gt_label_image = gt_labels[gt_seg_image]

    # Compute a pred segmentation image from the 1-NN of each pixel, wrt the
    # pred segment centers
    pred_xyz = torch.cat((pred_xy, torch.zeros_like(pred_xy[:, [0]])), dim=1).float()
    pred_nn = knn_2(pred_xyz, xyz.float(), 1, r_max=max(width, height))[0]
    pred_seg_image = pred_nn.view(height, width)
    pred_label_image = pred_labels[pred_seg_image]

    # Display the segment images
    if show:
        plt.subplot(2, 2, 1)
        plt.title('Ground truth instances')
        plt.imshow(gt_seg_image)
        plt.subplot(2, 2, 2)
        plt.title('Predicted instances')
        plt.imshow(pred_seg_image)
        plt.subplot(2, 2, 3)
        plt.title('Ground truth labels')
        plt.imshow(gt_label_image)
        plt.subplot(2, 2, 4)
        plt.title('Predicted labels')
        plt.imshow(pred_label_image)
        plt.show()

    # Organize the data into torchmetric-friendly format
    tm_targets = dict(
        masks=torch.stack([gt_seg_image == i_gt for i_gt in range(num_gt)]),
        labels=gt_labels)

    tm_preds = dict(
        masks=torch.stack([pred_seg_image == i_pred for i_pred in range(num_pred)]),
        labels=pred_labels,
        scores=pred_scores)

    tm_data = (tm_preds, tm_targets)

    # Organize the data into our custom format
    pred_idx = pred_seg_image.flatten()
    gt_idx = gt_seg_image.flatten()
    gt_y = gt_label_image.flatten()
    count = torch.ones_like(pred_idx)

    from src.data.instance import InstanceData
    instance_data = InstanceData(pred_idx, gt_idx, count, gt_y, dense=True)
    spt_data = (pred_scores, pred_labels, instance_data)

    return tm_data, spt_data


def generate_random_segment_data(
        num_img=2,
        num_gt_per_img=10,
        num_pred_per_img=14,
        num_classes=2,
        height=32,
        width=64,
        shift=5,
        random_pred_label=False,
        verbose=True):
    """Generate multiple images with random ground truth and predicted
    instance and semantic segmentation data. To make the images
    realistic, and to ensure that the instances form a PARTITION of the
    image, we rely on voronoi cells. Besides, to encourage a realistic
    overlap between the predicted and target instances, the prediction
    cell centers are sampled near the target samples.
    """
    tm_data = []
    spt_data = []

    for i_img in range(num_img):
        if verbose:
            print(f"\nImage {i_img + 1}/{num_img}")
        tm_data_, spt_data_ = generate_single_random_segment_image(
            num_gt=num_gt_per_img,
            num_pred=num_pred_per_img,
            num_classes=num_classes,
            height=height,
            width=width,
            shift=shift,
            random_pred_label=random_pred_label,
            show=verbose)
        tm_data.append(tm_data_)
        spt_data.append(spt_data_)

    return tm_data, spt_data


def _instance_cut_pursuit(
        node_x,
        node_logits,
        node_size,
        edge_index,
        edge_affinity_logits,
        loss_type='l2_kl',
        regularization=1e-2,
        x_weight=1,
        p_weight=1,
        cutoff=1,
        parallel=True,
        iterations=10,
        trim=False,
        discrepancy_epsilon=1e-4,
        temperature=1,
        dampening=0,
        verbose=False):
    """Partition an instance graph using cut-pursuit.

    :param node_x: Tensor of shape [num_nodes, num_dim]
        Node features
    :param node_logits: Tensor of shape [num_nodes, num_classes]
        Predicted classification logits for each node
    :param node_size: Tensor of shape [num_nodes]
        Size of each node
    :param edge_index: Tensor of shape [2, num_edges]
        Edges of the graph, in torch-geometric's format
    :param edge_affinity_logits: Tensor of shape [num_edges]
        Predicted affinity logits (ie in R+, before sigmoid) of each
        edge
    :param loss_type: str
        Rules the loss applied on the node features. Accepts one of
        'l2' (L2 loss on node features and probabilities),
        'l2_kl' (L2 loss on node features and Kullback-Leibler
        divergence on node probabilities)
    :param regularization: float
        Regularization parameter for the partition
    :param x_weight: float
        Weight used to mitigate the impact of the node position in the
        partition. The larger, the lesser features importance before
        the probabilities
    :param p_weight: float
        Weight used to mitigate the impact of the node probabilities in
        the partition. The larger, the lesser features importance before
        the features
    :param cutoff: float
        Minimum number of points in each cluster
    :param parallel: bool
        Whether cut-pursuit should run in parallel
    :param iterations: int
        Maximum number of iterations for each partition
    :param trim: bool
        Whether the input graph should be trimmed. See `to_trimmed()`
        documentation for more details on this operation
    :param discrepancy_epsilon: float
        Mitigates the maximum discrepancy. More precisely:
        `affinity=1 ⇒ discrepancy=1/discrepancy_epsilon`
    :param temperature: float
        Temperature used in the softmax when converting node logits to
        probabilities
    :param dampening: float
        Dampening applied to the node probabilities to mitigate the
        impact of near-zero probabilities in the Kullback-Leibler
        divergence
    :param verbose: bool
    :return:
    """

    # Sanity checks
    assert node_x.dim() == 2, \
        "`node_x` must have shape `[num_nodes, num_dim]`"
    assert node_logits.dim() == 2, \
        "`node_logits` must have shape `[num_nodes, num_classes]`"
    assert node_logits.shape[0] == node_x.shape[0], \
        "`node_logits` and `node_x` must have the same number of points"
    assert node_size.dim() == 1, \
        "`node_size` must have shape `[num_nodes]`"
    assert node_size.shape[0] == node_x.shape[0], \
        "`node_size` and `node_x` must have the same number of points"
    assert edge_index.dim() == 2 and edge_index.shape[0] == 2, \
        "`edge_index` must be of shape `[2, num_edges]`"
    edge_affinity_logits = edge_affinity_logits.squeeze()
    assert edge_affinity_logits.dim() == 1, \
        "`edge_affinity_logits` must be of shape `[num_edges]`"
    assert edge_affinity_logits.shape[0] == edge_index.shape[1], \
        "`edge_affinity_logits` and `edge_index` must have the same number " \
        "of edges"
    loss_type = loss_type.lower()
    assert loss_type in ['l2', 'l2_kl'], \
        "`loss_type` must be one of ['l2', 'l2_kl']"
    assert 0 < discrepancy_epsilon, \
        "`discrepancy_epsilon` must be strictly positive"
    assert 0 < temperature, "`temperature` must be strictly positive"
    assert 0 <= dampening <= 1, "`dampening` must be in [0, 1]"

    device = node_x.device
    num_nodes = node_x.shape[0]
    x_dim = node_x.shape[1]
    p_dim = node_logits.shape[1]
    dim = x_dim + p_dim
    num_edges = edge_affinity_logits.numel()

    assert num_nodes < np.iinfo(np.uint32).max, \
        "Too many nodes for `uint32` indices"
    assert num_edges < np.iinfo(np.uint32).max, \
        "Too many edges for `uint32` indices"

    # Initialize the number of threads used for parallel cut-pursuit
    num_threads = available_cpu_count() if parallel else 1

    # Exit if the graph contains only one node
    if num_nodes < 2:
        return torch.zeros(num_nodes, dtype=torch.long, device=device)

    # Trim the graph, if need be
    if trim:
        edge_index, edge_affinity_logits = to_trimmed(
            edge_index, edge_attr=edge_affinity_logits, reduce='mean')

    if verbose:
        print(
            f'Launching instance partition reg={regularization}, '
            f'cutoff={cutoff}')

    # User warning if the number of edges exceeds uint32 limits
    if num_edges > _MAX_NUM_EDGES and verbose:
        print(
            f"WARNING: number of edges {num_edges} exceeds the uint32 limit "
            f"{_MAX_NUM_EDGES}. Please update the cut-pursuit source code to "
            f"accept a larger data type for `index_t`.")

    # Convert affinity logits to discrepancies
    edge_affinity = edge_affinity_logits.sigmoid()
    edge_discrepancy = edge_affinity / (1 - edge_affinity + discrepancy_epsilon)

    # Convert edges to forward-star (or CSR) representation
    source_csr, target, reindex = edge_list_to_forward_star(
        num_nodes, edge_index.T.contiguous().cpu().numpy())
    source_csr = source_csr.astype('uint32')
    target = target.astype('uint32')
    edge_weights = edge_discrepancy.cpu().numpy()[reindex] * regularization \
        if edge_discrepancy is not None else regularization

    # Convert logits to class probabilities
    node_probas = torch.nn.functional.softmax(node_logits / temperature, dim=1)

    # Apply some dampening to the probability distributions. This brings
    # the distributions closer to a uniform distribution, limiting the
    # impact of near-zero probabilities in the Kullback-Leibler
    # divergence in the partition
    num_classes = node_probas.shape[1]
    node_probas = (1 - dampening) * node_probas + dampening / num_classes

    # Mean-center the node features, in case values have a very large
    # mean. This is optional, but favors maintaining values in a
    # reasonable float32 range
    node_x = node_x - node_x.mean(dim=0).view(1, -1)

    # Build the node features as the concatenation of positions and
    # class probabilities
    x = torch.cat((node_x, node_probas), dim=1)
    x = np.asfortranarray(x.cpu().numpy().T)
    node_size = node_size.float().cpu().numpy()

    # The `loss` term will decide which portion of `x` should be treated
    # with L2 loss and which should be treated with Kullback-Leibler
    # divergence
    l2_dim = dim if loss_type == 'l2' else x_dim

    # Weighting to apply on the features and probabilities
    coor_weights_dim = dim if loss_type == 'l2' else x_dim + 1
    coor_weights = np.ones(coor_weights_dim, dtype=np.float32)
    coor_weights[:x_dim] *= x_weight
    coor_weights[x_dim:] *= p_weight

    # Partition computation
    obj_index, x_c, cluster, edges, times = cp_d0_dist(
        l2_dim,
        x,
        source_csr,
        target,
        edge_weights=edge_weights,
        vert_weights=node_size,
        coor_weights=coor_weights,
        min_comp_weight=cutoff,
        cp_dif_tol=1e-2,
        K=4,
        cp_it_max=iterations,
        split_damp_ratio=0.7,
        verbose=verbose,
        max_num_threads=num_threads,
        balance_parallel_split=True,
        compute_Time=True,
        compute_List=True,
        compute_Graph=True)

    if verbose:
        delta_t = (times[1:] - times[:-1]).round(2)
        print(f'Instance partition times: {delta_t}')

    # Convert the obj_index to the input format
    obj_index = torch.from_numpy(obj_index.astype('int64')).to(device)

    return obj_index


def instance_cut_pursuit(
        batch,
        node_x,
        node_logits,
        stuff_classes,
        node_size,
        edge_index,
        edge_affinity_logits,
        loss_type='l2_kl',
        regularization=1e-2,
        x_weight=1,
        p_weight=1,
        cutoff=1,
        parallel=True,
        iterations=10,
        trim=False,
        discrepancy_epsilon=1e-4,
        temperature=1,
        dampening=0,
        verbose=False):
    """The forward step will compute the partition on the instance
    graph, based on the node features, node logits, and edge
    affinities. The partition segments will then be further merged
    so that there is at most one instance of each stuff class per
    batch item (ie per scene).

    :param batch: Tensor of shape [num_nodes]
        Batch index of each node
    :param node_x: Tensor of shape [num_nodes, num_dim]
        Predicted node embeddings
    :param node_logits: Tensor of shape [num_nodes, num_classes]
        Predicted classification logits for each node
    :param stuff_classes: List or Tensor
        List of 'stuff' class labels. These are used for merging
        stuff segments together to ensure there is at most one
        predicted instance of each 'stuff' class per batch item
    :param node_size: Tensor of shape [num_nodes]
        Size of each node
    :param edge_index: Tensor of shape [2, num_edges]
        Edges of the graph, in torch-geometric's format
    :param edge_affinity_logits: Tensor of shape [num_edges]
        Predicted affinity logits (ie in R+, before sigmoid) of each
        edge
    :param loss_type: str
        Rules the loss applied on the node features. Accepts one of
        'l2' (L2 loss on node features and probabilities),
        'l2_kl' (L2 loss on node features and Kullback-Leibler
        divergence on node probabilities)
    :param regularization: float
        Regularization parameter for the partition
    :param x_weight: float
        Weight used to mitigate the impact of the node position in the
        partition. The larger, the lesser features importance before
        the probabilities
    :param p_weight: float
        Weight used to mitigate the impact of the node probabilities in
        the partition. The larger, the lesser features importance before
        the features
    :param cutoff: float
        Minimum number of points in each cluster
    :param parallel: bool
        Whether cut-pursuit should run in parallel
    :param iterations: int
        Maximum number of iterations for each partition
    :param trim: bool
        Whether the input graph should be trimmed. See `to_trimmed()`
        documentation for more details on this operation
    :param discrepancy_epsilon: float
        Mitigates the maximum discrepancy. More precisely:
        `affinity=1 ⇒ discrepancy=1/discrepancy_epsilon`
    :param temperature: float
        Temperature used in the softmax when converting node logits to
        probabilities
    :param dampening: float
        Dampening applied to the node probabilities to mitigate the
        impact of near-zero probabilities in the Kullback-Leibler
        divergence
    :param verbose: bool

    :return: obj_index: Tensor of shape [num_nodes]
        Indicates which predicted instance each node belongs to
    """

    # Actual partition, returns a tensor indicating which predicted
    # object each node belongs to
    obj_index = _instance_cut_pursuit(
        node_x,
        node_logits,
        node_size,
        edge_index,
        edge_affinity_logits,
        loss_type=loss_type,
        regularization=regularization,
        x_weight=x_weight,
        p_weight=p_weight,
        cutoff=cutoff,
        parallel=parallel,
        iterations=iterations,
        trim=trim,
        discrepancy_epsilon=discrepancy_epsilon,
        temperature=temperature,
        dampening=dampening,
        verbose=verbose)

    # Compute the mean logits for each predicted object, weighted by
    # the node sizes
    obj_logits = scatter_mean_weighted(node_logits, obj_index, node_size)
    obj_y = obj_logits.argmax(dim=1)

    # Identify, out of the predicted objects, which are of type stuff.
    # These will need to be merged to ensure there as most one instance
    # of each stuff class in each scene
    obj_is_stuff = get_stuff_mask(obj_y, stuff_classes)

    # Distribute the object-wise labels to the nodes
    node_obj_y = obj_y[obj_index]
    node_is_stuff = obj_is_stuff[obj_index]

    # Since we only want at most one prediction of each stuff class
    # per batch item (ie per scene), we assign nodes predicted as a
    # stuff class to new indices. These new indices are built in
    # such a way that there can be only one instance of each stuff
    # class per batch item
    batch = batch if batch is not None else torch.zeros_like(obj_index)
    num_batch_items = batch.max() + 1
    final_obj_index = obj_index.clone()
    final_obj_index[node_is_stuff] = \
        obj_index.max() + 1 \
        + node_obj_y[node_is_stuff] * num_batch_items \
        + batch[node_is_stuff]
    final_obj_index, perm = consecutive_cluster(final_obj_index)

    return final_obj_index


def oracle_superpoint_clustering(
        nag,
        num_classes,
        stuff_classes,
        mode='pas',
        graph_kwargs=None,
        partition_kwargs=None):
    """Compute an oracle for superpoint clustering for instance and
    panoptic segmentation. This is a proxy for the highest achievable
    graph clustering performance with the superpoint partition at hand
    and the input clustering parameters.

    The output `InstanceData` can then be used to compute final
    segmentation metrics using:
      - `InstanceData.semantic_segmentation_oracle()`
      - `InstanceData.instance_segmentation_oracle()`
      - `InstanceData.panoptic_segmentation_oracle()`

    More precisely, for the optimal superpoint clustering:
      - build the instance graph on the input `NAG` `level`-partition
      - for each edge, the oracle perfectly predicts the affinity
      - for each node, the oracle perfectly predicts the offset
      - for each node, the oracle predicts the dominant label from its
        label histogram (excluding the 'void' label)
      - partition the instance graph using the oracle edge affinities,
        node offsets and node classes
      - merge superpoints if they are assigned to the same object
      - merge 'stuff' predictions together, so that there is at most 1
        prediction of each 'stuff' class per batch item

    :param nag: NAG object
    :param num_classes: int
        Number of classes in the dataset, allows differentiating between
        valid and void classes
    :param stuff_classes: List[int]
        List of labels for 'stuff' classes
    :param mode: str
        String characterizing whether edge affinities, node semantics,
        positions and offsets should be used in the graph clustering.
        'p': use node position.
        'o': use oracle offset.
        'a': use oracle edge affinities.
        's': use oracle node semantics.
        In contrast, not setting 'p', nor 'o' is equivalent to setting
        all nodes positions and offsets to 0.
        Similarly, not setting 'a' will set the same weight to all the
        edges.
        Finally, not setting 's' will set the same class to all the
        nodes.
    :param graph_kwargs: dict
        Dictionary of kwargs to be passed to the graph constructor
        `OnTheFlyInstanceGraph()`
    :param partition_kwargs: dict
        Dictionary of kwargs to be passed to the partition function
        `instance_cut_pursuit()`
    :return:
    """

    # TODO: maybe remove this function, redundant with grid_search_panoptic_partition

    # Local import to avoid import loop errors
    from src.transforms import OnTheFlyInstanceGraph
    from src.models.panoptic import PanopticSegmentationOutput
    from src.metrics import PanopticQuality3D

    # Instance graph computation
    graph_kwargs = {} if graph_kwargs is None else graph_kwargs
    graph_kwargs = dict(graph_kwargs, **dict(level=1, num_classes=num_classes))
    nag = OnTheFlyInstanceGraph(**graph_kwargs)(nag)

    # Get node target semantics, size and instance graph
    node_y = nag[1].y[:, :num_classes].argmax(dim=1)
    node_size = nag.get_sub_size(1)
    edge_index = nag[1].obj_edge_index

    # Prepare input for instance graph partition. If 's' is used, the
    # oracle will assign the target semantic label to each node
    # NB: we assign only to valid classes and ignore void
    # NB2: `instance_cut_pursuit()` expects logits, which it converts to
    # probabilities using a softmax, hence the `one_hot * 100`
    node_logits = one_hot(node_y, num_classes=num_classes).float() * 100

    # Otherwise, the nodes will all have the same logits and the
    # semantics will not influence the partition
    if 's' not in mode.lower():
        partition_kwargs['p_weight'] = 0

    # Prepare edge affinity logits. If affinities are not used, we set
    # all edge affinity logits to 0 (ie 0.5 sigmoid-ed weights)
    edge_affinity_logits = torch.special.logit(nag[1].obj_edge_affinity) \
        if 'a' in mode.lower() \
        else torch.zeros(edge_index.shape[1], device=nag.device)

    # Prepare node position features. If 'o' is used, the oracle
    # perfectly predicts the offset to the object center for each node,
    # except for stuff and void classes, whose offset is set to 0
    if 'o' in mode.lower():
        node_x = nag[1].obj_pos
        is_stuff = get_stuff_mask(node_y, stuff_classes)
        node_x[is_stuff] = nag[1].pos[is_stuff]

    # If 'p' only node positions are used
    elif 'p' in mode.lower():
        node_x = nag[1].pos

    # Otherwise, positions and offsets are not used in the partition
    else:
        partition_kwargs['x_weight'] = 0
        node_x = nag[1].pos * 0

    # For each node, recover the index of the batch item it belongs to
    batch = nag[1].batch if nag[1].batch is not None \
        else torch.zeros(nag[1].num_nodes, dtype=torch.long, device=nag.device)

    # Instance graph partition
    partition_kwargs = {} if partition_kwargs is None else partition_kwargs
    obj_index = instance_cut_pursuit(
        batch,
        node_x,
        node_logits,
        stuff_classes,
        node_size,
        edge_index,
        edge_affinity_logits,
        **partition_kwargs)

    # Gather results in an output object
    output = PanopticSegmentationOutput(
        node_logits,
        stuff_classes,
        edge_affinity_logits,
        # node_offset_pred,
        node_size)

    # Store the panoptic segmentation results in the output object
    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
    output.y_hist = nag[1].y
    output.obj_index_pred = obj_index

    # Create the metrics tracking objects
    panoptic_metrics = PanopticQuality3D(
        num_classes,
        ignore_unseen_classes=True,
        stuff_classes=stuff_classes,
        compute_on_cpu=True)

    # Recover the predicted instance score, semantic label and instance
    # partition
    obj_score, obj_y, instance_data = output.panoptic_pred()

    # Compute the metrics on the oracle partition
    panoptic_metrics.update(obj_y, instance_data.cpu())
    results = panoptic_metrics.compute()

    return results


def get_stuff_mask(y, stuff_classes):
    """Helper function producing a boolean mask of size `y.shape[0]`
    indicating which of the `y` (labels if 1D or logits/probabilities if
    2D) are among the `stuff_classes`.
    """
    # Get labels from y, in case y are logits
    labels = y.long() if y.dim() == 1 else y.argmax(dim=1)

    # Search the labels belonging to the set of stuff classes
    stuff_classes = torch.as_tensor(
        stuff_classes, dtype=labels.dtype, device=labels.device)
    return torch.isin(labels, stuff_classes)


def compute_panoptic_metrics(
        model,
        datamodule,
        stage='val',
        graph_kwargs=None,
        partition_kwargs=None,
        verbose=True):
    """Helper function to compute the semantic, instance, panoptic
    segmentation metrics of a model on a given dataset, for given
    instance graph and partition parameters.
    """
    # Local imports to avoid import loop errors
    from src.data import NAGBatch

    # Pick among train, val, and test datasets. It is important to note that
    # the train dataset produces augmented spherical samples of large
    # scenes, while the val and test dataset
    if stage == 'train':
        dataset = datamodule.train_dataset
        dataloader = datamodule.train_dataloader()
    elif stage == 'val':
        dataset = datamodule.val_dataset
        dataloader = datamodule.val_dataloader()
    elif stage == 'test':
        dataset = datamodule.test_dataset
        dataloader = datamodule.test_dataloader()
    else:
        raise ValueError(f"Unknown stage : {stage}")

    # Prevent `NAGAddKeysTo` from removing attributes to allow
    # visualizing them after model inference
    dataset = _set_attribute_preserving_transforms(dataset)

    # Set the instance graph construction parameters
    dataset = _set_graph_construction_parameters(dataset, graph_kwargs)

    # Set the partitioner parameters
    model, backup_kwargs = _set_partitioner_parameters(model, partition_kwargs)

    # Load a dataset item. This will return the hierarchical partition
    # of an entire tile, within a NAG object
    with torch.no_grad():
        enum = tqdm(dataloader) if verbose else dataloader
        for nag_list in enum:
            nag = NAGBatch.from_nag_list([nag.cuda() for nag in nag_list])

            # Apply on-device transforms on the NAG object. For the
            # train dataset, this will select a spherical sample of the
            # larger tile and apply some data augmentations. For the
            # validation and test datasets, this will prepare an entire
            # tile for inference
            nag = dataset.on_device_transform(nag)

            # NB: we use the "validation_step" protocol here, regardless
            # of the stage the data comes from
            model.validation_step(nag, None)

        # Actions taken from on_validation_epoch_end()
        # panoptic_results = model.val_panoptic.compute()
        # instance_miou = model.val_semantic.miou()
        # instance_oa = model.val_semantic.oa()
        # instance_macc = model.val_semantic.macc()
        panoptic = deepcopy(model.val_panoptic)
        instance = deepcopy(model.val_instance)
        semantic = deepcopy(model.val_semantic)
        model.val_affinity_oa.reset()
        model.val_affinity_f1.reset()
        model.val_panoptic.reset()
        model.val_semantic.reset()
        model.val_instance.reset()

    # Restore the partitioner initial kwargs
    model, _ = _set_partitioner_parameters(model, backup_kwargs)

    if not verbose:
        return panoptic, instance, semantic

    for k, v in panoptic.compute().items():
        print(f"{k:<22}: {v}")

    if not model.no_instance_metrics:
        for k, v in instance.compute().items():
            print(f"{k:<22}: {v}")

    print(f"mIoU                  : {semantic.miou().cpu().item()}")

    return panoptic, instance, semantic


def compute_panoptic_metrics_s3dis_6fold(
        fold_ckpt,
        experiment_config,
        stage='val',
        graph_kwargs=None,
        partition_kwargs=None,
        verbose=False):
    """Helper function to compute the semantic, instance, panoptic
    segmentation metrics of a model on a S3DIS 6-fold, for given
    instance graph and partition parameters.

    :param fold_ckpt: dict
        Dictionary with S3DIS fold numbers as keys and checkpoint paths
        as values
    :param experiment_config: str
        Experiment config to use for inference. For instance for S3DIS
        with stuff panoptic segmentation: 'panoptic/s3dis_with_stuff'
    :param stage: str
    :param graph_kwargs: dict
    :param partition_kwargs: dict
    :param verbose: bool
    :return:
    """
    # Local import to avoid import loop errors
    from src.metrics import PanopticQuality3D, MeanAveragePrecision3D, \
        ConfusionMatrix

    # Very ugly fix to ignore lightning's warning messages about the
    # trainer and modules not being connected
    import warnings
    warnings.filterwarnings("ignore")

    panoptic_list = []
    instance_list = []
    semantic_list = []
    no_instance_metrics = None
    min_instance_size = None
    num_classes = None
    stuff_classes = None

    for fold, ckpt_path in fold_ckpt.items():

        if verbose:
            print(f"\nFold {fold}")

        # Parse the configs using hydra
        cfg = init_config(overrides=[
            f"experiment={experiment_config}",
            f"datamodule.fold={fold}",
            f"ckpt_path={ckpt_path}"])

        # Instantiate the datamodule
        datamodule = hydra.utils.instantiate(cfg.datamodule)
        datamodule.prepare_data()
        datamodule.setup()

        # Instantiate the model
        model = hydra.utils.instantiate(cfg.model)

        # Load pretrained weights from a checkpoint file
        model = model._load_from_checkpoint(cfg.ckpt_path)
        model = model.eval().cuda()

        # Compute metrics on the fold
        panoptic, instance, semantic = compute_panoptic_metrics(
            model,
            datamodule,
            stage=stage,
            graph_kwargs=graph_kwargs,
            partition_kwargs=partition_kwargs,
            verbose=verbose)

        # Gather some details from the model and datamodule before
        # deleting them
        no_instance_metrics = model.no_instance_metrics
        min_instance_size = model.hparams.min_instance_size
        num_classes = datamodule.train_dataset.num_classes
        stuff_classes = datamodule.train_dataset.stuff_classes

        del model, datamodule

        # Store the metrics for each fold
        panoptic_list.append(panoptic)
        instance_list.append(instance)
        semantic_list.append(semantic)

    # Initialize the 6-fold metrics
    panoptic_6fold = PanopticQuality3D(
        num_classes,
        ignore_unseen_classes=True,
        stuff_classes=stuff_classes,
        compute_on_cpu=True)

    instance_6fold = MeanAveragePrecision3D(
        num_classes,
        stuff_classes=stuff_classes,
        min_size=min_instance_size,
        compute_on_cpu=True,
        remove_void=True)

    semantic_6fold = ConfusionMatrix(num_classes)

    # Group together per-fold panoptic and semantic results
    for i in range(len(panoptic_list)):

        panoptic_6fold.instance_data += panoptic_list[i].instance_data
        panoptic_6fold.prediction_semantic += panoptic_list[i].prediction_semantic

        if not no_instance_metrics:
            instance_6fold.prediction_score += instance_list[i].prediction_score
            instance_6fold.prediction_semantic += instance_list[i].prediction_semantic
            instance_6fold.instance_data += instance_list[i].instance_data

        semantic_6fold.confmat += semantic_list[i].confmat.cpu()

    # Print computed the metrics
    print(f"\n6-fold")
    for k, v in panoptic_6fold.compute().items():
        print(f"{k:<22}: {v}")

    if not no_instance_metrics:
        for k, v in instance_6fold.compute().items():
            print(f"{k:<22}: {v}")

    print(f"mIoU                  : {semantic_6fold.miou().cpu().item()}")

    return (panoptic_6fold, panoptic_list), (instance_6fold, instance_list), (semantic_6fold, semantic_list)


def _set_graph_construction_parameters(dataset, graph_kwargs):
    """Searches for the last occurrence of `OnTheFlyInstanceGraph` among
    the `on_device_transform` of the dataset and modifies the graph
    construction parameters passed in the `graph_kwargs` dictionary.
    """
    if graph_kwargs is None:
        return dataset

    # Local imports to avoid import loop errors
    from src.transforms import OnTheFlyInstanceGraph

    # Search for the `OnTheFlyInstanceGraph` instance graph construction
    # transform among the on-device transforms
    i_transform = None
    for i, transform in enumerate(dataset.on_device_transform.transforms):
        if isinstance(transform, OnTheFlyInstanceGraph):
            i_transform = i

    # Set OnTheFlyInstanceGraph parameters if need be
    if i_transform is not None and graph_kwargs is not None:
        for k, v in graph_kwargs.items():
            setattr(dataset.on_device_transform.transforms[i_transform], k, v)

    return dataset


def _set_partitioner_parameters(model, partition_kwargs):
    """Modifies the `model.partitioner` parameters with parameters
    passed in the `partition_kwargs` dictionary.
    """
    backup_kwargs = {}

    if partition_kwargs is None:
        return model, backup_kwargs

    # Set partitioner parameters if need be
    if partition_kwargs is not None:
        for k, v in partition_kwargs.items():
            backup_kwargs[k] = getattr(model.partitioner, k, None)
            setattr(model.partitioner, k, v)

    return model, backup_kwargs


def _forward_multi_partition(
        model,
        nag,
        partition_kwargs,
        mode='pas'):
    """Local helper to compute multiple instance partitions from the
    same input data, based on diverse partition parameter settings.
    """
    # Local import to avoid import loop errors
    from src.models.panoptic import PanopticSegmentationOutput

    # Make sure each element of `partition_kwargs` is a list,
    # to facilitate computing Cartesian product of the lists for
    # grid search
    partition_kwargs = {
        k: v if isinstance(v, list) else [v]
        for k, v in partition_kwargs.items()}

    with torch.no_grad():

        # Extract features
        x = model.net(nag)

        # Compute level-1 or multi-level semantic predictions
        semantic_pred = [head(x_) for head, x_ in zip(model.head, x)] \
            if model.multi_stage_loss else model.head(x)

        # Recover level-1 features only
        x = x[0] if model.multi_stage_loss else x

        # 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)
        edge_affinity_logits = model.edge_affinity_head(x_edge).squeeze()

        # Ignore predicted affinities (sets all edge affinity logits to
        # 0, which will set edge weights to 0.5 for the partition)
        if 'a' not in mode.lower():
            edge_affinity_logits = edge_affinity_logits * 0

        # Oracle edge affinities
        elif 'A' in mode:
            edge_affinity_logits = torch.special.logit(nag[1].obj_edge_affinity)

        # Ignore predicted semantic labels
        if 's' not in mode.lower():
            partition_kwargs['p_weight'] = [0]

        # Oracle node semantics predicts perfect semantic logits for
        # each node
        # NB: we assign only to valid classes and ignore void
        # NB2: `instance_cut_pursuit()` expects logits, which it
        # converts to probabilities using a softmax, hence the
        # `one_hot * 10`
        elif 'S' in mode:
            node_y = nag[1].y[:, :model.num_classes].argmax(dim=1)
            node_logits = one_hot(
                node_y, num_classes=model.num_classes).float() * 10
            if model.multi_stage_loss:
                semantic_pred[0] = node_logits
            else:
                semantic_pred = node_logits

        # Ignore positions and predicted offsets
        if 'p' not in mode.lower() and 'o' not in mode.lower():
            partition_kwargs['x_weight'] = [0]

        # Compute node offset predictions
        elif 'o' in mode:
            node_offset_pred = model.node_offset_head(x)

            # Forcefully set 0-offset for nodes with stuff predictions
            node_logits = semantic_pred[0] if model.multi_stage_loss \
                else semantic_pred
            is_stuff = get_stuff_mask(node_logits, model.stuff_classes)
            node_offset_pred[is_stuff] = 0

        # Oracle node offsets sets perfect offsets for all nodes and
        # keeps node centroid for nodes with target stuff label
        # (ie 0-offset)
        elif 'O' in mode:
            is_stuff = get_stuff_mask(nag[1].y, model.stuff_classes)
            nag[1].pos[~is_stuff] = nag[1].obj_pos[~is_stuff]

        # Compute the partition on the Cartesian product of parameters
        partition_keys = list(partition_kwargs.keys())
        enum = [
            {k: v for k, v in zip(partition_keys, values)}
            for values in product(*partition_kwargs.values())]
        partitions = {}
        for kwargs in tqdm(enum):

            # Apply the kwargs to the partitioner
            model, backup_kwargs = _set_partitioner_parameters(model, kwargs)

            # Gather results in an output object
            output = PanopticSegmentationOutput(
                semantic_pred,
                model.stuff_classes,
                edge_affinity_logits,
                # node_offset_pred,
                nag.get_sub_size(1))

            # Compute the panoptic partition
            output = model._forward_partition(nag, output)

            # Store the predicted partition wrt the parameter values
            # (can't directly store kwargs dict because unhashable)
            partitions[tuple(kwargs.values())] = output.obj_index_pred

            # Restore the initial partitioner kwargs
            model, _ = _set_partitioner_parameters(model, backup_kwargs)

        output = model.get_target(nag, output)

    return output, partitions, partition_keys


def grid_search_panoptic_partition(
        model,
        dataset,
        i_cloud=0,
        graph_kwargs=None,
        partition_kwargs=None,
        mode='pas',
        panoptic=True,
        instance=False):
    """Runs a grid search on the partition parameters to find the best
    setup on a given sample `dataset[i_cloud]`.

    :param model: PanopticSegmentationModule
    :param dataset: BaseDataset
    :param i_cloud: int
        The grid search will be computed on `dataset[i_cloud]`
    :param graph_kwargs: dict
        Dictionary of parameters to be passed to the instance graph
        constructor `OnTheFlyInstanceGraph`. NB: the grid search does
        not cover these parameters---only a single value can be passed
        for each of these parameters
    :param partition_kwargs: dict
        Dictionary of parameters to be passed to `model.partitioner`.
        Passing a list of values for a given parameter will trigger the
        grid search across these values. Beware of the combinatorial
        explosion !
    :param mode: str
        String characterizing whether edge affinities, node semantics,
        positions and offsets should be used in the graph clustering.
        'p': use node position.
        'o': use predicted node offset.
        'O': use oracle offset.
        'a': use predicted edge affinity.
        'A': use oracle edge affinities.
        's': use predicted node semantics.
        'S': use oracle node semantics.
        In contrast, not setting 'p', 'o', nor 'O' is equivalent to
        setting all node positions and offsets to 0.
        Similarly, not setting 'a' nor 'A' will set the same weight to
        all the edges.
        Finally, not setting 's', nor 'S' will set the same class to all
        the nodes.
    :param panoptic: bool
        Whether panoptic segmentation metrics should be computed
    :param instance: bool
        Whether instance segmentation metrics should be computed
    :return:
    """
    # TODO: grid search on the whole dataset rather than a single cloud

    # Local import to avoid import loop errors
    from src.metrics import PanopticQuality3D, MeanAveragePrecision3D

    assert panoptic or instance, \
        "At least 'panoptic' or 'instance' must be True"

    # Limit the column header size for printed tables
    max_len = 6

    # Prevent `NAGAddKeysTo` from removing attributes to allow
    # visualizing them after model inference
    dataset = _set_attribute_preserving_transforms(dataset)

    # Set the instance graph construction parameters
    dataset = _set_graph_construction_parameters(dataset, graph_kwargs)

    # Load a dataset item. This will return the hierarchical partition
    # of an entire tile, within a NAG object
    nag = dataset[i_cloud]

    # Apply on-device transforms on the NAG object. For the train
    # dataset, this will select a spherical sample of the larger tile
    # and apply some data augmentations. For the validation and test
    # datasets, this will prepare an entire tile for inference
    nag = dataset.on_device_transform(nag.cuda())

    # Compute the partition for each parameterization
    output, partitions, partition_keys = _forward_multi_partition(
        model,
        nag,
        partition_kwargs,
        mode=mode)

    # Get the target labels
    output = model.get_target(nag, output)

    # Create the metrics tracking objects
    instance_metrics = MeanAveragePrecision3D(
        model.num_classes,
        stuff_classes=model.stuff_classes,
        min_size=model.hparams.min_instance_size,
        compute_on_cpu=True,
        remove_void=True)

    panoptic_metrics = PanopticQuality3D(
        model.num_classes,
        ignore_unseen_classes=True,
        stuff_classes=model.stuff_classes,
        compute_on_cpu=True)

    # Compute and print metric results for each partition setup
    results = {}
    results_data = []
    best_pq = -1
    best_map = -1
    best_pq_params = None
    best_map_params = None
    for (kwargs_values), obj_index_pred in partitions.items():

        # Reconstruct the kwargs dict from the kwargs values
        kwargs = {k: v for k, v in zip(partition_keys, kwargs_values)}

        output.obj_index_pred = obj_index_pred

        obj_score, obj_y, instance_data = output.panoptic_pred()
        obj_score = obj_score.detach().cpu()
        obj_y = obj_y.detach().cpu()

        if panoptic:
            panoptic_metrics.update(obj_y, instance_data.cpu())
            panoptic_results = panoptic_metrics.compute()
            panoptic_metrics.reset()
            if panoptic_results.pq > best_pq:
                best_pq_params = tuple(kwargs.values())
                best_pq = panoptic_results.pq
        else:
            panoptic_results = None

        if instance:
            instance_metrics.update(obj_score, obj_y, instance_data.cpu())
            instance_results = instance_metrics.compute()
            instance_metrics.reset()
            if instance_results.map > best_map:
                best_map_params = tuple(kwargs.values())
                best_map = instance_results.map
        else:
            instance_results = None

        # Store the panoptic and instance metric results for the
        # parameters at hand
        results[tuple(kwargs.values())] = (panoptic_results, instance_results)

        # Track the results to build a global summary DataFrame
        current_results = [*kwargs.values()]
        if panoptic:
            current_results += [
                round(panoptic_results.pq.item() * 100, 2),
                round(panoptic_results.sq.item() * 100, 2),
                round(panoptic_results.rq.item() * 100, 2)]
        if instance:
            current_results += [
                round(instance_results.map.item() * 100, 2),
                round(instance_results.map_50.item() * 100, 2)]
        results_data.append(current_results)

    # Print a DataFrame summarizing the results
    metric_columns = []
    if panoptic:
        metric_columns += ['PQ', 'SQ', 'RQ']
    if instance:
        metric_columns += ['mAP', 'mAP 50']
    with pd.option_context('display.precision', 2):
        print(pd.DataFrame(
            data=results_data,
            columns=[
                *[
                    x[:max_len - 1] + '.' if len(x) > max_len else x
                    for x in partition_keys
                ],
                *metric_columns]))
    print()

    # Print more details about the best panoptic setup
    if panoptic and best_pq_params is not None:

        # Print global results
        print(f"\nBest panoptic setup: PQ={100 * best_pq:0.2f}")
        with pd.option_context('display.precision', 2):
            print(pd.DataFrame(
                data=[best_pq_params],
                columns=[
                    x[:max_len - 1] + '.' if len(x) > max_len else x
                    for x in partition_keys]))

        print()

        # Print per-class results
        res = results[best_pq_params][0]
        with pd.option_context('display.precision', 2):
            print(pd.DataFrame(
                data=torch.column_stack([
                    res.pq_per_class.mul(100),
                    res.sq_per_class.mul(100),
                    res.rq_per_class.mul(100),
                    res.precision_per_class.mul(100),
                    res.recall_per_class.mul(100),
                    res.tp_per_class,
                    res.fp_per_class,
                    res.fn_per_class]),
                index=dataset.class_names[:-1],
                columns=['PQ', 'SQ', 'RQ', 'PREC.', 'REC.', 'TP', 'FP', 'FN']))
        print()

        # Store the best panoptic partition indexing in the output
        output.obj_index_pred = partitions[best_pq_params]

    # Print more details about the best instance setup
    if instance and best_map_params is not None:

        # Print global results
        print(f"\nBest instance setup: mAP={100 * best_map:0.2f}")
        with pd.option_context('display.precision', 2):
            print(pd.DataFrame(
                data=[best_map_params],
                columns=[
                    x[:max_len - 1] + '.' if len(x) > max_len else x
                    for x in partition_keys]))
        print()

        # Print per-class results
        res = results[best_map_params][1]
        thing_class_names = [
            c for i, c in enumerate(dataset.class_names) if i in dataset.thing_classes]
        with pd.option_context('display.precision', 2):
            print(pd.DataFrame(
                data=torch.column_stack([res.map_per_class.mul(100)]),
                index=thing_class_names,
                columns=['mAP']))
        print()

    return output, partitions, results