| | from __future__ import print_function, division
|
| | import torch, os, glob
|
| | from torch.utils.data import Dataset, DataLoader
|
| | import numpy as np
|
| | from PIL import Image
|
| | import cv2
|
| |
|
| |
|
| | class LabDataset(Dataset):
|
| |
|
| | def __init__(self, rootdir=None, filelist=None, resize=None):
|
| |
|
| | if filelist:
|
| | self.file_list = filelist
|
| | else:
|
| | assert os.path.exists(rootdir), "@dir:'%s' NOT exist ..."%rootdir
|
| | self.file_list = glob.glob(os.path.join(rootdir, '*.*'))
|
| | self.file_list.sort()
|
| | self.resize = resize
|
| |
|
| | def __len__(self):
|
| | return len(self.file_list)
|
| |
|
| | def __getitem__(self, idx):
|
| | bgr_img = cv2.imread(self.file_list[idx], cv2.IMREAD_COLOR)
|
| | if self.resize:
|
| | bgr_img = cv2.resize(bgr_img, (self.resize,self.resize), interpolation=cv2.INTER_CUBIC)
|
| | bgr_img = np.array(bgr_img / 255., np.float32)
|
| | lab_img = cv2.cvtColor(bgr_img, cv2.COLOR_BGR2LAB)
|
| |
|
| |
|
| | lab_img = torch.from_numpy(lab_img.transpose((2, 0, 1)))
|
| | bgr_img = torch.from_numpy(bgr_img.transpose((2, 0, 1)))
|
| | gray_img = (lab_img[0:1,:,:]-50.) / 50.
|
| | color_map = lab_img[1:3,:,:] / 110.
|
| | bgr_img = bgr_img*2. - 1.
|
| | return {'gray': gray_img, 'color': color_map, 'BGR': bgr_img} |
