11

import os
import glob
import h5py
import random
from PIL import Image
from matplotlib import pyplot as plt

import torch.utils.data as data
from torch.utils.data import DataLoader
import torchvision.transforms as tfs
from torchvision.transforms import functional as FF

from Utils.metrics import *
from Utils.option import *


class train_Dataset(data.Dataset):
    def __init__(self, path, mode='train', size=opt.crop_size, **kwargs):
        super(train_Dataset,self).__init__()
        self.size=size
        self.mode=mode
        self.format=format
        self.haze_imgs_dir = os.path.join(path, 'low')  # haze_img dir path
        self.haze_imgs_list = os.listdir(self.haze_imgs_dir)  # haze_img name list
        self.haze_imgs = [os.path.join(self.haze_imgs_dir, img) for img in self.haze_imgs_list]  # haze_img path list
        self.clear_dir = os.path.join(path, 'high')  # clean_img dir path
        self.mask = False
        if 'restored_mask' in kwargs:
            self.restored_mask = kwargs['restored_mask'].long()
            self.mask = True

        self.length = len(self.haze_imgs_list)
    def __getitem__(self, index):
        haze = Image.open(self.haze_imgs[index])
        if isinstance(self.size, int):
            while haze.size[0]<self.size or haze.size[1]<self.size :
                index = random.randint(0,self.length)
                haze = Image.open(self.haze_imgs[index])
        haze_name = self.haze_imgs[index].split('/')[-1]
        id = haze_name.split('_')[0]
        clear_name = id
        clear = Image.open(os.path.join(self.clear_dir, clear_name))
        clear = tfs.CenterCrop(haze.size[::-1])(clear)

        if not isinstance(self.size,str):
            i,j,h,w=tfs.RandomCrop.get_params(haze,output_size=(self.size,self.size))
            haze = FF.crop(haze,i,j,h,w)
            clear = FF.crop(clear,i,j,h,w)
        rand_hor = random.randint(0, 1)
        rand_rot = random.randint(0, 3)
        haze =self.augData_haze(haze.convert("RGB"), rand_hor, rand_rot)
        clear = self.augData_clear(clear.convert("RGB"), rand_hor, rand_rot)
        mask_flag = self.restored_mask[index] if self.mask else -1

        return haze, clear, index, mask_flag

    def augData_haze(self, haze, rand_hor, rand_rot):
        haze=tfs.RandomHorizontalFlip(rand_hor)(haze)
        if rand_rot:
            haze=FF.rotate(haze,90*rand_rot)
        haze=tfs.ToTensor()(haze)
        # haze=tfs.Normalize(mean=[0.64, 0.6, 0.58],std=[0.14,0.15, 0.152])(haze)
        return haze

    def augData_clear(self, clear, rand_hor, rand_rot):
        clear=tfs.RandomHorizontalFlip(rand_hor)(clear)
        if rand_rot:
            clear=FF.rotate(clear,90*rand_rot)
        clear=tfs.ToTensor()(clear)
        return clear

    def __len__(self):
        return self.length


class test_Dataset(data.Dataset):
    def __init__(self, path, mode='test'):
        super(test_Dataset,self).__init__()
        self.mode=mode
        self.format=format
        self.haze_imgs_dir = os.path.join(path, 'low')  # haze_img dir path
        self.haze_imgs_list = os.listdir(self.haze_imgs_dir)  # haze_img name list
        self.haze_imgs = [os.path.join(self.haze_imgs_dir, img) for img in self.haze_imgs_list]  # haze_img path list
        self.clear_dir = os.path.join(path, 'high')  # clean_img dir path

        self.length = len(self.haze_imgs_list)
    def __getitem__(self, index):
        haze = Image.open(self.haze_imgs[index])
        haze_name = self.haze_imgs[index].split('/')[-1]
        id = haze_name.split('_')[0]
        clear_name = id
        clear = Image.open(os.path.join(self.clear_dir, clear_name))
        clear = tfs.CenterCrop(haze.size[::-1])(clear)

        haze = self.augData_haze(haze.convert("RGB"))
        clear = self.augData_clear(clear.convert("RGB"))
        return haze, clear, haze_name

    def augData_haze(self, haze):
        haze=tfs.ToTensor()(haze)
        # haze=tfs.Normalize(mean=[0.64, 0.6, 0.58],std=[0.14,0.15, 0.152])(haze)
        return haze

    def augData_clear(self, clear):
        clear=tfs.ToTensor()(clear)
        return clear

    def __len__(self):
        return self.length