def initialize(self, illuminants):
# init to zero if no illuminants are provided
if illuminants is None:
clusters = torch.FloatTensor(np.zeros((self.k, 3))).unsqueeze(0)
return clusters
else:
# convert to desired color space first
illuminants = np.array(illuminants)
if self.color_space == 'rg_bg':
illuminants_new = np.zeros((illuminants.shape[0], 2))
for i in range(illuminants.shape[0]):
illuminants_new[
i, 0] = illuminants[i, 0] / illuminants[i, 1] # r / g
illuminants_new[
i, 1] = illuminants[i, 2] / illuminants[i, 1] # b / g
illuminants = illuminants_new
# fit GMM
g = mixture.GMM(n_components=self.n_components)
g.fit(illuminants)
self._plot_gmm(g, illuminants)
# convert results to RGB
if self.color_space == 'rgb':
clusters = g.sample(self.k) # sample GMM
# normalize illuminants
for i in range(clusters.shape[0]):
clusters[i, :] = normalize_illuminant(clusters[i, :])
elif self.color_space == 'rg_bg':
clusters_rg_bg = g.sample(self.k) # sample GMM
clusters = np.zeros((clusters_rg_bg.shape[0], 3))
# convert to rgb
for i in range(clusters.shape[0]):
rgb_vector = np.array(
[clusters_rg_bg[i, 0], 1.0, clusters_rg_bg[i, 1]])
clusters[i, :] = normalize_illuminant(rgb_vector)
# pytorch candidates
clusters = torch.FloatTensor(clusters).unsqueeze(0)
return clusters
def initialize(self, illuminants):
if illuminants is None:
clusters = torch.FloatTensor(np.zeros(
(self.k * self.k, 3))).unsqueeze(0)
return clusters
else:
illuminants = np.array(illuminants)
min_rg = 100
max_rg = 0
min_bg = 100
max_bg = 0
for i in range(illuminants.shape[0]):
rg = illuminants[i, 0] / illuminants[i, 1] # r / g
bg = illuminants[i, 2] / illuminants[i, 1] # b / g
min_rg = min(min_rg, rg)
max_rg = max(max_rg, rg)
min_bg = min(min_bg, bg)
max_bg = max(max_bg, bg)
real_k = self.k * self.k
clusters_rg_bg = np.zeros((real_k, 2))
step_rg = (max_rg - min_rg) / self.k
step_bg = (max_bg - min_bg) / self.k
for i in range(self.k):
rg = min_rg + step_rg * i
for j in range(self.k):
bg = min_bg + step_bg * j
rg_bg = np.array([rg, bg])
#print('i: '+str(i)+' j: '+str(j)+' '+str(rg_bg))
clusters_rg_bg[i * self.k + j, :] = rg_bg
clusters = np.zeros((clusters_rg_bg.shape[0], 3))
# convert to rgb
for i in range(clusters.shape[0]):
rgb_vector = np.array(
[clusters_rg_bg[i, 0], 1.0, clusters_rg_bg[i, 1]])
clusters[i, :] = normalize_illuminant(rgb_vector)
clusters = torch.FloatTensor(clusters).unsqueeze(0)
return clusters
def initialize(self, illuminants):
if illuminants is None:
# if no illuminants are provided, init with zeros
clusters = torch.FloatTensor(np.zeros((self.k, 3))).unsqueeze(0)
return clusters
else:
# convert illuminants to the desired color space
illuminants = np.array(illuminants)
if self.color_space == 'rg_bg':
illuminants_new = np.zeros((illuminants.shape[0], 2))
for i in range(illuminants.shape[0]):
illuminants_new[
i, 0] = illuminants[i, 0] / illuminants[i, 1] # r / g
illuminants_new[
i, 1] = illuminants[i, 2] / illuminants[i, 1] # b / g
illuminants = illuminants_new
elif self.color_space == 'log_rg_bg':
illuminants_new = np.zeros((illuminants.shape[0], 2))
for i in range(illuminants.shape[0]):
illuminants_new[i, 0] = math.log(
illuminants[i, 0]) - math.log(
illuminants[i, 1]) # log(r / g)
illuminants_new[i, 1] = math.log(
illuminants[i, 2]) - math.log(
illuminants[i, 1]) # log(b / g)
illuminants = illuminants_new
# run K-Means
kmeans = sklearn.cluster.KMeans(n_clusters=self.k).fit(illuminants)
# convert back to RGB color space
if self.color_space == 'rgb':
clusters = np.copy(kmeans.cluster_centers_)
# normalize illuminants
for i in range(clusters.shape[0]):
clusters[i, :] = normalize_illuminant(clusters[i, :])
elif self.color_space == 'rg_bg':
clusters_rg_bg = np.copy(kmeans.cluster_centers_)
clusters = np.zeros((clusters_rg_bg.shape[0], 3))
# convert to rgb
for i in range(clusters.shape[0]):
rgb_vector = np.array(
[clusters_rg_bg[i, 0], 1.0, clusters_rg_bg[i, 1]])
clusters[i, :] = normalize_illuminant(rgb_vector)
elif self.color_space == 'log_rg_bg':
clusters_rg_bg = np.copy(kmeans.cluster_centers_)
clusters = np.zeros((clusters_rg_bg.shape[0], 3))
# convert to rgb
for i in range(clusters.shape[0]):
rgb_vector = np.array([
math.exp(clusters_rg_bg[i, 0]), 1.0,
math.exp(clusters_rg_bg[i, 1])
])
clusters[i, :] = normalize_illuminant(rgb_vector)
# output: pytorch tensor
clusters = torch.FloatTensor(clusters).unsqueeze(0)
return clusters