Upload models/basic.py

models/basic.py

@@ -0,0 +1,504 @@
+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)