使用 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