models/basic.py

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.nn.utils.spectral_norm as spectral_norm
from torch.autograd import Function
from utils import util, cielab
import cv2, math, random

def tensor2array(tensors):
    arrays = tensors.detach().to("cpu").numpy()
    return np.transpose(arrays, (0, 2, 3, 1))


def rgb2gray(color_batch):
    #! gray = 0.299*R+0.587*G+0.114*B
    gray_batch = color_batch[:, 0, ...] * 0.299 + color_batch[:, 1, ...] * 0.587 + color_batch[:, 2, ...] * 0.114
    gray_batch = gray_batch.unsqueeze_(1)
    return gray_batch


def getParamsAmount(model):
    params = list(model.parameters())
    count = 0
    for var in params:
        l = 1
        for j in var.size():
            l *= j
        count += l
    return count


def checkAverageGradient(model):
    meanGrad, cnt = 0.0, 0
    for name, parms in model.named_parameters():
        if parms.requires_grad:
            meanGrad += torch.mean(torch.abs(parms.grad))
            cnt += 1
    return meanGrad.item() / cnt


def get_random_mask(N, H, W, minNum, maxNum):
    binary_maps = np.zeros((N, H*W), np.float32)
    for i in range(N):
        locs = random.sample(range(0, H*W), random.randint(minNum,maxNum))
        binary_maps[i, locs] = 1
    return binary_maps.reshape(N,1,H,W)


def io_user_control(hint_mask, spix_colors, output=True):
    cache_dir = '/apdcephfs/private_richardxia'
    if output:
        print('--- data saving')
        mask_imgs = tensor2array(hint_mask) * 2.0 - 1.0
        util.save_images_from_batch(mask_imgs, cache_dir, ['mask.png'], -1)
        fake_gray = torch.zeros_like(spix_colors[:,[0],:,:])
        spix_labs = torch.cat((fake_gray,spix_colors), dim=1)
        spix_imgs = tensor2array(spix_labs)
        util.save_normLabs_from_batch(spix_imgs, cache_dir, ['color.png'], -1)
        return hint_mask, spix_colors
    else:
        print('--- data loading')
        mask_img = cv2.imread(cache_dir+'/mask.png', cv2.IMREAD_GRAYSCALE)
        mask_img = np.expand_dims(mask_img, axis=2) / 255.
        hint_mask = torch.from_numpy(mask_img.transpose((2, 0, 1)))
        hint_mask = hint_mask.unsqueeze(0).cuda()
        bgr_img = cv2.imread(cache_dir+'/color.png', cv2.IMREAD_COLOR)
        rgb_img = cv2.cvtColor(bgr_img, cv2.COLOR_BGR2RGB)
        rgb_img = np.array(rgb_img / 255., np.float32)
        lab_img = cv2.cvtColor(rgb_img, cv2.COLOR_RGB2LAB)
        lab_img = torch.from_numpy(lab_img.transpose((2, 0, 1)))
        ab_chans = lab_img[1:3,:,:] / 110.
        spix_colors = ab_chans.unsqueeze(0).cuda()
        return hint_mask.float(), spix_colors.float()


class Quantize(Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        y = x.round()
        return y

    @staticmethod
    def backward(ctx, grad_output):
        """
        In the backward pass we receive a Tensor containing the gradient of the loss
        with respect to the output, and we need to compute the gradient of the loss
        with respect to the input.
        """
        inputX = ctx.saved_tensors
        return grad_output


def mark_color_hints(input_grays, target_ABs, gate_maps, kernel_size=3, base_ABs=None):
    ## to highlight the seeds with 1-pixel margin
    binary_map = torch.where(gate_maps>0.7, torch.ones_like(gate_maps), torch.zeros_like(gate_maps))
    center_mask = dilate_seeds(binary_map, kernel_size=kernel_size)
    margin_mask = dilate_seeds(binary_map, kernel_size=kernel_size+2) - center_mask
    ## drop colors
    dilated_seeds = dilate_seeds(gate_maps, kernel_size=kernel_size+2)
    marked_grays = torch.where(margin_mask > 1e-5, torch.ones_like(gate_maps), input_grays)
    if base_ABs is None:
        marked_ABs = torch.where(center_mask < 1e-5, torch.zeros_like(target_ABs), target_ABs)
    else:
        marked_ABs = torch.where(margin_mask > 1e-5, torch.zeros_like(base_ABs), base_ABs)
        marked_ABs = torch.where(center_mask > 1e-5, target_ABs, marked_ABs)
    return torch.cat((marked_grays,marked_ABs), dim=1)

def dilate_seeds(gate_maps, kernel_size=3):
    N,C,H,W = gate_maps.shape
    input_unf = F.unfold(gate_maps, kernel_size, padding=kernel_size//2)
    #! Notice: differentiable? just like max pooling?
    dilated_seeds, _ = torch.max(input_unf, dim=1, keepdim=True)
    output = F.fold(dilated_seeds, output_size=(H,W), kernel_size=1)
    #print('-------', input_unf.shape)
    return output


class RebalanceLoss(Function):
    @staticmethod
    def forward(ctx, data_input, weights):
        ctx.save_for_backward(weights)
        return data_input.clone()

    @staticmethod
    def backward(ctx, grad_output):
        weights, = ctx.saved_tensors
        # reweigh gradient pixelwise so that rare colors get a chance to
        # contribute
        grad_input = grad_output * weights
        # second return value is None since we are not interested in the
        # gradient with respect to the weights
        return grad_input, None


class GetClassWeights:
    def __init__(self, cielab, lambda_=0.5, device='cuda'):
        prior = torch.from_numpy(cielab.gamut.prior).cuda()
        uniform = torch.zeros_like(prior)
        uniform[prior > 0] = 1 / (prior > 0).sum().type_as(uniform)
        self.weights = 1 / ((1 - lambda_) * prior + lambda_ * uniform)
        self.weights /= torch.sum(prior * self.weights)

    def __call__(self, ab_actual):
        return self.weights[ab_actual.argmax(dim=1, keepdim=True)]


class ColorLabel:
    def __init__(self, lambda_=0.5, device='cuda'):
        self.cielab = cielab.CIELAB()
        self.q_to_ab = torch.from_numpy(self.cielab.q_to_ab).to(device)
        prior = torch.from_numpy(self.cielab.gamut.prior).to(device)
        uniform = torch.zeros_like(prior)
        uniform[prior>0] = 1 / (prior>0).sum().type_as(uniform)
        self.weights = 1 / ((1-lambda_) * prior + lambda_ * uniform)
        self.weights /= torch.sum(prior * self.weights)

    def visualize_label(self, step=3):
        height, width = 200, 313*step
        label_lab = np.ones((height,width,3), np.float32)
        for x in range(313):
            ab = self.cielab.q_to_ab[x,:]
            label_lab[:,step*x:step*(x+1),1:] = ab / 110.
        label_lab[:,:,0] = np.zeros((height,width), np.float32)
        return label_lab

    @staticmethod
    def _gauss_eval(x, mu, sigma):
        norm = 1 / (2 * math.pi * sigma)
        return norm * torch.exp(-torch.sum((x - mu)**2, dim=0) / (2 * sigma**2))
    
    def get_classweights(self, batch_gt_indx):
        #return self.weights[batch_gt_q.argmax(dim=1, keepdim=True)]
        return self.weights[batch_gt_indx]

    def encode_ab2ind(self, batch_ab, neighbours=5, sigma=5.0):
        batch_ab = batch_ab * 110.
        n, _, h, w = batch_ab.shape
        m = n * h * w
        # find nearest neighbours
        ab_ = batch_ab.permute(1, 0, 2, 3).reshape(2, -1) # (2, n*h*w) 
        cdist = torch.cdist(self.q_to_ab, ab_.t())
        nns = cdist.argsort(dim=0)[:neighbours, :]
        # gaussian weighting
        nn_gauss = batch_ab.new_zeros(neighbours, m)
        for i in range(neighbours):
            nn_gauss[i, :] = self._gauss_eval(self.q_to_ab[nns[i, :], :].t(), ab_, sigma)
        nn_gauss /= nn_gauss.sum(dim=0, keepdim=True)
        # expand
        bins = self.cielab.gamut.EXPECTED_SIZE
        q = batch_ab.new_zeros(bins, m)
        q[nns, torch.arange(m).repeat(neighbours, 1)] = nn_gauss        
        return q.reshape(bins, n, h, w).permute(1, 0, 2, 3)

    def decode_ind2ab(self, batch_q, T=0.38):
        _, _, h, w = batch_q.shape
        batch_q = F.softmax(batch_q, dim=1)
        if T%1 == 0:
            # take the T-st probable index
            sorted_probs, batch_indexs = torch.sort(batch_q, dim=1, descending=True)
            #print('checking [index]', batch_indexs[:,0:5,5,5])
            #print('checking [probs]', sorted_probs[:,0:5,5,5])
            batch_indexs = batch_indexs[:,T:T+1,:,:]
            #batch_indexs = torch.where(sorted_probs[:,T:T+1,:,:] > 0.25, batch_indexs[:,T:T+1,:,:], batch_indexs[:,0:1,:,:])
            ab = torch.stack([
                self.q_to_ab.index_select(0, q_i.flatten()).reshape(h,w,2).permute(2,0,1)
                for q_i in batch_indexs])
        else:
            batch_q = torch.exp(batch_q / T)
            batch_q /= batch_q.sum(dim=1, keepdim=True)
            a = torch.tensordot(batch_q, self.q_to_ab[:,0], dims=((1,), (0,)))
            a = a.unsqueeze(dim=1)
            b = torch.tensordot(batch_q, self.q_to_ab[:,1], dims=((1,), (0,)))
            b = b.unsqueeze(dim=1)
            ab = torch.cat((a, b), dim=1)
        ab = ab / 110.
        return ab.type(batch_q.dtype)


def init_spixel_grid(img_height, img_width, spixel_size=16):
    # get spixel id for the final assignment
    n_spixl_h = int(np.floor(img_height/spixel_size))
    n_spixl_w = int(np.floor(img_width/spixel_size))
    spixel_height = int(img_height / (1. * n_spixl_h))
    spixel_width = int(img_width / (1. * n_spixl_w))
    spix_values = np.int32(np.arange(0, n_spixl_w * n_spixl_h).reshape((n_spixl_h, n_spixl_w)))
    
    def shift9pos(input, h_shift_unit=1, w_shift_unit=1):
        # input should be padding as (c, 1+ height+1, 1+width+1)
        input_pd = np.pad(input, ((h_shift_unit, h_shift_unit), (w_shift_unit, w_shift_unit)), mode='edge')
        input_pd = np.expand_dims(input_pd, axis=0)
        # assign to ...
        top     = input_pd[:, :-2 * h_shift_unit,          w_shift_unit:-w_shift_unit]
        bottom  = input_pd[:, 2 * h_shift_unit:,           w_shift_unit:-w_shift_unit]
        left    = input_pd[:, h_shift_unit:-h_shift_unit,  :-2 * w_shift_unit]
        right   = input_pd[:, h_shift_unit:-h_shift_unit,  2 * w_shift_unit:]
        center = input_pd[:,h_shift_unit:-h_shift_unit,w_shift_unit:-w_shift_unit]
        bottom_right    = input_pd[:, 2 * h_shift_unit:,   2 * w_shift_unit:]
        bottom_left     = input_pd[:, 2 * h_shift_unit:,   :-2 * w_shift_unit]
        top_right       = input_pd[:, :-2 * h_shift_unit,  2 * w_shift_unit:]
        top_left        = input_pd[:, :-2 * h_shift_unit,  :-2 * w_shift_unit]
        shift_tensor = np.concatenate([     top_left,    top,      top_right,
                                            left,        center,      right,
                                            bottom_left, bottom,    bottom_right], axis=0)
        return shift_tensor

    spix_idx_tensor_ = shift9pos(spix_values)
    spix_idx_tensor = np.repeat(
        np.repeat(spix_idx_tensor_, spixel_height, axis=1), spixel_width, axis=2)
    spixel_id_tensor = torch.from_numpy(spix_idx_tensor).type(torch.float)

    #! pixel coord feature maps
    all_h_coords = np.arange(0, img_height, 1)
    all_w_coords = np.arange(0, img_width, 1)
    curr_pxl_coord = np.array(np.meshgrid(all_h_coords, all_w_coords, indexing='ij'))
    coord_feat_tensor = np.concatenate([curr_pxl_coord[1:2, :, :], curr_pxl_coord[:1, :, :]])
    coord_feat_tensor = torch.from_numpy(coord_feat_tensor).type(torch.float)

    return spixel_id_tensor, coord_feat_tensor


def split_spixels(assign_map, spixel_ids):
    N,C,H,W = assign_map.shape
    spixel_id_map = spixel_ids.expand(N,-1,-1,-1)
    assig_max,_ = torch.max(assign_map, dim=1, keepdim=True)
    assignment_ = torch.where(assign_map == assig_max, torch.ones(assign_map.shape).cuda(),torch.zeros(assign_map.shape).cuda())
    ## winner take all
    new_spixl_map_ = spixel_id_map * assignment_
    new_spixl_map = torch.sum(new_spixl_map_,dim=1,keepdim=True).type(torch.int)
    return new_spixl_map


def poolfeat(input, prob, sp_h=2, sp_w=2, need_entry_prob=False):
    def feat_prob_sum(feat_sum, prob_sum, shift_feat):
        feat_sum += shift_feat[:, :-1, :, :]
        prob_sum += shift_feat[:, -1:, :, :]
        return feat_sum, prob_sum
    
    b, _, h, w = input.shape
    h_shift_unit = 1
    w_shift_unit = 1
    p2d = (w_shift_unit, w_shift_unit, h_shift_unit, h_shift_unit)
    feat_ = torch.cat([input, torch.ones([b, 1, h, w], device=input.device)], dim=1)  # b* (n+1) *h*w
    prob_feat = F.avg_pool2d(feat_ * prob.narrow(1, 0, 1), kernel_size=(sp_h, sp_w), stride=(sp_h, sp_w)) # b * (n+1) * h* w
    send_to_top_left =  F.pad(prob_feat, p2d, mode='constant', value=0)[:, :, 2 * h_shift_unit:, 2 * w_shift_unit:]
    feat_sum = send_to_top_left[:, :-1, :, :].clone()
    prob_sum = send_to_top_left[:, -1:, :, :].clone()

    prob_feat = F.avg_pool2d(feat_ * prob.narrow(1, 1, 1), kernel_size=(sp_h, sp_w), stride=(sp_h, sp_w))  # b * (n+1) * h* w
    top = F.pad(prob_feat, p2d, mode='constant', value=0)[:,  :, 2 * h_shift_unit:, w_shift_unit:-w_shift_unit]
    feat_sum, prob_sum = feat_prob_sum(feat_sum, prob_sum, top)

    prob_feat = F.avg_pool2d(feat_ * prob.narrow(1, 2, 1), kernel_size=(sp_h, sp_w), stride=(sp_h, sp_w))  # b * (n+1) * h* w
    top_right = F.pad(prob_feat, p2d, mode='constant', value=0)[:,  :, 2 * h_shift_unit:, :-2 * w_shift_unit]
    feat_sum, prob_sum = feat_prob_sum(feat_sum, prob_sum, top_right)

    prob_feat = F.avg_pool2d(feat_ * prob.narrow(1, 3, 1), kernel_size=(sp_h, sp_w), stride=(sp_h, sp_w))  # b * (n+1) * h* w
    left = F.pad(prob_feat, p2d, mode='constant', value=0)[:,  :, h_shift_unit:-h_shift_unit, 2 * w_shift_unit:]
    feat_sum, prob_sum = feat_prob_sum(feat_sum, prob_sum, left)

    prob_feat = F.avg_pool2d(feat_ * prob.narrow(1, 4, 1), kernel_size=(sp_h, sp_w), stride=(sp_h, sp_w))  # b * (n+1) * h* w
    center = F.pad(prob_feat, p2d, mode='constant', value=0)[:, :, h_shift_unit:-h_shift_unit, w_shift_unit:-w_shift_unit]
    feat_sum, prob_sum = feat_prob_sum(feat_sum, prob_sum, center)

    prob_feat = F.avg_pool2d(feat_ * prob.narrow(1, 5, 1), kernel_size=(sp_h, sp_w), stride=(sp_h, sp_w))  # b * (n+1) * h* w
    right = F.pad(prob_feat, p2d, mode='constant', value=0)[:,  :, h_shift_unit:-h_shift_unit, :-2 * w_shift_unit]
    feat_sum, prob_sum = feat_prob_sum(feat_sum, prob_sum, right)

    prob_feat = F.avg_pool2d(feat_ * prob.narrow(1, 6, 1), kernel_size=(sp_h, sp_w), stride=(sp_h, sp_w))  # b * (n+1) * h* w
    bottom_left = F.pad(prob_feat, p2d, mode='constant', value=0)[:,  :, :-2 * h_shift_unit, 2 * w_shift_unit:]
    feat_sum, prob_sum = feat_prob_sum(feat_sum, prob_sum, bottom_left)

    prob_feat = F.avg_pool2d(feat_ * prob.narrow(1, 7, 1), kernel_size=(sp_h, sp_w), stride=(sp_h, sp_w))  # b * (n+1) * h* w
    bottom = F.pad(prob_feat, p2d, mode='constant', value=0)[:, :, :-2 * h_shift_unit, w_shift_unit:-w_shift_unit]
    feat_sum, prob_sum = feat_prob_sum(feat_sum, prob_sum, bottom)

    prob_feat = F.avg_pool2d(feat_ * prob.narrow(1, 8, 1), kernel_size=(sp_h, sp_w), stride=(sp_h, sp_w))  # b * (n+1) * h* w
    bottom_right = F.pad(prob_feat, p2d, mode='constant', value=0)[:, :, :-2 * h_shift_unit, :-2 * w_shift_unit]
    feat_sum, prob_sum = feat_prob_sum(feat_sum, prob_sum, bottom_right)
    pooled_feat = feat_sum / (prob_sum + 1e-8)
    if need_entry_prob:
        return pooled_feat, prob_sum
    return pooled_feat


def get_spixel_size(affinity_map, sp_h=2, sp_w=2, elem_thres=25):
    N,C,H,W = affinity_map.shape
    device = affinity_map.device
    assign_max,_ = torch.max(affinity_map, dim=1, keepdim=True)
    assign_map = torch.where(affinity_map==assign_max, torch.ones(affinity_map.shape, device=device), torch.zeros(affinity_map.shape, device=device))
    ## one_map = (N,1,H,W)
    _, elem_num_maps = poolfeat(torch.ones(assign_max.shape, device=device), assign_map, sp_h, sp_w, True)
    #all_one_map = torch.ones(elem_num_maps.shape).cuda()
    #empty_mask = torch.where(elem_num_maps < elem_thres/256, all_one_map, 1-all_one_map)
    return elem_num_maps
    

def upfeat(input, prob, up_h=2, up_w=2):
    # input b*n*H*W  downsampled
    # prob b*9*h*w
    b, c, h, w = input.shape

    h_shift = 1
    w_shift = 1

    p2d = (w_shift, w_shift, h_shift, h_shift)
    feat_pd = F.pad(input, p2d, mode='constant', value=0)

    gt_frm_top_left = F.interpolate(feat_pd[:, :, :-2 * h_shift, :-2 * w_shift], size=(h * up_h, w * up_w),mode='nearest')
    feat_sum = gt_frm_top_left * prob.narrow(1,0,1)

    top = F.interpolate(feat_pd[:, :, :-2 * h_shift, w_shift:-w_shift], size=(h * up_h, w * up_w), mode='nearest')
    feat_sum += top * prob.narrow(1, 1, 1)

    top_right = F.interpolate(feat_pd[:, :, :-2 * h_shift, 2 * w_shift:], size=(h * up_h, w * up_w), mode='nearest')
    feat_sum += top_right * prob.narrow(1,2,1)

    left = F.interpolate(feat_pd[:, :, h_shift:-w_shift, :-2 * w_shift], size=(h * up_h, w * up_w), mode='nearest')
    feat_sum += left * prob.narrow(1, 3, 1)

    center = F.interpolate(input, (h * up_h, w * up_w), mode='nearest')
    feat_sum += center * prob.narrow(1, 4, 1)

    right = F.interpolate(feat_pd[:, :, h_shift:-w_shift, 2 * w_shift:], size=(h * up_h, w * up_w), mode='nearest')
    feat_sum += right * prob.narrow(1, 5, 1)

    bottom_left = F.interpolate(feat_pd[:, :, 2 * h_shift:, :-2 * w_shift], size=(h * up_h, w * up_w), mode='nearest')
    feat_sum += bottom_left * prob.narrow(1, 6, 1)

    bottom = F.interpolate(feat_pd[:, :, 2 * h_shift:, w_shift:-w_shift], size=(h * up_h, w * up_w), mode='nearest')
    feat_sum += bottom * prob.narrow(1, 7, 1)

    bottom_right =  F.interpolate(feat_pd[:, :, 2 * h_shift:, 2 * w_shift:], size=(h * up_h, w * up_w), mode='nearest')
    feat_sum += bottom_right * prob.narrow(1, 8, 1)

    return feat_sum


def suck_and_spread(self, base_maps, seg_layers):
    N,S,H,W = seg_layers.shape
    base_maps = base_maps.unsqueeze(1)
    seg_layers = seg_layers.unsqueeze(2)
    ## (N,S,C,1,1) = (N,1,C,H,W) * (N,S,1,H,W)
    mean_val_layers = (base_maps * seg_layers).sum(dim=(3,4), keepdim=True) / (1e-5 + seg_layers.sum(dim=(3,4), keepdim=True))
    ## normalized to be sum one
    weight_layers = seg_layers / (1e-5 + torch.sum(seg_layers, dim=1, keepdim=True))
    ## (N,S,C,H,W) = (N,S,C,1,1) * (N,S,1,H,W)
    recon_maps = mean_val_layers * weight_layers
    return recon_maps.sum(dim=1)


#! copy from Richard Zhang [SIGGRAPH2017]
# RGB grid points maps to Lab range: L[0,100], a[-86.183,98,233], b[-107.857,94.478]
#------------------------------------------------------------------------------
def rgb2xyz(rgb):  # rgb from [0,1]
    # xyz_from_rgb = np.array([[0.412453, 0.357580, 0.180423],
        #  [0.212671, 0.715160, 0.072169],
        #  [0.019334, 0.119193, 0.950227]])
    mask = (rgb > .04045).type(torch.FloatTensor)
    if(rgb.is_cuda):
        mask = mask.cuda()
    rgb = (((rgb+.055)/1.055)**2.4)*mask + rgb/12.92*(1-mask)
    x = .412453*rgb[:,0,:,:]+.357580*rgb[:,1,:,:]+.180423*rgb[:,2,:,:]
    y = .212671*rgb[:,0,:,:]+.715160*rgb[:,1,:,:]+.072169*rgb[:,2,:,:]
    z = .019334*rgb[:,0,:,:]+.119193*rgb[:,1,:,:]+.950227*rgb[:,2,:,:]
    out = torch.cat((x[:,None,:,:],y[:,None,:,:],z[:,None,:,:]),dim=1)
    return out

def xyz2rgb(xyz):
    # array([[ 3.24048134, -1.53715152, -0.49853633],
    #        [-0.96925495,  1.87599   ,  0.04155593],
    #        [ 0.05564664, -0.20404134,  1.05731107]])
    r = 3.24048134*xyz[:,0,:,:]-1.53715152*xyz[:,1,:,:]-0.49853633*xyz[:,2,:,:]
    g = -0.96925495*xyz[:,0,:,:]+1.87599*xyz[:,1,:,:]+.04155593*xyz[:,2,:,:]
    b = .05564664*xyz[:,0,:,:]-.20404134*xyz[:,1,:,:]+1.05731107*xyz[:,2,:,:]
    rgb = torch.cat((r[:,None,:,:],g[:,None,:,:],b[:,None,:,:]),dim=1)
    #！ sometimes reaches a small negative number, which causes NaNs
    rgb = torch.max(rgb,torch.zeros_like(rgb))
    mask = (rgb > .0031308).type(torch.FloatTensor)
    if(rgb.is_cuda):
        mask = mask.cuda()
    rgb = (1.055*(rgb**(1./2.4)) - 0.055)*mask + 12.92*rgb*(1-mask)
    return rgb

def xyz2lab(xyz):
    # 0.95047, 1., 1.08883 # white
    sc = torch.Tensor((0.95047, 1., 1.08883))[None,:,None,None]
    if(xyz.is_cuda):
        sc = sc.cuda()
    xyz_scale = xyz/sc
    mask = (xyz_scale > .008856).type(torch.FloatTensor)
    if(xyz_scale.is_cuda):
        mask = mask.cuda()
    xyz_int = xyz_scale**(1/3.)*mask + (7.787*xyz_scale + 16./116.)*(1-mask)
    L = 116.*xyz_int[:,1,:,:]-16.
    a = 500.*(xyz_int[:,0,:,:]-xyz_int[:,1,:,:])
    b = 200.*(xyz_int[:,1,:,:]-xyz_int[:,2,:,:])
    out = torch.cat((L[:,None,:,:],a[:,None,:,:],b[:,None,:,:]),dim=1)
    return out

def lab2xyz(lab):
    y_int = (lab[:,0,:,:]+16.)/116.
    x_int = (lab[:,1,:,:]/500.) + y_int
    z_int = y_int - (lab[:,2,:,:]/200.)
    if(z_int.is_cuda):
        z_int = torch.max(torch.Tensor((0,)).cuda(), z_int)
    else:
        z_int = torch.max(torch.Tensor((0,)), z_int)
    out = torch.cat((x_int[:,None,:,:],y_int[:,None,:,:],z_int[:,None,:,:]),dim=1)
    mask = (out > .2068966).type(torch.FloatTensor)
    if(out.is_cuda):
        mask = mask.cuda()
    out = (out**3.)*mask + (out - 16./116.)/7.787*(1-mask)
    sc = torch.Tensor((0.95047, 1., 1.08883))[None,:,None,None]
    sc = sc.to(out.device)
    out = out*sc
    return out

def rgb2lab(rgb, l_mean=50, l_norm=50, ab_norm=110):
    #! input rgb: [0,1]
    #! output lab: [-1,1]
    lab = xyz2lab(rgb2xyz(rgb))
    l_rs = (lab[:,[0],:,:]-l_mean) / l_norm
    ab_rs = lab[:,1:,:,:] / ab_norm
    out = torch.cat((l_rs,ab_rs),dim=1)
    return out
    
def lab2rgb(lab_rs, l_mean=50, l_norm=50, ab_norm=110):
    #! input lab: [-1,1]
    #! output rgb: [0,1]
    l_ = lab_rs[:,[0],:,:] * l_norm + l_mean
    ab = lab_rs[:,1:,:,:] * ab_norm
    lab = torch.cat((l_,ab), dim=1)
    out = xyz2rgb(lab2xyz(lab))
    return out


if __name__ == '__main__':
    minL, minA, minB = 999., 999., 999.
    maxL, maxA, maxB = 0., 0., 0.
    for r in range(256):
        print('h',r)
        for g in range(256):
            for b in range(256):
                rgb = np.array([r,g,b], np.float32).reshape(1,1,-1) / 255.0
                #lab_img = cv2.cvtColor(rgb, cv2.COLOR_RGB2LAB)
                rgb = torch.from_numpy(rgb.transpose((2, 0, 1)))
                rgb = rgb.reshape(1,3,1,1)
                lab = rgb2lab(rgb)
                lab[:,[0],:,:] = lab[:,[0],:,:] * 50 + 50
                lab[:,1:,:,:] = lab[:,1:,:,:] * 110
                lab = lab.squeeze()
                lab_float = lab.numpy()
                #print('zhang vs. cv2:', lab_float, lab_img.squeeze())
                minL = min(lab_float[0], minL)
                minA = min(lab_float[1], minA)
                minB = min(lab_float[2], minB)
                maxL = max(lab_float[0], maxL)
                maxA = max(lab_float[1], maxA)
                maxB = max(lab_float[2], maxB)
    print('L:', minL, maxL)
    print('A:', minA, maxA)
    print('B:', minB, maxB)