Image Classification Using Unity Barracuda and Tensorflow
Preface
In this article, I will develop a sample use case and test it End2End to show customized AI capabilities in Unity. And then I will introduce you to the Barracuda library. After this introduction, I will be presenting our use case(which is an Image classification task) using Barracuda, ONNX, VGG16 (layer), and Tensorflow Both from the Python Model training part to the Unity Model usage part.
Introduction
Note: You can find the Python Project in the References section
If you want to use AI capabilities with Unity, you might find a library called ML-Agents. But this ML-only library is using another library as a wrapper called Barracuda. That means if you want to use other than just ML (like Deep learning), you may want to use Barracuda.
An example of Barracuda inside Unity Editor:
Barracuda is a lightweight cross-platform neural network inference library developed by Unity, for Unity. Barracuda can run neural networks on both the GPU and CPU. If you want to check other products made with Barracuda inference, here.
So, How Does It Works?
It is pretty straightforward if you have a trained model ready.
- Train an AI model (we trained a basic CNN model using TF)
- Upload your pre-trained model into Unity
- Create a Worker (Unity uses this Worker as a wrapper for using models)
- Tensorize your input (Convert your dataset into tensor objects)
- Execute and check the output
here is a sample code snippet about image classification on different types of balls
using UnityEngine;
using Unity.Barracuda;
public class sampleCodeSnippet : MonoBehaviour
{
public NNModel onnxAsset;
public Texture2D imageToRecognise;
private IWorker worker;
void Start()
{
// Create a worker for executing the model
worker = onnxAsset.CreateWorker();
using (var input = new Tensor(imageToRecognise, channels: 3))
{
// execute neural network with specific input and get results back
var output = worker.Execute(input).PeekOutput();
// the following line will access values of the output tensor causing the main thread to block until neural network execution is done
var indexWithHighestProbability = output[0];
UnityEngine.Debug.Log($"Image was recognised as class number: " + output[0] + " " + output[1]);
}
}
}
onnxAsset takes our ONNX model, imageToRecognise is just an image for us to tensorize and execute with the model and worker is just here for creating a worker and executing the model with the given tensorized image.
Wait, What is ONNX?
As you can see in the above image, ONNX is an open format built to represent our models.
For example: Let’s say you built your model using Keras or Pytorch. Your model input can be changed to ONNX with 1 to 2 lines of Python code for extracting and using inside Unity. Example links: Keras (Tensorflow) Pythorch
Our Case-Study
In our case study, We gathered 80 to 90 beach-ball and football pictures from different resources. The aim was to build an image classification app using Barracuda in Unity.
We trained this model in Python than extracted this model in ONNX format, and then used this model inside Unity.
Below, you can see step-by-step clarification of the process.
Data Gathering
First of all, we would like to classify easy-to-find objects. So we agreed on 2 different ball images. beach-ball and football. We find the dataset mostly manually. We combine internet-found pictures with the pictures we screen-shotted in Unity, which makes 80–90 per category. Internet-found ones are diverse and have a distraction in them while the in-Unity ones are dedicated and just consist of the balls inside.
Model Training (Python Part)
You can find the Python project here.
The libraries that we used are Tensorflow, ImageDataGenerator from tensorflow.keras.preprocessing.image and Numpy.
Note: Tensorflow version is v2.9.2
We split our data for testing and training and then train a basic CNN model which includes convolutional and max-pooling layers. Then we extract it as ONNX for placing inside Unity and using it with the Barracuda library.
Then we converted out tensorflow model to ONNX using the tf2onnx package. tf2onnx supports the conversion of TensorFlow (tf-1.x or tf-2.x), Keras, tensorflow.js, and tflite models to ONNX via command line or Python API.
You can find the .ipynb file here -> https://github.com/MustafaCQN/Ball-Classification/blob/main/ball_classification.ipynb
Unity Project
In our Unity project, we are using this model to recognize given ball types. We are having 2 different balls as we described earlier, beach ball and football. B and F are the probability of the beach ball and the football. The recognized class prompt is the class with the highest probability.
Let’s first talk about the hierarchy. The main camera and the directional light are the components that every Unity project has, They will provide us the camera view and the light to see the objects in the Scene. Image canvas is the canvas for holding the photos and text outputs as the given Image(1.). “photo” is the first Image and “photo_2” is the second Image that we show in the Image(1.). beachText and footballText are the text outputs for the respected Images.
Let’s talk about the folder structure of the project. The Models folder is the folder that we put our trained models in, the Scenes folder is the Unity Scenes for the given scenario, Scripts are the folder that we reside our C# files inside, and Test Pictures are the pictures for us the try while testing the scene as we saw in the Image(1.), Lastly TextMesh Pro folder is the Unity-built folder that Unity uses for creating the TextMesh Pro instance in the hierarchy.
When we look inside the Scripts folder, we can see there are 4 different scripts in it. But we are only using one of them, so they are all doing the same work with different types and we want to let them rest for you to take code snippets for your specific use case when needed.
The project is using the sampleCodeSnippet.cs, Simplest is the most simple use-case for the given environment and others are just code examples for you to take snippets from as I mentioned (your specific use-case).
Lastly, I would like to show the Unity object in that we put our code in. You can change the location of the sampleCodeSnippet, but for ease find, I am showing the location (inside of the photo). As we see the Sample Code Snippet waits for us to give him the ONNX asset, images to recognize, photos for him to show, and the text place for him to print the outputs. Don’t forget to give them these objects from the Hierarchy or Directories.
Conclusion
In this article, we learned that we can train our models with Python and extract the model using ONNX for use inside Unity to execute our trained models. We learned what are the most important subcontexts inside Barracuda and ONNX.
Thank you for reading this article, feel free to ask if you have any questions.
