使用 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 datasetvar 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 pipelinevar 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 modelmlModel = trainingPipeline.Fit(trainingDataView);
}

现在，使用这个训练模型，我们可以尝试对一个新图像进行分类。通过为其中一个图像创建模型输入，然后将它传递到使用分类器构建的模型创建的预测引擎。

static ModelOutput Classify(string filePath)
{// Create input to classifyModelInput input = new ModelInput() { ImageSource = filePath };// Load model and predictvar 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}.");
}

结果如图所示。