Tea Sickness Classification - from EDA to model training and deployment

Key Features

• Full ML Pipeline (EDA -> Training -> API -> Deployment)
• Transfer Learning with EfficientNet-B0
• RESTful API with input validation
• Unit tests
• Deployed on GCP Cloud Run

Project Overview

The project's goal is to classify tea leaf images into 8 different classes, corresponding to various diseases, with one class being healthy. It was built to demonstrate the full ML workflow, from raw data to the deployed application on Google Cloud Platform. The demonstration is available here

Architecture & Tech Stack

• Frontend: Streamlit
• Backend: FastAPI
• Model: Fine-tuned EfficientNet-B0
• Deployment: Docker + Google Cloud Run

Architecture diagram:

User -> Streamlit UI -> FastAPI -> EfficientNet-B0 -> Prediction

Dataset

The dataset contained 885 images of tea leaves, classified into 8 different classes corresponding to 7 different diseases and one healthy class. It was obtained on HuggingFace at https://huggingface.co/datasets/yunusserhat/tea_sickness_dataset.

Citation:

@article{kimutai2022tea,
  title     = {Tea sickness dataset},
  author    = {Kimutai, Gibson and Förster, Anna},
  journal   = {Mendeley Data},
  volume    = {2},
  year      = {2022},
  doi       = {10.17632/j32xdt2ff5.2}
}

Model Development

• Architecture: EfficientNet-B0 (Transfer Learning)
• Training Strategy: Early layers were frozen. Only the last two layers were fine-tuned, and the final classification head was changed to account for the number of classes present in the dataset.
• Performance: Achieved 81% global accuracy on the test set. The goal of the project was more to demonstrate the end-to-end workflow rather than the search for performance.

Classification Report

               precision    recall  f1-score   support

  Anthracnose       0.62      0.80      0.70        10
   algal_leaf       1.00      1.00      1.00        11
bird_eye_spot       0.62      0.80      0.70        10
 brown_blight       0.82      0.75      0.78        12
   gray_light       1.00      0.40      0.57        10
      healthy       0.88      1.00      0.93         7
red_leaf_spot       0.94      1.00      0.97        15
   white_spot       0.85      0.79      0.81        14

     accuracy                           0.82        89
    macro avg       0.84      0.82      0.81        89
 weighted avg       0.84      0.82      0.81        89

Confusion Matrix

Project Structure

├── api/                    # FastAPI backend code
├── app/                    # Streamlit frontend code
├── models/                 # Saved model weights
├── notebooks/              # EDA and experiments
├── src/                    # Source code for training and data processing
├── tests/                  # Unit tests
├── Dockerfile.api          # API container
├── Dockerfile.ui           # UI container
├── docker-compose.yml      # Multi-container orchestration
├── requirements.txt        # Full dependencies (for local dev/training)
├── requirements-api.txt    # API-only dependencies
└── requirements-ui.txt     # UI-only dependencies

Installation and Usage

Local Development

Prerequisites: Python 3.13

Installation:

pip install -r requirements.txt

Running the API:

uvicorn api.main:app --reload

Running the Frontend (in a separate terminal):

streamlit run app/streamlit_app.py

Training:

python src/models/train.py --epochs 10

Running tests:

pytest tests/

Docker (Recommended)

Run both API and UI with a single command:

docker compose up --build

This will start:

• API: http://localhost:8000 (with interactive docs at /docs)
• UI: http://localhost:8501

To run in detached mode:

docker compose up -d --build

To stop:

docker compose down

Deployment

The app is deployed on Google Cloud Platform using Cloud Run. Frontend and Backend are deployed as separate services, each with its own Dockerfile and minimal dependencies:

• Dockerfile.api - Contains PyTorch and FastAPI
• Dockerfile.ui - Contains Streamlit only

This separation ensures faster builds and smaller container images.

Note on model weights: For this portfolio project, model weights are committed to the repository for ease of demonstration. In a production environment, weights would be stored in cloud storage (e.g., Google Cloud Storage) or a model registry (e.g., MLflow, Vertex AI) and downloaded at deployment time to avoid bloating the Git repository.

Future Improvements

• Implement class weights to optimize for business outcomes (e.g., penalize "sick -> healthy" misclassifications more heavily, as they lead to untreated diseases and crop loss)
• Improve model performance by testing different architectures and parameters
• Move model weights to cloud storage (GCS) instead of committing to Git, following production best practices for artifact management