How To Train YOLOv5 For Recognizing Game Objects In Real-Time

This article will look at how to train YOLOv5 for object detection and recognition with Python and PyTorch on a custom dataset. For this example, the dataset will contain objects from a game called Albion Online.

The other day, I watched a series on how to write game bots on the YouTube channel “Learn Code By Gaming.” In the video titled “Training a Cascade Classifier — OpenCV Object Detection in Games #8”, the author shows how he builds a simple object classifier.

I found the results from that video inspiring. Therefore, I thought it would be interesting to try to improve those results by using a more advanced classifier algorithm. This is how I started experimenting with the YOLOv5 object detection algorithm using Python and PyTorch. The content of this article is based on these experiments.

In what follows, I’ll show you how to take screenshots from a free game called Albion Online. These screenshots will also be edited for privacy, as multiple players will be visible. This will have no impact on the object detection and recognition later on.

Subsequently, I’ll show you how to label the game objects that represent basic resources in a specific area of the game world of Albion Online. Once labeled, both the labels and the screenshots will be used as training and validation data for the YOLOv5 algorithm. The result will be a new PyTorch model of our data.

Finally, I’ll show you how to export the trained PyTorch model as a TensorRT model optimized for fast inference. This is necessary for (quasi) real-time inference, so object detection and recognition can be performed while playing the game.

Take screenshots from your game

Let’s start by taking some screenshots from the Albion Online game. So make sure to have the game installed — it’s free. Also, note I’ll be working on Windows.