# Test code for MM24: DERO
from torchvision import transforms as T
from torchvision.datasets import ImageFolder
from torch.utils import data
from PIL import Image
import math
import torch
import os
import cv2
import numpy as np
import scipy.io as scio
import os
from PIL import Image
import numpy as np
import random
import torch
from torchvision import transforms
from torch.utils import data
from torch.utils.data import Dataset
import imageio
import imutils

class ImageLoader_for_test(data.Dataset):
	def __init__(self,data_dir,image_size,w_path,transform=None):
		super(ImageLoader_for_test,self).__init__()
		self.data_dir = data_dir
		self.transform = transform
		self.img_paths = os.listdir(data_dir)
		self.image_size = image_size
		W = scio.loadmat(w_path)
		self.w = W['w']
    
	def __len__(self):
		return len(self.img_paths)

	def __getitem__(self,index):
		imagesize = self.image_size
		curr_img_path = self.img_paths[index]
		nn = curr_img_path.index('.')
		num = int(curr_img_path[:nn])
		img = cv2.imread(self.data_dir + curr_img_path,1)
		data_img = cv2.resize(img,(imagesize,imagesize))

		m = self.w[num,:]

		Data_img = np.zeros_like(data_img)#[:,:,:]
		for i in range(3):
			Data_img[:,:,i] = data_img[:,:,2-i]

		Data_img = Data_img.transpose((2,0,1))
		Data_img = np.float32(Data_img/255*2-1)
		return  Data_img, m, num


def get_loader(image_dir, image_size=512, 
               batch_size=32, dataset='train', mode='train', num_workers=0, w_path = 'results/WatermarkMatrix/w.mat'):
    """Build and return a data loader."""
    transform = []

    if dataset in ['test_accuracy','test_embed']:
    	dataset = ImageLoader_for_test(image_dir,image_size,w_path)
    data_loader = data.DataLoader(dataset=dataset,
                                  batch_size=batch_size,
                                  shuffle=True,#(mode=='train')
                                  num_workers=num_workers,
                                  drop_last=True)
    return data_loader

