Pixels

######################### DO NOT CHANGE THIS CELL WHEN SUBMITTING ##########################
FINAL_EVALUATION_MODE = False # During grading on the grading system, the flag is automatically changed to True.

# ######################### DO NOT CHANGE THIS CELL WHEN SUBMITTING ##########################

import numpy as np
import torch
import torch.nn as nn
import math

if not FINAL_EVALUATION_MODE:
    import matplotlib.pyplot as plt

seed = 42
torch.manual_seed(seed)
np.random.seed(seed)

class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(1, 16, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(16, 32, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(32 * 7 * 7, 128),
            nn.ReLU(),
            nn.Linear(128, 10)
        )
    def forward(self, x):
        return self.net(x)


def load_trained_model(model_path="mnist_classifier.pth"):
    """Load a trained model."""
    model = SimpleCNN()
    model.load_state_dict(torch.load(model_path, map_location='cpu'))
    model.eval()
    return model


def predict(images, model):
    # normalize input!
    images = images.float()
    images = (images - 0.1307) / 0.3081
    with torch.no_grad():
        logits = model(images)
        probs = torch.softmax(logits, dim=1)
        preds = torch.argmax(probs, dim=1)
    return preds, probs

def round_half_up(number: float) -> int:
    return int(math.floor(number + 0.5))

def compute_score(acc_val: float):
    upper_limit = 0.6
    lower_limit = 0.35
    if acc_val > upper_limit:
        return 100
    elif acc_val < lower_limit:
        return 0
    else:
        return round_half_up((acc_val - lower_limit) / (upper_limit - lower_limit) * 100)

# ######################### DO NOT CHANGE THIS CELL WHEN SUBMITTING ##########################

model = load_trained_model("data/mnist_classifier.pth")
model.net.eval()

def load_split(path):
    d = np.load(path)
    return torch.from_numpy(d["images"]), torch.from_numpy(d["labels"]).long()

train_imgs, train_labels = load_split("data/train.npz")
val_imgs, val_labels = load_split("data/val.npz")

print(f"Train: {len(train_imgs)}, Val: {len(val_imgs)}")

# ######################### DO NOT CHANGE THIS CELL WHEN SUBMITTING ##########################

if not FINAL_EVALUATION_MODE:
    fig, axes = plt.subplots(2, 10, figsize=(20, 4))
    for i in range(20):
        row = (i // 10)
        col = i % 10
        predicted = predict(train_imgs[i:i+1], model)[0].item()
        axes[row, col].imshow(train_imgs[i, 0].numpy(), cmap='gray', vmin=0, vmax=1)
        axes[row, col].set_title(f"y={train_labels[i].item()}, pred={predicted}", fontsize=8)
        axes[row, col].axis('off')
    plt.tight_layout()
    plt.show()

def solution(img):
    """
    Classification of masked MNIST images.

    Args:
        img: tensor (1, 28, 28) - the masked image to classify

    Returns:
        the predicted label (0-9) for the given image
    """
    # TODO: write your solution here!

    # Example: direct model prediction on the masked images (baseline)
    img = img.unsqueeze(0)  # add the batch dimension
    predicted_label, _ = predict(img, model)
    return predicted_label.item()

######################### DO NOT CHANGE THIS CELL WHEN SUBMITTING ##########################

if not FINAL_EVALUATION_MODE:
    accs = []
    for img, label in zip(val_imgs, val_labels):
        pred = solution(img)
        accs.append(pred == label.item())
    acc = np.mean(accs)
    print(f"  Accuracy on the validation set: {acc*100:.2f}%")

    score = compute_score(acc)
    print(f"  Number of points: {score}")