Part 3: Developing the AI Model for Waste Sorting

1. Introducing AI Models You Can Use

Common AI models for waste sorting include:

MobileNetV2 / MobileNetV3: lightweight and optimized for edge devices like the Raspberry Pi.
EfficientNet-Lite: higher accuracy while still staying compact.
YOLOv8n: suited to detecting multiple objects in a single frame.
TensorFlow Lite INT8: optimized for running on a Raspberry Pi, but note that not every model can be quantized to INT8 without losing accuracy or needing an extra calibration step.

2. Gathering and Preparing Data

To train the model, you need to gather images of different waste types. You can use public data sources like TACO, or take photos yourself in a controlled environment.

Tip: make sure images are labeled accurately and consistently to improve the model's accuracy.

Split the data into training, validation, and test sets with a 70:15:15 ratio. The training set trains the model, the validation set tunes parameters, and the test set evaluates the model's final accuracy.

3. Creating and Training the AI Model with TensorFlow

Use TensorFlow to create and train the model:


import tensorflow as tf
from tensorflow.keras import layers, models

# Tạo mô hình
model = models.Sequential([
    layers.Input(shape=(224, 224, 3)),
    layers.Conv2D(32, (3, 3), activation='relu'),
    layers.MaxPooling2D(pool_size=(2, 2)),
    layers.Conv2D(64, (3, 3), activation='relu'),
    layers.MaxPooling2D(pool_size=(2, 2)),
    layers.Flatten(),
    layers.Dense(128, activation='relu'),
    layers.Dense(4, activation='softmax')  # 4 nhóm rác
])

# Biên dịch mô hình
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

# Huấn luyện mô hình
# Lưu ý: train_data, train_labels, val_data, val_labels cần được tiền xử lý và chia thành các lô dữ liệu
model.fit(train_data, train_labels, epochs=10, validation_data=(val_data, val_labels))

4. Exporting the Model to TensorFlow Lite Format

After training, export the model to TensorFlow Lite format to use on the Raspberry Pi:


# Xuất mô hình
converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()

# Lưu mô hình
with open('model.tflite', 'wb') as f:
    f.write(tflite_model)

5. Checking the Model's Accuracy

Evaluate the model's accuracy using the test dataset:


test_loss, test_acc = model.evaluate(test_data, test_labels)
print(f"Độ chính xác: {test_acc * 100:.2f}%")

6. Saving the Model and Preparing It for the Raspberry Pi

After testing, save the converted TensorFlow Lite model and transfer it to the Raspberry Pi. Use the TensorFlow Lite library to run the model on the device. For hardware acceleration like the Hailo-8L, you need to install the vendor's SDK and use their model-running API, not just plain TensorFlow Lite.

Warning: make sure the Raspberry Pi has enough resources to run the AI model efficiently. Prefer tflite_runtime over the full TensorFlow package due to performance and memory constraints.

References

TensorFlow Documentation