Battery_Defect_Detection_system

A deep learning-based defect detection system for battery cells, supporting both CNN​ and Vision Transformer​ models for multi-class classification of battery cell images.

基于深度学习的电池电芯缺陷检测系统，同时支持 CNN​ 和 Vision Transformer​ 模型，对电池电芯图像进行多分类缺陷识别。

A small side-project demo built in my spare time. The dataset format is already aligned with the code—with just a few tweaks, you can easily train your own model.

Dataset download: https://pan.quark.cn/s/3f2f3b464813

一个课余时间制作的小demo，数据集格式匹配，稍加修改代码，就可以训练出自己的模型。

数据集可以在下方链接中下载：https://pan.quark.cn/s/3f2f3b464813

Features

Multi-Class Classification

• No Defect
• Crack Only
• Inactive Only
• Crack + Inactive

Dual Model Architecture

• CNN (Convolutional Neural Network) for fast and stable training
• Vision Transformer with Patch Embedding for global feature capture

Complete Pipeline

• 1.Training from scratch / Loading pre-trained weights
• 2.Batch evaluation with accuracy, recall, and F1-score
• 3.Single-image prediction with confidence scores
• 4.Per-sample result inspection on test set

---

Project Structure

├── label.csv              # Dataset labels (semicolon-separated)
├── images/                # Battery cell images
├── cnn_model.pth          # Trained CNN weights (generated)
├── transformer_model.pth  # Trained Transformer weights (generated)
└── main.py                # Entry point with interactive menu

---

Quick Start

1. Install Dependencies

pip install torch torchvision pandas scikit-learn Pillow

2. Prepare Dataset

Ensure your label.csv contains the following columns:

Column	Description
filename	Image file path
poly_wafer	Poly wafer flag
crack	Crack defect label (0/1)
inactive	Inactive defect label (0/1)

3. Run

python main.py

---

Model Architecture

CNN Model

Input (1×128×128)
  → Conv2d(1→32) + ReLU + MaxPool2d
  → Conv2d(32→64) + ReLU + MaxPool2d
  → AdaptiveAvgPool2d(4×4)
  → Flatten → Linear(1024→128) → ReLU → Dropout
  → Linear(128→4)

Vision Transformer

Input (1×128×128)
  → Patch Embedding (16×16 patches → 128-dim tokens)
  → Positional Encoding
  → Transformer Encoder (4 layers, 8 heads)
  → Mean Pooling → Linear(128→4)

---

Evaluation Metrics

The system reports:

Metric	Description
Accuracy	Overall classification accuracy
Recall (Macro)	Per-class recall averaged equally
F1-Score (Macro)	Harmonic mean of precision and recall

---

Training Configuration

Parameter	Value
Image Size	128 × 128 (grayscale)
Batch Size	16
Optimizer	Adam (lr=0.001)
LR Scheduler	StepLR (step=10, γ=0.5)
Loss Function	CrossEntropyLoss
Data Augmentation	RandomHorizontalFlip + RandomRotation(10°)
Train/Test Split	80% / 20% (stratified)

---