from glob import glob

from torch.utils.data import Dataset
from PIL import Image


class MSCOCO(Dataset):
    def __init__(self, root, transform, img_list=None):
        assert root[-1] == '/', "root to COCO dataset should end with \'/\', not {}.".format(
            root)

        if img_list:
            self.image_paths = []
            with open(img_list, 'r') as r:
                lines = r.read().splitlines()
                for line in lines:
                    self.image_paths.append(root + line)
        else:
            self.image_paths = sorted(glob(root + "*.jpg"))
        self.transform = transform

    def __getitem__(self, index):
        """
        Args:
            index (int): Index
        Returns:
            object: image.
        """
        img_path = self.image_paths[index]

        img = Image.open(img_path).convert('RGB')

        if self.transform is not None:
            img = self.transform(img)

        return img

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

class Kodak(Dataset):
    def __init__(self, root, transform):

        assert root[-1] == '/', "root to Kodak dataset should end with \'/\', not {}.".format(
            root)

        self.image_paths = sorted(glob(root + "*.png"))
        self.transform = transform

    def __getitem__(self, index):
        """
        Args:
            index (int): Index
        Returns:
            object: image.
        """
        img_path = self.image_paths[index]

        img = Image.open(img_path).convert('RGB')

        if self.transform is not None:
            img = self.transform(img)

        return img

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