import torch

import numpy as np

import cv2

import os

from torch.utils.data import Dataset

class ToTensor(object):

    def __call__(self, sample):

        entry = {}

        for k in sample:

            if k == 'rect':

                entry[k] = torch.IntTensor(sample[k])

            else:

                entry[k] = torch.FloatTensor(sample[k])

        return entry

class InpaintingDataset(Dataset):

    def __init__(self, info_list, root_dir='', im_size=(256, 256), transform=None):

        self.filenames = open(info_list, 'rt').read().splitlines()

        self.root_dir = root_dir

        self.transform = transform

        self.im_size = im_size

        np.random.seed(2018)

    def __len__(self):

        return len(self.filenames)

    def read_image(self, filepath):

        image = cv2.imread(filepath)

        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        h, w, c = image.shape

        if h != self.im_size[0] or w != self.im_size[1]:

            ratio = max(1.0 * self.im_size[0] / h, 1.0 * self.im_size[1] / w)

            im_scaled = cv2.resize(image, None, fx=ratio, fy=ratio)

            h, w, _ = im_scaled.shape

            h_idx = (h - self.im_size[0]) // 2

            w_idx = (w - self.im_size[1]) // 2

            im_scaled = im_scaled[h_idx:h_idx + self.im_size[0], w_idx:w_idx + self.im_size[1], :]

            im_scaled = np.transpose(im_scaled, [2, 0, 1])

        else:

            im_scaled = np.transpose(image, [2, 0, 1])

        return im_scaled

    def __getitem__(self, idx):

        image = self.read_image(os.path.join(self.root_dir, self.filenames[idx]))

        sample = {'gt': image}

        if self.transform:

            sample = self.transform(sample)

        return sample