Inpainting with Implicit Neural Representations

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

# When your solution is checked, the value of the FINAL_EVALUATION_MODE flag will be changed to True
FINAL_EVALUATION_MODE = False

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

import json
import os
from pathlib import Path

import cv2
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
from skimage.metrics import peak_signal_noise_ratio
from torch.utils.data import DataLoader, Dataset
from tqdm import tqdm
import tempfile
import tarfile

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

# Setting the seed of the pseudo-random number generator to ensure deterministic results.

seed = 42
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

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

class ImageDataset(Dataset):
    """
    The ImageDataset class represents a dataset.
    It handles images and their corresponding masks, and also generates coordinates
    for every pixel of an image together with a time value (t_value).
    """

    # Number of frames in the video
    VIDEO_LENGTH = 79

    def __init__(
        self,
        img_dir: Path,
        mask_dir: Path,
        frame_names: list[str],
        mode: str,
        json_path: str | None = None,
    ):
        """
        Initialises an object of the ImageDataset class.

        Parameters:
        -----------
        img_dir : Path
            Path to the directory with the images.
        mask_dir : str
            Path to the directory with the masks.
        frame_names : list[str]
            List of image file names (without paths).
        mode : str
            Operating mode of the dataset ("train" or "val").
        json_path : str | None
            Path to the JSON file with the rectangle coordinates (default None).
        """
        self.img_dir = img_dir
        self.mask_dir = mask_dir
        self.frame_names = frame_names
        self.mode = mode

        if mode not in ["train", "val"]:
            raise ValueError("Invalid mode. Use 'train' or 'val'.")

        if mode == "val" and json_path is None:
            raise ValueError(
                "In 'val' mode you must provide the path to the JSON file with the rectangle coordinates."
            )

        if self.mode == "val":
            # Loading the coordinates of the missing rectangle from the JSON file
            with open(json_path, "r") as f:
                self.rect_coords = json.load(f)

    def __len__(self) -> int:
        """Returns the number of frames in the dataset.

        Returns:
            Number of frames: int
        """
        return len(self.frame_names)

    def __getitem__(self, idx: int) -> tuple:
        """
        Returns the data for the given index.

        Parameters
        ----------
        idx : int
            Frame index.

        Returns
        -------
        Dict[str, Any]
            A dictionary containing, depending on the mode, (x, y, t),
            images, masks and other data.
        """
        # Getting the frame name
        frame_name = self.frame_names[idx]
        img_name = frame_name
        mask_name = frame_name.replace(".jpg", ".png")

        # Loading the image and the mask
        img = cv2.imread(os.path.join(self.img_dir, img_name))
        mask = cv2.imread(os.path.join(self.mask_dir, mask_name), cv2.IMREAD_GRAYSCALE)

        if self.mode == "train":
            # Loading the rectangle coordinates and normalisation
            h, w = img.shape[:2]
            img = img.astype("float32") / 255.0
            mask = mask.astype("float32") / 255.0

            # Generating the (x, y) coordinate grid
            x = torch.linspace(-1, 1, w)
            y = torch.linspace(-1, 1, h)
            grid_y, grid_x = torch.meshgrid(y, x, indexing="ij")
            coords = torch.stack([grid_x, grid_y], dim=-1).reshape(h * w, 2)

            # Computing the time value (t_value)
            t_value = int(os.path.splitext(frame_name)[0])
            t_value = (t_value * 2) / (self.VIDEO_LENGTH - 1)
            t = torch.full((coords.shape[0], 1), t_value, dtype=torch.float32)
            coords = torch.cat([coords, t], dim=-1)

            # Flattening the image and the mask to a per-pixel format
            img = torch.tensor(img).view(-1, 3)
            mask = torch.tensor(mask).view(-1, 1)

            output = {
                "coordinates": coords,  # (x, y, t)
                "rgb": img,  # RGB values of the rectangle
                "mask": mask,  # Mask of the rectangle
            }

        elif self.mode == "val":
            # Getting the rectangle coordinates
            x1, y1, x2, y2 = (
                self.rect_coords[frame_name]["x1"],
                self.rect_coords[frame_name]["y1"],
                self.rect_coords[frame_name]["x2"],
                self.rect_coords[frame_name]["y2"],
            )

            # Dimensions of the image and the mask
            h, w = x2 - x1, y2 - y1

            # Normalisation of the image and the mask
            img = img.astype("float32") / 255.0
            mask = mask.astype("float32") / 255.0

            # Generating the (x, y) coordinate grid within the rectangle
            x = torch.linspace(x1 / img.shape[1] * 2 - 1, x2 / img.shape[1] * 2 - 1, w)
            y = torch.linspace(y1 / img.shape[0] * 2 - 1, y2 / img.shape[0] * 2 - 1, h)
            grid_y, grid_x = torch.meshgrid(y, x, indexing="ij")
            coords = torch.stack([grid_x, grid_y], dim=-1).reshape(h * w, 2)

            # Computing t_value
            t_value = int(os.path.splitext(frame_name)[0])
            t_value = (t_value * 2) / (self.VIDEO_LENGTH - 1)
            t = torch.full((coords.shape[0], 1), t_value, dtype=torch.float32)
            coords = torch.cat([coords, t], dim=-1)

            # Converting the image and the mask to tensor format
            img = torch.tensor(img).view(-1, 3)
            mask = torch.tensor(mask).view(-1, 1)

            output = {
                "coordinates": coords,  # (x, y, t)
                "original_image": img,  # Original image
                "original_mask": mask,  # Original mask
                "rectangle_coords": [x1, y1, x2, y2],  # Rectangle coordinates
            }

        return output

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


def create_dataloaders(
    dir_name: Path, batch_size: int = 1, num_workers: int = 0
) -> tuple[DataLoader, DataLoader]:
    """
    Creates DataLoader objects for the training and validation sets.

    Parameters
    ----------
    dir_name : Path
        Path to the main directory containing the data.
    batch_size : int, optional
        Batch size for the training loader, default 1.
    num_workers : int, optional
        Number of workers used for loading the data, default 0.

    Returns
    -------
    tuple[DataLoader, DataLoader]
        Returns a tuple containing the DataLoaders for the training set and the validation set.
    """
    # Path to the training set
    train_path = Path(dir_name / "train")
    frame_names = sorted(os.listdir(train_path / "images"))
    train_dataset = ImageDataset(
        Path(train_path / "images"),
        Path(train_path / "masks"),
        frame_names,
        mode="train",
    )

    # Path to the validation set
    val_path = Path(dir_name / "val")
    frame_names = sorted(os.listdir(val_path / "images"))
    val_dataset = ImageDataset(
        Path(val_path / "images"),
        Path(val_path / "masks"),
        frame_names,
        mode="val",
        json_path=val_path / "rectangles.json",
    )

    # Creating the DataLoaders
    train_loader = DataLoader(
        train_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers
    )
    val_loader = DataLoader(
        val_dataset, batch_size=1, shuffle=False, num_workers=num_workers
    )

    return train_loader, val_loader

######################### DO NOT CHANGE THIS CELL ##########################
tempdir = tempfile.TemporaryDirectory()
TMP_DIR = tempdir.name
DATA_PATH = Path(TMP_DIR) / Path("data")

def unpack_tar_gz(filename: str, path: Path = DATA_PATH) -> None:
    """ Unpacks a tar.gz archive """
    with tarfile.open(filename, "r:gz") as tar:
        tar.extractall(path=path)
    
unpack_tar_gz("./train.tar.gz", DATA_PATH / Path("train"))
unpack_tar_gz("./val.tar.gz", DATA_PATH / Path("val"))
train_loader, val_loader = create_dataloaders(
    DATA_PATH, batch_size=2, num_workers=0
)

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


def grade(
    model: torch.nn.Module,
    data_loader: torch.utils.data.DataLoader,
    device: torch.device = "cuda",
) -> tuple[float, float, float]:
    """Evaluates the model on the data from the DataLoader.

    The function computes the mean PSNR (Peak Signal-to-Noise Ratio) and the accuracy
    for the masks, and also returns the final score.

    Parameters
    ----------
    model : torch.nn.Module
        The model to evaluate.
    data_loader : torch.utils.data.DataLoader
        DataLoader containing the data for evaluation.
    device : torch.device
        The device on which the computations are performed (e.g. 'cuda' or 'cpu').

    Returns
    -------
    tuple[float, float, float]
        Returns a tuple containing:
        - The final points (float)
        - The mean PSNR (float)
        - The mean accuracy (float)
    """
    model.eval()  # Setting the model to evaluation mode
    psnr_list, acc_list = [], []  # Lists for storing the PSNR and accuracy results

    with torch.no_grad():  # Disabling gradients for evaluation
        for batch_idx, batch in enumerate(data_loader):
            # Getting the data from the batch
            coordinates = batch["coordinates"].to(device)
            original_image = (
                batch["original_image"][0].cpu().numpy().reshape(256, 256, 3)
            )
            original_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)
            original_mask = batch["original_mask"][0].cpu().numpy().reshape(256, 256)
            rect_coords = batch["rectangle_coords"]  # (x1, y1, x2, y2)

            # Moving the model to the GPU
            model = model.cuda()
            rgb_pred, mask_pred = model(coordinates)
            rgb_pred = rgb_pred.squeeze(0).detach().cpu().numpy()
            rgb_pred = rgb_pred[:, [2, 1, 0]]
            mask_pred = mask_pred.squeeze(0).detach().cpu().numpy().squeeze()

            # Creating the reconstructed images
            inpainted_img = original_image.copy()
            inpainted_mask = original_mask.copy()
            x1, y1, x2, y2 = map(int, rect_coords)
            idx = 0
            for y in range(y1, y2):
                for x in range(x1, x2):
                    inpainted_img[y, x, :] = rgb_pred[idx]
                    inpainted_mask[y, x] = mask_pred[idx]
                    idx += 1

            # Computing the PSNR for the rectangle region
            gt_region = original_image[y1:y2, x1:x2, :]
            pred_region = inpainted_img[y1:y2, x1:x2, :]
            psnr = peak_signal_noise_ratio(gt_region, pred_region, data_range=1.0)
            psnr_list.append(psnr)

            # Computing the accuracy for the mask region
            gt_mask_region = original_mask[y1:y2, x1:x2]
            pred_mask_region = inpainted_mask[y1:y2, x1:x2]
            pred_mask_region = (pred_mask_region > 0.5).astype(np.uint8)
            acc = np.mean(pred_mask_region == gt_mask_region)
            acc_list.append(acc)

            # Visualising the results for the first batch
            if batch_idx == 0:
                fig, axs = plt.subplots(1, 3, figsize=(15, 5))
                axs[0].imshow(original_image)
                axs[0].set_title("Original Image")
                axs[1].imshow(inpainted_img)
                axs[1].set_title("Inpainted Image")
                axs[2].imshow(inpainted_mask, cmap="gray")
                axs[2].set_title("Inpainted Mask")
                for ax in axs:
                    ax.axis("off")
                plt.tight_layout()
                plt.show()

    # Computing the mean PSNR and accuracy results
    psnr_score = np.mean(psnr_list)
    acc_score = np.mean(acc_list)

    # Computing the points based on PSNR
    if psnr_score < 15.5:
        p_psnr = 0
    elif psnr_score < 23.5:
        p_psnr = (psnr_score - 15.5) * 5/4
    else:
        p_psnr = 10

    # Computing the points based on accuracy
    if acc_score < 0.83:
        p_acc = 0
    elif acc_score < 0.98:
        p_acc = (acc_score - 0.83) * 20/3
    else:
        p_acc = 1

    # Total points
    points = 7 * p_psnr + 30 * p_acc

    return int(round(points, 0)), psnr_score, acc_score

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

class DummyINR(nn.Module):
    """
    DummyINR is a simple neural network model that generates pink RGB values
    and a zero mask for the given input coordinates.
    """

    def __init__(self):
        """Initialises an instance of the DummyINR class."""
        super(DummyINR, self).__init__()

    def forward(self, coords):
        """
        Computes pink RGB values and a zero mask for the input coordinates.

        Parameters:
        -----------
        coords : torch.Tensor
            Input tensor of shape (B, N, 3), where B is the batch size,
            N is the number of points, and 3 are the coordinates (x, y, t).

        Returns:
        --------
        tuple:
            Two values containing:
                - rgb : torch.Tensor
                    Tensor of pink RGB values in the range [0, 1] of shape (B, N, 3).
                - mask : torch.Tensor
                    Tensor of the zero mask of shape (B, N, 1).
        """
        batch_size, num_points, _ = coords.shape

        pink_rgb = torch.tensor(
            [1, 0, 1], dtype=torch.float32, device=coords.device
        )
        rgb = pink_rgb.unsqueeze(0).unsqueeze(0).expand(batch_size, num_points, -1)  # Shape: (B, N, 3)

        # Mask values in the range [0, 1]
        mask = torch.zeros(
            (batch_size, num_points, 1), dtype=torch.float32, device=coords.device
        )  # Shape: (B, N, 1)

        return rgb, mask

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

if not FINAL_EVALUATION_MODE:
    dummy_model = DummyINR()
    points, psnr, accuracy = grade(dummy_model, val_loader)
    print(f"The example solution scored {points} pts on the validation set.")
    print(f"PSNR: {psnr:.2f}")
    print(f"Accuracy: {accuracy:.2f}")

class YourSolution(nn.Module):
    """The YourSolution class

    The class implements an INR neural network model that processes the input
    coordinates (x, y, t) and returns RGB values and a mask.

    Attributes:
    -----------
    No attributes to initialise in the constructor.
    The model should be defined in the `__init__` method.
    """

    def __init__(self):
        """Initialises the YourSolution model.

        Define all layers and components of the model here.
        """
        super(YourSolution, self).__init__()
        # Model initialisation
        pass

    def forward(self, coords):
        """Processes the input coordinates and returns the model outputs.

        Parameters:
        -----------
        coords : torch.Tensor
            Input tensor of shape (B, N, 3), where each column
            corresponds to the coordinates (x, y, t).

        Returns:
        --------
        tuple
            Two output values:
                - "rgb" : torch.Tensor
                    Tensor with RGB values in the range [0, 1].
                - "mask" : torch.Tensor
                    Tensor with mask values in the range [0, 1].
        """
        # Model implementation
        pass

def train(
    model,
    train_loader,
    epochs=1,
    lr=1000,
    device="cuda" if torch.cuda.is_available() else "cpu",
):
    """Function that trains the model on the given dataset.

    Parameters
    ----------
    model : torch.nn.Module
        The model to be trained.
    train_loader : torch.utils.data.DataLoader
        DataLoader containing the training data.
    epochs : int, optional
        Number of training epochs (default 1).
    lr : float, optional
        Learning rate (default 1000).
    device : str, optional
        The device on which the model is to be trained (default "cuda", if available).

    Notes
    -----
    The function requires further implementation of the training loop and the learning logic.
    """
    # Build the training loop
    model = model.to(device)
    optimizer = torch.optim.Adam(
        model.parameters(), lr=lr
    )  # You may change the optimiser
    pass

# Enter the number of epochs you want to run and the chosen learning rate
EPOCHS = 40
LEARNING_RATE = 1e-4

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

model = YourSolution()
train(model, train_loader, epochs=EPOCHS, lr=LEARNING_RATE)

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

if not FINAL_EVALUATION_MODE:
    points, psnr, accuracy = grade(model, val_loader)
    print(f"Your solution scored {points} pts on the validation set.")
    print(f"PSNR: {psnr:.2f}")
    print(f"Classification accuracy: {accuracy:.2f}%")

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

if FINAL_EVALUATION_MODE:
    import cloudpickle

    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(model, f)