utils/util.py

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))