Pytorch：如何计算用于语义分割的 IoU（Jaccard 指数）

Question

有人可以提供一个关于如何在 pytorch 中为语义分割计算 IoU（联合交叉）的玩具示例吗？

Answer 1

截至 2021 年，无需实现您自己的 IoU，因为torchmetrics它已配备 -这是链接。它被命名torchmetrics.JaccardIndex（以前torchmetrics.IoU）并计算您想要的内容。它可以与 PyTorch 和 PyTorch Lightning 配合使用，也可以进行分布式训练。

从文档中：

torchmetrics.JaccardIndex(num_classes, ignore_index=None, absent_score=0.0, threshold=0.5, multilabel=False, reduction='elementwise_mean', compute_on_step=None, **kwargs)

通过并集计算交集，或 Jaccard 索引计算：

J(A,B) = \frac{|A\cap B|}{|A\cup B|}

其中：A和B都是大小相同的张量，包含整数类值。它们可能会受到输入数据的转换（参见下面的描述）。请注意，它与 box IoU 不同。

适用于二进制、多类和多标签数据。接受模型输出的概率或预测中的整数类值。适用于多维预测和目标。

转发接受

• preds（浮点或长张量）：(N, ...)或(N, C, ...)其中 C 是类别数
• target（长张量）：(N, ...)如果 preds 和 target 具有相同的形状，并且 preds 是浮点张量，我们使用 self.threshold参数将其转换为整数标签。二元概率和多标签概率就是这种情况。

如果 preds 有一个额外的维度，就像多类分数的情况一样，我们对 执行 argmax dim=1。

官方示例：

from torchmetrics import JaccardIndex
target = torch.randint(0, 2, (10, 25, 25))
pred = torch.tensor(target)
pred[2:5, 7:13, 9:15] = 1 - pred[2:5, 7:13, 9:15]
jaccard = JaccardIndex(task='multiclass', num_classes=2)
jaccard(pred, target)

返回tensor(0.9660)。

Answer 2

我在某处找到了它并为我改编了它。如果我能再次找到它，我会发布链接。抱歉，如果这是重复的。
这里的关键函数是名为 的函数iou。包装函数evaluate_performance不是通用的，但它表明需要在计算之前迭代所有结果IoU。

import torch 
import pandas as pd  # For filelist reading
import myPytorchDatasetClass  # Custom dataset class, inherited from torch.utils.data.dataset


def iou(pred, target, n_classes = 12):
  ious = []
  pred = pred.view(-1)
  target = target.view(-1)

  # Ignore IoU for background class ("0")
  for cls in xrange(1, n_classes):  # This goes from 1:n_classes-1 -> class "0" is ignored
    pred_inds = pred == cls
    target_inds = target == cls
    intersection = (pred_inds[target_inds]).long().sum().data.cpu()[0]  # Cast to long to prevent overflows
    union = pred_inds.long().sum().data.cpu()[0] + target_inds.long().sum().data.cpu()[0] - intersection
    if union == 0:
      ious.append(float('nan'))  # If there is no ground truth, do not include in evaluation
    else:
      ious.append(float(intersection) / float(max(union, 1)))
  return np.array(ious)


def evaluate_performance(net):
    # Dataloader for test data
    batch_size = 1  
    filelist_name_test = '/path/to/my/test/filelist.txt'
    data_root_test = '/path/to/my/data/'
    dset_test = myPytorchDatasetClass.CustomDataset(filelist_name_test, data_root_test)
    test_loader = torch.utils.data.DataLoader(dataset=dset_test,  
                                              batch_size=batch_size,
                                              shuffle=False,
                                              pin_memory=True)
    data_info = pd.read_csv(filelist_name_test, header=None)
    num_test_files = data_info.shape[0]  
    sample_size = num_test_files

    # Containers for results
    preds = Variable(torch.zeros((sample_size, 60, 36, 60)))
    gts = Variable(torch.zeros((sample_size, 60, 36, 60)))

    dataiter = iter(test_loader) 
    for i in xrange(sample_size):
        images, labels, filename = dataiter.next()
        images = Variable(images).cuda()
        labels = Variable(labels)
        gts[i:i+batch_size, :, :, :] = labels
        outputs = net(images)
        outputs = outputs.permute(0, 2, 3, 4, 1).contiguous()
        val, pred = torch.max(outputs, 4)
        preds[i:i+batch_size, :, :, :] = pred.cpu()
    acc = iou(preds, gts)
    return acc

Answer 3

假设您的输出形状为[32, 256, 256]#32 是小批量大小，256x256 是图像的高度和宽度，并且标签也是相同的形状。

jaccard_similarity_score然后你可以在一些重塑后使用sklearn 。

如果两者都是火炬张量，则：

lbl = labels.cpu().numpy().reshape(-1)
target = output.cpu().numpy().reshape(-1)

现在：

from sklearn.metrics import jaccard_similarity_score as jsc
print(jsc(target,lbl))