Multi-Label Classification

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

FINAL_EVALUATION_MODE: bool = False  # We will set this flag to True during grading.

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

import os

import numpy as np

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as fun
import torchvision

from torch.utils.data import DataLoader, TensorDataset
from torchvision import transforms

from tqdm import tqdm

from sklearn.metrics import f1_score

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

def seed_everything(seed: int):
    """Sets the seed for reproducibility of results in Python, NumPy and PyTorch."""

    os.environ["PYTHONHASHSEED"] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

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

# number of classes to classify
N_CLASSES: int = 10

# mapping of the class index to its name
LABEL_NAMES: dict[int, str] = {
    0: "T-shirt/top",
    1: "Trouser",
    2: "Pullover",
    3: "Dress",
    4: "Coat",
    5: "Sandal",
    6: "Shirt",
    7: "Sneaker",
    8: "Bag",
    9: "Ankle boot"
}

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

if not FINAL_EVALUATION_MODE:
    FILES: list[str] = [
        "runway-mnist/train-x.npz",
        "runway-mnist/train-y.npz",
        "runway-mnist/val-x.npz",
        "runway-mnist/val-y.npz"
    ]

    # download the data again if anything is missing
    if not all(os.path.exists(file) for file in FILES):
        import gzip
        import tarfile
        import shutil

        if not os.path.exists("runway-mnist"):
            os.mkdir("runway-mnist")

        COMPRESSED_ARCHIVE = "runway-mnist.tar.gz"
        TAR_ARCHIVE  = COMPRESSED_ARCHIVE.rstrip(".gz")
        DOWNLOAD_URL = "https://drive.google.com/uc?id=1oNAFYdJyCVe3Po90KLUPxAG9HGuL7NSw"

        try:
            import gdown
        except ImportError as err:
            raise RuntimeError("To download the dataset, you need a local installation of the gdown package: `pip install gdown`") from err

        gdown.download(DOWNLOAD_URL, str(COMPRESSED_ARCHIVE), quiet=False)

        with gzip.open(COMPRESSED_ARCHIVE, "rb") as compressed:
            with open(TAR_ARCHIVE, "wb") as archive:
                shutil.copyfileobj(compressed, archive)

        os.remove(COMPRESSED_ARCHIVE)
        print(f"Decompressed: {TAR_ARCHIVE}")

        with tarfile.open(TAR_ARCHIVE, "r") as tar:
            tar.extractall("runway-mnist")

        os.remove(TAR_ARCHIVE)
        print(f"Extracted: {TAR_ARCHIVE}")

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

SEED: int = 42
seed_everything(SEED)

def load_x(usage) -> torch.Tensor:
    path = f"runway-mnist/{usage}-x.npz"
    return torch.tensor(np.load(path)["images"], dtype=torch.float32).unsqueeze(1)

def load_y(usage) -> torch.Tensor:
    path = f"runway-mnist/{usage}-y.npz"
    return torch.tensor(np.load(path)["labels"], dtype=torch.long)

train_dataset = TensorDataset(load_x("train"), load_y("train"))
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

val_dataset = TensorDataset(load_x("val"), load_y("val"))
val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False)

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

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)

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

def compute_score(f1: float) -> int:
    """Computes the score in points from the value of the F1 metric."""
    lower_bound = 0.57
    upper_bound = 0.87

    if f1 <= lower_bound:
        return 0
    elif lower_bound < f1 < upper_bound:
        return int(round(100 * (f1 - lower_bound) / (upper_bound - lower_bound)))
    else:
        return 100

def evaluate_algorithm(model, predict, loader) -> float:
    """Computes the metrics and scores your solution on their basis. Returns the computed value of the F1 metric."""
    preds = []
    labels = []
    model.eval()
    with torch.no_grad():
        for x, y in loader:
            prediction = predict(model, x.to(device)).cpu()
            preds.append(prediction)
            labels.append(y)

    predictions = torch.cat(preds)
    labels = torch.cat(labels)

    f1 = f1_score(labels.numpy(), predictions.numpy(), average="macro")
    points = compute_score(f1)
    print(f"Your F1 score: {f1:.3f}, which gives {points} points.")
    return f1

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

class NaiveSolution(nn.Module):
    """Naive solution."""

    def __init__(self):
        super().__init__()

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        """This naive model predicts that all classes are present in the image."""
        BATCH_SIZE = input.size(0)
        return torch.Tensor([1] * BATCH_SIZE * N_CLASSES).reshape(BATCH_SIZE, -1)

def train_naive(_: NaiveSolution):
    """The model does not require training."""
    pass

def predict_naive(model: NaiveSolution, input: torch.Tensor) -> torch.Tensor:
    """Our model returns the model's predictions directly; we do not process them further."""
    return model(input).to(torch.long)

######################### DO NOT CHANGE THIS CELL WHEN SUBMITTING #############################
if not FINAL_EVALUATION_MODE:
    naive_solution = NaiveSolution().to(device)
    naive_solution.train()
    train_naive(naive_solution)
    evaluate_algorithm(naive_solution, predict_naive, val_loader)

class Solution(nn.Module):
    """Your solution."""

    def __init__(self):
        super().__init__()

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        """Model inference"""
        BATCH_SIZE = input.size(0)
        return torch.rand(BATCH_SIZE * N_CLASSES).reshape(BATCH_SIZE, -1)

def train_solution(_: Solution):
    """Training loop for your model."""
    pass

def predict_solution(model: Solution, input: torch.Tensor) -> torch.Tensor:
    """Classification using the model.
    This function is kept separate to make it easy to post-process the model's outputs."""
    predictions = model(input).round().to(torch.long)
    return predictions

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

solution = Solution().to(device)
solution.train()

train_solution(solution)

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

if not FINAL_EVALUATION_MODE:
    evaluate_algorithm(solution, predict_solution, val_loader)