Upload 6 files

utils/cielab.py

@@ -0,0 +1,71 @@
+from functools import partial
+import numpy as np
+
+class ABGamut:
+    RESOURCE_POINTS = "./utils/gamut_pts.npy"
+    RESOURCE_PRIOR = "./utils/gamut_probs.npy"
+    DTYPE = np.float32
+    EXPECTED_SIZE = 313
+    def __init__(self):
+        self.points = np.load(self.RESOURCE_POINTS).astype(self.DTYPE)
+        self.prior = np.load(self.RESOURCE_PRIOR).astype(self.DTYPE)
+        assert self.points.shape == (self.EXPECTED_SIZE, 2)
+        assert self.prior.shape == (self.EXPECTED_SIZE,)
+
+
+class CIELAB:
+    L_MEAN = 50
+    AB_BINSIZE = 10
+    AB_RANGE = [-110 - AB_BINSIZE // 2, 110 + AB_BINSIZE // 2, AB_BINSIZE]
+    AB_DTYPE = np.float32
+    Q_DTYPE = np.int64
+
+    RGB_RESOLUTION = 101
+    RGB_RANGE = [0, 1, RGB_RESOLUTION]
+    RGB_DTYPE = np.float64
+
+    def __init__(self, gamut=None):
+        self.gamut = gamut if gamut is not None else ABGamut()
+        a, b, self.ab = self._get_ab()
+        self.ab_gamut_mask = self._get_ab_gamut_mask(
+            a, b, self.ab, self.gamut)
+
+        self.ab_to_q = self._get_ab_to_q(self.ab_gamut_mask)
+        self.q_to_ab = self._get_q_to_ab(self.ab, self.ab_gamut_mask)
+
+    @classmethod
+    def _get_ab(cls):
+        a = np.arange(*cls.AB_RANGE, dtype=cls.AB_DTYPE)
+        b = np.arange(*cls.AB_RANGE, dtype=cls.AB_DTYPE)
+        b_, a_ = np.meshgrid(a, b)
+        ab = np.dstack((a_, b_))
+        return a, b, ab
+
+    @classmethod
+    def _get_ab_gamut_mask(cls, a, b, ab, gamut):
+        ab_gamut_mask = np.full(ab.shape[:-1], False, dtype=bool)
+        a = np.digitize(gamut.points[:, 0], a) - 1
+        b = np.digitize(gamut.points[:, 1], b) - 1
+        for a_, b_ in zip(a, b):
+            ab_gamut_mask[a_, b_] = True
+
+        return ab_gamut_mask
+
+    @classmethod
+    def _get_ab_to_q(cls, ab_gamut_mask):
+        ab_to_q = np.full(ab_gamut_mask.shape, -1, dtype=cls.Q_DTYPE)
+        ab_to_q[ab_gamut_mask] = np.arange(np.count_nonzero(ab_gamut_mask))
+
+        return ab_to_q
+
+    @classmethod
+    def _get_q_to_ab(cls, ab, ab_gamut_mask):
+        return ab[ab_gamut_mask] + cls.AB_BINSIZE / 2
+
+    def bin_ab(self, ab):
+        ab_discrete = ((ab + 110) / self.AB_RANGE[2]).astype(int)
+
+        a, b = np.hsplit(ab_discrete.reshape(-1, 2), 2)
+
+        return self.ab_to_q[a, b].reshape(*ab.shape[:2])
+

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}

utils/gamut_probs.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:19d00659c7d6f6ee47456fd2c19c86a073f7124875e3d5ab9d601864e062b56c
+size 2584

utils/gamut_pts.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b5dec01315c34f43f1c8c089e84c45ae35d1838d8e77ed0e7ca930f79ffa450e
+size 5088

utils/util.py

@@ -0,0 +1,178 @@
+from __future__ import division
+from __future__ import print_function
+import os, glob, shutil, math, json
+from queue import Queue
+from threading import Thread
+from skimage.segmentation import mark_boundaries
+import numpy as np
+from PIL import Image
+import cv2, torch
+
+def get_gauss_kernel(size, sigma):
+    '''Function to mimic the 'fspecial' gaussian MATLAB function'''
+    x, y = np.mgrid[-size//2 + 1:size//2 + 1, -size//2 + 1:size//2 + 1]
+    g = np.exp(-((x**2 + y**2)/(2.0*sigma**2)))
+    return g/g.sum()
+
+
+def batchGray2Colormap(gray_batch):
+    colormap = plt.get_cmap('viridis')
+    heatmap_batch = []
+    for i in range(gray_batch.shape[0]):
+        # quantize [-1,1] to {0,1}
+        gray_map = gray_batch[i, :, :, 0]
+        heatmap = (colormap(gray_map) * 2**16).astype(np.uint16)[:,:,:3]
+        heatmap_batch.append(heatmap/127.5-1.0)
+    return np.array(heatmap_batch)
+
+
+class PlotterThread():
+    '''log tensorboard data in a background thread to save time'''
+    def __init__(self, writer):
+        self.writer = writer
+        self.task_queue = Queue(maxsize=0)
+        worker = Thread(target=self.do_work, args=(self.task_queue,))
+        worker.setDaemon(True)
+        worker.start()
+
+    def do_work(self, q):
+        while True:
+            content = q.get()
+            if content[-1] == 'image':
+                self.writer.add_image(*content[:-1])
+            elif content[-1] == 'scalar':
+                self.writer.add_scalar(*content[:-1])
+            else:
+                raise ValueError
+            q.task_done()
+
+    def add_data(self, name, value, step, data_type='scalar'):
+        self.task_queue.put([name, value, step, data_type])
+
+    def __len__(self):
+        return self.task_queue.qsize()
+
+
+def save_images_from_batch(img_batch, save_dir, filename_list, batch_no=-1, suffix=None):
+    N,H,W,C = img_batch.shape
+    if C == 3:
+        #! rgb color image
+        for i in range(N):
+            # [-1,1] >>> [0,255]
+            image = Image.fromarray((127.5*(img_batch[i,:,:,:]+1.)).astype(np.uint8))
+            save_name = filename_list[i] if batch_no==-1 else '%05d.png' % (batch_no*N+i)
+            save_name = save_name.replace('.png', '-%s.png'%suffix) if suffix else save_name
+            image.save(os.path.join(save_dir, save_name), 'PNG')
+    elif C == 1:
+        #! single-channel gray image
+        for i in range(N):
+            # [-1,1] >>> [0,255]
+            image = Image.fromarray((127.5*(img_batch[i,:,:,0]+1.)).astype(np.uint8))
+            save_name = filename_list[i] if batch_no==-1 else '%05d.png' % (batch_no*img_batch.shape[0]+i)
+            save_name = save_name.replace('.png', '-%s.png'%suffix) if suffix else save_name
+            image.save(os.path.join(save_dir, save_name), 'PNG')
+    else:
+        #! multi-channel: save each channel as a single image
+        for i in range(N):
+            # [-1,1] >>> [0,255]
+            for j in range(C):
+                image = Image.fromarray((127.5*(img_batch[i,:,:,j]+1.)).astype(np.uint8))
+                if batch_no == -1:
+                    _, file_name = os.path.split(filename_list[i])
+                    name_only, _ = os.path.os.path.splitext(file_name)
+                    save_name = name_only + '_c%d.png' % j
+                else:
+                    save_name = '%05d_c%d.png' % (batch_no*N+i, j)
+                save_name = save_name.replace('.png', '-%s.png'%suffix) if suffix else save_name
+                image.save(os.path.join(save_dir, save_name), 'PNG')
+    return None
+
+
+def save_normLabs_from_batch(img_batch, save_dir, filename_list, batch_no=-1, suffix=None):
+    N,H,W,C = img_batch.shape
+    if C != 3:
+        print('@Warning:the Lab images are NOT in 3 channels!')
+        return None
+    # denormalization: L: (L+1.0)*50.0 | a: a*110.0| b: b*110.0
+    img_batch[:,:,:,0] = img_batch[:,:,:,0] * 50.0 + 50.0
+    img_batch[:,:,:,1:3] = img_batch[:,:,:,1:3] * 110.0
+    #! convert into RGB color image
+    for i in range(N):
+        rgb_img = cv2.cvtColor(img_batch[i,:,:,:], cv2.COLOR_LAB2RGB)
+        image = Image.fromarray((rgb_img*255.0).astype(np.uint8))
+        save_name = filename_list[i] if batch_no==-1 else '%05d.png' % (batch_no*N+i)
+        save_name = save_name.replace('.png', '-%s.png'%suffix) if suffix else save_name
+        image.save(os.path.join(save_dir, save_name), 'PNG')
+    return None
+
+
+def save_markedSP_from_batch(img_batch, spix_batch, save_dir, filename_list, batch_no=-1, suffix=None):
+    N,H,W,C = img_batch.shape
+    #! img_batch: BGR nd-array (range:0~1)
+    #! map_batch: single-channel spixel map
+    #print('----------', img_batch.shape, spix_batch.shape)
+    for i in range(N):
+        norm_image = img_batch[i,:,:,:]*0.5+0.5
+        spixel_bd_image = mark_boundaries(norm_image, spix_batch[i,:,:,0].astype(int), color=(1,1,1))
+        #spixel_bd_image = cv2.cvtColor(spixel_bd_image, cv2.COLOR_BGR2RGB)
+        image = Image.fromarray((spixel_bd_image*255.0).astype(np.uint8))
+        save_name = filename_list[i] if batch_no==-1 else '%05d.png' % (batch_no*N+i)
+        save_name = save_name.replace('.png', '-%s.png'%suffix) if suffix else save_name
+        image.save(os.path.join(save_dir, save_name), 'PNG')
+    return None
+
+
+def get_filelist(data_dir):
+    file_list = glob.glob(os.path.join(data_dir, '*.*'))
+    file_list.sort()
+    return file_list
+    
+
+def collect_filenames(data_dir):
+    file_list = get_filelist(data_dir)
+    name_list = []
+    for file_path in file_list:
+        _, file_name = os.path.split(file_path)
+        name_list.append(file_name)
+    name_list.sort()
+    return name_list
+
+
+def exists_or_mkdir(path, need_remove=False):
+    if not os.path.exists(path):
+        os.makedirs(path)
+    elif need_remove:
+        shutil.rmtree(path)
+        os.makedirs(path)
+    return None
+
+
+def save_list(save_path, data_list, append_mode=False):
+    n = len(data_list)
+    if append_mode:
+        with open(save_path, 'a') as f:
+            f.writelines([str(data_list[i]) + '\n' for i in range(n-1,n)])
+    else:
+        with open(save_path, 'w') as f:
+            f.writelines([str(data_list[i]) + '\n' for i in range(n)])
+    return None
+    
+    
+def save_dict(save_path, dict):
+    json.dumps(dict, open(save_path,"w"))
+    return None
+
+
+if __name__ == '__main__':
+    data_dir = '../PolyNet/PolyNet/cache/'
+    #visualizeLossCurves(data_dir)
+    clbar = GamutIndex()
+    ab, ab_gamut_mask = clbar._get_gamut_mask()
+    ab2q = clbar._get_ab_to_q(ab_gamut_mask)
+    q2ab = clbar._get_q_to_ab(ab, ab_gamut_mask)
+    maps = ab_gamut_mask*255.0
+    image = Image.fromarray(maps.astype(np.uint8))
+    image.save('gamut.png', 'PNG')
+    print(ab2q.shape)
+    print(q2ab.shape)
+    print('label range:', np.min(ab2q), np.max(ab2q))