import os
import torch
import numpy as np
from torch.utils.data import Dataset
from PIL import Image
from skimage import color


def load_img(img_path):
	out_np = np.asarray(Image.open(img_path))
	if(out_np.ndim==2):
		out_np = np.tile(out_np[:,:,None],3)
	return out_np


class ColDataset(Dataset):
    """Colored Postcards dataset"""

    def __init__(self, image_dir, transform=None):
        """
        :param iamge_dir (string): Directory with all the images
        :param transform (callable, optional): Optional transform to be applied on a sample
        """
        filepaths = os.listdir(image_dir)
        filepaths = [os.path.join(image_dir, f) for f in filepaths]
        self.filepaths = filepaths
        self.transform = transform

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

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()

        img_path = self.filepaths[idx]
        image = load_img(img_path)
        image_lab = color.rgb2lab(image)
        tens_l = torch.Tensor(image_lab[:, :, 0])[None, :, :]
        tens_ab = torch.Tensor(image_lab[:, :, 1:])[:, :].transpose(2, 1).transpose(1, 0)
        sample = {'L': tens_l, 'ab': tens_ab}

        if self.transform:
            sample = self.transform(sample)

        return sample

