LeafNet - Leaf Disease Classification Neural Network

Computer vision deep neural network for leaf disease classification using transfer learning with TensorFlow/Keras (MobileNetV2).

Key scripts

• scripts/load_data.sh: download and organize dataset
• Distribution.py: quick dataset exploration and plots (optional)
• Augmentation.py: image augmentations and dataset balancing (optional)
• Transformation.py: PlantCV-based preprocessing/transforms (optional)
• train.py: train a MobileNetV2-based classifier with modular hyperparameter tuning and export as ZIP artifact
• predict.py: predict on a single image importing the trained model from the ZIP artifact

Installation

This project uses uv for Python and dependency management.

Prerequisites

• Python (managed by uv)
• uv installed (pipx install uv or see https://docs.astral.sh/uv/)

Setup

# Install the project Python and dependencies
uv python install
uv sync

# Optionally work in the venv via uv
uv run python --version

Data download

# Download and arrange the dataset
./scripts/load_data.sh

# Resulting structure (examples)
images/
  Apple/
    image (1).JPG ...
  Grape/
    image (1).JPG ...

The dataset is a folder-structured image corpus (JPEG), with one subfolder per class (e.g., Apple/Apple_Black_rot, Apple/Apple_healthy, Apple/Apple_rust, Apple/Apple_scab).

Project structure and workflow

Part 1 - Distribution (optional)

Distribution.py inspects a labeled image directory (one subfolder per class) and produces simple visualizations (pie/bar charts) of the class counts. Use it to verify dataset balance and spot missing/empty classes before training.

Example:

uv run python Distribution.py --data_dir images/Apple

Part 2 - Augmentation (optional)

Augmentation.py provides basic image augmentations (flip, crop, distortion, etc.) and a utility to balance a dataset by synthesizing new images for minority classes. In training, class balancing is applied on the training split to reduce class imbalance effects.

CLI examples:

# Augment a single image and preview
uv run python Augmentation.py path/to/image.jpg --display

# Balance a dataset in-place (adds augmented copies to smaller classes)
uv run python Augmentation.py images/Apple

Part 3 - Transformation (optional)

Transformation.py uses PlantCV and OpenCV to produce preprocessing outputs such as masks and ROI-restricted overlays. This helps focus the model on the leaf area and can generate interpretable intermediate representations (e.g., masked leaves, pseudo-landmarks). While the current train.py uses Keras preprocessing, you can batch-generate transformed datasets here if you want a transformation-first pipeline.

Examples:

# Process all images from a source folder into a destination folder
uv run python Transformation.py -src images/Apple -dst transformed/Apple

# Process a single image and visualize all steps
uv run python Transformation.py images/Apple/some_image.JPG

Part 4 - Model (training and inference)

The modeling pipeline centers on transfer learning with MobileNetV2, keeping things fast and lightweight while achieving solid accuracy.

• Running training

uv run python train.py images/Apple

This will print dataset stats, training progress, final validation accuracy, and create leaffliction.zip.

• Dataset handling in train.py:
  • Input expects one folder per class under a root directory (e.g., images/Apple or images/Grape).
  • train.py performs an 80/20 deterministic split per class and saves the train and validation set into data/
  • On the training split, in-place class balancing is applied to mitigate class imbalance. The final train set is saved in train_aug

Example: data dir after running training on images/Apple:

data/
├── train_aug              # Training dataset balanced with data augmentation
│   ├── Apple_Black_rot
│   ├── Apple_healthy
│   ├── Apple_rust
│   └── Apple_scab
└── val                    # Validation dataset
    ├── Apple_Black_rot
    ├── Apple_healthy
    ├── Apple_rust
    └── Apple_scab

• Architecture

  • Base: MobileNetV2 (ImageNet weights) with the base frozen initially
  • Head: GlobalAveragePooling2D → Dropout(0.3) → Dense(num_classes, softmax)
  • Optimizers/Loss: Adam with categorical cross-entropy

• Training schedule

  1. Transfer-learning stage (frozen base) for a few epochs
  2. Fine-tuning: unfreeze the last N layers of the base and train with a lower LR
  3. Class weights computed from the training generator to handle imbalance
  4. EarlyStopping, ReduceLROnPlateau, and ModelCheckpoint callbacks

• Outputs and portability

  • Final model saved as best_transfer_final.keras
  • Exported ZIP artifact leaffliction.zip includes:
    • model.keras (the model)
    • label_encoder.pkl (class mapping)
    • training_info.json (e.g., validation accuracy, class names)
  • This ZIP is all you need for inference with predict.py.

• Running predictions

uv run python predict.py path/to/image.jpg --model_zip leaffliction.zip

It will load the ZIP, preprocess the image to 256×256 with MobileNetV2 preprocessing, predict the class, display inputs/outputs, and report confidence along with the validation accuracy stored in the ZIP.

Tips and notes

• Image formats: .jpg/.jpeg are supported throughout.
• GPU: if a GPU is present, TensorFlow is configured for memory growth.
• Reproducibility: the split is deterministic via sorted filenames; control your data ordering if you want a specific split.
• Extensibility: swap MobileNetV2 for a different backbone in train.py and adjust preprocessing accordingly.