Upload utils/dataset_lab.py

utils/dataset_lab.py

@@ -0,0 +1,37 @@
+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)
+        #print('--------L:', np.min(lab_img[:,:,0]), np.max(lab_img[:,:,0]))
+        #print('--------ab:', np.min(lab_img[:,:,1:3]), np.max(lab_img[:,:,1:3]))
+        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}