使用 ML.NET 識別樂高顏色塊

每一個樂高迷都擁有很多的顏色塊，需要進行排序和按類型分揀，按照《Organizing your LEGO Bricks》或許有所幫助，但這不是一個簡單的任務，因為有很多顏色塊有非常微妙的差異。如果換作一個典型的程序員可以做什么來解決這個問題呢？你猜對了 - 建立一個程序使用 ML.NET 來識別樂高的顏色塊。

首先，我們將創建一個控制台應用並添加所需的包

> dotnet new console
> dotnet add package Microsoft.ML
> dotnet add package Microsoft.ML.Vision
> dotnet add package Microsoft.ML.ImageAnalytics
> dotnet add package SciSharp.TensorFlow.Redist

在項目文件夾的根目錄中，我將創建一個名為 pieces 的子文件夾，並在此文件夾中創建一些顏色分類的子文件夾，放置訓練集中的每種顏色的圖片。

使用時，我們需要定義輸入和輸出模型（分類器提供分類結果）。

public class ModelInput
{
    public string Label { get; set; }
    public string ImageSource { get; set; }
}
 
public class ModelOutput
{
    public String PredictedLabel { get; set; }
}

為了訓練模型，我們首先創建一個由目錄中的圖像組成的輸入數據集，並將其作為標簽分配它們位於的目錄的名稱。在此之后，我們創建訓練管道，最后，使用數據進行訓練以創建模型。

static void TrainModel()
{
    // Create the input dataset
    var inputs = new List<ModelInput>();
    foreach (var subDir in Directory.GetDirectories(inputDataDirectoryPath))
    {
        foreach (var file in Directory.GetFiles(subDir))
        {
            inputs.Add(new ModelInput() { Label = subDir.Split("\\").Last(), ImageSource = file });
        }
    }
    var trainingDataView = mlContext.Data.LoadFromEnumerable<ModelInput>(inputs);
    // Create training pipeline
    var dataProcessPipeline = mlContext.Transforms.Conversion.MapValueToKey("Label", "Label")
                                .Append(mlContext.Transforms.LoadRawImageBytes("ImageSource_featurized", null, "ImageSource"))
                                .Append(mlContext.Transforms.CopyColumns("Features", "ImageSource_featurized"));
    var trainer = mlContext.MulticlassClassification.Trainers.ImageClassification(new ImageClassificationTrainer.Options() { LabelColumnName = "Label", FeatureColumnName = "Features" })
                                .Append(mlContext.Transforms.Conversion.MapKeyToValue("PredictedLabel", "PredictedLabel"));
    IEstimator<ITransformer> trainingPipeline = dataProcessPipeline.Append(trainer);
    // Create the model
    mlModel = trainingPipeline.Fit(trainingDataView);
}

現在，使用這個訓練模型，我們可以嘗試對一個新圖像進行分類。通過為其中一個圖像創建模型輸入，然后將它傳遞到使用分類器構建的模型創建的預測引擎。

static ModelOutput Classify(string filePath)
{
    // Create input to classify
    ModelInput input = new ModelInput() { ImageSource = filePath };
    // Load model and predict
    var predEngine = mlContext.Model.CreatePredictionEngine<ModelInput, ModelOutput>(mlModel);
    return predEngine.Predict(input);
}

最后讓我們用4種不同的顏色來測試這一點。

static void Main()
{
    TrainModel();
 
    var result = Classify(Environment.CurrentDirectory + Path.DirectorySeparatorChar + "Black.jpg");
    Console.WriteLine($"Testing with black piece. Prediction: {result.PredictedLabel}.");
    result = Classify(Environment.CurrentDirectory + Path.DirectorySeparatorChar + "Blue.jpg");
    Console.WriteLine($"Testing with blue piece. Prediction: {result.PredictedLabel}.");
    result = Classify(Environment.CurrentDirectory + Path.DirectorySeparatorChar + "Green.jpg");
    Console.WriteLine($"Testing with green piece. Prediction: {result.PredictedLabel}.");
    result = Classify(Environment.CurrentDirectory + Path.DirectorySeparatorChar + "Yellow.jpg");
    Console.WriteLine($"Testing with yellow piece. Prediction: {result.PredictedLabel}.");
}

結果如圖所示。

4張圖片對了3個！略微有點令人失望。但這是一個很好的開始，因為它給了我們機會去深入，並試圖了解如何改進分類，使其更准確。也許它需要更多的訓練數據，也許有更好的分類算法我們可以使用！

項目完整示例代碼和訓練數據在GIthub上：https://github.com/BeanHsiang/Vainosamples/tree/master/CSharp/ML/LegoColorIdentifier