Samia-Nawaz/ViT-with-Token-Level-Supervision-for-Wheat-Disease-Classification

# 🌾 Vision Transformer for Wheat Leaf Disease Classification This repository implements a Vision Transformer (ViT) model for fine-grained classification of wheat leaf diseases using token-level supervision, as described in the research paper: > **Vision Transformer-Based Fine-Grained Classification of Wheat Leaf Diseases with Token-Level Supervision** --- ## 🧠 Key Features - Implements Vision Transformer with patch-based embeddings. - Incorporates **token-level supervision** combining class and patch token loss. - Preprocessing with data augmentation to increase generalization. - Custom training loop using TensorFlow and Keras. - Visualization of attention via **Grad-CAM** for model interpretability. - Benchmarking against other CNN architectures. --- ## 📁 Directory Structure ``` . ├── data_loader.py # Load and split dataset into train/val/test ├── preprocessing.py # Resize and augment images ├── patch_embedding.py # Patch embedding layer for ViT ├── vit_model.py # Vision Transformer architecture ├── token_supervision.py # Custom loss: class token + patch token ├── train.py # Custom training loop using token-level loss ├── evaluate.py # Evaluate model and generate Grad-CAM ├── README.md # Project documentation └── vit_token_supervised_model.h5 # (Generated after training) ``` --- ## 📦 Requirements Install the dependencies: ```bash pip install tensorflow pandas scikit-learn matplotlib seaborn ``` --- ## 🗂️ Dataset Use the **Wheat Disease Images Small Dataset** from: 📎 Zenodo Dataset - doi.org Organize it in the following structure: ``` dataset/ ├── Yellow Rust/ ├── Brown Rust/ ├── Septoria/ ├── Mildew/ └── Healthy/ ``` Update the dataset path in `train.py` and `evaluate.py`: ```python dataset_path = "path_to_your_dataset" ``` --- ## 🚀 Training Run the training script: ```bash python train.py ``` The model will be …