Multispectral Data Segmentation

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

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

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

import os
import random
from tqdm import tqdm

import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, TensorDataset, DataLoader

# Additional libraries that you may use in your solution
import xgboost
import sklearn

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
DATA_DIR = "./data"
TRAIN_DATA_PATH = os.path.join(DATA_DIR, "train.npz")
VALID_DATA_PATH = os.path.join(DATA_DIR, "valid.npz")

assert torch.cuda.is_available(), "No graphics card (GPU) was found!"

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

seed = 12345

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

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

class BaseDataset(Dataset):
    """
    Multispectral dataset class.
    """
    def __init__(self, data_path: str):
        data = np.load(data_path)
        self.bands = data["bands"]
        self.segmentations = data["segmentations"]

    def __len__(self):
        """Returns the number of samples in the dataset."""
        return self.bands.shape[0]

    def __getitem__(self, idx):
        """Returns the sample with index idx - its multispectral bands and its segmentation."""
        bands = self.bands[idx]
        segmentations = self.segmentations[idx]

        bands = torch.from_numpy(bands).float()
        segmentations = torch.from_numpy(segmentations).long()

        return bands, segmentations

def setup_data():
    """
    Downloads the datasets used in the task.
    """
    import gdown
    os.makedirs(DATA_DIR, exist_ok=True)

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

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

if not FINAL_EVALUATION_MODE:
    setup_data()

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

def calculate_miou(y_true, y_pred):
    """
    Computes the mIoU (mean Intersection over Union) metric used to evaluate the model.
    Helper function used when evaluating the solution.
    """
    assert y_true.shape == y_pred.shape
    assert y_true.device == y_pred.device

    num_classes = 4
    device = y_true.device
    batch_size = y_true.shape[0]
    y_true_flat = y_true.reshape(batch_size, y_true.shape[1] * y_true.shape[2])
    y_pred_flat = y_pred.reshape(batch_size, y_pred.shape[1] * y_pred.shape[2])

    # The final mean is computed only over the classes present in the reference data (ground truth)
    intersections = torch.zeros((batch_size, num_classes), dtype=torch.float32, device=device)
    unions = torch.zeros((batch_size, num_classes), dtype=torch.float32, device=device)
    true_present = torch.zeros((batch_size, num_classes), dtype=torch.bool, device=device)

    for cls in range(num_classes):
        y_true_c = (y_true_flat == cls)
        y_pred_c = (y_pred_flat == cls)
        true_present[:, cls] = torch.any(y_true_c, dim=1)

        intersections[:, cls] = torch.sum(y_true_c & y_pred_c, dim=1).to(torch.float32)
        unions[:, cls] = torch.sum(y_true_c | y_pred_c, dim=1).to(torch.float32)

    # Computes the number of relevant classes for each sample
    num_present_classes = torch.sum(true_present, dim=1).to(torch.float32)

    # Although the computations are performed for all classes, we sum only the classes that actually occur (in the reference data)
    iou_per_class = torch.nan_to_num(intersections / unions, nan=0.0)
    sum_iou = torch.sum(iou_per_class * true_present, dim=1)
    miou_scores = torch.nan_to_num(sum_iou / num_present_classes, nan=0.0).tolist()
    assert len(miou_scores) == batch_size

    return miou_scores


def calculate_channel_count(dataloader):
    """
    Helper function used when evaluating the solution.
    Computes the number of bands used in training and evaluating the model.
    """
    tensors = dataloader.dataset.tensors

    if len(tensors) != 2:
        raise ValueError("The dataset should contain only labels and images (__getitem__ should return a tuple of size 2)")

    x, _ = tensors

    # x.shape [dataset size, channels, height, width]
    if x.ndim != 4:
        raise ValueError("The processed data must have 4 dimensions [batch size, channels, height, width]")

    if x.shape[0] < 15: #the validation set is the smallest and has 15 samples
        raise ValueError(f"The first dimension is reserved for the dataset size; your code should not change it!")

    if x.shape[2] != 30 or x.shape[3] != 30:
        raise ValueError("The processed data must have size 30x30 in the 3rd and 4th dimensions (.shape[2] == .shape[3] == 30)")

    channels_count = x.shape[1]
    return channels_count

def transform_dataset(preprocessing, dataset:BaseDataset) -> TensorDataset:
    """
    Processes the given dataset and keeps it in the computer's memory (RAM).
    Calls the .transform function, implemented by the participant, of the data preparation class.
    """
    processed_labels = []
    processed_images = []

    for raw_image, label in dataset:
        processed_labels.append(label.clone())

        transformed_image = preprocessing.transform(raw_image.clone())
        processed_images.append(transformed_image)

    labels_tensor = torch.stack(processed_labels)
    images_tensor = torch.stack(processed_images)
    memory_dataset = TensorDataset(images_tensor, labels_tensor)
    return memory_dataset

def evaluate(train_model, preprocessing, data_path: str):
    """
    Main function that evaluates the task.
    The same function will be called on the Contest Platform.

    1. Fits the data processing class on the training set.
    2. Processes the given dataset using the fitted class.
    3. Evaluates the model with the mIoU metric on the given processed dataset.
    4. Evaluates the number of bands used in evaluating the model - checks the shape of the data.
    """

    # Loads the training set and fits the data preparation class on it
    train_ds = BaseDataset(data_path=TRAIN_DATA_PATH)
    preprocessing = preprocessing()
    assert hasattr(preprocessing, "fit"), "The data preparation function must implement the .fit method"
    preprocessing.fit(train_ds)

    # Loads the target dataset and processes it using the data preparation class
    target_ds = BaseDataset(data_path=data_path)
    assert hasattr(preprocessing, "transform"), "The data preparation function must implement the .transform method"
    transformed_dataset = transform_dataset(preprocessing=preprocessing, dataset=target_ds)
    dataloader = DataLoader(transformed_dataset, batch_size=8, shuffle=False)

    model = train_model()

    if hasattr(model, "to") and callable(getattr(model, "to", None)):
        model = model.to(DEVICE)

    if hasattr(model, "eval") and callable(getattr(model, "eval", None)):
        model.eval()

    mious = []

    # Evaluates the model with the mIoU metric
    with torch.no_grad():

        for x, y in dataloader:
            x = x.to(DEVICE)
            y = y.to(DEVICE)

            y_pred = model(x)
            if hasattr(y_pred, "to") and callable(getattr(y_pred, "to", None)):
                y_pred = y_pred.to(DEVICE)
            if not torch.is_tensor(y_pred):
                y_pred = torch.tensor(y_pred)
            y_pred = torch.argmax(y_pred, dim=1)

            assert y_pred.shape == y.shape
            assert y_pred.max() <= 4 and y_pred.min() >= 0

            miou = calculate_miou(y, y_pred)
            mious.extend(miou)

    # Evaluates the number of bands used in evaluating the model
    channels_count = calculate_channel_count(dataloader=dataloader)

    # Computes the mean mIoU over all samples
    miou = sum(mious) / len(mious)

    return miou, channels_count

def compute_score(miou, channels_count):
    """
    Computes the score for the task: calculates the final number of points for the task using the formula given in the task statement.
    The same function will be called on the Contest Platform.
    """
    band_score = max(0, min(100, 100 * ((12 - channels_count) / (12 - 3))**2))
    miou_score = max(0, min(100, 100 * (miou - 0.6) / (0.8 - 0.6)))

    print(f"Number of channels: {channels_count}")
    print(f"mIoU: {miou:.3f} \n")

    print(f"Channel score: {band_score:.3f}")
    print(f"mIoU score: {miou_score:.3f}")

    if miou < 0.5:
        total_score = int(0)
    else:
        total_score = 0.25 * band_score + 0.75 * miou_score
        total_score = int(round(total_score))
    print(f"Estimated number of points for the task: {total_score}")
    return total_score

######################### DO NOT CHANGE THIS CELL ###########################
class BasicPreprocessing():
  """
  Example implementation of a dataset processing class.

  This is only an example solution, and it does not reduce the number of channels.
  The fewer channels you use, the more points you will receive for this part of the task.
  """

  def __init__(self):
    self.std_per_channel = None

  def fit(self, dataset:BaseDataset):

    # Collects all images into a single matrix [N_samples, 12, 30, 30]
    images = []
    for item in dataset:
        image_tensor, _ = item
        images.append(image_tensor)
    images_tensor = torch.stack(images)  # shape: [N, 12, 30, 30]

    # Computes the std (standard deviation) over the dimensions: (0: sample, 2: H, 3: W)
    self.std_per_channel = images_tensor.std(dim=(0, 2, 3))

    return self

  def transform(self, image_tensor: torch.Tensor) -> torch.Tensor:
      # Reshapes the mean to [12, 1, 1] to allow broadcasting against [12, 30, 30].
      # Remember, this is only an example solution; merely dividing by the std is not the most effective solution.
      std_reshaped = self.std_per_channel.view(-1, 1, 1)

      return image_tensor / std_reshaped

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

class BasicModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(12, 4)

    def forward(self, bands):
        # bands.shape [b, c, h, w]
        return self.linear(bands.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)

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

def train_basic_model():

    epochs = 40
    lr = 0.001
    batch_size = 8
    model = BasicModel()
    model = model.to(DEVICE)

    # Prepares the data preprocessing parameters using only
    # the training set.

    raw_train_ds = BaseDataset(data_path=TRAIN_DATA_PATH)
    raw_valid_ds = BaseDataset(data_path=VALID_DATA_PATH)

    preprocessing = BasicPreprocessing()
    preprocessing.fit(raw_train_ds)

    train_ds = transform_dataset(preprocessing=preprocessing, dataset=raw_train_ds)
    valid_ds = transform_dataset(preprocessing=preprocessing, dataset=raw_valid_ds)
    train_dataloader = DataLoader(train_ds, batch_size=batch_size, shuffle=True)
    valid_dataloader = DataLoader(valid_ds, batch_size=batch_size, shuffle=False)

    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()

    for epoch in range(epochs):
        model.train()
        for x, y in tqdm(train_dataloader, total=len(train_dataloader), desc="Training"):
            x = x.to(DEVICE)
            y = y.to(DEVICE)

            y_pred = model(x)
            loss = criterion(y_pred, y)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        model.eval()
        with torch.no_grad():
            valid_loss = 0
            mious = []
            for x, y in tqdm(valid_dataloader, total=len(valid_dataloader), desc="Validation"):
                x = x.to(DEVICE)
                y = y.to(DEVICE)

                y_pred = model(x)
                loss = criterion(y_pred, y)

                valid_loss += loss.item()

                y_pred = torch.argmax(y_pred, dim=1)
                miou = calculate_miou(y, y_pred)
                mious.extend(miou)

            valid_loss = valid_loss / len(valid_dataloader)
            print(f"Epoch {epoch+1} loss: {valid_loss}, mIoU: {sum(mious) / len(mious)}")

    return model

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

if not FINAL_EVALUATION_MODE:
    miou, channels_count = evaluate(train_basic_model, BasicPreprocessing, VALID_DATA_PATH)
    print("-"*50)
    compute_score(miou, channels_count)

# Do not change the name of the class
# This class may only process the samples; it must not use the data labels.
# No method may perform segmentation either.

class YourPreprocessing():
  def __init__(self):
    pass

  def fit(self, dataset: BaseDataset):
    return dataset

  def transform(self, image_tensor: torch.Tensor) -> torch.Tensor:
    return image_tensor

class YourModel(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, bands):
        # random predictions
        segmentation = torch.randint(0, 4, (bands.shape[0], 1, 30, 30))
        # Performs one-hot encoding because evaluate uses torch.argmax
        one_hot = torch.zeros(bands.shape[0], 4, 30, 30)
        one_hot.scatter_(1, segmentation, 1)
        return one_hot

# Do not change the name of the function
def train_your_model():
    return YourModel()

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

if not FINAL_EVALUATION_MODE:
    miou, channels_count = evaluate(train_your_model, YourPreprocessing, VALID_DATA_PATH)
    print("-"*50)
    compute_score(miou, channels_count)