Decision Localisation

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

import os
import random
import numpy as np
import cv2
import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, models

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

seed = 42
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
np.random.seed(seed)
random.seed(seed)

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

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

def download_data():
    """Downloads the dataset from Google Drive and saves it in the 'data' folder."""
    import shutil
    import gdown

    # Create or reset the 'data' folder
    if not os.path.exists('data'):
        os.makedirs('data')
    else:
        shutil.rmtree('data')
        os.makedirs('data')

    # Download the file from Google Drive and save it in the 'data' folder
    url_model = "https://drive.google.com/file/d/1b6WhjG_GDihBNKLIXtFcaclJKkTWbZZk/view?usp=drive_link"
    gdown.download(url_model, f'data/model.pth', quiet=True, fuzzy=True)

    url_data = "https://drive.google.com/file/d/19jSIGHXp_BY_gkyGdYKKc7ls-QU7_Rcn/view?usp=share_link"
    gdown.download(url_data, f'data/val_data.npz', quiet=True, fuzzy=True)


if not FINAL_EVALUATION_MODE:
    download_data()

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

# Definition of the model architecture.
class ResNet18Classifier(nn.Module):
    def __init__(self, pretrained=False, num_classes=1):
        super().__init__()

        if pretrained:
            weights = models.ResNet18_Weights.DEFAULT
        else:
            weights = None

        resnet = models.resnet18(weights=weights)

        self.encoder = nn.Sequential(
            resnet.conv1,
            resnet.bn1,
            resnet.relu,
            resnet.maxpool,
            resnet.layer1,
            resnet.layer2,
            resnet.layer3,
            resnet.layer4,
        )

        self.pool = nn.AdaptiveAvgPool2d((1, 1))
        self.head = nn.Linear(512, num_classes)

    def forward(self, x):
        x = self.encoder(x)
        x = self.pool(x)
        x = torch.flatten(x, 1)
        x = self.head(x)
        return x

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

# Model initialisation
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model_A = ResNet18Classifier(num_classes=1).to(device)

# Loading the trained weights
model_A.load_state_dict(torch.load("data/model.pth", map_location=device))
model_A.eval()

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

class CocoSegmentationDataset(Dataset):
    def __init__(self, npz_path, mode="val"):
        """
        npz_path: Path to the .npz file (e.g. 'data/val_data.npz')
        mode: Key inside the npz
        """
        if not os.path.exists(npz_path):
            raise FileNotFoundError(f"File not found: {npz_path}")

        # Loading the data from the NPZ
        data = np.load(npz_path, allow_pickle=True)

        # Key mapping
        key = f"{mode}_ds"
        if key not in data:
            raise KeyError(f"The file {npz_path} does not contain the key {key}")

        self.samples = data[key].tolist()

        # Definition of the ImageNet transformation
        self.normalize = transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        )

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

    def __getitem__(self, idx):
        sample = self.samples[idx]

        # Image conversion (uint8 [H,W,C] -> float32 [C,H,W])
        # 'image' is your JPG from the original dataset
        img_np = sample["image"]
        img_tensor = torch.from_numpy(img_np).to(torch.float32) / 255.0
        img_tensor = img_tensor.permute(2, 0, 1) # Change to [C, H, W]
        img_tensor = self.normalize(img_tensor)

        # 4. Mask conversion (uint8 [H,W] -> float32 [1,H,W])
        # 'mask' is your PNG mask
        mask_np = sample["mask"]
        mask_tensor = torch.tensor(mask_np / 255.0, dtype=torch.float32)

        # Adding a channel dimension if the mask is 2D
        if mask_tensor.ndim == 2:
            mask_tensor = mask_tensor.unsqueeze(0)

        return img_tensor, mask_tensor

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

# Initialisation
val_dataset = CocoSegmentationDataset("data/val_data.npz")
val_loader = DataLoader(val_dataset, batch_size=1, shuffle=True)

# Your solution - make changes only here.
# Do not change the function arguments or the number of returned values.
# The function has been filled in with an example solution that is far
# from optimal - in particular, it does not use the given models at all.
# It is meant to serve as a starting point for building a solution.

def your_solution(img, model):
    _, H, W = img.shape

    # Randomly generated mask.
    mask = torch.randint(0, 2, (H, W), dtype=torch.float32)

    return mask

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

# Displaying an example image from the dataset.
if not FINAL_EVALUATION_MODE:
    import matplotlib.pyplot as plt

    def plot_example(dataset, model, your_solution):
        img, mask = dataset[0]
        pred = your_solution(img, model)
        if torch.is_tensor(pred):
            pred = pred.detach().cpu().numpy()
        mask_to_show = mask.squeeze().cpu().numpy()

        # Channel permutation
        img_display = img.permute(1, 2, 0).numpy()
        # Reversing the normalisation for the image
        img_display = img_display * np.array([0.229, 0.224, 0.225]) + np.array([0.485, 0.456, 0.406])
        # Clipping the values to the range [0, 1] for correct visualisation
        img_display = np.clip(img_display, 0, 1)

        plt.figure(figsize=(12, 4))
        plt.subplot(1, 3, 1)
        plt.title("Original")
        plt.imshow(img_display)

        plt.subplot(1, 3, 2)
        plt.title("Mask (GT)")
        plt.imshow(mask_to_show, cmap='gray')

        plt.subplot(1, 3, 3)
        plt.title("Prediction")
        plt.imshow(pred, cmap='jet', alpha=0.5)
        plt.imshow(img_display, alpha=0.5)
        plt.show()

    # Running the function that visualises an example image from the dataset together with the GT mask and the prediction.
    plot_example(val_dataset, model_A, your_solution)

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

# Computes the IoU metric between two masks.

def compute_iou(pred, target):
    # We make sure the tensors are binary (0 or 1)
    pred = (pred > 0.5).float()
    target = (target > 0.5).float()

    intersection = (pred * target).sum()
    union = (pred + target).clamp(0, 1).sum()
    return (intersection + 1e-8) / (union + 1e-8)

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

# Computing the score for participants (public, uses val_loader)
def compute_score(loader, model):
    model.eval()
    ious = []
    device = next(model.parameters()).device

    with torch.no_grad():
        for img, mask in loader:
            img, mask = img.to(device), mask.to(device)

            # We pass a single image (C, H, W) to the participant's solution
            # We assume that your_solution returns a mask of shape (H, W) or (1, H, W)
            pred = your_solution(img[0], model)
            mask = mask.to(pred.device)

            # Matching the shapes for IoU
            ious.append(compute_iou(pred.squeeze(), mask.squeeze()))

    mIoU = sum(ious) / len(ious)

    # Scaling the score: 0.191 -> 0 pts, 0.245 -> 100 pts
    raw_score = 100 * (mIoU - 0.191) / (0.245 - 0.191)
    score = int(torch.clamp(raw_score, min=0, max=100).round().item())

    return mIoU, score

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

# Evaluates your solution on the validation set by computing the mean IoU.
# It should obtain a similar value on the secret test set.

if not FINAL_EVALUATION_MODE:
    mIoU, score = compute_score(val_loader, model_A)
    print(f"mIoU: {mIoU:.4f} | Estimated number of points: {score}/100")