Tag: lego

There comes a time in every LEGO^® fan’s life where you have so many pieces that sorting stuff by the type of piece is not good enough (following the 6 steps of lego organization). But this is not a simple task since there are many colors, some with very subtle differences. So what can a typical programmer do to solve this problem? You guessed it – build a program that can identify the color of a LEGO piece :-). Since I’m writing in C# most of the time, I decided the most natural thing to do was play with ML.NET and see where that takes me. In this first try, I’ve created a console app that I train on a labeled set of LEGO pieces and then give new pieces to see if it can identify their color correctly. Disclaimer: I’m not an ML expert (yet :-)) and this project is just a way to start learning, so don’t expect a deep explanation of why things work. Hopefully, I’ll have more time in the future to learn deeper and answer those questions, but for now, I’m treating ML.NET as a black box and doing what I call monkey coding. So having said that, let’s get started!

First, we’ll create a console app and add the required packages

> 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

In the root of the project folder, I’m going to create a subfolder called “pieces” and in this folder a subfolder for each of the colors I have in my training set. In the color folder, I’m putting the pictures for that color. Like this:

When working with ML.NET we need to define the input model, which is what the trainer/classifier receives, and an output model, which is what the classifier gives as the classification result. In our case, this is pretty simple:

public class ModelInput
{
    public string Label { get; set; }
    public string ImageSource { get; set; }
}

public class ModelOutput
{
    public String PredictedLabel { get; set; }
}

To train a model, we first create an input data set consisting of the images in the pieces directory and assigning them as a label the name of the directory where they are located. After this, we create the training pipeline and finally, train with the data to create the model.

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);
}

Now using this trained model, we can try and classify a new image. This is done by creating a model input for one of the images and passing it to a prediction engine that is created with the model built by the classifier.

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);
}

So putting it all together, let’s test this using 4 different pieces of those colors.

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}.");
}

And the results… (drum roll!!!)

Three out of four! That is… a bit disappointing :-(. But it is a great start because it gives us the opportunity to go deeper and try to understand how we can improve the classification to make it more accurate. Maybe it needs more training data? Maybe there is a better classification algorithm we can use? So many things to learn!

The full code for this tutorial (and much more) can be found on GitHub: https://github.com/vainolo/Vainosamples/tree/v6/CSharp/ML/LegoColorIdentifier.

Keep safe, enjoy life, and until next time, happy coding!