Imbalanced Classification

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

FINAL_EVALUATION_MODE = False
# While checking your solution, we will change this value to True
# The value of this flag M U S T be set to False in the solution you send us!

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

import abc
import os

import glob
import gdown
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import matplotlib.pyplot as plt
import zipfile

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

GDRIVE_DATA = [
    ("1bR87z7ZI3gLK0vAGkyr_cnVGZ9P9bO7A", "train_data.zip"),
    ("1TA0lWnjJCv3lyRMML4JNHsJz3RJ-TUwZ", "valid_data.zip"),
]

def download_data():
    for file_id, zip_name in GDRIVE_DATA:
        folder_name = zip_name.split(".")[0]
        if not os.path.exists(folder_name):
            url = f'https://drive.google.com/uc?id={file_id}'
            gdown.download(url, output=zip_name, quiet=True)
            with zipfile.ZipFile(zip_name, 'r') as zip_ref:
                    zip_ref.extractall(folder_name)
            os.remove(zip_name)

download_data()

class ImageDataset(torch.utils.data.Dataset):
    """Implementation of torch's dataset abstraction."""
    def __init__(self, dataset_type: str):
        self.filelist = glob.glob(f"{dataset_type}_data/*")
        self.labels   = [0 if "normal" in path else 1 for path in self.filelist]

    def __len__(self):
        return len(self.filelist)

    def __getitem__(self, idx) -> tuple[torch.Tensor, int]:
        if torch.is_tensor(idx):
            idx = idx.tolist()
        image = torchvision.transforms.functional.to_tensor(plt.imread(self.filelist[idx])[:,:,0])
        label = self.labels[idx]
        return image, label
    
    def loader(self, **kwargs) -> torch.utils.data.DataLoader:
        """
        Create a `DataLoader` for the current dataset.

        All `**kwargs` will be passed to the constructor of `torch.utils.data.DataLoader`.
        In short, `DataLoader`s are a data-loading abstraction that provides a convenient interface.
        You can learn more about them here: https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader
        """
        return torch.utils.data.DataLoader(self, **kwargs)
    
train_dataset: ImageDataset = ImageDataset("train")
valid_dataset: ImageDataset = ImageDataset("valid")

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

def accuracy_to_points(accuracy: float) -> float:
    """Compute the score from the prediction accuracy."""
    return (round(accuracy, 2) - 0.5) * 2 if accuracy > 0.5 else 0.0

def grade(model):
    """Assess how many points the current task will receive."""
    model.eval()
    test_loader = valid_dataset.loader()
    correct = 0
    total = 0
    with torch.no_grad():
        for [images, labels] in test_loader:
            outputs = model(images).squeeze()
            incorrect_indices = torch.where((outputs > 0.5).int() != labels)[0]
            correct += len(labels) - len(incorrect_indices)
            total += len(labels)
        accuracy = correct / total if total != 0 else 0
        if not FINAL_EVALUATION_MODE:
            print(f"Accuracy: {int(round(accuracy, 2) * 100)}%")
        return accuracy_to_points(accuracy)

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

class CnnClassifier(torch.nn.Module, abc.ABC):
    MODEL_PATH: str = "cnn-classifier.pth"
    
    @classmethod
    def load(cls):
        """Load the model from a file."""
        model = cls()
        model.load_state_dict(torch.load(cls.MODEL_PATH))
        return model


    @classmethod
    @abc.abstractmethod
    def create_with_training(cls):
        """Save the model to a file."""
        pass

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

if not FINAL_EVALUATION_MODE:
    class DummyCnnClassifier(CnnClassifier):
        def forward(self, x):
            batch_size, *_ = x.shape
            return torch.zeros(batch_size)
        
        @classmethod
        def create_with_training(cls):
            return cls()
    
    dummy_model = DummyCnnClassifier.create_with_training()
    print(f"DummyCnnClassifier -- Score: {grade(dummy_model)} pts")

    del dummy_model
    del DummyCnnClassifier

class YourCnnClassifier(CnnClassifier):
    def forward(self, x):
        batch_size, *_ = x.shape
        return torch.zeros(batch_size)

    @classmethod
    def create_with_training(cls):
        return cls()

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

your_model = YourCnnClassifier.load() if FINAL_EVALUATION_MODE else YourCnnClassifier.create_with_training()

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

def evaluate_model(model):
    """Assess how many points the current task will receive."""
    return grade(model)

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

if not FINAL_EVALUATION_MODE:
    print(f"YourCnnClassifier -- Score: {evaluate_model(your_model):.2f} pts")