-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathutil.py
132 lines (102 loc) · 3.7 KB
/
util.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
"""
By Richard Zhang
"""
from __future__ import print_function
import torch
import numpy as np
from PIL import Image
import os
from collections import OrderedDict
l_cent = 50.
l_norm = 100.
ab_norm = 110.
# Color conversion code
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)
# if(torch.sum(torch.isnan(out))>0):
# print('rgb2xyz')
# embed()
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)
rgb = torch.max(rgb,torch.zeros_like(rgb)) # sometimes reaches a small negative number, which causes NaNs
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)
# if(torch.sum(torch.isnan(rgb))>0):
# print('xyz2rgb')
# embed()
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)
# if(torch.sum(torch.isnan(out))>0):
# print('xyz2lab')
# embed()
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
# if(torch.sum(torch.isnan(out))>0):
# print('lab2xyz')
# embed()
return out
def rgb2lab(rgb):
lab = xyz2lab(rgb2xyz(rgb))
l_rs = (lab[:,[0],:,:]-l_cent)/l_norm
ab_rs = lab[:,1:,:,:]/ab_norm
out = torch.cat((l_rs,ab_rs),dim=1)
# if(torch.sum(torch.isnan(out))>0):
# print('rgb2lab')
# embed()
return out
def lab2rgb(lab_rs):
l = lab_rs[:,[0],:,:]*l_norm + l_cent
ab = lab_rs[:,1:,:,:]*ab_norm
lab = torch.cat((l,ab),dim=1)
out = xyz2rgb(lab2xyz(lab))
# if(torch.sum(torch.isnan(out))>0):
# print('lab2rgb')
# embed()
return out