QC4: 4. MNIST Handwritten Character Detection using PyTorch, Keras and Tensorflow

!pip install torch torchvision tensorflow matplotlib pillow -q

# TensorFlow
import tensorflow as tf
from tensorflow import keras

# PyTorch
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms

# Common
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image

# TensorFlow dataset
(X_train, y_train), (X_test, y_test) = keras.datasets.mnist.load_data()

print("Train shape:", X_train.shape)

# Normalize
X_train = X_train / 255.0
X_test = X_test / 255.0

# Reshape for CNN
X_train_tf = X_train.reshape(-1,28,28,1)
X_test_tf = X_test.reshape(-1,28,28,1)

model_tf = keras.Sequential([
    keras.layers.Conv2D(32,(3,3),activation='relu',input_shape=(28,28,1)),
    keras.layers.MaxPooling2D((2,2)),

    keras.layers.Conv2D(64,(3,3),activation='relu'),
    keras.layers.MaxPooling2D((2,2)),

    keras.layers.Flatten(),
    keras.layers.Dense(64,activation='relu'),
    keras.layers.Dense(10,activation='softmax')
])

model_tf.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)

model_tf.summary()

history_tf = model_tf.fit(
    X_train_tf, y_train,
    epochs=5,
    batch_size=64,
    validation_split=0.1
)

loss, acc = model_tf.evaluate(X_test_tf, y_test)
print("TensorFlow Test Accuracy:", acc)

transform = transforms.Compose([
    transforms.ToTensor()
])

train_dataset = torchvision.datasets.MNIST(
    root='./data', train=True, download=True, transform=transform
)

test_dataset = torchvision.datasets.MNIST(
    root='./data', train=False, download=True, transform=transform
)

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=64)

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(1,32,3)
        self.conv2 = nn.Conv2d(32,64,3)
        self.pool = nn.MaxPool2d(2,2)

        self.fc1 = nn.Linear(64*5*5,64)
        self.fc2 = nn.Linear(64,10)

    def forward(self,x):
        x = self.pool(torch.relu(self.conv1(x)))
        x = self.pool(torch.relu(self.conv2(x)))
        x = x.view(-1,64*5*5)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model_pt = CNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model_pt.parameters(), lr=0.001)

for epoch in range(5):
    for images, labels in train_loader:
        outputs = model_pt(images)
        loss = criterion(outputs, labels)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    print(f"Epoch {epoch+1} done")

correct = 0
total = 0

with torch.no_grad():
    for images, labels in test_loader:
        outputs = model_pt(images)
        _, predicted = torch.max(outputs.data,1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print("PyTorch Test Accuracy:", correct/total)

from PIL import Image, ImageOps
import numpy as np
import matplotlib.pyplot as plt
import torch

img_path = "test.jpg"

img = Image.open(img_path).convert('L')
img = ImageOps.invert(img)
img = img.resize((28,28))
img_array = np.array(img)

img_array = img_array / 255.0

plt.imshow(img_array, cmap='gray')
plt.title("Processed Image")
plt.axis('off')
plt.show()

img_tf = img_array.reshape(1,28,28,1)
pred_tf = model_tf.predict(img_tf)
print("TensorFlow Prediction:", np.argmax(pred_tf))

img_pt = torch.tensor(img_array).float().unsqueeze(0).unsqueeze(0)
output = model_pt(img_pt)
_, pred_pt = torch.max(output,1)
print("PyTorch Prediction:", pred_pt.item())