Add application files

app.py

@@ -0,0 +1,25 @@
+import gradio as gr
+from predict import predict_image
+# Custom CSS to make the label text bigger
+def predict(image_dict):
+    # Extract the "composite" key from the dictionary
+    composite_image = image_dict["composite"]
+    # composite_image.save("sketchpad_output.png")  # Save as PNG
+    predicted = predict_image(composite_image)
+    # print(predicted)
+    return predicted #, composite_image  # Directly return the PIL image
+css = """
+.big-label {
+    font-size: 24px; /* Adjust this value to make the label bigger */
+    font-weight: bold; /* Optional: to make it bold */
+}
+"""
+demo = gr.Interface(
+    fn=predict,
+    inputs=gr.Sketchpad(type="pil", brush=gr.Brush(default_size=20)),  # Ensure it returns a PIL image
+    outputs=[gr.Label(num_top_classes=3, label="Predicted number is:", elem_classes=["big-label"])],
+    css=css
+)
+
+if __name__ == "__main__":
+    demo.launch()

device.py

@@ -0,0 +1,15 @@
+import torch
+
+def get_device():
+    if torch.cuda.is_available:
+        # print('cuda is available')
+        return 'cuda'
+    elif torch.backends.mps.is_available:
+        # print('mps is available')
+        return 'mps'
+    else:
+        # print('using cpu')
+        return 'cpu'
+
+device = get_device()
+# print(device)

predict.py

@@ -0,0 +1,105 @@
+import os
+from PIL import Image
+import numpy as np
+import torch
+from torchvision import transforms
+from device import get_device
+from simple_nn import SimpleNN
+from simple_cnn import SimpleCNN
+from view_image import view_image, view_tensor_image
+
+transform = transforms.Compose([
+    transforms.ToTensor(),
+    transforms.Normalize((0.5,), (0.5,))
+])
+def predict_image(image):
+    
+    model = SimpleCNN()
+    with open('mnist_simple_cnn.pht', 'rb') as f:
+        state_dict = torch.load(f, weights_only=True)
+    model.load_state_dict(state_dict)
+    model.eval()
+
+    image = image.convert('RGBA')
+    grayscale_image = Image.new("L", image.size, 255)  # Create a white background
+    grayscale_image.paste(image.convert("L"), mask=image.split()[3])  # Use alpha channel as mask
+    grayscale_image = grayscale_image.resize((28, 28))  # Resize to 28x28 pixels
+    
+    grayscale_image.save("processed_image.png")
+
+    image_np = np.array(grayscale_image)
+    image_np = 255 - image_np  # Invert colors (MNIST has white digits on black)
+
+    # Normalize to range [0, 1]
+    image_np = image_np / 255.0
+
+    image_tensor = transform(image_np)  # Add batch and channel dimensions
+    image_tensor = image_tensor.unsqueeze(0)
+    image_tensor = image_tensor.to(torch.float32)
+
+    # image_tensor = transform(grayscale_image).unsqueeze(0)  # Add batch and channel dimensions
+    
+    with torch.no_grad():
+        output = model(image_tensor)
+        #_, predicted = torch.max(output.data, 1)
+        probabilities = torch.softmax(output, dim=1) 
+    # Convert probabilities to a list of (class, probability)
+    class_probabilities = {
+        str(class_index): prob.item() for class_index, prob in enumerate(probabilities[0])
+    }
+    print(class_probabilities)
+    # class_probabilities = {}
+    return class_probabilities
+
+def predict(model_path, image_path):
+    model = SimpleCNN()
+    with open(model_path, 'rb') as f:
+        state_dict = torch.load(f, weights_only=True)
+    model.load_state_dict(state_dict)
+    model.eval()
+
+    # Load and preprocess the image
+    image = Image.open(image_path).convert("L")  # Convert to grayscale
+    # view_image(image=image)
+    # Resize to 28x28
+    image = image.resize((28, 28))
+
+    # Convert to NumPy array and invert colors if needed
+    image_np = np.array(image)
+    image_np = 255 - image_np  # Invert colors (MNIST has white digits on black)
+
+    # Normalize to range [0, 1]
+    image_np = image_np / 255.0
+
+    # Convert to tensor
+    
+    image_tensor = transform(image_np)  # Add batch and channel dimensions
+    image_tensor = image_tensor.unsqueeze(0)
+    image_tensor = image_tensor.to(torch.float32)
+    # Ensure the tensor is in the correct dtype
+    
+    # view_tensor_image(image_tensor=image_tensor)
+    
+    with torch.no_grad():
+        output = model(image_tensor)
+        #_, predicted = torch.max(output.data, 1)
+        probabilities = torch.softmax(output, dim=1) 
+    # Convert probabilities to a list of (class, probability)
+    class_probabilities = {
+        str(class_index): prob.item() for class_index, prob in enumerate(probabilities[0])
+    }
+    # return predicted.item()
+    return class_probabilities
+if __name__ == "__main__":
+    device = get_device()
+    model_path = "trained_model/mnist_simple_cnn.pht"
+   
+    # Loop through all files in the test folder
+    test_folder = "test/"
+    for filename in os.listdir(test_folder):
+        if filename.endswith(".png"):  # Only process .png files (you can add more extensions if needed)
+            image_path = os.path.join(test_folder, filename)
+            predicted = predict(model_path = "mnist_model.pht",image_path=image_path)
+            print(F"[INFO] The predicted results of the image {image_path} are: {predicted}")
+            print()
+           

simple_cnn.py

@@ -0,0 +1,28 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class SimpleCNN(nn.Module):
+    def __init__(self):
+        super(SimpleCNN, self).__init__()
+        # Convolutional layers
+        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)  # 28x28 -> 28x28
+        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)  # 28x28 -> 28x28
+        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)  # 28x28 -> 14x14
+
+        self.bn1 = nn.BatchNorm2d(32)
+        self.bn2 = nn.BatchNorm2d(64)
+
+        # Fully connected layers
+        self.fc1 = nn.Linear(64 * 14 * 14, 128)
+        self.dropout = nn.Dropout(0.5)
+        self.fc2 = nn.Linear(128, 10)
+
+    def forward(self, x):
+        x = F.relu(self.bn1(self.conv1(x)))  # Apply first convolution and ReLU
+        x = self.pool(F.relu(self.bn2(self.conv2(x))))  # Apply second convolution, ReLU, and pooling
+        x = torch.flatten(x, 1) # Flatten the feature maps
+        x = F.relu(self.fc1(x))  # Fully connected layer with ReLU
+        x = self.dropout(x)
+        x = self.fc2(x)  # Output layer
+        return x

simple_nn.py

@@ -0,0 +1,17 @@
+import torch
+import torch.nn as nn
+
+class SimpleNN(nn.Module):
+    def __init__(self):
+        super(SimpleNN, self).__init__()
+        self.fc1 = nn.Linear(28*28, 128)
+        self.fc2 = nn.Linear(128, 64)
+        self.fc3 = nn.Linear(64, 10)
+
+   
+    def forward(self, x):
+        x = x.view(-1, 28*28)
+        x = torch.relu(self.fc1(x))
+        x = torch.relu(self.fc2(x))
+        x = self.fc3(x)
+        return x

test.py

@@ -0,0 +1,130 @@
+import os
+from PIL import Image
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torchvision import datasets, transforms
+from torch.utils.data import DataLoader
+from torch.utils.tensorboard import SummaryWriter
+from simple_nn import SimpleNN
+from simple_cnn import SimpleCNN
+from device import get_device
+from view_image import view_image,view_batch_images, save_batch_images
+from tqdm import tqdm  # Import tqdm for the progress bar
+
+root_path = os.path.expanduser('data')
+
+# Define transforms for training and testing
+transforms = {
+    'train': transforms.Compose([
+        transforms.RandomRotation(10),
+        transforms.RandomHorizontalFlip(),
+        transforms.ToTensor(),
+        transforms.Normalize((0.5,), (0.5,))  # Normalize for MNIST
+    ]),
+    'valid_test': transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize((0.5,), (0.5,))
+    ])
+}
+
+# Define dataset and dataloader
+train_dataset = datasets.MNIST(root=root_path, download=True, train=True, transform=transforms['train'])
+test_dataset = datasets.MNIST(root=root_path, download=True, train=False, transform=transforms['valid_test'])
+
+train_loader = DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)
+test_loader = DataLoader(dataset=test_dataset, batch_size=64, shuffle=False)
+
+model = SimpleCNN()
+device = get_device()
+
+model.to(device=device)
+criterion = nn.CrossEntropyLoss()
+optimizer = optim.Adam(model.parameters(), lr=0.001)
+epochs = 50
+for epoch in range(epochs):
+    model.train()  # Set the model to training mode
+    epoch_loss = 0  # Initialize epoch loss
+    correct = 0  # Track number of correct predictions
+    total = 0  # Track total predictions
+
+    with tqdm(total=len(train_loader), desc=f"Epoch {epoch+1}/{epochs}", unit="batch") as pbar:
+        for batch_idx, (data, target) in enumerate(train_loader):
+            data, target = data.to(device), target.to(device)  # Move data to device
+            
+            # Forward pass
+            output = model(data)
+            loss = criterion(output, target)
+
+            # Backward pass and optimization
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+            epoch_loss += loss.item()  # Add current batch loss to epoch loss
+
+            # Calculate accuracy for this batch
+            _, predicted = torch.max(output, 1)
+            total += target.size(0)
+            correct += (predicted == target).sum().item()
+
+            # Update the progress bar
+            pbar.set_postfix(loss=loss.item(), accuracy=100. * correct / total)
+            pbar.update(1)
+
+    # Calculate the average loss and accuracy for this epoch
+    avg_loss = epoch_loss / len(train_loader)
+    accuracy = 100. * correct / total
+
+
+    print(f"Epoch {epoch+1}/{epochs}, Loss: {avg_loss:.4f}, Accuracy: {accuracy:.2f}%")
+
+    # Now validation (on validation set)
+    correct_val = 0
+    total_val = 0
+    model.eval()  # Set the model to evaluation mode
+    with torch.no_grad():
+        for data, target in test_loader:  # Assuming `val_loader` is your validation data loader
+            data, target = data.to(device), target.to(device)
+            output = model(data)
+            _, predicted = torch.max(output.data, 1)
+            total_val += target.size(0)
+            correct_val += (predicted == target).sum().item()
+
+    accuracy_val = 100 * correct_val / total_val
+    print(f'Validation Accuracy: {accuracy_val:.2f}%')
+
+
+# Save model after training
+torch.save(model.state_dict(), "mnist_simple_cnn.pht")
+print(f'Model saved')
+
+
+model.eval()
+
+test_folder = "test/"
+
+# Loop through all files in the test folder
+for filename in os.listdir(test_folder):
+    if filename.endswith(".png"):  # Only process .png files (you can add more extensions if needed)
+        image_path = os.path.join(test_folder, filename)
+
+        # Load and preprocess the image
+        image = Image.open(image_path).convert("L")  # Convert to grayscale
+        image = image.resize((28, 28))  # Resize to 28x28 pixels
+        image_tensor = transforms['valid_test'](image).unsqueeze(0)  # Add batch and channel dimensions
+        image_tensor = image_tensor.to(device)
+
+        # Make prediction
+        with torch.no_grad():
+            output = model(image_tensor)
+            probabilities = torch.softmax(output, dim=1)
+
+        # Convert probabilities to a list of (class, probability)
+        class_probabilities = {
+            str(class_index): prob.item() for class_index, prob in enumerate(probabilities[0])
+        }
+
+        # Print or store the predictions for the current image
+        print(f"Predictions for {filename}: {class_probabilities}")
+        print()  # Line break for separation

train.py

@@ -0,0 +1,109 @@
+import os
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torchvision import datasets, transforms
+from torch.utils.data import DataLoader
+from torch.utils.tensorboard import SummaryWriter
+from simple_nn import SimpleNN
+from simple_cnn import SimpleCNN
+from device import get_device
+from view_image import view_batch_images, save_batch_images
+from tqdm import tqdm  # Import tqdm for the progress bar
+
+
+def train(model, device, train_loader, test_loader, criterion, optimizer, epochs = 5):
+    
+    # Initialize TensorBoard writer
+    # writer = SummaryWriter('runs/mnist_experiment')
+
+    # Train the model
+    for epoch in range(epochs):
+        model.train()
+        epoch_loss = 0
+        with tqdm(total=len(train_loader), desc=f"Epoch {epoch+1}/{epochs}", unit="batch") as pbar:
+            for batch_idx, (data, target) in enumerate(train_loader):
+                # Forward pass
+                data, target = data.to(device), target.to(device)
+                output = model(data)
+                loss = criterion(output, target)
+
+                # Backward pass and optimization
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+
+                epoch_loss += loss.item()
+                
+                pbar.set_postfix(loss=loss.item())
+                pbar.update(1)
+
+        # Log the average loss for this epoch to TensorBoard
+        # avg_loss = epoch_loss / len(train_loader)
+        # writer.add_scalar('Loss/train', avg_loss, epoch)
+
+        print(f'Epoch {epoch+1}/{epochs}, Loss: {loss.item():.4f}')
+        validation(model,device,test_loader)
+    # After training, visualize with TensorBoard
+    # writer.close()
+
+# Test the model
+def validation(model, device, data_loader):
+    correct = 0
+    total = 0
+    model.eval()  # Set the model to evaluation mode
+    with torch.no_grad():
+        for data, target in data_loader:
+            data, target = data.to(device), target.to(device)
+            output = model(data)
+            _, predicted = torch.max(output.data, 1)
+            total += target.size(0)
+            correct += (predicted == target).sum().item()
+
+    accuracy = 100 * correct / total
+    print(f'Accuracy on test set: {accuracy:.2f}%')
+
+# Save the trained model
+def save_model(model, model_save_path):
+    torch.save(model.state_dict(), model_save_path)
+    print(f'[INFO] Model saved to {model_save_path}')
+
+
+if __name__ == "__main__":
+    root_path = os.path.expanduser('data')
+    # Load the dataset
+    transforms = {
+        'train': transforms.Compose([
+            transforms.RandomRotation(10),
+            transforms.RandomHorizontalFlip(),
+            transforms.ToTensor(),
+            transforms.Normalize((0.5,), (0.5,))  # Normalize for MNIST
+        ]),
+        'valid_test' : transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize((0.5,), (0.5,))
+        ])
+    }
+    train_dataset = datasets.MNIST(root=root_path, download=False, train=True, transform=transforms['train'])
+    test_dataset = datasets.MNIST(root=root_path, download=True, train=False, transform=transforms['valid_test'])
+    train_loader = DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)
+    test_loader = DataLoader(dataset=test_dataset, batch_size=64, shuffle=False)
+    # view_batch_images(train_loader=train_loader)
+
+    # Get a batch of images
+    # data_iter = iter(train_loader)
+    # images, labels = next(data_iter)
+    # Save the images
+    # save_batch_images(images, save_dir="output_images", prefix="mnist_image", file_format="png")
+
+    # Initialize the model, loss function, and optimizer
+    device = get_device()
+    # model = SimpleNN()
+    model = SimpleCNN()
+    model.to(device=device)
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.Adam(model.parameters(), lr=0.001)
+
+    epochs = 20
+    train(model=model,device=device,train_loader=train_loader,test_loader=test_loader,criterion=criterion,optimizer=optimizer,epochs = epochs)
+    save_model(model=model,model_save_path="mnist_model.pht")

view_image.py

@@ -0,0 +1,63 @@
+import os
+from PIL import Image
+import matplotlib.pyplot as plt
+import torchvision
+from torchvision.transforms import ToPILImage
+
+def view_image(image):
+    plt.imshow(image, cmap="gray")
+    plt.title("Grayscale Image")
+    plt.axis("off")  # Hide axes for better visualization
+    plt.show()
+
+def view_tensor_image(image_tensor, title="Image"):
+    image_np = image_tensor.squeeze().numpy()
+    plt.imshow(image_np)
+    plt.title(title)
+    plt.axis('off')
+    plt.show()
+
+def view_batch_images(train_loader, num_images=8):
+    """
+    Display a batch of images from the train_loader.
+    
+    Parameters:
+        train_loader (DataLoader): The DataLoader containing the images.
+        num_images (int): Number of images to display from the batch.
+    """
+    data_iter = iter(train_loader)
+    images, labels = next(data_iter)  # Get a batch of images and labels
+    
+    # Make a grid of images
+    img_grid = torchvision.utils.make_grid(images[:num_images], nrow=num_images, normalize=True)
+    img_np = img_grid.numpy().transpose((1, 2, 0))  # Rearrange dimensions for plotting
+    
+    plt.figure(figsize=(12, 6))
+    plt.imshow(img_np, cmap="gray")
+    plt.title("Batch of Images")
+    plt.axis("off")
+    plt.show()
+
+
+def save_batch_images(images, save_dir, prefix="image", file_format="png", unnormalize=None):
+    """
+    Save each image in a batch to a specified directory.
+    
+    Parameters:
+        images (torch.Tensor): Batch of images with shape (B, C, H, W).
+        save_dir (str): Directory to save the images.
+        prefix (str): Prefix for the saved image filenames.
+        file_format (str): File format for the saved images (e.g., "png", "jpg").
+        unnormalize (callable, optional): Function to unnormalize the images before saving.
+    """
+    os.makedirs(save_dir, exist_ok=True)  # Create the directory if it doesn't exist
+    to_pil = ToPILImage()  # Converts tensors to PIL images
+    
+    for idx, image in enumerate(images):
+        if unnormalize:
+            image = unnormalize(image)  # Apply unnormalization if provided
+        
+        pil_image = to_pil(image)  # Convert to PIL Image
+        filename = os.path.join(save_dir, f"{prefix}_{idx}.{file_format}")
+        pil_image.save(filename)
+        print(f"Saved: {filename}")

view_model_information.py

@@ -0,0 +1,26 @@
+import os
+from PIL import Image
+import numpy as np
+import torch
+from torchvision import transforms
+from device import get_device
+from simple_nn import SimpleNN
+from simple_cnn import SimpleCNN
+from view_image import view_image, view_tensor_image
+
+model_path = "mnist_model.pht"
+# Load model
+model = SimpleCNN()
+with open(model_path, 'rb') as f:
+    state_dict = torch.load(f, weights_only=True)
+model.load_state_dict(state_dict)
+
+# View model information
+print(model)  # Display the model architecture
+
+# For more detailed information about the model's parameters:
+print(f"Model summary: {model}")
+
+# You can also view the parameters' details (e.g., number of parameters, layers, etc.)
+for name, param in model.named_parameters():
+    print(f"Parameter: {name}, Shape: {param.shape}")