Teach AI to Read Handwritten Numbers — Build a Real CNN

# Project 8 — Teach AI to Read Handwritten Numbers

Real PyTorch CNN trained from scratch on scikit-learn's offline **1,797 handwritten 8×8 digit images**, with a Streamlit upload/predict app. The dataset is built into scikit-learn — no Kaggle login, API key or external dataset download.

## Beginner setup

1. Install Python 3.12 and VS Code. Open the folder `projects/digit-recognizer` in VS Code using File → Open Folder.
2. Choose Terminal → New Terminal. Windows: `py -3.12 -m venv .venv` and `.venv\\Scripts\\activate`. macOS/Linux: `python3.12 -m venv .venv && source .venv/bin/activate`.
3. Run `python -m pip install -r requirements.txt` (PyTorch is a large library; installation can take time).
4. Run `python -m pytest -q`. Tests include real training, image validation and model reloading.
5. Train from scratch: `python train.py`. This saves a model and JSON metrics under `artifacts/` and prints held-out test accuracy.
6. Start the app: `python -m streamlit run app.py`. Open the displayed localhost address, commonly http://localhost:8501.
7. Draw a single digit with a dark pen on white paper, photograph/crop it closely or create a simple PNG in Paint, upload it and observe the original and preprocessed 8×8 image.
8. Inspect all ten model probabilities; look at errors and note that accuracy on the held-out built-in dataset does not guarantee accuracy on handwritten camera images.

## Architecture calculation

Input `1×8×8` → convolution `16×8×8` → max pooling `16×4×4` → convolution `32×4×4` → max pooling `32×2×2` → flatten `32·2·2=128` values → dense 64 → 10 class logits.

A 3×3 convolution with 1 input channel and 16 outputs has `16 × (1 × 3 × 3 + 1) = 160` learned weights including biases. The second has `32 × (16 × 3 × 3 + 1) = 4,640` learned parameters. Explain ReLU, pooling, cross-entropy, gradient descent and the 80/20 split with a further validation split.

## Limits

The dataset contains small normalized greyscale handwritten digits. External pen-on-paper images can differ substantially in stroke thickness, lighting, rotation and framing. The model does not perform OCR of sentences or multiple digits. Test and validation data must never be used as additional training labels.

## Files

`src/digits.py` contains dataset preparation, CNN, training, saving and inference; `train.py` trains; `app.py` runs Streamlit; `tests/test_digits.py` checks the critical stages. Generated artifacts are excluded from Git.

torch==2.8.0
scikit-learn==1.7.2
numpy==2.3.3
Pillow==11.3.0
streamlit==1.50.0
pytest==8.4.2
matplotlib==3.10.7

"""Reproducible digit classification using scikit-learn's built-in 8x8 images.

A small PyTorch CNN is trained from scratch; the model is not pretrained.
The test split remains untouched until final evaluation.
"""
from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
import json

import numpy as np
import torch
from PIL import Image, ImageOps
from sklearn.datasets import load_digits
from sklearn.metrics import accuracy_score, confusion_matrix, classification_report
from sklearn.model_selection import train_test_split
from torch import nn

ROOT = Path(__file__).resolve().parents[1]
MODEL_PATH = ROOT / "artifacts" / "digits-cnn.pt"
METRICS_PATH = ROOT / "artifacts" / "metrics.json"
SEED = 42

class DigitCNN(nn.Module):
    """8x8 greyscale image -> 10 handwritten-digit probabilities."""
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2),
            nn.Conv2d(16, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2),
        )
        self.classifier = nn.Sequential(nn.Flatten(), nn.Linear(32 * 2 * 2, 64),
                                        nn.ReLU(), nn.Linear(64, 10))

    def forward(self, images: torch.Tensor) -> torch.Tensor:
        return self.classifier(self.features(images))


@dataclass
class Splits:
    train_x: np.ndarray
    val_x: np.ndarray
    test_x: np.ndarray
    train_y: np.ndarray
    val_y: np.ndarray
    test_y: np.ndarray


def load_splits(seed: int = SEED) -> Splits:
    images, targets = load_digits(return_X_y=False).images, load_digits(return_X_y=False).target
    images = (images / 16.0).astype(np.float32)
    first_x, test_x, first_y, test_y = train_test_split(
        images, targets, test_size=0.20, random_state=seed, stratify=targets
    )
    train_x, val_x, train_y, val_y = train_test_split(
        first_x, first_y, test_size=0.20, random_state=seed, stratify=first_y
    )
    return Splits(train_x, val_x, test_x, train_y, val_y, test_y)


def to_tensor(images: np.ndarray) -> torch.Tensor:
    if images.ndim != 3 or images.shape[1:] != (8, 8):
        raise ValueError("Expected (N, 8, 8) greyscale images")
    if not np.isfinite(images).all() or images.min() < 0 or images.max() > 1:
        raise ValueError("Images must be finite pixel intensities from 0 to 1")
    return torch.from_numpy(images.astype(np.float32)).unsqueeze(1)


def fit(epochs: int = 15, learning_rate: float = 0.003, seed: int = SEED) -> dict:
    if epochs < 1 or epochs > 100:
        raise ValueError("Choose between 1 and 100 epochs")
    torch.manual_seed(seed)
    np.random.seed(seed)
    torch.set_num_threads(2)
    parts = load_splits(seed)
    model = DigitCNN()
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
    criterion = nn.CrossEntropyLoss()
    x, y = to_tensor(parts.train_x), torch.from_numpy(parts.train_y).long()
    history = []
    for epoch in range(epochs):
        model.train()
        order = torch.randperm(len(x))
        for indices in order.split(64):
            optimizer.zero_grad()
            logits = model(x[indices])
            loss = criterion(logits, y[indices])
            loss.backward()
            optimizer.step()
        model.eval()
        with torch.no_grad():
            val_logits = model(to_tensor(parts.val_x))
            val_accuracy = (val_logits.argmax(dim=1).numpy() == parts.val_y).mean()
        history.append({"epoch": epoch + 1, "train_loss": float(loss.item()),
                        "val_accuracy": round(float(val_accuracy), 6)})
    model.eval()
    with torch.no_grad():
        predictions = model(to_tensor(parts.test_x)).argmax(dim=1).numpy()
    accuracy = float(accuracy_score(parts.test_y, predictions))
    confusion = confusion_matrix(parts.test_y, predictions, labels=list(range(10)))
    metrics = {"dataset": "scikit-learn digits (1,797 8x8 images)",
               "seed": seed, "epochs": epochs, "test_accuracy": accuracy,
               "history": history, "confusion_matrix": confusion.tolist(),
               "classification_report": classification_report(parts.test_y, predictions,
                                                               output_dict=True, zero_division=0)}
    MODEL_PATH.parent.mkdir(parents=True, exist_ok=True)
    torch.save(model.state_dict(), MODEL_PATH)
    METRICS_PATH.write_text(json.dumps(metrics, indent=2), encoding="utf-8")
    return metrics


def load_model(model_path: Path = MODEL_PATH) -> DigitCNN:
    model = DigitCNN()
    model.load_state_dict(torch.load(model_path, map_location="cpu", weights_only=True))
    model.eval()
    return model


def prepare_uploaded_image(raw: bytes) -> np.ndarray:
    """Convert a learner-uploaded drawing into the dataset's bright-ink 8x8 convention."""
    from io import BytesIO
    if not raw or len(raw) > 4_000_000:
        raise ValueError("Upload a nonempty image smaller than 4 MB")
    try:
        with Image.open(BytesIO(raw)) as image:
            if image.width > 4096 or image.height > 4096:
                raise ValueError("Image dimensions are too large")
            gray = image.convert("L")
            # Convention: dark background with light strokes as in sklearn digits.
            if float(np.asarray(gray).mean()) > 127:
                gray = ImageOps.invert(gray)
            gray = ImageOps.autocontrast(gray)
            gray = ImageOps.pad(gray, (8, 8), color=0)
            arr = np.asarray(gray, dtype=np.float32) / 255.0
    except (OSError, SyntaxError) as exc:
        raise ValueError("Cannot decode this image") from exc
    if arr.max() < 0.05:
        raise ValueError("No visible digit found")
    return arr


def predict(model: DigitCNN, image: np.ndarray) -> tuple[int, list[float]]:
    tensor = to_tensor(image[None])
    with torch.no_grad():
        probabilities = torch.softmax(model(tensor), dim=1)[0].numpy()
    return int(probabilities.argmax()), probabilities.tolist()

from src.digits import fit

if __name__ == "__main__":
    results = fit()
    print(f"Test accuracy on held-out sklearn digits: {results['test_accuracy']:.2%}")
    print("Saved artifacts/digits-cnn.pt and artifacts/metrics.json")

"""Streamlit app; train first: python train.py"""
from __future__ import annotations
import io
from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
import streamlit as st
from PIL import Image
from src.digits import MODEL_PATH, METRICS_PATH, load_model, prepare_uploaded_image, predict

st.set_page_config(page_title="Handwritten Digit Recognizer | LearnMLAcademy", page_icon="🔢", layout="wide")
st.title("Teach AI to Read Handwritten Numbers")
st.caption("A real 8×8 convolutional neural network trained on 1,797 built-in handwritten digit images.")
if not MODEL_PATH.exists() or not METRICS_PATH.exists():
    st.warning("Model not trained yet. In the project terminal run: python train.py")
    st.stop()

@st.cache_resource
def cached_model():
    return load_model()

import json
metrics = json.loads(METRICS_PATH.read_text(encoding="utf-8"))
st.metric("Held-out test accuracy", f'{metrics["test_accuracy"]:.1%}')
st.caption("This measures only the built-in 8×8 dataset. Phone-camera images can perform worse.")
upload = st.file_uploader("Upload a clear handwritten digit (PNG/JPG)", type=["png", "jpg", "jpeg"])
if upload is None:
    st.info("Upload a dark-ink drawing on a light background, or white ink on black. The app will convert it to 8×8 pixels.")
else:
    try:
        image = prepare_uploaded_image(upload.getvalue())
        col1, col2 = st.columns(2)
        with col1:
            st.image(upload.getvalue(), caption="Your original drawing", width=220)
        with col2:
            st.image(image, caption="Actual 8×8 input after preprocessing", width=220, clamp=True)
        prediction, probabilities = predict(cached_model(), image)
        st.subheader(f"Predicted digit: {prediction}")
        st.caption("Softmax probabilities are relative model scores, not guaranteed real-world correctness.")
        st.bar_chart({str(i): float(score) for i, score in enumerate(probabilities)})
    except (ValueError, OSError) as exc:
        st.error(f"Could not read the digit: {exc}")

with st.expander("How does the neural network learn?"):
    st.markdown("Each 3×3 convolution scans local pixel patterns; ReLU keeps positive activations; "
                "2×2 max pooling reduces width and height. The dense layers choose among 10 digits.")
    st.code("Input: 1×8×8 → Conv: 16×8×8 → Pool: 16×4×4\n"
            "→ Conv: 32×4×4 → Pool: 32×2×2\n"
            "→ Flatten: 128 → Dense: 64 → Logits: 10", language="text")
    st.markdown("A training epoch predicts labels, calculates cross-entropy loss, propagates gradients "
                "backward, then updates weights using Adam. The test set is held out until final evaluation.")
st.divider()
st.caption("Educational project. Avoid interpreting scores as calibrated confidence.")

from io import BytesIO
import numpy as np
import pytest
import torch
from PIL import Image
from src.digits import DigitCNN, fit, load_model, load_splits, predict, prepare_uploaded_image, to_tensor

def test_dataset_has_stratified_disjoint_splits():
    parts = load_splits()
    assert len(parts.train_x) + len(parts.val_x) + len(parts.test_x) == 1797
    assert parts.train_x.shape[1:] == (8, 8)
    assert set(parts.train_y) == set(range(10))
    assert set(parts.test_y) == set(range(10))

def test_model_shapes_and_normalized_output():
    model = DigitCNN()
    logits = model(torch.zeros(5, 1, 8, 8))
    assert logits.shape == (5, 10)
    value, probs = predict(model, np.ones((8, 8), dtype=np.float32))
    assert value in range(10)
    assert abs(sum(probs) - 1) < 1e-5

def test_bad_inputs_are_rejected():
    with pytest.raises(ValueError): to_tensor(np.ones((5, 9, 9), dtype=np.float32))
    with pytest.raises(ValueError): to_tensor(np.ones((5, 8, 8), dtype=np.float32) * 2)
    with pytest.raises(ValueError): prepare_uploaded_image(b"not-an-image")
    with pytest.raises(ValueError): prepare_uploaded_image(b"")

def test_uploaded_black_ink_preprocessing():
    arr = np.ones((64, 64), dtype=np.uint8) * 255
    arr[10:54, 25:39] = 0
    buffer = BytesIO()
    Image.fromarray(arr).save(buffer, format="PNG")
    result = prepare_uploaded_image(buffer.getvalue())
    assert result.shape == (8, 8)
    assert result.min() >= 0 and result.max() <= 1
    assert result.max() > 0.05

def test_training_is_reproducible_and_model_reload_works(tmp_path, monkeypatch):
    import src.digits as d
    monkeypatch.setattr(d, "MODEL_PATH", tmp_path / "digit.pt")
    monkeypatch.setattr(d, "METRICS_PATH", tmp_path / "metrics.json")
    result = d.fit(epochs=3)
    assert 0.3 < result["test_accuracy"] <= 1
    assert len(result["confusion_matrix"]) == 10
    model = d.load_model(tmp_path / "digit.pt")
    parts = load_splits()
    number, scores = d.predict(model, parts.test_x[0])
    assert number in range(10)
    assert abs(sum(scores) - 1) < 1e-5