zhanghang1989 / PyTorch-Encoding

A CV toolkit for my papers.
https://hangzhang.org/PyTorch-Encoding/
MIT License
2.04k stars 450 forks source link

Which index matches which class in ADE20k? #415

Open mohwald opened 1 year ago

mohwald commented 1 year ago

When I output with any model, I got probabilities per class on each pixel. E.g. when I run

model = encoding.models.get_model('DeepLab_ResNeSt101_ADE', pretrained=True).eval()
img = torch.randn(1,3,256,256)
output = model.evaluate(img)
print(output.shape)

I got torch.Size([1, 150, 256, 256]), since there are 150 classes in ADE20k dataset.

However, I did not find anywhere in the code, which prediction of the output[0, :, x, y] corresponds to which of the 150 classes in ADE dataset. Where can I find list of classes?