Anomaly Detection

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

FINAL_EVALUATION_MODE = False
# While grading 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 random
import warnings
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
from torch.utils.data import DataLoader, Dataset

warnings.filterwarnings("ignore", category=UserWarning)

np.random.seed(0)
random.seed(0)
torch.manual_seed(0)

if not FINAL_EVALUATION_MODE:
    ! gdown https://drive.google.com/uc?id=1QoTW4eZctWvzmjrNrwmai9zhL8RWnhBU
    # ! gdown https://drive.google.com/uc?id=1bD38bZf8pcUyinuvYDXO56RPlXUluzxn
    # ! gdown https://drive.google.com/uc?id=108d2ERztZzq5_ZXyizxeWDOCctG5xhIy

    ! unzip anomaly.zip
    ! unzip train.zip
    ! unzip valid.zip

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

TRAIN_DIR: Path = Path("./train")
VALID_DIR: Path = Path("./valid")

TRAIN_CSV: Path = Path("./train.csv")
VALID_CSV: Path = Path("./valid.csv")

# Set your batch size -- we will require this variable to be present during grading and will use it during testing
BATCH_SIZE: int = 32

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

class ImageDataset(Dataset):
    def __init__(self, dir: Path, csv: Path):
        self.dir: Path = dir
        self.csv = pd.read_csv(csv)

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

    def __getitem__(self, idx) -> tuple[torch.Tensor, int]:
        if torch.is_tensor(idx):
            idx = idx.tolist()
        path, label = self.csv.iloc[idx]
        img = plt.imread(self.dir / path)

        to_tensor = transforms.ToTensor()
        return to_tensor(img), label


def train_dataloader() -> DataLoader:
    """Create a Dataloader with the training data."""
    return DataLoader(ImageDataset(TRAIN_DIR, TRAIN_CSV), batch_size=BATCH_SIZE, shuffle=True)


def valid_dataloader() -> DataLoader:
    """Create a Dataloader with the validation data."""
    return DataLoader(ImageDataset(VALID_DIR, VALID_CSV), batch_size=BATCH_SIZE, shuffle=False)

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return torch.rand(x.shape[0])

    def predict(self, batch: torch.Tensor) -> torch.Tensor:
        """
        Assign 0 or 1 to each image in the batch.
        Do not change or add arguments in this method!
        """
        predictions = torch.ones(batch.shape[0])
        return predictions


def train_model() -> Model:
    """Create and train the model"""
    return Model()

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

model = train_model()

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

def grade_solution(model):
    dataloader = valid_dataloader() # During grading we will replace this with test_dataloader

    predictions = np.concatenate([
        model.predict(images).cpu().to(dtype=torch.int32).numpy() for images, _ in dataloader
    ], axis=0)

    labels = np.concatenate([label for _, label in dataloader], axis=0, dtype=np.int32)

    accuracy = sum(labels == predictions) / len(labels)

    score = min(max(accuracy - 0.6, 0.0), 0.3) / 0.3

    print(f"Accuracy: {accuracy}")
    print(f'Your score is {score} pts')
    return score

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

if not FINAL_EVALUATION_MODE:
    grade_solution(model)