Ниже приведена реализация среднего значения IoU (пересечение по объединению) в PyTorch.
def mIOU(label, pred, num_classes=19):
pred = F.softmax(pred, dim=1)
pred = torch.argmax(pred, dim=1).squeeze(1)
iou_list = list()
present_iou_list = list()
pred = pred.view(-1)
label = label.view(-1)
# Note: Following for loop goes from 0 to (num_classes-1)
# and ignore_index is num_classes, thus ignore_index is
# not considered in computation of IoU.
for sem_class in range(num_classes):
pred_inds = (pred == sem_class)
target_inds = (label == sem_class)
if target_inds.long().sum().item() == 0:
iou_now = float('nan')
else:
intersection_now = (pred_inds[target_inds]).long().sum().item()
union_now = pred_inds.long().sum().item() + target_inds.long().sum().item() - intersection_now
iou_now = float(intersection_now) / float(union_now)
present_iou_list.append(iou_now)
iou_list.append(iou_now)
return np.mean(present_iou_list)
Прогноз вашей модели будет в горячей форме, поэтому сначала возьмите softmax (если ваша модель еще не работает), а затем argmax
, чтобы получить индекс с наибольшей вероятностью для каждого пикселя . Затем мы вычисляем IoU для каждого класса (и в конце берем среднее значение).
Мы можем изменить и предсказание, и метку как одномерные векторы (я читал, что это ускоряет вычисления). Для каждого класса мы сначала определяем индексы этого класса, используя pred_inds = (pred == sem_class)
и target_inds = (label == sem_class)
. Результирующие pred_inds
и target_inds
будут иметь 1 пиксель, помеченный как этот конкретный класс, а 0 для любого другого класса.
Тогда существует вероятность, что цель вообще не содержит этот конкретный класс. Это сделает вычисление IoU этого класса недействительным, поскольку его нет в цели. Итак, вы назначаете таким классам NaN IoU (чтобы вы могли идентифицировать их позже) и не включаете их в вычисление среднего.
Если конкретный класс присутствует в цели, то pred_inds[target_inds]
даст вектор из единиц и нулей, где индексы с единицей - это те, в которых прогноз и цель равны, в противном случае - ноль. Сумма всех элементов этого даст нам пересечение.
Если мы сложим все элементы pred_inds
и target_inds
, мы получим объединение + пересечение пикселей этого конкретного класса. Итак, мы вычитаем уже рассчитанное пересечение, чтобы получить объединение. Затем мы можем разделить пересечение и объединение, чтобы получить IoU этого конкретного класса и добавить его в список допустимых IoU.
В конце вы берете среднее значение всего списка, чтобы получить mIoU. Если вам нужен коэффициент игральных костей, вы можете рассчитать его аналогичным образом.