Labeling class Name

**The model is 79.82897162437439 % certain that the image has a predicted class of 2

How to make the output show the real name if the class instead of the index ?

def predict(image, model):
    # Pass the image through our model
    image = image.to(device)
    output = model.forward(image)
    
    # Reverse the log function in our output
    output = torch.exp(output)
    
    # Get the top predicted class, and the output percentage for
    # that class
    probs, classes = output.topk(1, dim=1)
    
   
    
    return probs.item(), classes.item()

You could create a mapping between the class indices and the corresponding names using a dict:

idx_to_name = {
  0: 'class0',
  1: 'class1',
  ...
}

print(idx_to_name[target])
1 Like