Data Prototypes

######################### DO NOT CHANGE THIS CELL ##########################

FINAL_EVALUATION_MODE = False  # During checking, we will set this flag to True.

######################### DO NOT CHANGE THIS CELL ##########################
# The libraries needed to run this notebook are imported below.
# In your solution you will be able to use the other libraries listed
# above. The place for importing them is further below, in the section
# for your solution.

import os
import torch
import pickle
import numpy as np
import math
import torch.nn as nn
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader
from torchvision.models import mobilenetv3, MobileNet_V3_Small_Weights
import torchinfo
import tarfile
from tqdm import tqdm, trange
from typing import Tuple, Type

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

if not FINAL_EVALUATION_MODE:
    import gdown
    import os

    GDRIVE_DATA = [
        ("15v0YZx6frCzdVT9zPcoD7Vr-N0mzFKsN", "models/mobilenet_v3_small.pth"),
        ("1gHeU8Of1BYbrFWG3Roi1esYq_pevhSAo", "data/val_embeddings.pkl"),
        ("1p0CHP5APgjx18Iq1Fj5rKq9nRt9ZlP_X", "data/Cifar100/train_val.gz")
    ]
    
    for file_id, output in GDRIVE_DATA:        
        url = f'https://drive.google.com/uc?id={file_id}'
        os.makedirs(os.path.dirname(output), exist_ok=True)
        gdown.download(url, output, quiet=False)
        
        print(f"Downloaded: {output}")

######################### DO NOT CHANGE THIS CELL ##########################
class Net(nn.Module):
    """
    This is a class representing the classifier that will be used to
    transform images into their embeddings.
    """
    def __init__(self, verbose=False):
        super().__init__()

        self.net = self.get_net().to(DEVICE)

        self.IMAGE_HEIGHT = 224
        self.IMAGE_WIDTH = 224
        self.INPUT_CHANNELS = 3

        if verbose:
            self.summary()

    @staticmethod
    def get_net() -> nn.Module:
        mobilenet = mobilenetv3.mobilenet_v3_small(weights=None)
        mobilenet.load_state_dict(torch.load('models/mobilenet_v3_small.pth'))
        return mobilenet

    def summary(self) -> None:
        """
        Prints a description and a summary of the network using the
        torchinfo library.
        """
        print(
            torchinfo.summary(
                self.net,
                input_size=(
                    1,
                    self.INPUT_CHANNELS,
                    self.IMAGE_HEIGHT,
                    self.IMAGE_WIDTH,
                ),
            )
        )

    def get_embedding(self, x: torch.Tensor) -> torch.Tensor: 
        """
        For the given image `x`, returns its embedding
        """
        x = self.net(x)
        return torch.flatten(x, 1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)

# Network initialisation
model = Net(verbose=True)

######################### DO NOT CHANGE THIS CELL ##########################
class Dataset:
    """
    Class representing the CIFAR-100 dataset. This dataset is divided into exactly
    100 classes. The class provides the methods `get_trian_loader` and `get_val_loader`.
    """
    def __init__(self, train_dataset, val_dataset, verbose: bool = False):
        self.dataset_name = "CIFAR-100"
        self.mean = (0.485, 0.456, 0.406)
        self.std = (0.229, 0.224, 0.225)
        self.NUM_OF_CLASSES = 100

        self.train_dataset = train_dataset

        self.val_dataset = val_dataset

        if verbose:
            self.summary()

    def get_train_loader(self, batch_size: int = 128) -> DataLoader:
        """Returns the training data loader."""
        return DataLoader(self.train_dataset, batch_size=batch_size, shuffle=True)

    def get_val_loader(self, batch_size: int = 128) -> DataLoader:
        """Returns the validation data loader."""
        return DataLoader(self.val_dataset, batch_size=batch_size, shuffle=False)

    def summary(self, show_image: bool = True):
        """Displays a short summary of the dataset and an example image."""
        print(f"\n📊 Summary of {self.dataset_name} Dataset")
        print(f"🔢 Number of classes: {self.NUM_OF_CLASSES}")
        print(f"📁 Number of training examples: {len(self.train_dataset)}")
        print(f"📁 Number of validation examples: {len(self.val_dataset)}")
        print(
            f"📦 Examples per class in train (approx): {len(self.train_dataset) // self.NUM_OF_CLASSES}"
        )
        print(
            f"📦 Examples per class in val (approx): {len(self.val_dataset) // self.NUM_OF_CLASSES}"
        )

        image, label = self.train_dataset[2]
        print(f"🖼️ Example image shape: {image.shape}")
        print(f"🏷️ Example label: {label} ({self.train_dataset.classes[label]})")

        if show_image:
            img_np = image.permute(1, 2, 0)
            # unnormalize
            img_np = img_np * torch.tensor(self.std) + torch.tensor(self.mean)
            img_np = img_np.clip(0, 1).numpy()

            plt.imshow(img_np)
            plt.title(f"Class: {label}")
            plt.axis("off")
            plt.show()

######################### DO NOT CHANGE THIS CELL ##########################
READONLY_DATA_PATH = "data/Cifar100"
OUTPUT_PATH = "file_output"
os.makedirs(OUTPUT_PATH, exist_ok=True)
def unpack_tar_gz(filename: str, path: str = OUTPUT_PATH) -> None:
    """ Unpacks a tar.gz archive """
    with tarfile.open(filename, "r:gz") as tar:
        tar.extractall(path=path)
        print(f"Extracted {filename} to {path}")


def load_dataset_from_pkl(filename: str) -> Dataset:
    """ Loads a dataset from a .pkl file """
    with open(filename, "rb") as f:
        dataset = pickle.load(f)
    return dataset

######################### DO NOT CHANGE THIS CELL ##########################

unpack_tar_gz(
    os.path.join(READONLY_DATA_PATH, "train_val.gz"), os.path.join(OUTPUT_PATH)
)


loaded_data = load_dataset_from_pkl(os.path.join(OUTPUT_PATH, "train_val.pkl"))

# Extract the training and validation sets
train_dataset = loaded_data["train"]
val_dataset = loaded_data["val"]
classes = loaded_data["classes"]
train_dataset.classes = classes
val_dataset.classes = classes

######################### DO NOT CHANGE THIS CELL #########################

ds = Dataset(train_dataset, val_dataset, verbose=True)

BATCH_SIZE = 128
train_loader = ds.get_train_loader(BATCH_SIZE)
val_loader = ds.get_val_loader(BATCH_SIZE)

######################### DO NOT CHANGE THIS CELL ##########################
def get_embeddings(model, dataloader: DataLoader) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Extracts the embeddings and labels from the model for the data from the dataloader.
    """
    model.eval()
    embeddings = []
    labels = []
    with torch.no_grad():
        for images, batch_labels in tqdm(dataloader):
            images, batch_labels = images.to(DEVICE), batch_labels.to(DEVICE)
            batch_embeddings = model.get_embedding(images)
            embeddings.append(batch_embeddings)
            labels.append(batch_labels)
    return torch.cat(embeddings), torch.cat(labels)


def save_embeddings(embeddings: torch.Tensor, labels: torch.Tensor, filename: str) -> None:
    """
    Saves the embeddings and labels to a file.
    """
    data = {"embeddings": embeddings, "labels": labels}
    with open(filename, "wb") as f:
        pickle.dump(data, f)


def load_embeddings(filename: str) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Loads the embeddings and labels from a file.
    """
    with open(filename, "rb") as f:
        data = pickle.load(f)
    return data["embeddings"], data["labels"]


embedding_path = os.path.join(OUTPUT_PATH, "train_embeddings.pkl")
val_embedding_path = os.path.join(OUTPUT_PATH, "val_embeddings.pkl")

# Compute the embeddings or load previously computed ones from a file
if os.path.exists(embedding_path):
    print(f"Loading cached training embeddings from: {embedding_path}")
    train_embeddings, train_labels = load_embeddings(embedding_path)
else:
    print("Generating new training embeddings...")
    train_embeddings, train_labels = get_embeddings(model, train_loader)
    save_embeddings(train_embeddings, train_labels, embedding_path)
    print(f"Saved training embeddings to: {embedding_path}")

if os.path.exists(val_embedding_path):
    print(f"Loading cached validation embeddings from: {val_embedding_path}")
    val_embeddings, val_labels = load_embeddings(val_embedding_path)
else:
    print("Generating new validation embeddings...")
    val_embeddings, val_labels = get_embeddings(model, val_loader)
    save_embeddings(val_embeddings, val_labels, val_embedding_path)
    print(f"Saved validation embeddings to: {val_embedding_path}")

train_embeddings = train_embeddings.to(DEVICE)
val_embeddings = val_embeddings.to(DEVICE)

train_labels = train_labels.to(DEVICE)
val_labels = val_labels.to(DEVICE)

######################### DO NOT CHANGE THIS CELL ##########################
def accuracy(
    prototypes: torch.Tensor,
    prototypes_labels: torch.Tensor,
    embeddings_val: torch.Tensor,
    labels_val: torch.Tensor
) -> float:
    """
    Computes the classification accuracy based on the nearest prototypes.
    """
    dist = torch.cdist(embeddings_val, prototypes)
    idx = dist.argmin(dim=1).to(prototypes.device)
    predictions = prototypes_labels[idx]
    predicted_class = predictions.clone().detach().to(labels_val.device)
    return (predicted_class == labels_val).to(torch.float32).mean().item()


class SolutionHolder():    
    """
    Class responsible for registering and evaluating prototype solutions
    based on the embedding and label data.
    """
    def __init__(
        self,
        train_embeddings: torch.Tensor,
        train_labels: torch.Tensor,
        val_embeddings: torch.Tensor,
        val_labels: torch.Tensor
    ):
        self.train_embeddings = train_embeddings
        self.train_labels = train_labels
        self.val_embeddings = val_embeddings
        self.val_labels = val_labels
        self.sol_list = []
    
    def assert_solution(self, prototypes: torch.Tensor, prototypes_labels: torch.Tensor):
        """
        Checks whether the solution meets the required conditions.
        """
        # Constant required conditions
        REQUIRED_NUM_PROTOTYPES = 150
        MAX_LABEL = 100
        DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        # Checking the types
        assert isinstance(prototypes, torch.Tensor), "prototypes must be of type torch.Tensor"
        assert isinstance(prototypes_labels, torch.Tensor), "prototypes_labels must be of type torch.Tensor"

        # Checking the data type
        assert prototypes.dtype == torch.float32, "prototypes.dtype must be torch.float32"
        assert prototypes_labels.dtype == torch.long, "prototypes_labels.dtype must be torch.long"

        # Checking the shapes
        assert prototypes.shape[0] == REQUIRED_NUM_PROTOTYPES, f"prototypes should have 150 rows, got {prototypes.shape[0]}"
        assert prototypes_labels.shape[0] == REQUIRED_NUM_PROTOTYPES, f"prototypes_labels should have 150 elements, got {prototypes_labels.shape[0]}"
        assert len(prototypes_labels.shape) == 1, "prototypes_labels must be a 1D vector"

        # Checking the label values
        assert torch.all((prototypes_labels >= 0) & (prototypes_labels < MAX_LABEL)), "All labels in prototypes_labels must be in the range [0, 100)"

    
    def register_class_solution(self, solution_class: Type, name: str) -> None:
        """
        Registers a solution based on a class containing the get_prototypes method.
        """
        solution = solution_class()
        prototypes, prototypes_labels = solution.get_prototypes(self.train_embeddings, self.train_labels)
        self.assert_solution(prototypes, prototypes_labels)
        self.register_solution(prototypes, prototypes_labels, name)

    def register_solution(self, prototypes: torch.Tensor, prototypes_labels: torch.Tensor, name: str) -> None:
        """
        Adds a solution to the list of solutions.
        """
        self.sol_list.append({
            "name": name,
            "prototypes": prototypes.to(DEVICE),
            "prototypes_labels": prototypes_labels.to(DEVICE),
        })

    def points_from_accuracy(self, acc: float) -> int:
        """
        Converts the accuracy into a number of points (0–100).
        """
        result = (acc - 0.4) / (0.68 - 0.4) * 100
        result = int(round(result))
        return min(100, max(0, result))

    def _print_one_solution(
        self,
        prototypes: torch.Tensor,
        prototypes_labels: torch.Tensor,
        name: str
    ) -> None:
        """
        Displays the accuracy and the number of points for one solution.
        """
        acc_val = accuracy(prototypes, prototypes_labels, self.val_embeddings, self.val_labels)
        points = self.points_from_accuracy(acc_val)
        print(f"{name} ({prototypes.shape[0]} prototypes) Validation Accuracy: {acc_val * 100:.2f}% | (Points {points})")

    def print_solutions(self) -> None:
        """
        Displays the results for all registered solutions.
        """
        for sol_dict in self.sol_list:
            name = sol_dict['name']
            prototypes = sol_dict['prototypes']
            prototypes_labels = sol_dict['prototypes_labels']
            
            self._print_one_solution(
                prototypes,
                prototypes_labels,
                name
            )

def get_prototypes(self, train_embeddings, train_labels) -> Tuple[torch.Tensor, torch.Tensor]:
    ...
    return prototypes, prototypes_labels

# Your imports ...

class YourSolution(): 
    def get_random_150_samples(self, unique_labels: np.ndarray) -> np.ndarray:
        """
        Randomly allocates 1–2 samples per class, 150 samples in total.
        """
        num_samples_per_class = np.ones_like(unique_labels)
        double_samples = np.random.choice(
            num_samples_per_class.shape[0],
            50,
            replace = False
        )
        num_samples_per_class[double_samples] = num_samples_per_class[double_samples] + 1

        return num_samples_per_class

    def get_random_samples_per_class(
        self,
        embeddings: torch.Tensor,
        labels: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Randomly selects embedding samples from each class according to the distribution.
        """
        embeddings = embeddings.numpy(force=True)
        labels = labels.numpy(force=True)

        unique_labels = np.unique(labels)
        num_samples_per_class = self.get_random_150_samples(unique_labels)

        random_samples = []
        for i, label in enumerate(unique_labels):
            class_indices = np.where(labels == label)[0]
            selected_indices = np.random.choice(class_indices, size=num_samples_per_class[i], replace=False)
            random_samples.extend(selected_indices)

        return torch.tensor(embeddings[random_samples], device=DEVICE), \
                torch.tensor(labels[random_samples], device=DEVICE)
        
    def get_prototypes(
        self,
        train_embeddings: torch.Tensor,
        train_labels: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Returns 150 random prototypes and their corresponding labels.

        [IMPORTANT] Remember to implement the `get_prototypes` method in your class
        as described in the `Constraints` section. This method can serve you as a template.
        """
        prototypes, prototypes_labels = self.get_random_samples_per_class(
            train_embeddings, train_labels
        )
        prototypes = prototypes.to(DEVICE)
        prototypes_labels = prototypes_labels.to(DEVICE)
        return prototypes, prototypes_labels

######################### DO NOT CHANGE THIS CELL ##########################
if not FINAL_EVALUATION_MODE:   
    solution_holder = SolutionHolder(
        train_embeddings=train_embeddings,
        train_labels=train_labels,
        val_embeddings=val_embeddings,
        val_labels=val_labels,
    )

    solution_holder.register_class_solution(YourSolution, "Your solution")
    solution_holder.print_solutions()

if not FINAL_EVALUATION_MODE:   
    solution_holder.register_class_solution(YourSolution, "Different solution")
    solution_holder.print_solutions()

######################### DO NOT CHANGE THIS CELL ##########################
if FINAL_EVALUATION_MODE: 
    solution = YourSolution()
    prototypes, prototypes_labels = solution.get_prototypes(
        train_embeddings, 
        train_labels
    )
     
    os.makedirs(OUTPUT_PATH, exist_ok=True)
    torch.save(prototypes, os.path.join(OUTPUT_PATH, "prototypes.pt"),)
    torch.save(prototypes_labels, os.path.join(OUTPUT_PATH, "prototypes_labels.pt"))