Lego Sorting System

An automated pipeline that detects Lego brick shape (and, optionally, color) from a live video feed, as the first stage of an automated sorting system. Physical sorting hardware (Arduino/Raspberry Pi actuation) is a later phase -- this covers everything up through live detection.

---

1. Project Overview

The pipeline has four stages:

1. Capture (capture/) -- guided image capture from a VDO.Ninja stream or USB webcam, walking through each brick class and a checklist of orientation variations.
2. Labeling (labeling/) -- a local, offline bounding-box annotation tool (no cloud account needed), plus a script that splits labeled images into a YOLO-ready train/val/test dataset.
3. Training (training/) -- fine-tunes a pretrained YOLOv8 model on your dataset. Includes a grayscale-conversion step so the shape model can't use color as a classification shortcut (see Section 5).
4. Inference (inference/) -- runs the trained model against a live feed, with an optional second stage (color_classification/) that classifies the color of each detected piece separately.

Project structure

lego_sorter/
├── requirements.txt
├── setup.bat
├── README.md                      <- this file
├── capture/                        # Phase 1: guided image capture
│   ├── config.py                   #   BRICK_CLASSES, VARIATION_PROMPTS
│   ├── base.py, vdo_ninja_source.py, webcam_source.py
│   ├── session.py, main.py
├── labeling/                       # Phase 1.5: local annotation
│   ├── annotate.py, split_dataset.py
├── training/                       # Phase 2: model training
│   ├── make_grayscale_dataset.py
│   ├── train.py, evaluate.py
├── inference/                      # Phase 3: live detection
│   └── live_infer.py
├── color_classification/           # Phase 3b: color stage (optional)
│   ├── config.py, capture_reference.py, classify.py
├── data/
│   ├── raw/                        # captured + labeled images, by class
│   └── color_references/           # saved color reference swatches
├── dataset/                        # generated split (color)
└── dataset_gray/                   # generated split (grayscale, for training)

---

2. Setup

1. Ensure Python is installed and on your PATH.
2. From the project root, run:

   setup.bat
   

   This installs everything in requirements.txt, then runs playwright install chromium (a required separate download for Playwright's browser binary -- the pip package alone isn't enough), then verifies the key imports.
3. (Optional, recommended as the project grows) Use a virtual environment:

   python -m venv venv
   venv\Scripts\activate
   setup.bat
   

   Activate it manually each session afterward -- a batch script can't leave a venv active in the terminal window that launched it.

---

3. Usage, phase by phase

3.1 Capture

python -m capture.main --source vdoninja --url "https://vdo.ninja/?view=AgWSsA9y"
python -m capture.main --source webcam --device 0

Walks through each class in capture/config.py's BRICK_CLASSES, then each variation prompt, capturing IMAGES_PER_VARIATION shots per variation (s = capture, r = redo, n = skip variation, q = quit).

3.2 Labeling

python -m labeling.annotate

Draws one bounding box per image (class is inferred from the folder). Controls: drag to draw, s save + next, u undo, x exclude a bad photo, q quit (safe to resume later -- it skips anything already labeled).

3.3 Split into train/val/test

python -m labeling.split_dataset

Produces dataset/ in YOLO format, split per-class so small datasets still get proportional representation, with data.yaml included.

3.4 Convert to grayscale (for shape training)

python -m training.make_grayscale_dataset

Produces dataset_gray/ -- same split and labels, grayscale (3-channel) images. See Section 5 for why this matters.

3.5 Train

python -m training.train --data dataset_gray/data.yaml

Fine-tunes YOLOv8n from COCO-pretrained weights. Best weights land at runs/train/lego_shape_mvp/weights/best.pt.

3.6 Evaluate

python -m training.evaluate --weights runs/detect/runs/train/lego_shape_mvp/weights/best.pt

Prints overall + per-class mAP50 and mAP50-95 on the held-out test split.

3.7 (Optional) Capture color references

python -m color_classification.capture_reference --source vdoninja --url "https://vdo.ninja/?view=AgWSsA9y"

One photo per color you want to detect (see Section 6).

3.8 Live inference

python -m inference.live_infer --source vdoninja --url "https://vdo.ninja/?view=AgWSsA9y" --weights "runs\detect\runs\train\lego_shape_mvp\weights\best.pt"

Runs shape detection (on a grayscale copy of each frame) and, if color references exist, classifies the color of each detection from the original color frame. Press q to quit. Add --no-color to skip color entirely.

---

4. Adding new brick types later

Say you're currently detecting 2x2, 2x3, 2x4 and want to add 2x5, 2x6, 2x7. The pipeline is designed so this doesn't require starting over:

1. Append the new classes to the end of BRICK_CLASSES in capture/config.py:

   BRICK_CLASSES = [
       "2x2", "2x3", "2x4",   # existing -- do not reorder or remove these
       "2x5", "2x6", "2x7",   # new -- always add to the END
   ]

   This ordering matters. Each class's position in this list becomes its numeric class ID in every YOLO label file already on disk. If you insert a new class in the middle or reorder existing ones, every already-labeled image's class ID silently points to the wrong class. Appending to the end keeps all your existing labels valid.

2. Capture data for the new classes only:

   python -m capture.main --source vdoninja --url "..."
   

   The capture tool walks through every class currently in BRICK_CLASSES each run -- it doesn't know which ones you've already captured. The simplest way to capture only the new classes: temporarily comment out the existing class names in BRICK_CLASSES for this capture session, then restore the full list (existing + new) afterward before labeling. Aim for the same volume you used for the original classes.

3. Label the new images:

   python -m labeling.annotate
   

   Since already-labeled images are automatically skipped, this will only prompt you for the new, unlabeled images.

4. Re-run the split (rebuilds dataset/ from all labeled images, old and new):

   python -m labeling.split_dataset
   

5. Regenerate the grayscale dataset:

   python -m training.make_grayscale_dataset
   

6. Retrain. With datasets this size, a full retrain from the COCO-pretrained base (yolov8n.pt) is fast (minutes) and safer than incrementally fine-tuning your previous best.pt -- incremental fine-tuning on a tiny new batch risks the model drifting on the classes it already knew (catastrophic forgetting):

   python -m training.train --data dataset_gray/data.yaml
   

7. Re-evaluate all classes, not just the new ones -- check the confusion matrix for cross-confusion between old and new classes (e.g. 2x4 vs 2x5 is a similarly subtle difference to your earlier 2x2 vs 2x3 confusion, and is exactly the kind of thing to check for).

---

5. Why grayscale training (and how it works)

If a shape class is represented by only one or two colors during training, the model can learn to use color as a shortcut for classification instead of the actual geometric shape -- it's an easier signal to pick up on than subtle stud-count differences. That shortcut then breaks the moment you test with a color the model didn't associate with that class.

Converting training images to grayscale removes color information from the model's input entirely, so this shortcut is physically unavailable -- the model is forced to learn the actual silhouette (outline, stud pattern, aspect ratio).

This must be applied consistently: training/make_grayscale_dataset.py converts the training data, and inference/live_infer.py converts each live frame to grayscale the same way before running shape detection. Training and inference need matching preprocessing, or you've recreated a different version of the same domain-shift problem.

One upside: since color no longer matters for the shape model, you don't need every color of every shape in your training data -- orientation variety matters far more than color variety now.

---

6. Color classification

Color is handled as a separate, much simpler problem from shape -- you don't need a deep model or images of every shape in every color:

1. Capture one reference photo per color you want to detect:

   python -m color_classification.capture_reference --source vdoninja --url "..."
   

   Point the camera at a solid-color piece centered in frame for each color listed in color_classification/config.py's COLOR_NAMES. The tool samples the center of the frame (avoiding background) and stores its average color in LAB space (a color space where Euclidean distance corresponds well to perceived color difference) to data/color_references/references.json.

2. Classification at inference time: once the shape model detects a piece and its bounding box (on the grayscale frame), live_infer.py crops that same region from the original color frame and calls classify_color(), which compares the crop's average color (background pixels excluded via a brightness threshold) against your saved references and returns the closest match.

3. Limitation to keep in mind: this is a simple nearest-reference color match, not a trained model -- it's lighting-sensitive in the same way shape detection was before you fixed that with grayscale. Capture your color references under lighting similar to your deployment setup, and expect to recapture references if you change your rig's lighting meaningfully.

To add a new color later: add its name to COLOR_NAMES in color_classification/config.py and re-run capture_reference.py -- it will only prompt for colors not already in the references file.

---

7. Known limitations / next steps

• Color classification accuracy hasn't been validated yet the way shape detection has -- treat the first live test as diagnostic, same as the shape model's first run.
• Physical sorting mechanism (Arduino/Raspberry Pi actuation) is not yet built -- live_infer.py's per-frame detection output (class + confidence
  • color) is the hook point for that once you're ready.
• For scaling to a much larger piece catalog later, see the discussion of embedding-based / few-shot retrieval recognition from earlier in this project -- worth revisiting once class count grows into the dozens.