Pytorch：如何为语义分割计算 IoU（Jaccard 索引）

Question

有人可以提供一个玩具示例，说明如何在 pytorch 中计算 IoU（联合交集）以进行语义分割吗？

Answer 1

我在某个地方找到了这个并为我改编了它。如果我能再次找到它，我会发布链接。抱歉，如果这是重复的。
这里的关键函数是名为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 2

从 2021 年开始，您无需实施自己的 IoU，因为 torchmetrics 配备了它 - here's the link。它适用于 PyTorch 和 PyTorch Lightning，也适用于分布式训练。

来自文档：

torchmetrics.IoU(num_classes, ignore_index=None, absent_score=0.0, threshold=0.5, reduction='elementwise_mean', compute_on_step=True, dist_sync_on_step=False, process_group=None)

计算交集而不是联合，或 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 有一个额外的维度，就像多类分数的情况一样，我们会在 dim=1 上执行 argmax。

官方示例：

from torchmetrics import IoU
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]
iou = IoU(num_classes=2)
iou(pred, target)

Answer 3

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

然后你就可以使用 sklearn 的jaccard_similarity_score 进行一些整形了。

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

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))