Training a CNN for LEGO Bricks Recognition
https://medium.com/rocket-science-team/training-a-cnn-for-lego-bricks-recognition-f285ffab3327
This project focuses on building a convolutional neural network (CNN) model to classify LEGO bricks. The model is trained on a dataset of images, aiming to recognize 20 different types of LEGO bricks. It was developed as part of a Data Science bootcamp and is designed to help LEGO collectors and enthusiasts identify pieces quickly.
The full project description can be found in PROJECT_DESCRIPTION.md.
- Aleksandra Baran - LinkedIn | GitHub
- Dariusz Balcer - LinkedIn | GitHub
- Grzegorz Gomza - LinkedIn | GitHub
LEGO is a globally recognized brand, and with this project, we aim to create a system that can recognize individual LEGO bricks from photos. This recognition could be helpful for collectors, enthusiasts, and even in app development for various LEGO-related functionalities.
We used convolutional neural networks (CNNs) for this image classification project. Two models were implemented:
- RocketNet: A custom CNN designed for this project.
- AlexNet: A scaled-down version of the well-known AlexNet architecture.
Both models were trained on a subset of the B200C LEGO Classification Dataset from Kaggle, which contains approximately 4,000 images per class for 20 different LEGO brick types.
-
RocketNet:
- Simple architecture with three convolutional layers, dropout, and dense layers.
- Achieved an accuracy of 88.9% on the validation set.
-
AlexNet:
- More complex architecture with five convolutional layers and two dense layers.
- Achieved an accuracy of 90.9% on the validation set but had some concentrated errors.
We used Grad-CAM (Gradient-weighted Class Activation Mapping) to visualize which parts of the images were being focused on by the models during classification. This helped in understanding the decision-making process of the models.
The trained model has several potential applications, such as:
- LEGO Building App: To recognize available bricks and suggest models to build.
- Assistive Technology for Visually Impaired: LEGO recognition with voice prompts for easy building.
- Collection Management App: For sorting and categorizing bricks.
- LEGO Shopping and Exchange App: To help identify and catalog LEGO pieces for easy transactions.
The project demonstrates the effectiveness of CNNs in classifying LEGO bricks. Despite hardware limitations, the models showed good performance, with the simpler RocketNet providing comparable results to the more complex AlexNet. This approach could be further improved and expanded for real-world applications.