File size: 5,970 Bytes
d9bb75c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This software may be used and distributed in accordance with
# the terms of the DINOv3 License Agreement.

import logging
import numpy as np
import pandas as pd
import torch


pd.set_option("display.max_rows", 200)

logger = logging.getLogger("dinov3")


SEGMENTATION_METRICS = ["mIoU", "acc", "aAcc", "dice", "fscore", "precision", "recall"]


def calculate_segmentation_metrics(
    pre_eval_results,
    metrics=["mIoU"],
    beta=1,
):
    """Calculate the segmentation metrics after aggregating all the intermediate results.

    Args:
        pre_eval_results (list): Lists of (area_intersect, area_union, area_pred_label, and area_label).
            These are intermediate results to compute the final metrics such as iou, fscore, etc.
        metrics (list): Metrics to compute. Defaults to ["mIoU"].
        beta (int): Parameter for computing F-score. Defaults to 1 (for computing F1-score).

    Returns:
        Dictionary of final metrics.
    """
    pre_eval_results = tuple(zip(*pre_eval_results))
    assert len(pre_eval_results) == 4
    total_area_intersect = sum(pre_eval_results[0])
    total_area_union = sum(pre_eval_results[1])
    total_area_pred_label = sum(pre_eval_results[2])
    total_area_label = sum(pre_eval_results[3])
    metrics_dict = total_area_to_metrics(
        total_area_intersect,
        total_area_union,
        total_area_pred_label,
        total_area_label,
        metrics=metrics,
        beta=beta,
    )
    df = pd.DataFrame(
        {
            "Class Index": np.arange(len(metrics_dict["mIoU"])),
            "mIoU": 100 * metrics_dict["mIoU"].cpu().numpy(),
        }
    )
    logger.info(f"mIoU per class:\n{df.to_string(index=False)}")
    return {
        "mIoU": metrics_dict["mIoU"].nanmean(),
        "acc": metrics_dict["acc"].nanmean(),
        "aAcc": metrics_dict["aAcc"].nanmean(),
        "dice": metrics_dict["dice"].nanmean(),
        "fscore": metrics_dict["fscore"].nanmean(),
        "precision": metrics_dict["precision"].nanmean(),
        "recall": metrics_dict["recall"].nanmean(),
    }


def preprocess_nonzero_labels(label, ignore_index=255):
    label_new = label.clone()
    label_new[label_new == ignore_index] += 1
    label_new -= 1
    label_new[label_new == -1] = ignore_index
    return label_new


def calculate_intersect_and_union(pred_label, label, num_classes, ignore_index=255, reduce_zero_label=False):
    """Calculate intersection and Union.
    Args:
        pred_label (torch.Tensor): Prediction segmentation map
        label (torch.Tensor): Ground truth segmentation map
        num_classes (int): Number of categories.
        ignore_index (int): Index that will be ignored in evaluation.
        reduce_zero_label (bool): Indicates whether or not label 0 is to be ignored.
    """
    pred_label = pred_label.float()  # Enables float tensor operations
    if reduce_zero_label:
        label = preprocess_nonzero_labels(label, ignore_index=ignore_index)

    mask = label != ignore_index
    pred_label = pred_label[mask]
    label = label[mask]
    intersect = pred_label[pred_label == label]
    area_intersect = torch.histc(intersect.float(), bins=(num_classes), min=0, max=num_classes - 1)
    area_pred_label = torch.histc(pred_label.float(), bins=(num_classes), min=0, max=num_classes - 1)
    area_label = torch.histc(label.float(), bins=(num_classes), min=0, max=num_classes - 1)
    area_union = area_pred_label + area_label - area_intersect

    return torch.stack([area_intersect, area_union, area_pred_label, area_label])


def total_area_to_metrics(
    total_area_intersect,
    total_area_union,
    total_area_pred_label,
    total_area_label,
    metrics=["mIoU"],
    beta=1,
):
    """Calculate evaluation metrics
    Args:
        total_area_intersect (torch.Tensor): The intersection of prediction and
            ground truth histogram on all classes.
        total_area_union (torch.Tensor): The union of prediction and ground truth
            histogram on all classes.
        total_area_pred_label (torch.Tensor): The prediction histogram on all
            classes.
        total_area_label (torch.Tensor): The ground truth histogram on all classes.
        metrics (list[str] | str): Metrics to be evaluated,
            can be 'mIoU', 'mDice', or 'mFscore'.
        beta (int): Parameter for computing F-score. Defaults to 1 (for computing F1-score).
    Returns:
        float: Overall accuracy on all images.
        ndarray: Per category accuracy, shape (num_classes, ).
        ndarray: Per category evaluation metrics, shape (num_classes, ).
    """

    def f_score(precision, recall, beta=1):
        score = (1 + beta**2) * (precision * recall) / ((beta**2 * precision) + recall)
        return score

    if isinstance(metrics, str):
        metrics = [metrics]
    allowed_metrics = ["mIoU", "dice", "fscore"]
    if not set(metrics).issubset(set(allowed_metrics)):
        raise KeyError(f"metrics {metrics} is not supported")

    all_acc = total_area_intersect.sum() / total_area_label.sum()
    ret_metrics = dict({"aAcc": all_acc})
    for metric in metrics:
        if metric == "mIoU":
            ret_metrics["mIoU"] = total_area_intersect / total_area_union
            ret_metrics["acc"] = total_area_intersect / total_area_label
        elif metric == "dice":
            ret_metrics["dice"] = 2 * total_area_intersect / (total_area_pred_label + total_area_label)
            ret_metrics["acc"] = total_area_intersect / total_area_label
        elif metric == "fscore":
            precision = total_area_intersect / total_area_pred_label
            recall = total_area_intersect / total_area_label
            f_value = torch.tensor([f_score(x[0], x[1], beta) for x in zip(precision, recall)])
            ret_metrics["fscore"] = f_value
            ret_metrics["precision"] = precision
            ret_metrics["recall"] = recall
    return ret_metrics