Coin Counting Machine

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

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

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

import os
import torch
import pickle
import gdown
import numpy as np
import torch.nn as nn
import matplotlib.pyplot as plt
import seaborn as sns
import torchvision.transforms.v2 as T
from collections.abc import Callable
from matplotlib import patches
from matplotlib.collections import PatchCollection
from torchvision.models import resnet18
from torchvision.ops import box_iou
from torch.utils.data import Dataset
from tqdm import tqdm
from torchmetrics.detection.mean_ap import MeanAveragePrecision 

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

assert torch.cuda.is_available(), "CUDA is not available!"

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

seed = 12345

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

######################### DO NOT CHANGE THIS CELL ###########################
# Cell containing helper functions for preparing the data.

class CoinsDataset(Dataset):
    """
    Coin dataset loaded from a pickle file.
    
    Args:
        pickle_file (str): Path to the pickle file containing the data.
        transform (callable, optional): Transformations applied to the images and labels.
    """
    def __init__(
            self, 
            pickle_file: str, 
            transform: Callable | None = None
        ):
        self.transform = transform

        with open(pickle_file, 'rb') as f:
            self.data = pickle.load(f)

    def __len__(self) -> int:
        """Returns the number of samples in the dataset."""
        return len(self.data)
    
    def __getitem__(
            self, 
            idx: int
        ) -> dict:
        """
        Retrieves a data sample by its index.
        
        Args:
            idx (int): Index of the sample.
        
        Returns:
            dict: A dictionary containing the image and its corresponding target objects (boxes, labels).
        """
        sample = self.data[idx]
        image = sample['image']
        target = {
            'boxes': sample['boxes'],
            'labels': sample['labels']
        }

        if self.transform:
            image, target = self.transform(image, target)

        return {
            'image': image,
            **target
        }


def setup_data(
        train_transform: Callable | None = None, 
        val_transform: Callable | None = None, 
        root: str = 'data/'
    ) -> tuple:
    """
    Prepares the training and validation datasets, downloading them if necessary.
    
    Args:
        train_transform (callable, optional): Augmentations for the training set.
        val_transform (callable, optional): Augmentations for the validation set.
        root (str, optional): Base directory for the data files.
    
    Returns:
        tuple: The datasets (train_ds, val_ds).
    """
    if train_transform is None:
        train_transform = T.Compose([T.ToImage(), T.ToDtype(torch.float32, scale=True)])
    if val_transform is None:
        val_transform = T.Compose([T.ToImage(), T.ToDtype(torch.float32, scale=True)])

    train_file = root+'train.pkl'
    val_file = root+'val.pkl'

    if not os.path.exists(root):
        os.makedirs(root)

    if not os.path.exists(train_file):
        url = "https://drive.google.com/uc?id=1KC8FBlCuwh9ITUt0CcPeRJCpqy5j4WBp"
        gdown.download(url, train_file, quiet=True)

    if not os.path.exists(val_file):
        url = "https://drive.google.com/uc?id=1Oza4UjnmAUeae2cA8YDwMWxVHOb7SKdP"
        gdown.download(url, val_file, quiet=True)

    train_ds = CoinsDataset(root+'train.pkl', transform=train_transform)
    val_ds = CoinsDataset(root+'val.pkl', transform=val_transform)

    return train_ds, val_ds

######################### DO NOT CHANGE THIS CELL ###########################
# Cell containing helper functions for computing the metrics that assess the quality of the model.

def predict_all_bounding_boxes(
        model: nn.Module, 
        ds: Dataset
    ) -> list:
    """
    Function that predicts all bounding boxes for a dataset using the model.

    Args:
        model: Object detection model.
        ds: Dataset.

    Returns:
        A list containing all predicted bounding boxes for each sample in the dataset.
    """
    all_pred_bboxes = []

    for sample in ds:
        img = sample["image"].to(DEVICE)
        pred_bboxes = model(img)
        all_pred_bboxes.append(pred_bboxes)

    return all_pred_bboxes


def calculate_map(
        predictions: list, 
        ds: Dataset, 
        return_all: bool = False
    ) -> dict | float:
    """
    Function that computes the mean average precision (mAP) of the predicted bounding boxes for the whole dataset.

    Args:
        predictions: A list containing the predicted bounding boxes for each sample in the dataset.
        ds: Dataset containing the ground-truth bounding boxes.
        return_all (bool, optional): Whether to return all metrics or only mAP:0.5:0.95:0.05.

    Returns:
        The mAP value, or all metrics in the form of a dictionary.
    """
    meta = []

    for img_meta in predictions:
        entry = {
            "boxes": [],
            "labels": [],
            "scores": []
        }

        for box in img_meta:
            entry["boxes"].append(box[:4])
            entry["labels"].append(box[4])
            entry["scores"].append(box[5])

        meta.append(entry)

    for i in range(len(meta)):
        meta[i]['boxes'] = torch.tensor(meta[i]['boxes'])
        meta[i]['labels'] = torch.tensor(meta[i]['labels']).view(-1)
        meta[i]['scores'] = torch.tensor(meta[i]['scores'])

    mAP = MeanAveragePrecision()

    GT = [{
        "boxes": sample["boxes"],
        "labels": sample["labels"]
    } for sample in ds]

    output = mAP(meta, GT)

    if return_all:
        return output

    return mAP(meta, GT)['map'].item()

def compute_confusion_matrix(
        predictions: list, 
        ds: Dataset, 
        iou_threshold: float = 0.5
    ) -> np.ndarray:
    """
    Function that computes the confusion matrix of the predicted bounding boxes for the whole dataset.

    Args:
        predictions: A list containing the predicted bounding boxes for each sample in the dataset.
        ds: Dataset containing the ground-truth bounding boxes.
        iou_threshold (float, optional): IoU threshold for assigning a prediction to an object.

    Returns:
        The confusion matrix.
    """
    num_classes = 10  # 9 coin classes + 1 background
    conf_matrix = np.zeros((num_classes, num_classes), dtype=int)

    for pred_boxes, item in zip(predictions, ds):
        # Ground truth
        gt_boxes = item['boxes']
        gt_labels = item['labels']

        # predictions
        pred_boxes_tensor = torch.tensor([p[:4] for p in pred_boxes]) if pred_boxes else torch.empty((0, 4))
        pred_labels = torch.tensor([p[4] for p in pred_boxes]) if pred_boxes else torch.empty((0,), dtype=torch.long)

        # IoU values for all pairs of predictions and ground truth
        iou_matrix = box_iou(pred_boxes_tensor, gt_boxes) if pred_boxes else torch.empty((0, gt_boxes.shape[0]))

        # assign the predictions to the ground truth based on IoU
        matched_gt = set()
        for pred_idx, ious in enumerate(iou_matrix):
            max_iou, gt_idx = torch.max(ious, dim=0)
            if max_iou >= iou_threshold and gt_idx.item() not in matched_gt:  
                conf_matrix[gt_labels[gt_idx].item(), pred_labels[pred_idx].item()] += 1
                matched_gt.add(gt_idx.item())
            else:
                conf_matrix[-1, pred_labels[pred_idx].item()] += 1  # False positive

        # Add all unassigned ground-truth objects as False negatives
        for gt_idx in range(len(gt_boxes)):
            if gt_idx not in matched_gt:
                conf_matrix[gt_labels[gt_idx].item(), -1] += 1

    return conf_matrix

######################### DO NOT CHANGE THIS CELL ###########################
# Cell containing helper functions for visualising the results

colors = ["red", "green", "blue", "yellow", "black", "purple", "orange", "brown", "pink"]
label_names = ['1 grosz', '2 grosze', '5 groszy', '10 groszy', '20 groszy', '50 groszy', '1 złotych', '2 złote', '5 złotych']  # Polish coin names (1 złoty = 100 groszy); "grosze"/"groszy" and "złote"/"złotych" are plural forms

def show_sample(
        sample: dict
    ):
    """
    Function that displays an image with the bounding boxes of the objects.

    Args:
        sample: A dictionary containing the image and the labels.
    """
    image = sample['image']
    meta = sample

    plt.figure(figsize=(9, 5))
    plt.imshow(image.permute((1, 2, 0)))
    plt.xticks([])
    plt.yticks([])

    if meta is not None:
        patches_list = []
        legend_labels = []

        for bbox, label in zip(meta['boxes'], meta['labels']):
            points = np.array(bbox)
            points = points.astype(int)

            # Draw the rectangle that encloses the object
            rect = patches.Rectangle(
                (points[0], points[1]), 
                points[2] - points[0], 
                points[3] - points[1], 
                linewidth=2, 
                edgecolor=colors[label.item()],
                facecolor='none'
            )
            patches_list.append(rect)

            # legend with unique labels
            if label_names[label.item()] not in legend_labels:
                legend_labels.append(label_names[label.item()])

        patch_collection = PatchCollection(patches_list, match_original=True)
        plt.gca().add_collection(patch_collection)

        # Add a legend with unique labels
        handles = [patches.Patch(color=colors[i], label=label) for i, label in enumerate(label_names) if label in legend_labels]
        plt.legend(handles=handles, loc="upper right")

    plt.show()


def plot_detection_results_grid(
        predictions: list, 
        ds: Dataset, 
        size: tuple = (2, 3)
    ):
    """
    Function that displays the detection results on selected validation photos.

    Args:
        predictions: A list containing the predicted bounding boxes for each sample in the dataset.
        ds: Dataset.
        size (tuple, optional): Size of the image grid.
    """

    fig, axes = plt.subplots(size[0], size[1], figsize=(size[1]*6, size[0]*5))

    for i, (sample, pred_meta) in enumerate(zip(ds, predictions)):
        if i >= size[0] * size[1]:
            break

        img = sample["image"]

        ax = axes[i // size[1], i % size[1]]
        ax.imshow(img.permute(1, 2, 0))
        ax.axis('off')

        for x1, y1, x2, y2, label, _ in pred_meta:
            rect = patches.Rectangle(
                (x1, y1),
                x2 - x1,
                y2 - y1,
                linewidth=2,
                edgecolor=colors[label],
                facecolor='none',
            )
            ax.add_patch(rect)
    
    handles = [patches.Patch(color=colors[i], label=label) for i, label in enumerate(label_names)]
    plt.legend(handles=handles, loc='upper right')
    plt.suptitle("Coin detection on six selected validation photos")
    plt.tight_layout()
    plt.show()


def plot_confusion_matrix(
        predictions: list, 
        ds: Dataset, 
        iou_threshold: float = 0.5
    ):
    """
    Function that displays the confusion matrix for coin detection.

    Args:
        predictions: The coin detection results.
        ds: Dataset.
        iou_threshold float: IoU threshold for assigning a prediction to an object
    """
    conf_matrix = compute_confusion_matrix(predictions, ds, iou_threshold=iou_threshold)

    labels = ["1 grosz", "2 grosze", "5 groszy", "10 groszy", "20 groszy", "50 groszy", "1 złotych", "2 złote", "5 złotych", "none (background)"]
    plt.figure(figsize=(10, 8))
    sns.heatmap(conf_matrix, annot=True, fmt="d", cmap="Blues", xticklabels=labels, yticklabels=labels, cbar=False)

    plt.xlabel("The model predicted", labelpad=15)
    plt.ylabel("whereas it should have predicted", labelpad=15)
    plt.xticks(rotation=45)

    plt.title("Confusion matrix for coin detection using IOU={}".format(iou_threshold))
    plt.show()

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

train_ds, val_ds = setup_data(root='./data/')

if not FINAL_EVALUATION_MODE:
    print("Number of photos in the training set:", len(train_ds), ", number of photos in the validation set:", len(val_ds))

    sample = train_ds[0]

    print("Each sample contains:", list(sample.keys()))
    print("Each photo has the dimensions:", list(sample["image"].shape))
    print("This sample contains", sample['labels'].shape[0], "objects")
    print("Each rectangle is described by", sample['boxes'].shape[1], "values (x1, y1, x2, y2)")

    show_sample(sample) # display an example photo with the objects marked

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

class BasicCNNClassifier(nn.Module):
    def __init__(self):
        super(BasicCNNClassifier, self).__init__()
        self.net = resnet18(weights=None, num_classes=10)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Function that classifies the image x into one of 10 classes: 9 coins + background.

        Args:
            x (torch.Tensor): The image to classify.

        Returns:
            torch.Tensor: The model's predictions.
        """
        return self.net(x)

######################### DO NOT CHANGE THIS CELL ###########################
# code for converting a photo containing many coins into training data for the classifier
# in 90% of cases, we randomly pick an object from the image and crop its surroundings
# in 10% of cases, we choose a completely random fragment of the image

class ClassificationDataPreprocessor:
    def __init__(
            self, 
            num_of_crops_per_image: int = 4, 
            box_size: int = 64
        ):
        self.num_of_crops_per_image = num_of_crops_per_image
        self.box_size = box_size

    def get_label_of_crop(
            self, 
            crop_box: tuple, 
            boxes: torch.Tensor, 
            labels: torch.Tensor,
            iou_threshold: float = 0.5
        ) -> int:
        """
        Function that returns the label for an image crop, depending on whether the crop overlaps a coin with an IoU value > 0.5.

        Args:
            crop_box (tuple): Coordinates of the top-left and bottom-right corners of the image crop (x1, y1, x2, y2).
            boxes (torch.Tensor): Tensor containing the bounding boxes of the coins.
            labels (torch.Tensor): Tensor containing the labels of the coins.
            iou_threshold (float): IoU threshold for assigning a prediction to an object.
        """
        for box, label in zip(boxes, labels):
            if box_iou(torch.tensor(crop_box).view(1, 4), box.view(1, 4)) > iou_threshold:
                return label
        return 9 # values 0-8 are the coin labels, so 9 will be the background

    def __call__(
            self, 
            batch: list
        ) -> dict:
        """
        Function that converts images containing many coins into training data for the classifier, consisting of image crops with their labels.

        Args:
            batch (list): A list of dictionaries containing the images, the bounding boxes of the coins and their labels.

        Returns:
            dict: A dictionary containing the image crops and their labels.
        """
        crops, labels = [], []

        for sample in batch:
            img = sample['image']

            for _ in range(self.num_of_crops_per_image):
                if torch.rand(1) < 0.1: # in one draw out of 10, we sample the background (this is how we try to balance the dataset)
                    central_x = torch.randint(0, img.shape[2], (1,)).item()
                    central_y = torch.randint(0, img.shape[1], (1,)).item()
                else: # in the remaining 90%, we randomly pick an object from the image and crop its surroundings
                    idx_gt = torch.randint(0, sample['labels'].shape[0], (1,)).item()
                    central_x = (sample['boxes'][idx_gt, 0] + sample['boxes'][idx_gt, 2]) // 2 + torch.randint(-10, 10, (1,)).item()
                    central_y = (sample['boxes'][idx_gt, 1] + sample['boxes'][idx_gt, 3]) // 2 + torch.randint(-10, 10, (1,)).item()

                x1 = np.clip(central_x - self.box_size // 2, 0, img.shape[2] - self.box_size)
                y1 = np.clip(central_y - self.box_size // 2, 0, img.shape[1] - self.box_size)
                x2, y2 = x1 + self.box_size, y1 + self.box_size

                crop = img[:, y1:y2, x1:x2]
                label = self.get_label_of_crop((x1, y1, x2, y2), sample['boxes'], sample['labels'])

                crops.append(crop)
                labels.append(label)

        return {
            "crops": torch.stack(crops, dim=0),
            "labels": torch.tensor(labels)
        }

######################### DO NOT CHANGE THIS CELL ###########################
# code for validation using the sliding-window method

class SlidingWindowDetector(nn.Module):
    """
    Class implementing an object detector based on the sliding-window method. It is an example solution to the coin detection problem.

    Args:
        classifier (nn.Module): The object classifier model. For each position of the sliding window,
                    the classifier returns a prediction for the given image crop.
        crop_size: Size of the sliding window
        stride: the distance by which the window is moved in each iteration.
    """
    def __init__(
            self, 
            classifier: nn.Module, 
            crop_size: int = 64, 
            stride: int = 32
        ):
        super(SlidingWindowDetector, self).__init__()
        self.classifier = classifier
        self.crop_size = crop_size
        self.stride = stride

    def forward(
            self, 
            image: torch.Tensor
        ) -> list:
        """
        Function that predicts the objects in an image using the sliding-window method.

        Args:
            image: The image to process.

        Returns:
            a list of the objects found, each in the form of a tuple (x1, y1, x2, y2, label, confidence).
        """
        # move the image to the appropriate device (the same one as the model)
        device = next(self.parameters()).device
        image = image.to(device)

        found_objects = [] # list of objects, where each object is a tuple (x1, y1, x2, y2, label, confidence)

        # slide the window over the image
        for y in range(0, image.shape[1] - self.crop_size, self.stride):
            for x in range(0, image.shape[2] - self.crop_size, self.stride):
                crop = image[:, y:y+self.crop_size, x:x+self.crop_size]
                pred = self.classifier(crop.unsqueeze(0))[0]

                if pred.argmax() != 9: # we skip the background
                    label = pred.argmax().item()
                    confidence = torch.softmax(pred, dim=0)[label].item()
                    found_objects.append((x, y, x+self.crop_size, y+self.crop_size, label, confidence))

        return found_objects

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

def train_basic_detector(
        train_ds: Dataset
    ) -> SlidingWindowDetector:
    """
    Function that trains a coin classifier by creating image crops and classifying them.
    Once the classifier has been trained, we create an object detector based on the sliding-window method (SlidingWindowDetector).

    Args:
        train_ds (Dataset): The training dataset.

    Returns:
        SlidingWindowDetector: A trained object detector that uses the sliding-window method.
    """

    # what should the size of the sliding window be? the answer to this question can be found in the data...
    # let us compute the average size of all objects in the training set
    total = 0
    boxes = 0

    for sample in train_ds:
        total += sum(sample["boxes"][:, 2] - sample["boxes"][:, 0]) # x2 - x1 (we assume that the objects are square, so we ignore y)
        boxes += len(sample["boxes"])

    box_size = (total / boxes).item()

    print("Average object size in the training set:", box_size) # the result should be about 62

    # this value is close to 64, which is a power of 2; this will make the network architecture easier for us
    box_size = 64
    preprocess = ClassificationDataPreprocessor(num_of_crops_per_image=128, box_size=box_size)

    # let us prepare the model, the optimiser and the loss function (thanks to the sliding-window technique, the problem of training our model has become a classification problem)
    model = BasicCNNClassifier()
    model.to(DEVICE)
    optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=0.001)
    criterion = nn.CrossEntropyLoss()

    # let us prepare the dataloaders
    train_dl = torch.utils.data.DataLoader(train_ds, batch_size=2, shuffle=True, collate_fn=preprocess)

    # let us train the model for the chosen number of epochs
    epochs = 30

    pbar = tqdm(range(epochs), desc="Training", total=epochs)
    for _ in pbar:
        epoch_losses = []

        for batch in train_dl:
            X = batch["crops"].to(DEVICE)
            y = batch["labels"].to(DEVICE)

            optimizer.zero_grad()
            output = model(X)
            loss = criterion(output, y)
            loss.backward()
            optimizer.step()

            epoch_losses.append(loss.detach().cpu().item())

        avg_loss = np.mean(epoch_losses)
        pbar.set_postfix({'train loss': avg_loss})

    model.eval()

    # we construct an object detector based on the sliding-window method, which uses the trained classifier to make predictions for the crops
    return SlidingWindowDetector(model)

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

if not FINAL_EVALUATION_MODE:
    train_ds, val_ds = setup_data(root='data/')

    model = train_basic_detector(train_ds)

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

if not FINAL_EVALUATION_MODE:
    out = predict_all_bounding_boxes(model, val_ds)

    map_val = calculate_map(out, val_ds, return_all=True)

    print(f"mAP using the IoU=0.5 threshold on the validation set: {map_val['map_50'].item():.2f}")
    print(f"mAP over many IoU values on the validation set:  {map_val['map']:.2f}, this is the metric that is scored in the contest. Your task is to maximise it.")

    plot_detection_results_grid(out, val_ds)

    plot_confusion_matrix(out, val_ds, iou_threshold=0.5)
    plot_confusion_matrix(out, val_ds, iou_threshold=0.8)

# here you can implement your detector, which will return a list of the objects found in the format (x1, y1, x2, y2, label, confidence)

class YourDetector(nn.Module):
    def __init__(self):
        # here you can initialise your model
        super(YourDetector, self).__init__()
        pass

    def forward(
            self, 
            img: torch.Tensor
        ) -> list:
        """
        Your function for detecting coins in an image.

        Args:
            img: The image to process.

        Returns:
            a list of the objects found, each in the form of a tuple (x1, y1, x2, y2, label, confidence).
        """
        # implement the logic of your detector here
        return []  # an example model that returns an empty list

# definitions of the augmentations for the training and validation sets. By default, None means no augmentation
train_transform = None
val_transform = None

train_ds, val_ds = setup_data(train_transform, val_transform, root='data/')

# train your model here
# ... 

your_model = YourDetector()

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

if not FINAL_EVALUATION_MODE:
    your_out = predict_all_bounding_boxes(your_model, val_ds)

    your_map_val = calculate_map(your_out, val_ds, return_all=False) # we return only the averaged mAP (the main scoring criterion)

    score = (np.clip(your_map_val, 0.2, 0.85) - 0.2) / 0.65 * 100
    score = int(round(score))

    print(f"mAP on the validation set: {your_map_val:.2f}")
    print(f"Estimated number of points for the task: {score}")

    plot_detection_results_grid(your_out, val_ds)

    plot_confusion_matrix(your_out, val_ds, iou_threshold=0.5)
    plot_confusion_matrix(your_out, val_ds, iou_threshold=0.8)

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

if FINAL_EVALUATION_MODE:
    import cloudpickle

    # If the model has parameters, set it to evaluation mode and move it to the CPU
    if list(your_model.parameters()):
        your_model.eval()
        your_model.cpu()

    OUTPUT_PATH = "file_output"
    FUNCTION_FILENAME = "your_model.pkl"
    FUNCTION_OUTPUT_PATH = os.path.join(OUTPUT_PATH, FUNCTION_FILENAME)

    if not os.path.exists(OUTPUT_PATH):
        os.makedirs(OUTPUT_PATH)

    with open(FUNCTION_OUTPUT_PATH, "wb") as f:
        cloudpickle.dump(your_model, f)