Press "Enter" to skip to content

使用 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个!略微有点令人失望。但这是一个很好的开始,因为它给了我们机会去深入,并试图了解如何改进分类,使其更准确。也许它需要更多的训练数据,也许有更好的分类算法我们可以使用!

Be First to Comment

发表评论

您的电子邮箱地址不会被公开。 必填项已用*标注