Déployer un CNN MNIST sur Hugging Face avec Streamlit & Docker

"""
MNIST CNN Training Script
Author: Dr. Papa-Sega WADE

Usage:
    python train_model.py              # train with defaults
    python train_model.py --epochs 5   # fast training (99%+ from epoch 4)
    python train_model.py --output models/mnist_v2.pth
"""

import os
from pathlib import Path

import click
import torch
import torch.nn as nn
import torch.nn.functional as F
from loguru import logger
from torchvision import datasets, transforms


# ======================================
# MODEL DEFINITION
# ======================================
class MNISTNet(nn.Module):
    """Compact CNN for MNIST digit recognition.

    Architecture:
        Conv2D(32) -> MaxPool -> Conv2D(64) -> MaxPool -> Conv2D(64) -> FC(128) -> Dropout -> FC(10)

    Input:  (B, 1, 28, 28) float32, normalized with MNIST mu/sigma (0.1307, 0.3081)
    Output: (B, 10) raw logits
    Params: ~130k
    """

    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, 3)   # 28x28 -> 26x26
        self.conv2 = nn.Conv2d(32, 64, 3)  # 13x13 -> 11x11
        self.conv3 = nn.Conv2d(64, 64, 3)  # 5x5   -> 3x3
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(64 * 3 * 3, 128)
        self.dropout = nn.Dropout(0.5)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.pool(F.relu(self.conv1(x)))  # -> (B, 32, 13, 13)
        x = self.pool(F.relu(self.conv2(x)))  # -> (B, 64,  5,  5)
        x = F.relu(self.conv3(x))             # -> (B, 64,  3,  3)
        x = x.view(-1, 64 * 3 * 3)
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        return self.fc2(x)

    def count_parameters(self) -> int:
        return sum(p.numel() for p in self.parameters())


# ======================================
# TRAINING
# ======================================
def train(model, train_loader, optimizer, criterion, device, epoch, total_epochs):
    """Run one training epoch."""
    model.train()
    total_loss = 0.0
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        loss = criterion(model(data), target)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
        if batch_idx % 100 == 0:
            pct = 100.0 * batch_idx / len(train_loader)
            logger.info(
                "Epoch {epoch}/{total} [{pct:.0f}%]  loss={loss:.4f}",
                epoch=epoch, total=total_epochs, pct=pct, loss=loss.item(),
            )
    return total_loss / len(train_loader)


def evaluate(model, test_loader, device):
    """Evaluate model on test set. Returns accuracy in [0, 1]."""
    model.eval()
    correct = total = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            pred = model(data).argmax(dim=1)
            correct += pred.eq(target).sum().item()
            total += target.size(0)
    return correct / total


# ======================================
# CLI ENTRYPOINT
# ======================================
@click.command()
@click.option("--epochs",     default=5,                  show_default=True)
@click.option("--batch-size", default=128,                show_default=True)
@click.option("--lr",         default=1e-3,               show_default=True)
@click.option("--data-dir",   default="./data",           show_default=True)
@click.option("--output",     default="mnist_cnn_v1.pth", show_default=True)
def main(epochs, batch_size, lr, data_dir, output):
    """Train a compact CNN on MNIST and save weights to a .pth file."""
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    logger.info("Device: {device}", device=device)

    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,)),
    ])
    train_dataset = datasets.MNIST(data_dir, train=True,  download=True, transform=transform)
    test_dataset  = datasets.MNIST(data_dir, train=False, download=True, transform=transform)
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    test_loader  = torch.utils.data.DataLoader(test_dataset,  batch_size=1000, shuffle=False)

    model = MNISTNet().to(device)
    logger.info("Parameters: {n:,}", n=model.count_parameters())
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()

    best_acc = 0.0
    output_path = Path(output)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    for epoch in range(1, epochs + 1):
        avg_loss = train(model, train_loader, optimizer, criterion, device, epoch, epochs)
        acc = evaluate(model, test_loader, device)
        logger.info("Epoch {e}/{t}  loss={l:.4f}  acc={a:.4f}", e=epoch, t=epochs, l=avg_loss, a=acc)
        if acc > best_acc:
            best_acc = acc
            torch.save(model.state_dict(), output_path)
            logger.success("Saved best model -> {p}  (acc={a:.2f}%)", p=output_path, a=best_acc*100)

    logger.success("Training complete. Best: {a:.2f}%  -> {p}", a=best_acc*100, p=output_path)

if __name__ == "__main__":
    main()

"""
EUREKA MNIST Demo - Streamlit App (bundled model, no training at runtime)
Author: Dr. Papa-Sega WADE

Run locally:
    python train_model.py --epochs 5 --output mnist_cnn_v1.pth
    streamlit run src/app_with_model.py

On HuggingFace Space:
    Push mnist_cnn_v1.pth alongside src/ — app loads it instantly, zero training wait.
"""

import math, os
from pathlib import Path

# ── Inject dark theme at runtime ──
_CONFIG_DIR = Path(__file__).parent / ".streamlit"
_CONFIG_FILE = _CONFIG_DIR / "config.toml"
if not _CONFIG_FILE.exists():
    _CONFIG_DIR.mkdir(parents=True, exist_ok=True)
    _CONFIG_FILE.write_text(
        '[theme]\nbase = "dark"\nbackgroundColor = "#1C0F0F"\n'
        'secondaryBackgroundColor = "#2A1A0E"\nprimaryColor = "#E04E00"\n'
        'textColor = "#E0D0B0"\nfont = "sans serif"\n'
    )

import cv2, numpy as np, streamlit as st
import torch, torch.nn as nn, torch.nn.functional as F
from scipy.ndimage import center_of_mass
from streamlit_drawable_canvas import st_canvas
from torchvision import datasets, transforms

st.set_page_config(
    page_title="MNIST EUREKA - De la Terminale a la Revolution IA",
    page_icon="🧠", layout="wide",
)

# ── CNN Model (same architecture as train_model.py) ──
class MNISTNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, 3)
        self.conv2 = nn.Conv2d(32, 64, 3)
        self.conv3 = nn.Conv2d(64, 64, 3)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(64 * 3 * 3, 128)
        self.dropout = nn.Dropout(0.5)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = F.relu(self.conv3(x))
        x = x.view(-1, 64 * 3 * 3)
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        return self.fc2(x)

    def count_parameters(self):
        return sum(p.numel() for p in self.parameters())

# ── Load pre-trained weights ──
MODEL_PATH = Path("mnist_cnn_v1.pth")

@st.cache_resource
def load_model():
    if not MODEL_PATH.exists():
        st.error("Model file not found. Run: python train_model.py --epochs 5")
        st.stop()
    model = MNISTNet()
    model.load_state_dict(torch.load(MODEL_PATH, map_location="cpu", weights_only=True))
    model.eval()
    return model

model = load_model()

# ── Preprocessing (MNIST-faithful pipeline) ──
def preprocess_canvas(img_array):
    """Convert raw canvas RGBA to normalized 28x28 MNIST tensor."""
    if img_array is None:
        return None
    gray = cv2.cvtColor(img_array[:,:,:3].astype(np.uint8), cv2.COLOR_RGB2GRAY)
    ink_mask = (gray > 50).astype(np.uint8) * 255
    if np.sum(ink_mask) == 0:
        return None

    # Tight crop
    rows_ink = np.where(ink_mask.any(axis=1))[0]
    cols_ink = np.where(ink_mask.any(axis=0))[0]
    if len(rows_ink) == 0:
        return None
    cropped = ink_mask[rows_ink[0]:rows_ink[-1]+1, cols_ink[0]:cols_ink[-1]+1]

    # Proportional resize to max 20px
    h, w = cropped.shape
    if h >= w:
        new_h, new_w = 20, max(1, int(round(w * 20.0 / h)))
    else:
        new_w, new_h = 20, max(1, int(round(h * 20.0 / w)))
    resized = cv2.resize(cropped, (new_w, new_h), interpolation=cv2.INTER_AREA)
    _, resized = cv2.threshold(resized, 64, 255, cv2.THRESH_BINARY)

    # Pad to 28x28
    pad_t = int(math.ceil((28 - new_h) / 2.0))
    pad_b = int(math.floor((28 - new_h) / 2.0))
    pad_l = int(math.ceil((28 - new_w) / 2.0))
    pad_r = int(math.floor((28 - new_w) / 2.0))
    padded = np.pad(resized, ((pad_t, pad_b), (pad_l, pad_r)), "constant")

    # Recenter by center of mass
    cy, cx = center_of_mass(padded)
    M = np.float32([[1,0,int(round(14.0-cx))],[0,1,int(round(14.0-cy))]])
    centered = cv2.warpAffine(padded, M, (28, 28))

    # MNIST normalization
    img_f = centered.astype(np.float32) / 255.0
    img_f = (img_f - 0.1307) / 0.3081
    return img_f.reshape(1, 1, 28, 28)

# ── UI : canvas + prediction (simplified for article) ──
canvas_result = st_canvas(
    stroke_width=22, stroke_color="#FFFFFF", background_color="#1A1A1A",
    height=400, width=400, drawing_mode="freedraw", update_streamlit=True, key="canvas",
)

if canvas_result.image_data is not None:
    processed = preprocess_canvas(canvas_result.image_data)
    if processed is not None:
        tensor_in = torch.FloatTensor(processed)
        with torch.no_grad():
            probs = F.softmax(model(tensor_in), dim=1)[0].numpy()
        predicted = int(np.argmax(probs))
        confidence = float(probs[predicted]) * 100
        st.markdown(f"### Prédit : **{predicted}**  (confiance : {confidence:.1f}%)")

FROM python:3.11-slim

WORKDIR /app

RUN apt-get update && apt-get install -y \
    build-essential \
    curl \
    git \
    libgl1 \
    libglib2.0-0 \
    && rm -rf /var/lib/apt/lists/*

COPY requirements.txt ./
RUN pip install --no-cache-dir -r requirements.txt

# Injection explicite du modèle pré-entraîné
COPY mnist_cnn_v1.pth ./
COPY src/ ./src/

EXPOSE 8501
HEALTHCHECK CMD curl --fail http://localhost:8501/_stcore/health

ENTRYPOINT ["streamlit", "run", "src/app_with_model.py", \
    "--server.port=8501", \
    "--server.address=0.0.0.0", \
    "--server.headless=true"]