A computer vision system for identifying and classifying LEGO parts using deep learning techniques. The project combines object detection and metric learning to create a robust pipeline for LEGO part recognition from photographs.
Future implementation: The program will suggest LEGO sets the parts could be from, allowing the users to recreate original sets from a chaotic mix of lego parts.
FindMyBrick addresses the challenge of automatically identifying LEGO parts from images, which is particularly useful for sorting large collections of LEGO pieces or verifying inventory. The system employs a two-stage approach: first detecting individual LEGO pieces in images, then classifying each detected piece to determine its specific part number.
The project uses YOLOv8 for detecting individual LEGO pieces within images. The detection model is trained on custom labeled data to identify LEGO bricks with high accuracy, even in cluttered scenes with multiple pieces.
For classification, the system implements a Vision Transformer (ViT) pretrain on imagenet. Model is fine tuned with triplet loss and semi-hard negative mining. This metric learning approach creates embeddings that group similar LEGO parts together in the feature space while separating different parts.
Part matching is done through calculating distances between input and anchor embeddings of each known part. Last hidden state of ViT is used to extract embedding of each part, instead of using cls token to output class prediction, to ensure scaling, because additional LEGO parts will be added in the future.
Key Features:
- Semi-hard triplet mining: Automatically selects semi-hard negative examples during training to improve model robustness
- Image Hashing: Identifies and removes duplicate images from the crawled training dataset
- Multi-threaded data collection: Efficiently crawls LEGO part images from BrickLink catalog pages
- Validation framework: Testing pipeline to evaluate model performance
- Deep Learning: PyTorch, Transformers (Hugging Face)
- Computer Vision: YOLOv8 (Ultralytics), Pillow
- Data Processing: NumPy, pandas
- Web Scraping: Selenium, Multithreading
- Model Architecture: Vision Transformer (ViT-base-patch16-224)
- Training Strategy: Triplet loss with cosine distance metric
The training data consists of LEGO part images collected from Rebrickable and BrickLink, organized by part categories including:
- Animal parts
- Basic bricks
- Baseplates
- Specialized components
Images are preprocessed and deduplicated using hashing to ensure dataset quality. The system maintains separate anchor images (highest quality reference) and sample images for each part type.
- Data Preparation: Images are organized by part ID with quality-based ranking
- Triplet Generation: Semi-hard negative mining creates challenging training examples
- Model Training: ViT model learns discriminative embeddings using triplet loss
- Validation: Performance is evaluated using cosine similarity matching against anchor embeddings
The trained model can classify LEGO parts by:
- Processing input images through the detection pipeline to isolate individual pieces
- Generating embeddings for detected pieces using the trained ViT model
- Matching embeddings against the anchor database using cosine similarity
- Returning the closest matching part ID
The system achieves competitive accuracy on part classification tasks, with the metric learning approach proving particularly effective for handling the large number of similar-looking LEGO variants. The semi-hard mining strategy significantly improves the model's ability to distinguish between closely related parts.
- Expansion to additional LEGO part categories
- Real-time inference optimization
- Integration with existing LEGO sets
- Enhanced data augmentation techniques for improved robustness