computer-vision - 如何将 35 类城市景观数据集转换为 19 类？

Question

以下是我的代码的一小段。使用它，我可以在 cityscapes 数据集上训练我的名为“lolnet”的模型。但是数据集包含 35 个类/标签 [0-34]。

imports ***

trainloader = torch.utils.data.DataLoader(
    datasets.Cityscapes('/media/farshid/DataStore/temp/cityscapes/', split='train', mode='fine',
                    target_type='semantic', target_transform =trans,
                    transform=input_transform ), batch_size = batch_size, num_workers = 2)

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net = lolNet()
criterion = CrossEntropyLoss2d()

net.to(device)
num_of_classes = 34

for epoch in range(int(0), 200000):

    lr = 0.0001

    for batch, data in enumerate(trainloader, 0):

        inputs, labels = data
        labels = labels.long()
        inputs, labels = inputs.to(device), labels.to(device)

        labels = labels.view([-1, ])

        optimizer = optim.Adam(net.parameters(), lr=lr)

        optimizer.zero_grad()
        outputs = net(inputs)

        outputs = outputs.view(-1, num_of_class)


        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()

        outputs = outputs.to('cpu')
        outputs = outputs.data.numpy()
        outputs = outputs.reshape([-1, num_of_class])

        mask = np.zeros([outputs.shape[0]])
        #
        for i in range(len(outputs)):
            mask[i] = np.argmax(outputs[i])

        mask = mask.reshape([-1, 1])

        IoU = jaccard_score(labels.to('cpu').data, mask, average='micro')

但我只想在 19 个类(class)上训练我的模型。找到这 19 个类 here .要训​​练的标签存储为“ignoreInEval”= True。此 pytorch Dataloader helper对于这个数据集没有提供任何线索。

所以我的问题是如何使用 pytorch 的“datasets.Cityscapes”api 在这个数据集的所需 19 个类上训练我的模型。

Answer 1

已经有一段时间了，但留下一个对其他人有用的答案:
首先创建一个映射到 19 个类 + 背景。背景与具有忽略标志的不那么重要的类有关，如所述 here .

# Mapping of ignore categories and valid ones (numbered from 1-19)
    mapping_20 = { 
        0: 0,
        1: 0,
        2: 0,
        3: 0,
        4: 0,
        5: 0,
        6: 0,
        7: 1,
        8: 2,
        9: 0,
        10: 0,
        11: 3,
        12: 4,
        13: 5,
        14: 0,
        15: 0,
        16: 0,
        17: 6,
        18: 0,
        19: 7,
        20: 8,
        21: 9,
        22: 10,
        23: 11,
        24: 12,
        25: 13,
        26: 14,
        27: 15,
        28: 16,
        29: 0,
        30: 0,
        31: 17,
        32: 18,
        33: 19,
        -1: 0
    }

然后对于您为训练加载的每个标签图像(每个像素包含一个类的灰色图像，其模式为“{city}__{number}_{number}_gtFine_labelIds.png”)，运行下面的函数。
它将根据上面的映射转换每个像素，您的标签图像(蒙版)现在只有 20 个(19 个类 + 1 个背景)不同的值，而不是 35 个。

def encode_labels(mask):
    label_mask = np.zeros_like(mask)
    for k in mapping_20:
        label_mask[mask == k] = mapping_20[k]
    return label_mask

然后，您可以使用这些新数量的类正常训练您的模型。