Handwritten Digit Recogniser with a Full MLOps Workflow

This project demonstrates a complete end-to-end workflow for building, training, fine-tuning, and deploying a deep learning model to recognise handwritten digits. The application is built using PyTorch and containerised with Docker, with all experiments tracked using MLflow.

---

🚀 Live Demo

You can try out the live application, deployed on Hugging Face Spaces, at the link below. Draw a digit from 0 to 9 and see the model's prediction in real-time.

➡️ Live Demo on Hugging Face Spaces

---

Features

• Deep Learning Model: A Convolutional Neural Network (CNN) built with PyTorch for image classification.
• Experiment Tracking: All training and fine-tuning runs are tracked with MLflow, logging parameters, metrics, and model artifacts.
• Robust Training: K-Fold cross-validation is used to ensure the model is robust and to select the best performing base model.
• Transfer Learning: The base model is fine-tuned on a small, custom dataset of user-drawn digits to adapt to a specific drawing style.
• Containerisation: The entire application is containerised using Docker, ensuring a consistent and reproducible environment.
• Web Application: A simple web interface built with Flask allows users to draw digits and receive predictions.
• Cloud Deployment: The final containerised application is deployed on Hugging Face Spaces for public access.

---

📂 Project Structure

/Digit_Recogniser
|
├─── data/
│    └── train/              # Original MNIST training (<https://www.kaggle.com/datasets/hojjatk/mnist-dataset>)
|
├─── my_drawings/
│    └── 0/                  # Custom drawings for fine-tuning
│    └── 1/
│    └── ...
|
├─── models/
│    └── best_model_fold_2.pt    # Best base model saved from training
│    └── fine_tuned_best_model_1.pt     # Final fine-tuned model
|
├─── static/
│    └── index.html          # Frontend for the Flask app
|
├─── mlruns/                   # Folder auto-generated by MLflow to store experiment data
|       
|─── scripts/
│    └── train.py    # Script for training the base model
│    └── fine_tune.py    # Script for fine-tuning on custom data
│    └── engine.py   # The core training/validation loop function
│    └── interact.py   # Flask application for prediction
│    └── model.py    # CNN model definition
│    └── data_loader.py    # PyTorch Dataset/DataLoader classes
|
├─── requirements.txt          # Python dependencies
└─── Dockerfile                # Instructions to build the application container

---

Methodology

The project was executed in two main stages to build a robust and personalised model.

1. Base Model Training

First, a CNN was trained on the standard MNIST dataset. To ensure the model was robust and to select the best possible version, a 5-fold cross-validation strategy was used. Each fold was logged as a separate run in MLflow. I tracked key parameters (learning rate, epochs, etc.) and metrics (accuracy, F1-score) to compare the folds. The model from the best-performing fold (best_model_fold_2.pt) was saved and selected as the base model for the next stage.

2. Fine-Tuning on Custom Data

After training a general-purpose digit recogniser, a small dataset of 300 custom drawings (30 for each digit) was created. The best base model was then loaded, and its early convolutional layers were frozen. Only the final, fully-connected classification layers were retrained on this new dataset. This transfer learning approach was used to adapt the model to the characteristics of the inference module (e.g. drawing width). This fine-tuning process was also tracked as a separate experiment in MLflow.

---

Technologies Used

• Python
• PyTorch: For building and training the neural network.
• MLflow: For experiment tracking and model management.
• Docker: For containerising the application.
• Flask: For serving the web application and prediction endpoint.
• Gunicorn: As the production web server.
• Hugging Face Spaces: For hosting the final application.
• Scikit-learn: For K-Fold cross-validation and metrics.
• NumPy & Pillow: For data manipulation and image processing.
• Albumentations: For data augmentation.

---

⚙️ Setup and Local Usage

To run this project on your local machine, follow these steps:

1. Clone the repository:

   git clone https://github.com/muanderson/Handwritten-digit-recogniser
   cd Handwritten-digit-recogniser

2. Create and activate the Conda environment:

   conda create --name ENVIRONMENT_NAME python=3.9
   conda activate ENVIRONMENT_NAME

3. Install dependencies:

   pip install -r requirements.txt

4. Run the MLflow UI: Open a separate terminal, activate the environment, and run:

   mlflow ui

   You can view the tracking server at http://localhost:5000.

5. Run the training and fine-tuning scripts:

   # Train the base model on MNIST
   python train_with_mlflow.py
   
   # Fine-tune the best model on your custom drawings
   python fine_tune_with_mlflow.py

6. Run the prediction app:

   # Make sure you have the final fine-tuned model saved
   python interact.py

   The application will be available at http://127.0.0.1:5000 (or the port specified in the script).