Edge Image-gen for Predictive Maintenance

This project dives into how edge image gen model could be use in fields like robotics to perform anaomaly detection.

Comment

Inspiration

Traditional anomaly detection models (like simple classifiers or one-class SVMs) are designed to identify faults as soon as they happen (a binary alert: Fault/No Fault). However, we saw the potential of Variational Autoencoders (VAEs) to be used for more than just anomaly detection. We use the VAE's ability to model normalcy to:

  • Quantify Degradation: Provide a continuous, interpretable health score (the reconstruction error) instead of a simple binary alert.
  • Track Progression: Use this score to track the progressive degradation (corrosion) of industrial assets over time, enabling true Predictive Maintenance.
  • Edge Intelligence: Prove that this advanced Generative AI technique could be deployed directly onto a resource-constrained ARM device, eliminating reliance on costly, high-latency cloud processing.

🛠️ What it Does (VAEs for Predictive Maintenance)

Our system uses a Convolutional Variational Autoencoder (C-VAE) to act as an on-device "expert" in what "normal" and "healthy" assets look like.

  • Training Phase: The C-VAE is trained exclusively on images of non-corroded, healthy components. It learns the statistical distribution of "normal" pixels and shapes.
  • The Latent Space: The VAE maps the image to a concise, meaningful latent space, forcing the model to learn the fundamental features of the healthy component.
  • Anomaly Scoring: When a new image (from an inspection robot) is passed through the VAE, the model tries to reconstruct it.
    • Healthy Image: The C-VAE reconstructs the image accurately, resulting in a low Reconstruction Error (MSE).
    • Corroded Image (Anomaly): The C-VAE cannot plausibly reconstruct the corrosion, causing the Reconstruction Error to spike significantly.
  • Predictive Maintenance: This continuous error score is the degradation metric. By tracking the rising error over weeks or months, we can predict when a repair will be critically necessary, moving beyond simple reactive alerts.

🏗️ How We Built It

  1. Dataset Preparation: We utilized the MNIST dataset for initial validation (training on digit '0' as "healthy") and structured our workflow to support real-world industrial imaging.
  2. C-VAE Architecture: We designed a custom Convolutional VAE (C-VAE) using TensorFlow/Keras, specifically employing Conv2D layers for the Encoder and Conv2DTranspose layers for the Decoder to efficiently handle image feature extraction and generation.
  3. Custom Loss Function: We implemented the core VAE loss, which combines the Reconstruction Loss (Mean Squared Error) with the Kullback-Leibler (KL) Divergence Loss to ensure a smooth, probabilistic latent space.
  4. Training and Thresholding: We trained the VAE solely on the "healthy" images. The anomaly threshold was then set based on the 95th percentile of the reconstruction errors seen during the training of normal data, minimizing false positives.
  5. Edge Deployment: We structured the inference code to be ready for model quantization (reducing model size and computational complexity) and deployment onto a low-power ARM processor ( Raspberry Pi ).

🚧 Challenges We Ran Into

  • VAE Training Instability: VAEs, while more stable than GANs, still require careful tuning. Balancing the weight between the Reconstruction Loss and the KL Divergence Loss was challenging to prevent the VAE from simply memorizing the images (identity function).
  • Computational Constraints: Designing a C-VAE deep enough to learn complex image features but shallow enough to run inference with ultra-low latency (milliseconds) on the ARM target required rigorous testing of layer counts and filter sizes.

📚 What We Learned

  • ** Quantization:** We reinforced the necessity of model quantization and careful architecture design when deploying deep learning models to low-power Edge/ARM environments.
  • Metrics for Anomaly Detection: In anomaly detection, Recall (catching all defects) and a low

What's Next for Edge Image-gen for Predictive Maintenance

  • Real-World Calibration: Integrate the C-VAE with a real-world progressive corrosion dataset to fine-tune the model for specific material degradation patterns.
  • Time-Series Integration: Implement an LSTM layer on top of the VAE's output (the sequence of error scores) to predict when the threshold will be crossed, providing the exact time-to-failure.

To explore the full technical documentation, detailed methodology, and complete Convolutional VAE (C-VAE) code for the Edge Image-gen for Predictive Maintenance project, please visit our repository:

GitHub Link: https://github.com/Ryan-py/VAEs_on_ARM.git

We welcome contributions and feedback!

Built With

Share this project:

Updates