Draw a digit in your browser and a PyTorch CNN predicts what it is. A tiny, fully self-contained demo of image classification with PyTorch + Streamlit.
| File | Purpose |
|---|---|
model.py |
The CNN (2 conv + 2 fc layers, ~99% test accuracy) |
train.py |
Downloads MNIST and trains the model |
digit.py |
Turns your drawing into a centered 28×28 MNIST-style input |
app.py |
The Streamlit app: draw, predict, show class probabilities |
mnist_cnn.pt |
Pre-trained weights (committed, so it runs out of the box) |
Requires uv (curl -LsSf https://astral.sh/uv/install.sh | sh).
git clone [email protected]:cyber-shuttle/demo.git cybershuttle-demo
cd cybershuttle-demo
uv run streamlit run app.pyuv run creates the virtualenv and installs everything on first launch, then
opens the app at http://localhost:8501. Draw a digit — the prediction and a
probability bar chart update when you finish a stroke.
The repo ships with trained weights. To train from scratch:
uv run python train.py # ~1-2 min on CPU, writes mnist_cnn.ptA model trained on MNIST only classifies well if drawn digits are prepared the
same way MNIST's were. digit.preprocess mirrors that pipeline: crop to the
digit's bounding box, scale its longest side to 20px, center it in a 28×28 frame
by center of mass, then normalize. Skipping this (a naive resize of the whole
canvas) tanks accuracy.
Sanity-check that step on its own:
uv run python digit.py # runs digit.preprocess self-check