Convolutional Filters

Kernel (3x3 matrix)	Effect on the image
Identity The one in the centre makes the new pixel an exact copy of the original one. $$\begin{bmatrix} 0 & 0 & 0 \ 0 & 1 & 0 \ 0 & 0 & 0 \end{bmatrix}$$
Blurring Averages the value of a pixel with its neighbours, which smooths the image. $$\begin{bmatrix} \frac{1}{9} & \frac{1}{9} & \frac{1}{9} \ \frac{1}{9} & \frac{1}{9} & \frac{1}{9} \ \frac{1}{9} & \frac{1}{9} & \frac{1}{9} \end{bmatrix}$$
Sharpening Emphasises the differences between a pixel and its surroundings, bringing out details. $$\begin{bmatrix} 0 & -1 & 0 \ -1 & 5 & -1 \ 0 & -1 & 0 \end{bmatrix}$$
Vertical edge detection Strengthens places where there are dark pixels on the left and bright pixels on the right. $$\begin{bmatrix} -1 & 0 & 1 \ -2 & 0 & 2 \ -1 & 0 & 1 \end{bmatrix}$$

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

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

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

if not FINAL_EVALUATION_MODE:
    import matplotlib.pyplot as plt

import os
import torch
import numpy as np
import math
import shutil
from typing import Callable, Optional

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

######################### 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.

def download_data() -> None:
    """Downloads the dataset from Google Drive and saves it in the 'data' folder."""
    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 = "https://drive.google.com/file/d/1wzYqBQ56IMXpM7_gv05vzercvZSZy-Ps/view?usp=drive_link"
    output = f'data/train_val.npz'
    gdown.download(url, output, quiet=True, fuzzy=True)

def setup_data_from_npz(npz_path: str) -> tuple[list[dict], list[dict]]:
    """
    Loads triples of images (original, corrupted and downscaled (downsampled)
    corrupted) from an NPZ file and returns them as two lists of dictionaries.

    Args:
        npz_path (str): Path to the .npz file with the data.

    Returns:
        tuple[list[dict], list[dict]]:
            - train_ds: a list of dictionaries {'image_original': tensor,
                'image_corrupted': tensor, 'image_corrupted_downsampled': tensor}
            - val_ds:   a list of dictionaries in the same format, but for validation.
    """
    # Loading the data from the file
    data = np.load(npz_path, allow_pickle=True)
    train_list = data['train_ds'].tolist()
    val_list   = data['val_ds'].tolist()

    def to_tensor_triplet(sample: dict) -> dict:
        # We convert the loaded data (uint8) to float32
        # and scale it back to the range [0, 1] by dividing by 255.0
        return {
            "image_original": torch.from_numpy(sample["image_original"]).to(torch.float32) / 255.0,
            "image_corrupted": torch.from_numpy(sample["image_corrupted"]).to(torch.float32) / 255.0,
            "image_corrupted_downsampled": torch.from_numpy(sample["image_corrupted_downsampled"]).to(torch.float32) / 255.0,
        }

    train_ds = [to_tensor_triplet(s) for s in train_list]
    val_ds   = [to_tensor_triplet(s) for s in val_list]

    return train_ds, val_ds

######################### DO NOT CHANGE THIS CELL ###########################
# Cell containing helper functions for evaluating the solution.

def _pad_replicate(x: torch.Tensor, pad: int) -> torch.Tensor:
    """
    "Padding by replication" — adds a border made of the image's edge pixels.
    Useful for convolution operations (the kernel is in the centre, the edge does not "escape").

    Args:
        x (torch.Tensor): Image of shape (C, H, W).
        pad (int): Number of pixels to add on each side.

    Returns:
        torch.Tensor: Image of enlarged size (C, H+2*pad, W+2*pad).
    """
    if pad == 0:
        return x
    # Vertical padding: we duplicate the first and the last row.
    top    = x[:, :1, :].expand(-1, pad, -1)
    bottom = x[:, -1:, :].expand(-1, pad, -1)
    xvb = torch.cat([top, x, bottom], dim=1)          # (C, H+2p, W)
    # Horizontal padding: we duplicate the first and the last column.
    left  = xvb[:, :, :1].expand(-1, -1, pad)
    right = xvb[:, :, -1:].expand(-1, -1, pad)
    xpad = torch.cat([left, xvb, right], dim=2)       # (C, H+2p, W+2p)
    return xpad

def apply_kernel(image: torch.Tensor, kernel: torch.Tensor) -> torch.Tensor:
    """
    Applies the same filter (kernel) to each channel component of the image.

    Args:
        image (torch.Tensor): Image of shape (C, H, W) in the range [0,1].
        kernel (torch.Tensor): 2D filter of shape (K, K).

    Returns:
        torch.Tensor: New image (C, H, W) after applying the kernel.
    """
    assert image.dim() == 3, "image must have dimensions (C,H,W)"
    assert kernel.dim() == 2 and kernel.size(0) == kernel.size(1), "kernel must be a square (K,K)"
    assert kernel.size(0) <= 10, "the kernel size must not exceed 10x10"
    C, H, W = image.shape
    K = kernel.size(0)
    pad = K // 2

    # Padding by replication (without in-place modification — important for autograd).
    xpad = _pad_replicate(image, pad)                 # (C, H+2p, W+2p)

    # We create "windows" of neighbouring pixels for each location in the image: (C, H, W, K, K)
    sC, sH, sW = xpad.stride()
    windows = xpad.as_strided(
        size=(C, H, W, K, K),
        stride=(sC, sH, sW, sH, sW)
    )

    # We compute the correlation: element-wise multiplication and summation over the KxK neighbourhood
    out = (windows * kernel.view(1, 1, 1, K, K)).sum(dim=(-1, -2))   # (C,H,W)
    return out

def evaluate_solution(
    kernel: torch.Tensor,
    ds: list[dict],
    upsample_fun: Optional[Callable] = None
) -> float:
    """
    Computes the mean MSE (Mean Squared Error) on the given dataset.

    Args:
        kernel (torch.Tensor): 2D filter (e.g. 5x5) to be evaluated.
        ds (list[dict]): List of samples with the keys 'image_original',
            'image_corrupted' and 'image_corrupted_downsampled'.
        upsample_fun (Optional[Callable]): (Optional) Upsampling function
            (torch.Tensor -> torch.Tensor) to be evaluated.

    Returns:
        float: Mean MSE value for the whole set.
    """
    total_mse = 0.0
    kernel = kernel.to(DEVICE)

    # We do not need gradients during evaluation — this speeds things up and saves memory.
    with torch.no_grad():
        for sample in ds:
            original = sample['image_original'].to(DEVICE)

            # Choosing the appropriate "corrupted" version:
            # - if an upsampling function is given, we enlarge the downscaled version,
            # - otherwise we take the full-size version.
            if upsample_fun is not None:
                corrupted = sample['image_corrupted_downsampled'].to(DEVICE)
                corrupted = upsample_fun(corrupted)  # we expect (C, H, W)
            else:
                corrupted = sample['image_corrupted'].to(DEVICE)  # already (C, H, W)

            # Reconstructing the image by applying the kernel to the "corrupted" image.
            reconstructed = apply_kernel(corrupted, kernel)
            # MSE: the mean of the squared pixel-by-pixel differences.
            mse = torch.mean((reconstructed - original)**2)
            total_mse += mse.item()

    return total_mse / len(ds)

def round_half_up(number: float) -> int:
    """
    Rounds to the nearest integer (.5 always rounds up).
    """
    return int(math.floor(number + 0.5))

def compute_score(kernel_mse: float, upsampling_mse: float) -> tuple[int, int, int]:
    """
    Converts the two MSE results (for the kernel and for upsampling) into a single score of 0–100.
    Each MSE is mapped separately onto a point scale based on thresholds,
    and the results are then combined using weights.

    Args:
        kernel_mse (float): Mean MSE for part 1 (kernel).
        upsampling_mse (float): Mean MSE for part 2 (upsampling).

    Returns:
        tuple[int, int, int]: A tuple containing the final score (0–100),
            the score for part 1 (kernel) and the score for part 2 (upsampling).
    """

    # Quality thresholds: the lower the MSE, the better.
    # 100 pts are awarded for MSE <= lower_bound,
    # 0 pts for MSE >= upper_bound,
    # otherwise — between 100 and 0 (the score is scaled linearly).

    KERNEL_LOWER = 0.0077
    KERNEL_UPPER = 0.0285
    UPS_LOWER    = 0.0350
    UPS_UPPER    = 0.0542

    # Weights of the parts of the task.
    KERNEL_W = 0.7
    UPS_W    = 0.3

    def linear_score(mse: float, lower: float, upper: float) -> float:
        """
        Maps the MSE onto the interval [0,100] according to the thresholds.
        - mse <= lower  -> 100
        - mse >= upper  -> 0
        - in between    -> linearly between 100 and 0
        """
        if mse <= lower:
            return 100.0
        if mse >= upper:
            return 0.0
        # the closer to lower, the more points
        return 100.0 * (upper - mse) / (upper - lower)

    kernel_score = round_half_up(KERNEL_W * linear_score(kernel_mse, KERNEL_LOWER, KERNEL_UPPER))
    upsampling_score = round_half_up(UPS_W * linear_score(upsampling_mse, UPS_LOWER, UPS_UPPER))

    # We combine the scores and clip the total to [0,100] just in case.
    final_score = kernel_score + upsampling_score
    final_score = max(0, min(100, final_score))

    return final_score, kernel_score, upsampling_score

######################### DO NOT CHANGE THIS CELL ###########################
# Cell containing helper functions for visualisation.

def plot_results(
    original: torch.Tensor,
    corrupted: torch.Tensor,
    reconstructed: torch.Tensor,
    title: str = "Reconstruction results",
    labels: tuple = ("Original", "Corrupted", "Reconstructed"),
    center_smaller: bool = True,
    background: float = 1.0,
    show_shapes: bool = True,
    clip_for_display: bool = True
) -> None:
    """
    Displays three images side by side. If one of the images is smaller
    than the largest of the others and `center_smaller=True`, it is only
    **centred** on a canvas the size of the largest image (by means of
    padding with the background colour), so that the panel sizes are visually consistent.

    Args:
        original (torch.Tensor): Reference image in (C, H, W) format, values in [0, 1].
        corrupted (torch.Tensor): "Corrupted" image in (C, h, w) format; it may be smaller.
        reconstructed (torch.Tensor): Image after reconstruction in (C, H, W) format.
        title (str): Title of the whole figure.
        labels (tuple): Captions above the panels (from the left): ("Original", "Corrupted", "Reconstructed").
        center_smaller (bool): If True, smaller images are centred with padding (without scaling).
        background (float): Background (padding) colour in [0, 1]; 1.0 = white, 0.0 = black.
        show_shapes (bool): If True, the sizes (height×width) are appended to the panel titles.
        clip_for_display (bool): If True, clips only a copy of the data to [0,1]
            for imshow (eliminates warnings and "burnt-out" areas on screen).

    Returns:
        None
    """
    import numpy as np

    def to_hwc_np(t: torch.Tensor) -> np.ndarray:
        """
        Converts a (C, H, W) tensor into a NumPy array (H, W, C) in [0, 1],
        without modifying the values.
        """
        return t.detach().cpu().permute(1, 2, 0).numpy()

    def pad_center(img_hwc: np.ndarray, target_hw: tuple, bg: float) -> np.ndarray:
        """
        Pastes an image (H, W, C) into the centre of a canvas of size target_hw (H_t, W_t).

        Args:
            img_hwc (np.ndarray): Input image (H, W, C) in [0, 1].
            target_hw (tuple[int,int]): Target canvas size (H_t, W_t).
            bg (float): Background colour in [0, 1].

        Returns:
            np.ndarray: Image (H_t, W_t, C) with the input image centred.
        """
        Ht, Wt = target_hw
        h, w, c = img_hwc.shape
        # If the sizes already match, or we do not want to centre, we return the image unchanged
        if (h, w) == (Ht, Wt) or not center_smaller:
            return img_hwc

        # We create a background canvas and paste the image in the middle
        canvas = np.full((Ht, Wt, c), bg, dtype=img_hwc.dtype)
        top = (Ht - h) // 2
        left = (Wt - w) // 2
        canvas[top:top + h, left:left + w, :] = img_hwc
        return canvas

    # Conversion to HWC numpy
    img_o = to_hwc_np(original)
    img_c = to_hwc_np(corrupted)
    img_r = to_hwc_np(reconstructed)

    # Common frame = the maximum size among the inputs
    Hs = [img_o.shape[0], img_c.shape[0], img_r.shape[0]]
    Ws = [img_o.shape[1], img_c.shape[1], img_r.shape[1]]
    target_hw = (max(Hs), max(Ws))

    # Padding only if needed
    img_o_pad = pad_center(img_o, target_hw, background)
    img_c_pad = pad_center(img_c, target_hw, background)
    img_r_pad = pad_center(img_r, target_hw, background)

    # Optionally append the sizes to the titles
    titles = list(labels)
    if show_shapes:
        titles[0] = f"{titles[0]} ({img_o.shape[0]}×{img_o.shape[1]})"
        titles[1] = f"{titles[1]} ({img_c.shape[0]}×{img_c.shape[1]})"
        titles[2] = f"{titles[2]} ({img_r.shape[0]}×{img_r.shape[1]})"

    # Optional clipping for display only
    if clip_for_display:
        def clip01(a: np.ndarray) -> np.ndarray:
            return np.clip(a, 0.0, 1.0, out=a.copy())
        show_o = clip01(img_o_pad)
        show_c = clip01(img_c_pad)
        show_r = clip01(img_r_pad)
    else:
        show_o, show_c, show_r = img_o_pad, img_c_pad, img_r_pad

    # Drawing
    fig, axes = plt.subplots(1, 3, figsize=(15, 5))
    for ax, img, lab in zip(axes, (show_o, show_c, show_r), titles):
        ax.imshow(img, origin="upper")
        ax.set_title(lab)
        ax.axis("off")

    fig.suptitle(title)
    plt.tight_layout()
    plt.show()

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

if not FINAL_EVALUATION_MODE:
    download_data()

    # Path to the single NPZ file with the data (training + validation)
    NPZ_PATH = os.path.join('data', 'train_val.npz')

    # Load the sets from the NPZ file
    train_ds, val_ds = setup_data_from_npz(NPZ_PATH) # These variables will also be available during the final evaluation

    if train_ds is not None and val_ds is not None:
        print(f"Number of images in the training set: {len(train_ds)}")
        print(f"Number of images in the validation set: {len(val_ds)}")

        # Preview of example data
        sample_for_show = val_ds[0]
        plot_results(sample_for_show['image_original'], sample_for_show['image_corrupted'], sample_for_show['image_corrupted_downsampled'], title="Example data from the NPZ file", labels=("Original", "Corrupted", "Corrupted (downsampled)"))

def compute_kernel() -> torch.Tensor:
    """
    Creates a blurring kernel of size 9x9.
    """
    kernel = torch.ones(9, 9, device=DEVICE) / 81
    return kernel

def upsample(image: torch.Tensor) -> torch.Tensor:
    """
    Simple upsampling using the "nearest neighbour" method (nearest neighbor)
    to a size of 256x256.

    Idea:
    Each pixel of the smaller image is simply replicated several times
    vertically and horizontally. As a result, the image becomes larger, but no new
    details appear – we only see the "stretched" original pixels.

    Args:
        image (torch.Tensor): Image in (C, h, w) format, where:
            - C is the number of channels (e.g. 3 for RGB),
            - h, w are the height and width of the image (≤ 256).

    Returns:
        torch.Tensor: Image of shape (C, 256, 256) obtained by
        repeating the rows and columns of the input as many times as needed
        to cover the whole target size, and then cropping to the exact
        dimensions of 256x256.
    """
    assert image.dim() == 3, "Expecting a tensor of shape (C, h, w)."
    C, h, w = image.shape

    # Target dimensions.
    Ht, Wt = (256, 256)

    # How many times we must repeat the rows and columns to reach >= 256.
    rh = max(1, math.ceil(Ht / h))
    rw = max(1, math.ceil(Wt / w))

    # repeat_interleave replicates successive rows/columns the appropriate number of times.
    y = image.repeat_interleave(rh, dim=1).repeat_interleave(rw, dim=2)

    # It may happen that there are "too many" repetitions – in that case we crop to 256x256.
    return y[:, :Ht, :Wt]

def compute_kernel_upsampling(kernel_task1: torch.Tensor) -> torch.Tensor:
    """
    During evaluation, the kernel from the first part will be passed
    as an argument. You may use it or ignore it.
    """
    return kernel_task1

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

if not FINAL_EVALUATION_MODE:

    # ============== 1) First part of the task ==============
    # -- Training
    your_kernel = compute_kernel()
    # -- Evaluation of the solution
    your_kernel_mse = evaluate_solution(your_kernel, val_ds)

    # ============== 2) Second part of the task ==============
    # -- Training
    your_kernel_upsampling = compute_kernel_upsampling(your_kernel)
    # -- Evaluation of the solution
    your_upsampling_mse = evaluate_solution(your_kernel_upsampling, val_ds, upsample)

    score, kernel_score, upsampling_score = compute_score(your_kernel_mse, your_upsampling_mse)

    print(f"MSE on the validation set (kernel): {your_kernel_mse:.4f}. Number of points (kernel): {kernel_score}")
    print(f"MSE on the validation set (upsampling): {your_upsampling_mse:.4f}. Number of points (upsampling): {upsampling_score}")
    print(f"Estimated number of points for the task: {score}")

    # Visualisation of the effect of the kernel on a single validation image.
    sample_for_show = val_ds[0]
    reconstructed = apply_kernel(sample_for_show['image_corrupted'].to(DEVICE), your_kernel.to(DEVICE))
    plot_results(sample_for_show['image_original'], sample_for_show['image_corrupted'], reconstructed, title="Visualisation of the effect of your kernel")

    # Preview of the kernel itself (colour map).
    plt.figure(figsize=(5, 5))
    plt.imshow(your_kernel.cpu().detach().numpy(), cmap='viridis')
    plt.colorbar()
    plt.title("Your kernel")
    plt.show()

    # Visualisation of the effect of upsampling with the kernel on a single validation image.
    reconstructed = apply_kernel(upsample(sample_for_show['image_corrupted_downsampled'].to(DEVICE)), your_kernel_upsampling)
    plot_results(sample_for_show['image_original'], sample_for_show['image_corrupted_downsampled'], reconstructed, title="Visualisation of the effect of your upsampling method with the kernel", labels=("Original", "Corrupted (downsampled)", "Reconstructed"))