• Kesiena Berezi
• Ball Bearing Classification Project
• MEE206 - Mechanical Student Design Project
• Description: This project implements a complete computer vision and robotics system

•          for automatically classifying and sorting ball bearings by material type
  

•          (Brass, Nylon, Steel) and size (1/2" and 1/4"). The system combines
  

•          machine learning, real-time inference on Raspberry Pi, and mechanical
  

•          sorting using Arduino-controlled servos.
  

---

Ball Bearing Classification Project

Overview

This project implements a complete computer vision and robotics system for automatically classifying and sorting ball bearings by material type (Brass, Nylon, Steel) and size (1/2" and 1/4"). The system combines machine learning, real-time inference on Raspberry Pi, and mechanical sorting using Arduino-controlled servos.

Features

• Machine Learning Pipeline: VGG16-based CNN for material classification with transfer learning
• Real-time Inference: TensorFlow Lite deployment on Raspberry Pi with camera integration
• Mechanical Sorting: Arduino-controlled servo system for automated bearing routing
• Serial Communication: Seamless communication between Pi and Arduino for classification codes
• Data Visualization: Tools for exploring dataset and analyzing model performance
• Quantized Models: Optimized models for fast inference on edge devices

System Architecture

The system consists of three main components:

1. Raspberry Pi: Captures images and performs real-time classification
2. Arduino: Controls servo motors for mechanical sorting
3. Machine Learning Model: Classifies bearings into 6 categories (3 materials × 2 sizes)

File Descriptions

Machine Learning (ML/)

VGG16Train (2).py

• Trains a VGG16-based convolutional neural network for ball bearing classification
• Uses transfer learning with ImageNet weights and data augmentation
• Supports 3 material classes: Brass, Nylon, Steel
• Usage:

  python VGG16Train.py [train_dir] [validation_dir] [model_save_path]

VGG16ID (1).py

• Evaluates trained VGG16 model on test data
• Generates confusion matrix, classification reports, and visualizes misclassified samples
• Provides comprehensive model performance analysis
• Usage:

  python VGG16ID.py [model_path] [test_data_dir]

ReorgData_materials.py

• Organizes ball bearing image dataset into train/validation/test splits
• Creates structured directories for machine learning pipeline
• Supports 70/15/15 split ratio for Brass, Nylon, and Steel classes
• Usage:

  python ReorgData_materials.py

explore_images.py

• Visualizes sample images from ball bearing dataset
• Generates sample visualization grid for each material class
• Helps understand data distribution and image characteristics
• Usage:

  python explore_images.py

Raspberry Pi Scripts (Raspberry Scripts/)

raspberry_pi_inference.py

• Real-time ball bearing classification using TensorFlow Lite on Raspberry Pi
• Captures images from PiCamera2, performs inference, and displays results
• Implements voting system for stability with confidence scores and FPS display
• Usage:

  python raspberry_pi_inference.py

rasp_script.py

• Serial communication script for Raspberry Pi bearing classifier
• Listens for "CAPTURE" commands via serial, performs real-time classification
• Returns classification codes (0-6) to Arduino for mechanical sorting control
• Note: Uses quantized model for faster inference
• Usage:

  python rasp_script.py

Arduino Firmware (arduino firmware/)

main.cpp

• Arduino firmware for mechanical ball bearing sorting system
• Controls servo motors to sort bearings into 6 categories based on material and size
• Receives classification codes from Raspberry Pi via serial communication
• Positions servos accordingly for automated sorting
• Usage: Upload to Arduino Uno/Mega via Arduino IDE

Hardware Requirements

• Raspberry Pi 4 with PiCamera2 module
• Arduino Uno/Mega with servo motors
• Mechanical sorting mechanism with drop zones
• Ball bearings for testing (Brass, Nylon, Steel in 1/2" and 1/4" sizes)

Software Requirements

Python Dependencies

tensorflow>=2.5.0
numpy>=1.19.5
matplotlib>=3.4.2
pillow>=8.2.0
scikit-learn>=0.24.2
seaborn>=0.11.1
opencv-python>=4.5.2
tflite_runtime
picamera2
pyserial

Arduino Dependencies

• Arduino IDE
• Servo.h library

Installation & Setup

1. Clone the Repository

git clone <repository-url>
cd Ball_Bearing_Classification

2. Install Python Dependencies

pip install -r requirements.txt

3. Prepare Dataset

cd ML
python ReorgData_materials.py

4. Train Model

python "VGG16Train (2).py"

5. Evaluate Model

python "VGG16ID (1).py"

6. Setup Raspberry Pi

• Install PiCamera2: sudo apt install -y python3-picamera2
• Install TFLite Runtime for ARM
• Upload Arduino firmware

Usage Workflow

Training Phase

1. Organize Data: Run ReorgData_materials.py to split dataset
2. Train Model: Run VGG16Train (2).py to train VGG16 classifier
3. Evaluate Performance: Run VGG16ID (1).py to analyze results
4. Convert to TFLite: Export model for Raspberry Pi deployment

Deployment Phase

1. Start Arduino: Upload main.cpp and power on Arduino
2. Start Raspberry Pi: Run rasp_script.py for serial communication
3. Test System: Place bearings in inspection area
4. Monitor Sorting: Watch automated classification and sorting

Model Performance Notes

• Classification Accuracy: ~85% on test data
• Inference Time: ~30ms on Raspberry Pi 4
• Known Issues: Model occasionally misclassifies Brass as Steel (~15% error rate)
• Solution: Model was overfit to training setup for optimal performance
• Alternative: Consider using Google's Teachable Machine for faster training and deployment

Important Notes

• Quantized Model: The rasp_script.py uses a quantized TensorFlow Lite model for faster inference on the Raspberry Pi
• Overfitting Issue: The model was intentionally overfit to the specific training setup for optimal performance. Any changes to the physical setup (lighting, camera position, etc.) may cause misclassifications
• Teachable Machine Alternative: Due to the overfitting challenges, Google's Teachable Machine was found to be faster for training and deployment, with better performance on the Pi even when our custom model was quantized
• Model Limitations: The current model works well for the specific setup it was trained on, but may not generalize well to different environments

Troubleshooting

• Serial Communication Issues: Check port settings and baud rate
• Camera Not Found: Ensure PiCamera2 is properly installed
• Servo Movement Problems: Verify servo connections and power supply
• Model Loading Errors: Check TFLite runtime installation

Project Structure

Ball_Bearing_Classification/
├── ML/                    # Machine learning scripts
│   ├── VGG16Train (2).py
│   ├── VGG16ID (1).py
│   ├── ReorgData_materials.py
│   └── explore_images.py
├── Raspberry Scripts/     # Pi camera and inference
│   ├── raspberry_pi_inference.py
│   └── rasp_script.py
├── arduino firmware/      # Hardware control
│   └── main.cpp
├── models/               # Trained models
├── requirements.txt      # Python dependencies
└── README.md

License

This project is licensed under the MIT License.

Authors

Kesiena Berezi Developed for MEE206 - Mechanical Student Design Project.

Acknowledgments

• VGG16 architecture for transfer learning
• Google Teachable Machine
• TensorFlow Lite for edge deployment
• PiCamera2 for Raspberry Pi camera integration
• Arduino Servo library for mechanical control