Get Labels from loaded ML.NET model

Question

I have followed the ML.NET tutorial for image classification and already created my first Model. (https://docs.microsoft.com/en-us/dotnet/machine-learning/tutorials/image-classification)

When I load the Saved model I write the following:

trainedModel = mlContext.Model.Load(modelRelativePath & ModelName, modelSchema)

Once I run the model with a picture, it returns if the picture is a cat or dog. (The labels used are "CAT" and "DOG")

Is there a way to read the whole list of labels from the loaded model? I would like to display it once the model is loaded.

I have searched the trainedmodel tree but couldn't find them. Thanks

Answer 1

You can use the following method to get the current model labels as a string array

public string[] GetSlotNames(string name) {
    PredictionEngine<ImageData, ImagePrediction> predictionEngine = mlContext.Model.CreatePredictionEngine<ImageData, ImagePrediction>(mlModel);
    DataViewSchema.Column? column = predictionEngine.OutputSchema.GetColumnOrNull(name);

    VBuffer<ReadOnlyMemory<char>> slotNames = new VBuffer<ReadOnlyMemory<char>>();
    column.Value.GetSlotNames(ref slotNames);
    string[]? names = new string[slotNames.Length];
    int num = 0;

    foreach (var denseValue in slotNames.DenseValues()) {
        names[num++] = denseValue.ToString();
    }

    return names;
}

To use it provide it with your Score column from your ImagePrediction model

string[] labels = GetSlotNames("Score");