DEV Community

Cover image for Training an image classification model to identify goodies and baddies using ML.NET
Lee Englestone 💡🧠👨‍💻
Lee Englestone 💡🧠👨‍💻

Posted on • Edited on

Training an image classification model to identify goodies and baddies using ML.NET

TLDR; I created an image classification model to determine whether an image is of a goody or a baddy using images of good and bad cartoon characters.

Introduction

If you spent your childhood anything like me, you will have watched a lot of cartoon shows ranging from well known shows like the Teenage Mutant Ninja Turtles and Transformers to lesser known series such as Visionaries and Mask. You learn at an early age to determine who is a baddy and who is a goody just by their appearance.

Premise

I have wanted to play around with the Machine Learning capability in ML.NET and its ability to create image classification models for a while now. I had previously seen such examples where machine learning models were trained and created to identify things like species of plant to breed of dog from an image. I wanted to see if I could create one to identify good or bad cartoon characters based on their appearance. Something simple and not too ambitious.

My premise is that, most people at a glance are able to determine if a cartoon character is good or evil solely by their appearance, I wanted to see if I could create a machine learning model to do the same.

But first i'd need some images..

Building and Training the Model

So I set about downloading a lot of images of cartoon characters that are known to be good or bad and put them in separate folders named wait for it.. "good" and "bad". This would be my training data.

I then found a few code samples that handled image classification scenarios such as determining whether an image of a surface contained a crack or not. I extended these code samples to suit my purposes.

https://docs.microsoft.com/en-us/dotnet/machine-learning/tutorials/image-classification-api-transfer-learning

Bias

Unconscious bias is something that I always try and be conscious of every day in every situation.

I am not a data scientist, but I am aware of the problem of bias in machine learning datasets. For example, regarding the images that I have scoured the internet to include in my training images.. what if I unintentionally chose more good women than bad, or always chose smiling goodies and dark images of baddies?

I quickly realised that in a perfect world, for my experiment to be better, it would have to choose images in a lot more random manner than my manual curation.

That said, this is a throw away experiment, not a scientific paper so let's continue..

Problems Encountered

So I downloaded a bunch of images of goodies and baddies, and ran the same model training/creation sample code that I found online but was hitting a roadblock..

It just didn't seem to work! The code was populating the training set, the validation set but the test set was always empty and resulted in a empty sequence exception later on in the code..

IDataView trainSet = trainSplit.TrainSet;
//var trainSetPreview = trainSet.Preview();

IDataView validationSet = validationTestSplit.TrainSet;
//var validationSetPreview = validationSet.Preview();

IDataView testSet = validationTestSplit.TestSet;
//var testSetPreview = testSet.Preview(); // Empty when there aren't sufficient images!
Enter fullscreen mode Exit fullscreen mode

I then realised that the machine learning model (ImageClassificationTrainer.Architecture.ResnetV2101) expects the images in dimensions of 224px x 224px so spent a LOT of time cropping and resizing images, resulting in 70 images of goodies and 70 images of baddies.

But it still didn't seem to work. Same error..

Then, out of desperation, I wondered if the number of training images I was supplying was insufficient so I copied and pasted them all, creating a ..- copy.jpg version of each. Thereby doubling the number of training images (but had no idea of how this would effect the model).

I ran the code again and voila! this time it worked. It would seem that there is a minimum number of images needed before a useful model can be created?

I'm pretty sure that my workaround isn't a valid real world solution for training a model with a sufficient number of images, however for the purpose of being able to proceed.. I continued. I managed to create the model and save it to a .zip file for use later on.

Testing the Model

So here are the results of the first 6 test images I ran through the model (no these images weren't part of the training images).

Test images

Baddy 1
Expected: Bad
Result: Bad

Baddy 2
Expected: Bad
Result: Bad

Baddy 3
Expected: Bad
Result: Bad

Goody 1
Expected: Good
Result: Good

Goody 2
Expected: Good
Result: Good

Goody 3
Expected: Good
Result: Good

So I guess it kind of works!? Success!

I'm sure there are a lot more tests you could do to put the model through its paces but for the purpose of my little experiment, I am happy with the result.

I now have an image classification model that can determine whether a cartoon character is likely to be a goody or a baddy.

GoodyOrBaddy.com

oh, also I may have fed my domain registration habit and registered GoodyOrBaddy.com to put up a web UI front end for uploading and testing images of goodies and baddies. This isn't live yet and wont be in the near future (other priorities!).

Summary

As I have said, I am not a data scientist or machine learning expert, however this has been one of my more enjoyable experiences experimenting with some very powerful technology and starting to learn ML.NET and machine learning.

I'd love to hear feedback (positive criticism) as to how I could have better approached the experiment. I've quickly made the code, training and test images available on GitHub in case anyone wants to play around with this in the same way I have.

It does make me wonder how many other use cases, this image machine learning image classification functionality could be used in..

Thanks for reading, and if you have enjoyed, please share!

Follow me on Twitter (@LeeEnglestone) to see my future experiments.

-- Lee

Update

I now have now published https://www.GoodOrBaddy.com which calls the trained model which is hosted in an Azure Function. The model isn't 100% accurate but pretty accurate. Perhaps further training is required.

Top comments (0)