# AUTOGENERATED! DO NOT EDIT! File to edit: 02_utils.ipynb (unless otherwise specified).

__all__ = ['load_camvid_dataset', 'display_segmentation', 'display_segmentation_from_file', 'CamvidDataset']


# Cell
import matplotlib.pyplot as plt
import os
import torch
import torchvision.transforms.functional as tf
from PIL import Image


# Cell
def load_camvid_dataset(data_directory):
    with open(os.path.join(data_directory, "valid.txt"), "r") as f:
        val_names = [line.strip() for line in f]
    with open(os.path.join(data_directory, "codes.txt"), "r") as f:
        label_mapping = {l.strip(): i for i, l in enumerate(f)}
    data = []
    image_index_mapping = {}
    for im_f in os.listdir(os.path.join(data_directory, "images")):
        if im_f.split('.')[-1] != 'png':
            continue
        image_index_mapping[im_f] = len(data)
        fp = os.path.join(data_directory, "images", im_f)
        data.append(fp)
    for label_f in os.listdir(os.path.join(data_directory, "labels")):
        im_f = label_f.split('.')
        im_f[0] = '_'.join(im_f[0].split('_')[:-1])
        im_f = '.'.join(im_f)
        index = image_index_mapping[im_f]
        fp = os.path.join(data_directory, "labels", label_f)
        data[index] = (data[index], fp)
    val_indices = [image_index_mapping[name] for name in val_names]
    return data, val_indices, label_mapping


# Cell
def display_segmentation(image, target, ax=None):
    if ax:
        ax.imshow(image, cmap='gray')
    else:
        plt.imshow(image, cmap='gray')
    if ax:
        ax.imshow(target, cmap='jet', alpha=0.5)
    else:
        plt.imshow(target, cmap='jet', alpha=0.5)
    plt.show()

def display_segmentation_from_file(im_f, label_f):
    im, label = Image.open(im_f), Image.open(label_f)
    display_segmentation(im, label)


# Cell
class CamvidDataset(torch.utils.data.Dataset):
    def __init__(self, data, resize_shape=(360, 480), is_train=True):
        self.images, self.labels = [tpl[0] for tpl in data], \
                                   [tpl[1] for tpl in data]
        self.resize_shape = resize_shape
        self.is_train = is_train

    def transform(self, index):
        input, target = map(
            Image.open, (self.images[index], self.labels[index]))
        input, target = (
            tf.resize(input, self.resize_shape),
            tf.resize(target, self.resize_shape, interpolation=Image.NEAREST)
        )
        if self.is_train:
            horizontal_draw = torch.rand(1).item()
            vertical_draw = torch.rand(1).item()
            if horizontal_draw > 0.5:
                input, target = tf.hflip(input), tf.hflip(target)
            if vertical_draw > 0.5:
                input, target = tf.vflip(input), tf.vflip(target)

        input, target = map(tf.to_tensor, (input, target))
        torch.clamp((255 * target), 0, 32, out=target)
        return tf.normalize(input, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), target.long()

    def __getitem__(self, index):
        return self.transform(index)

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