import os
import glob
import torch.utils.data
import numpy as np
from PIL import Image
import LJY_utils

class Dataloader(torch.utils.data.Dataset):
    def __init__(self, path, transform):
        super().__init__()
        self.transform = transform

        assert os.path.exists(path)
        self.base_path = path

        #self.mean_image = self.get_mean_image()

        cur_file_paths = glob.glob(self.base_path + '/*.*')
        cur_file_paths.sort()
        self.file_paths = cur_file_paths

    def pil_loader(self,path):
        # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)
        with open(path, 'rb') as f:
            with Image.open(f) as img:
                return img.convert('RGB')

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

    def __getitem__(self, item):
        path = self.file_paths[item]
        img = self.pil_loader(path)
        if self.transform is not None:
            img = self.transform(img)


        return img


class fold_Dataloader(torch.utils.data.Dataset):
    def __init__(self,fold, path, transform, type ='train'):
        super().__init__()
        self.transform = transform

        self.type = type
        train_path, val_path = LJY_utils.fold_loader(fold, path)
        if self.type == 'train':
            self.file_paths = train_path[0]
        # self.Semantic_base_path = Semantic_path
        elif self.type == 'validation':
            self.file_paths = val_path[0]

    def pil_loader(self,path):
        # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)
        with open(path, 'rb') as f:
            with Image.open(f) as img:
                return img.convert('RGB')

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

    def __getitem__(self, item):
        path = self.file_paths[item]
        img = self.pil_loader(path)
        if self.transform is not None:
            img = self.transform(img)

        return img