menghanxia/disco

1from __future__ import print_function, division
2import torch, os, glob
3from torch.utils.data import Dataset, DataLoader
4import numpy as np
5from PIL import Image
6import cv2
7
8
9class LabDataset(Dataset):
10
11    def __init__(self, rootdir=None, filelist=None, resize=None):
12
13        if filelist:
14            self.file_list = filelist
15        else:
16            assert os.path.exists(rootdir), "@dir:'%s' NOT exist ..."%rootdir
17            self.file_list = glob.glob(os.path.join(rootdir, '*.*'))
18            self.file_list.sort()
19        self.resize = resize
20
21    def __len__(self):
22        return len(self.file_list)
23
24    def __getitem__(self, idx):
25        bgr_img = cv2.imread(self.file_list[idx], cv2.IMREAD_COLOR)
26        if self.resize:
27            bgr_img = cv2.resize(bgr_img, (self.resize,self.resize), interpolation=cv2.INTER_CUBIC)
28        bgr_img = np.array(bgr_img / 255., np.float32)
29        lab_img = cv2.cvtColor(bgr_img, cv2.COLOR_BGR2LAB)
30        #print('--------L:', np.min(lab_img[:,:,0]), np.max(lab_img[:,:,0]))
31        #print('--------ab:', np.min(lab_img[:,:,1:3]), np.max(lab_img[:,:,1:3]))
32        lab_img = torch.from_numpy(lab_img.transpose((2, 0, 1)))
33        bgr_img = torch.from_numpy(bgr_img.transpose((2, 0, 1)))
34        gray_img = (lab_img[0:1,:,:]-50.) / 50.
35        color_map = lab_img[1:3,:,:] / 110.
36        bgr_img = bgr_img*2. - 1.
37        return {'gray': gray_img, 'color': color_map, 'BGR': bgr_img}