from __future__ import annotations import gzip import json import random import struct from pathlib import Path from urllib.request import urlretrieve import numpy as np import torch from PIL import Image, ImageFilter from torch import nn from torch.utils.data import DataLoader, Dataset from mnist_model import DATA_DIR, MNIST_MEAN, MNIST_STD, MODEL_PATH, DigitCNN, center_digit, shift_pixels METRICS_PATH = Path("model/metrics.json") MNIST_BASE_URL = "https://storage.googleapis.com/cvdf-datasets/mnist/" MNIST_FILES = { "train_images": "train-images-idx3-ubyte.gz", "train_labels": "train-labels-idx1-ubyte.gz", "test_images": "t10k-images-idx3-ubyte.gz", "test_labels": "t10k-labels-idx1-ubyte.gz", } def download_real_mnist() -> None: DATA_DIR.mkdir(parents=True, exist_ok=True) for filename in MNIST_FILES.values(): target = DATA_DIR / filename if target.exists(): continue print(f"downloading {filename}") urlretrieve(MNIST_BASE_URL + filename, target) def read_images(path: Path) -> np.ndarray: with gzip.open(path, "rb") as f: magic, count, rows, cols = struct.unpack(">IIII", f.read(16)) if magic != 2051 or rows != 28 or cols != 28: raise ValueError(f"bad image file: {path}") data = np.frombuffer(f.read(), dtype=np.uint8) return data.reshape(count, 28, 28) def read_labels(path: Path) -> np.ndarray: with gzip.open(path, "rb") as f: magic, count = struct.unpack(">II", f.read(8)) if magic != 2049: raise ValueError(f"bad label file: {path}") data = np.frombuffer(f.read(), dtype=np.uint8) return data.astype(np.int64) def load_real_mnist() -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: download_real_mnist() x_train = read_images(DATA_DIR / MNIST_FILES["train_images"]) y_train = read_labels(DATA_DIR / MNIST_FILES["train_labels"]) x_test = read_images(DATA_DIR / MNIST_FILES["test_images"]) y_test = read_labels(DATA_DIR / MNIST_FILES["test_labels"]) return x_train, y_train, x_test, y_test def augment_canvas_style(image: Image.Image) -> Image.Image: angle = random.uniform(-14, 14) image = image.rotate(angle, resample=Image.Resampling.BILINEAR, fillcolor=0) target_size = random.choice([18, 19, 20, 21, 22, 23, 24]) image = center_digit(image, target_size=target_size) if random.random() < 0.25: image = image.filter(ImageFilter.MaxFilter(3)) if random.random() < 0.15: image = image.filter(ImageFilter.GaussianBlur(radius=0.45)) dy = random.randint(-2, 2) dx = random.randint(-2, 2) pixels = shift_pixels(np.array(image, dtype=np.uint8), dy, dx) return Image.fromarray(pixels, mode="L") def image_to_tensor(image: Image.Image) -> torch.Tensor: pixels = np.array(image, dtype=np.float32) / 255.0 tensor = torch.from_numpy(pixels).unsqueeze(0) return (tensor - MNIST_MEAN) / MNIST_STD class RealMNISTDataset(Dataset): def __init__(self, images: np.ndarray, labels: np.ndarray, augment: bool) -> None: self.images = images self.labels = labels self.augment = augment def __len__(self) -> int: return len(self.labels) def __getitem__(self, index: int) -> tuple[torch.Tensor, torch.Tensor]: image = Image.fromarray(self.images[index], mode="L") if self.augment: image = augment_canvas_style(image) else: image = center_digit(image, target_size=20) return image_to_tensor(image), torch.tensor(self.labels[index], dtype=torch.long) def evaluate(model: DigitCNN, loader: DataLoader, device: torch.device) -> float: model.eval() correct = 0 total = 0 with torch.no_grad(): for images, labels in loader: images = images.to(device) labels = labels.to(device) predictions = model(images).argmax(dim=1) correct += int((predictions == labels).sum().item()) total += int(labels.numel()) return correct / total def train_model(epochs: int = 5, batch_size: int = 128, learning_rate: float = 1e-3) -> None: random.seed(42) np.random.seed(42) torch.manual_seed(42) x_train, y_train, x_test, y_test = load_real_mnist() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"device={device}") print(f"train_images={len(y_train)} test_images={len(y_test)}") train_loader = DataLoader( RealMNISTDataset(x_train, y_train, augment=True), batch_size=batch_size, shuffle=True, num_workers=0, ) test_loader = DataLoader( RealMNISTDataset(x_test, y_test, augment=False), batch_size=batch_size, shuffle=False, num_workers=0, ) model = DigitCNN().to(device) optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=1e-4) criterion = nn.CrossEntropyLoss() history = [] best_accuracy = 0.0 MODEL_PATH.parent.mkdir(parents=True, exist_ok=True) for epoch in range(1, epochs + 1): model.train() running_loss = 0.0 seen = 0 for images, labels in train_loader: images = images.to(device) labels = labels.to(device) optimizer.zero_grad() loss = criterion(model(images), labels) loss.backward() optimizer.step() running_loss += float(loss.item()) * int(labels.numel()) seen += int(labels.numel()) test_accuracy = evaluate(model, test_loader, device) train_loss = running_loss / seen history.append({"epoch": epoch, "train_loss": train_loss, "test_accuracy": test_accuracy}) print(f"epoch={epoch} train_loss={train_loss:.4f} test_accuracy={test_accuracy:.4f}") if test_accuracy > best_accuracy: best_accuracy = test_accuracy torch.save(model.cpu().state_dict(), MODEL_PATH) model.to(device) metrics = { "model_type": "PyTorch DigitCNN", "training_data": "Real MNIST IDX files downloaded from cvdf-datasets", "optimization_strategy": [ "deeper convolutional neural network", "canvas-style augmentation: rotation, scaling, translation, thick strokes, light blur", "prediction-time test-time augmentation over multiple target sizes and shifts", ], "epochs": epochs, "batch_size": batch_size, "best_test_accuracy": best_accuracy, "history": history, } with open(METRICS_PATH, "w", encoding="utf-8") as f: json.dump(metrics, f, indent=2) print(f"saved model: {MODEL_PATH}") print(f"saved metrics: {METRICS_PATH}") print(f"best_test_accuracy={best_accuracy:.4f}") if __name__ == "__main__": train_model()