Introduction

Using TensorFlow to identify LEGO bricks is a practical image-classification problem with clear applications in inventory tools, educational games, and robotics. The model architecture matters, but data quality and labeling decisions matter more in early iterations. A reliable pipeline starts with clear class definitions, balanced image collection, and repeatable evaluation.

Define the Prediction Target Clearly

Before building the model, decide what each class represents:

• brick family such as plate, tile, brick
• exact part number
• part number plus color

These choices drastically change class count and dataset complexity. A common strategy is staged modeling:

1. predict shape family first
2. predict color as a separate output
3. combine predictions in post-processing

This reduces confusion and allows targeted error analysis.

Build a Dataset with Realistic Variation

Images should include multiple backgrounds, lighting conditions, and camera angles. If all training photos use one clean background, deployment performance will drop quickly in real scenes.

Simple directory layout:

1data/
2  train/
3    brick_1x2/
4    brick_2x4/
5    plate_2x2/
6  val/
7    brick_1x2/
8    brick_2x4/
9    plate_2x2/

Load datasets with TensorFlow utilities:

1import tensorflow as tf
2
3IMG_SIZE = (224, 224)
4BATCH_SIZE = 32
5
6train_ds = tf.keras.utils.image_dataset_from_directory(
7    "data/train",
8    image_size=IMG_SIZE,
9    batch_size=BATCH_SIZE,
10)
11
12val_ds = tf.keras.utils.image_dataset_from_directory(
13    "data/val",
14    image_size=IMG_SIZE,
15    batch_size=BATCH_SIZE,
16)
17
18AUTOTUNE = tf.data.AUTOTUNE
19train_ds = train_ds.cache().shuffle(1000).prefetch(AUTOTUNE)
20val_ds = val_ds.cache().prefetch(AUTOTUNE)

Start with Transfer Learning

For limited datasets, transfer learning is usually better than training from scratch.

1import tensorflow as tf
2
3num_classes = 3
4base = tf.keras.applications.MobileNetV2(
5    input_shape=(224, 224, 3),
6    include_top=False,
7    weights="imagenet",
8)
9base.trainable = False
10
11model = tf.keras.Sequential([
12    tf.keras.layers.Rescaling(1.0 / 255),
13    base,
14    tf.keras.layers.GlobalAveragePooling2D(),
15    tf.keras.layers.Dropout(0.2),
16    tf.keras.layers.Dense(num_classes, activation="softmax"),
17])
18
19model.compile(
20    optimizer=tf.keras.optimizers.Adam(1e-3),
21    loss="sparse_categorical_crossentropy",
22    metrics=["accuracy"],
23)
24
25history = model.fit(train_ds, validation_data=val_ds, epochs=8)

After baseline stability, unfreeze top layers for fine-tuning with a smaller learning rate.

Use Data Augmentation for Robustness

Augmentation helps the model tolerate orientation and lighting differences.

1augment = tf.keras.Sequential([
2    tf.keras.layers.RandomFlip("horizontal_and_vertical"),
3    tf.keras.layers.RandomRotation(0.08),
4    tf.keras.layers.RandomContrast(0.15),
5])

Apply augmentation in training only, not validation.

Evaluate with Class-Level Metrics

Overall accuracy can hide confusion between similar bricks. Use confusion matrix and per-class metrics.

1import numpy as np
2from sklearn.metrics import confusion_matrix, classification_report
3
4true_labels = []
5pred_labels = []
6
7for x_batch, y_batch in val_ds:
8    probs = model.predict(x_batch, verbose=0)
9    pred = np.argmax(probs, axis=1)
10    true_labels.extend(y_batch.numpy().tolist())
11    pred_labels.extend(pred.tolist())
12
13print(confusion_matrix(true_labels, pred_labels))
14print(classification_report(true_labels, pred_labels))

This quickly shows where more data is needed.

Deploy and Benchmark Early

If you plan edge deployment, convert to TensorFlow Lite and measure latency on target hardware.

1converter = tf.lite.TFLiteConverter.from_keras_model(model)
2tflite_model = converter.convert()
3
4with open("lego_classifier.tflite", "wb") as f:
5    f.write(tflite_model)

Keep preprocessing consistent between training and inference to avoid quality drops.

Include Hard Negatives

Add non-brick images, partial occlusions, and cluttered scenes as hard negatives. Without them, the model tends to over-predict known classes when input is ambiguous.

Hard negatives improve confidence calibration and reduce false positives in real-world use.

Common Pitfalls

A common pitfall is defining labels too aggressively, such as exact part plus color from day one, which makes data collection unmanageable.

Another pitfall is training on studio-clean images and deploying in cluttered environments.

A third pitfall is monitoring only top-line accuracy and ignoring class-level confusion.

Teams also forget to lock preprocessing behavior, causing train-inference mismatch after deployment.

Summary

• LEGO brick recognition is mainly a data and label-design problem, not only a model-design problem.
• Start with clear class scope and realistic image variation.
• Use transfer learning for strong early performance with limited data.
• Evaluate with confusion matrices, not just aggregate accuracy.
• Prepare deployment early by keeping preprocessing and runtime constraints explicit.