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)