-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathhist_loss.py
More file actions
176 lines (163 loc) · 6.5 KB
/
hist_loss.py
File metadata and controls
176 lines (163 loc) · 6.5 KB
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
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
"""
Copyright 2021 Mahmoud Afifi.
Mahmoud Afifi, Marcus A. Brubaker, and Michael S. Brown. "HistoGAN:
Controlling Colors of GAN-Generated and Real Images via Color Histograms."
In CVPR, 2021.
@inproceedings{afifi2021histogan,
title={Histo{GAN}: Controlling Colors of {GAN}-Generated and Real Images via
Color Histograms},
author={Afifi, Mahmoud and Brubaker, Marcus A. and Brown, Michael S.},
booktitle={CVPR},
year={2021}
}
"""
import torch
import torch.nn as nn
from PIL import Image
import matplotlib.pyplot as plt
import torch.nn.functional as F
import torchvision.transforms as transforms
import numpy as np
EPS = 1e-6
class RGBuvHistBlock(nn.Module):
def __init__(self, h=64, insz=150, resizing='interpolation',
method='inverse-quadratic', sigma=0.02, intensity_scale=True, device='cuda'):
super(RGBuvHistBlock, self).__init__()
self.h = h
self.insz = insz
self.device = device
self.resizing = resizing
self.method = method
self.intensity_scale = intensity_scale
if self.method == 'thresholding':
self.eps = 6.0 / h
else:
self.sigma = sigma
def forward(self, x):
x = torch.clamp(x, 0, 1)
if x.shape[2] > self.insz or x.shape[3] > self.insz:
if self.resizing == 'interpolation':
x_sampled = F.interpolate(x, size=(self.insz, self.insz),
mode='bilinear', align_corners=False)
elif self.resizing == 'sampling':
inds_1 = torch.LongTensor(
np.linspace(0, x.shape[2], self.h, endpoint=False)).cuda()
inds_2 = torch.LongTensor(
np.linspace(0, x.shape[3], self.h, endpoint=False)).cuda()
x_sampled = x.index_select(2, inds_1)
x_sampled = x_sampled.index_select(3, inds_2)
# else:
# raise Exception(
# f'Wrong resizing method. It should be: interpolation or sampling. '
# f'But the given value is {self.resizing}.')
else:
x_sampled = x
L = x_sampled.shape[0] # size of mini-batch
if x_sampled.shape[1] > 3:
x_sampled = x_sampled[:, :3, :, :]
X = torch.unbind(x_sampled, dim=0)
hists = torch.zeros((x_sampled.shape[0], 3, self.h, self.h)).cuda()
for l in range(L):
I = torch.t(torch.reshape(X[l], (3, -1)))
II = torch.pow(I, 2)
if self.intensity_scale:
Iy = torch.unsqueeze(torch.sqrt(II[:, 0] + II[:, 1] + II[:, 2] + EPS), dim=1)
else:
Iy = 1
Iu0 = torch.unsqueeze(torch.log(I[:, 0] + EPS) - torch.log(I[:, 1] + EPS),
dim=1)
Iv0 = torch.unsqueeze(torch.log(I[:, 0] + EPS) - torch.log(I[:, 2] + EPS),
dim=1)
diff_u0 = abs(Iu0 - torch.unsqueeze(torch.tensor(np.linspace(-3, 3, num=self.h)), dim=0).cuda())
diff_v0 = abs(Iv0 - torch.unsqueeze(torch.tensor(np.linspace(-3, 3, num=self.h)), dim=0).cuda())
if self.method == 'thresholding':
diff_u0 = torch.reshape(diff_u0, (-1, self.h)) <= self.eps / 2
diff_v0 = torch.reshape(diff_v0, (-1, self.h)) <= self.eps / 2
elif self.method == 'RBF':
diff_u0 = torch.pow(torch.reshape(diff_u0, (-1, self.h)),
2) / self.sigma ** 2
diff_v0 = torch.pow(torch.reshape(diff_v0, (-1, self.h)),
2) / self.sigma ** 2
diff_u0 = torch.exp(-diff_u0) # Radial basis function
diff_v0 = torch.exp(-diff_v0)
elif self.method == 'inverse-quadratic':
diff_u0 = torch.pow(torch.reshape(diff_u0, (-1, self.h)),
2) / self.sigma ** 2
diff_v0 = torch.pow(torch.reshape(diff_v0, (-1, self.h)),
2) / self.sigma ** 2
diff_u0 = 1 / (1 + diff_u0) # Inverse quadratic
diff_v0 = 1 / (1 + diff_v0)
# else:
# raise Exception(
# f'Wrong kernel method. It should be either thresholding, RBF,'
# f' inverse-quadratic. But the given value is {self.method}.')
diff_u0 = diff_u0.type(torch.float32)
diff_v0 = diff_v0.type(torch.float32)
a = torch.t(Iy * diff_u0)
hists[l, 0, :, :] = torch.mm(a, diff_v0)
Iu1 = torch.unsqueeze(torch.log(I[:, 1] + EPS) - torch.log(I[:, 0] + EPS),
dim=1)
Iv1 = torch.unsqueeze(torch.log(I[:, 1] + EPS) - torch.log(I[:, 2] + EPS),
dim=1)
diff_u1 = abs(
Iu1 - torch.unsqueeze(torch.tensor(np.linspace(-3, 3, num=self.h)),
dim=0).cuda())
diff_v1 = abs(
Iv1 - torch.unsqueeze(torch.tensor(np.linspace(-3, 3, num=self.h)),
dim=0).cuda())
if self.method == 'thresholding':
diff_u1 = torch.reshape(diff_u1, (-1, self.h)) <= self.eps / 2
diff_v1 = torch.reshape(diff_v1, (-1, self.h)) <= self.eps / 2
elif self.method == 'RBF':
diff_u1 = torch.pow(torch.reshape(diff_u1, (-1, self.h)),
2) / self.sigma ** 2
diff_v1 = torch.pow(torch.reshape(diff_v1, (-1, self.h)),
2) / self.sigma ** 2
diff_u1 = torch.exp(-diff_u1) # Gaussian
diff_v1 = torch.exp(-diff_v1)
elif self.method == 'inverse-quadratic':
diff_u1 = torch.pow(torch.reshape(diff_u1, (-1, self.h)),
2) / self.sigma ** 2
diff_v1 = torch.pow(torch.reshape(diff_v1, (-1, self.h)),
2) / self.sigma ** 2
diff_u1 = 1 / (1 + diff_u1) # Inverse quadratic
diff_v1 = 1 / (1 + diff_v1)
diff_u1 = diff_u1.type(torch.float32)
diff_v1 = diff_v1.type(torch.float32)
a = torch.t(Iy * diff_u1)
hists[l, 1, :, :] = torch.mm(a, diff_v1)
Iu2 = torch.unsqueeze(torch.log(I[:, 2] + EPS) - torch.log(I[:, 0] + EPS),
dim=1)
Iv2 = torch.unsqueeze(torch.log(I[:, 2] + EPS) - torch.log(I[:, 1] + EPS),
dim=1)
diff_u2 = abs(
Iu2 - torch.unsqueeze(torch.tensor(np.linspace(-3, 3, num=self.h)),
dim=0).cuda())
diff_v2 = abs(
Iv2 - torch.unsqueeze(torch.tensor(np.linspace(-3, 3, num=self.h)),
dim=0).cuda())
if self.method == 'thresholding':
diff_u2 = torch.reshape(diff_u2, (-1, self.h)) <= self.eps / 2
diff_v2 = torch.reshape(diff_v2, (-1, self.h)) <= self.eps / 2
elif self.method == 'RBF':
diff_u2 = torch.pow(torch.reshape(diff_u2, (-1, self.h)),
2) / self.sigma ** 2
diff_v2 = torch.pow(torch.reshape(diff_v2, (-1, self.h)),
2) / self.sigma ** 2
diff_u2 = torch.exp(-diff_u2) # Gaussian
diff_v2 = torch.exp(-diff_v2)
elif self.method == 'inverse-quadratic':
diff_u2 = torch.pow(torch.reshape(diff_u2, (-1, self.h)),
2) / self.sigma ** 2
diff_v2 = torch.pow(torch.reshape(diff_v2, (-1, self.h)),
2) / self.sigma ** 2
diff_u2 = 1 / (1 + diff_u2) # Inverse quadratic
diff_v2 = 1 / (1 + diff_v2)
diff_u2 = diff_u2.type(torch.float32)
diff_v2 = diff_v2.type(torch.float32)
a = torch.t(Iy * diff_u2)
hists[l, 2, :, :] = torch.mm(a, diff_v2)
# normalization
hists_normalized = hists / (
((hists.sum(dim=1)).sum(dim=1)).sum(dim=1).view(-1, 1, 1, 1) + EPS)
return hists_normalized