Skip to content

LEGO AI Classifier v1.0 – Pretrained Model (Work in Progress)

Pre-release
Pre-release

Choose a tag to compare

@Kane-Pan Kane-Pan released this 28 Mar 07:05
· 8 commits to main since this release
4287d0d

(Project is still ongoing)

This release includes a pretrained PyTorch model for classifying LEGO pieces using computer vision. The model was trained on over 600,000 images across 447 official LEGO brick classes, using a combination of real photographs (10%) and rendered images (90%). The model achieves 99.49% validation accuracy.

Files Included:

train_lego_classifier.py - Script to train or resume training the LEGO classifier model on image datasets using PyTorch.

lego-live_classifier.py - Real-time classification script that uses a connected camera or image input to predict LEGO brick classes using the pre-trained model.

AI_Lego_Classifier_Model.pth - Pre-trained PyTorch checkpoint containing model weights, optimizer/scheduler states, and training metadata

Dataset Source

This model was trained using the open dataset:
LEGO Bricks for Training Classification Network (v1.1)
Gdańsk University of Technology (2021)
447 LEGO part classes
Rendered + photographed images
Licensed under CC BY 4.0
Dataset page & download