simple_MNIST_reduced_for_CNN

train_and_inference.py

import zipfile
import os
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader
from PIL import Image
import matplotlib.pyplot as plt
import pandas as pd

# Paths
zip_file_path = '/path/to/archive.zip'  # Update this path
extract_path = '/path/to/extracted_data'  # Update this path

# Extract the archive
with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
    zip_ref.extractall(extract_path)

# List the directories to identify the structure
data_path = os.path.join(extract_path, 'Reduced MNIST Data')
train_dir = os.path.join(data_path, 'Reduced Trainging data')
test_dir = os.path.join(data_path, 'Reduced Testing data')

# Transform to normalize the data
transform = transforms.Compose([
    transforms.Grayscale(num_output_channels=1),
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

# Load the datasets with ImageFolder
train_dataset = ImageFolder(train_dir, transform=transform)
test_dataset = ImageFolder(test_dir, transform=transform)

# Create data loaders
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)

# Define a simple neural network model
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 16, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
        self.fc1 = nn.Linear(16 * 14 * 14, 128)
        self.fc2 = nn.Linear(128, 64)
        self.fc3 = nn.Linear(64, 10)
    
    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))
        x = x.view(-1, 16 * 14 * 14)
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# Ensure you have a device (GPU if available)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model = SimpleCNN().to(device)

# Define loss function and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# Train the model for 3 epochs
num_epochs = 3
losses = []
for epoch in range(num_epochs):
    model.train()
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)
        
        # Zero the gradients
        optimizer.zero_grad()
        
        # Forward pass
        outputs = model(images)
        loss = criterion(outputs, labels)
        losses.append(loss.item())
        
        # Backward pass and optimization
        loss.backward()
        optimizer.step()
    
    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')

# Perform inferences on all images in each label folder (0-9)
label_folders = range(10)
inference_results_all = []

for label in label_folders:
    label_dir = os.path.join(test_dir, str(label))
    for image_file in os.listdir(label_dir):
        image_path = os.path.join(label_dir, image_file)
        image = Image.open(image_path).convert('L')
        image_tensor = transform(image).unsqueeze(0).to(device)
        
        with torch.no_grad():
            sample_output = model(image_tensor)
            predicted_label = torch.argmax(sample_output, 1).item()
            inference_results_all.append((image_file, predicted_label, label))

# Convert results to a DataFrame for better visualization
results_df_all = pd.DataFrame(inference_results_all, columns=['Image', 'Predicted Label', 'Actual Label'])

# Calculate the accuracy for each label and the overall accuracy
results_df_all['Correct'] = results_df_all['Predicted Label'] == results_df_all['Actual Label']

accuracy_per_label = results_df_all.groupby('Actual Label')['Correct'].mean() * 100
overall_accuracy = results_df_all['Correct'].mean() * 100

# Print overall accuracy
print(f"Overall Accuracy: {overall_accuracy:.2f}%")

# Plot the accuracy distribution
plt.figure(figsize=(10, 6))
plt.bar(accuracy_per_label.index, accuracy_per_label.values, color='skyblue')
plt.xlabel('Label')
plt.ylabel('Accuracy (%)')
plt.title('Accuracy Distribution for Each Label')
plt.xticks(accuracy_per_label.index)
plt.ylim(0, 100)

# Display the plot
plt.show()

And here is the dataset links: https://www.kaggle.com/datasets/mohamedgamal07/reduced-mnist/code