Integrating TensorFlow Models with Unity

Oct 6, 2024

Importing TensorFlow Models into Unity

Introduction

  • Overview of the tutorial: Importing a Python model generated using TensorFlow into Unity.
  • Standard method: Using ONNX format with Barracuda library.
  • Alternative method: Running a Python server that communicates with Unity via WebSocket.

Reasons for Using the Alternative Method

  • Issues with converting TensorFlow models (h5, saved models) into ONNX format.
  • Challenges such as version mismatches and optimizers complicating the conversion.
  • Simplicity of the new method despite needing a Python server running.

How the Alternative Method Works

  • A Python server runs TensorFlow model.
  • Unity communicates through WebSocket for model input and output.
  • Example model trained on the EMNIST dataset (letters).
    • Ability to draw letters and get predictions.

User Interface Overview

  • UI features:
    • Two sliders controlling the radius of the input area and prediction strength.
    • Option for constant prediction (may stretch WebSocket).
    • Reset button to clear the drawing table.

Code Structure

Overview of Classes

  1. Pixel Class: Controls individual pixels (black, white, shades of gray).
  2. Pixels Handler Class: Generates a 21x21 pixel table as input and handles drawing.
    • Uses physics overlap circle to modify pixels based on mouse input.
  3. Prediction Function: Reads pixel data and sends it to the Python server for predictions.
  4. Error Handling: Placeholder for managing exceptions and connection issues.

Python Server Setup

  • Python code to load the model and set up a connection using ZeroMQ.
  • Continuously listens for incoming bytes of pixel data from Unity.
  • Converts incoming bytes to float array and makes predictions using TensorFlow.
  • Sends predictions back to Unity.

Important Notes

  • Ensure the Python server runs whenever Unity is launched.
  • Consider implementing scripts to manage server restarts.
  • Potential to run the Python server externally for better performance.

Conclusion

  • This method allows executing a model without full conversion to Unity.
  • Useful for testing models quickly.
  • Encouragement to leave feedback and suggestions for future tutorials.

Call to Action

  • Like and subscribe for more similar content.
  • Share experiences in the comments regarding TensorFlow models and Unity integration.