def testAgainstRefImpl(self):
batch_size = 100
pred_illum_rgb = np.random.uniform(low=sys.float_info.epsilon,
size=(batch_size, 3))
pred_illum_rgb /= np.sum(pred_illum_rgb, axis=1, keepdims=True)
true_illum_rgb = np.random.uniform(low=sys.float_info.epsilon,
size=(batch_size, 3))
true_illum_rgb /= np.sum(true_illum_rgb, axis=1, keepdims=True)
true_wb_rgb = 1.0 / true_illum_rgb
true_wb_rgb /= np.expand_dims(true_wb_rgb[:, 1], axis=1)
true_scene_rgb = np.random.uniform(low=sys.float_info.epsilon,
size=(batch_size, 3))
input_scene_rgb = true_scene_rgb / true_wb_rgb
height = 64
width = 48
rgb_stats = np.tile(input_scene_rgb[:, np.newaxis, np.newaxis, :],
reps=[1, height, width, 1])
with self.test_session() as sess:
actual_result = sess.run(
losses.anisotropic_reproduction_loss(
ops.rgb_to_uv(tf.constant(pred_illum_rgb)),
ops.rgb_to_uv(tf.constant(true_illum_rgb)),
tf.constant(rgb_stats)))
expected_result = sess.run(
losses.anisotropic_reproduction_error(
tf.constant(pred_illum_rgb), tf.constant(true_illum_rgb),
tf.constant(true_scene_rgb)))
np.testing.assert_almost_equal(actual_result, expected_result)
def testRgbToUvAndBack(self):
batch_size = 100
rgb = np.random.uniform(size=(batch_size, 3))
# normalized RGB values to unit vectors
rgb /= np.linalg.norm(rgb, axis=1, keepdims=True)
np.testing.assert_almost_equal(
rgb, self._eval(ops.uv_to_rgb(ops.rgb_to_uv(tf.constant(rgb)))))
def testBasicCorrectness(self):
"""Test with ground truth values equal to predictions."""
batch_size = 100
# Generate random covariance matrices. Add a small epsilon on the diagonal
# to make them definite positive.
sigma = tf.eye(2, batch_shape=(batch_size, ))
rgb = np.random.uniform(low=sys.float_info.epsilon,
size=(batch_size, 3)).astype('float32')
uv = ops.rgb_to_uv(rgb)
# If ground truth values are equal to predictions, and the matrix 'sigma'
# is constant (over the batch size), the result will be a constant as well.
actual_result = losses.gaussian_negative_log_likelihood(uv, sigma, uv)
expected_result = tf.repeat(actual_result[0], batch_size)
self.assertAllClose(expected_result, actual_result)
def testApplyWb(self):
batch_size = 100
width = 64
height = 48
channels = 3
rgbs = np.random.uniform(size=(batch_size, height, width, channels))
wb_gains = np.random.uniform(low=0.1,
high=1.0,
size=(batch_size, channels))
wb_gains /= np.expand_dims(wb_gains[:, 1], axis=1)
expected_wb_rgbs = rgbs * wb_gains[:, np.newaxis, np.newaxis, :]
np.testing.assert_almost_equal(
expected_wb_rgbs,
self._eval(
ops.apply_wb(tf.constant(rgbs),
ops.rgb_to_uv(tf.constant(1.0 / wb_gains)))))
def _model_fn(features, labels, mode, params):
"""Model function for tf.Estimator.
Args:
features (dict): A dictionary maps to a batch of training features.
The following keys are expected:
{'name': an unique id for this example.
'rgb': the RGB image.
'extended_feature': a (float) number representing some unique feature
of this example.
labels (dict): A dictionary map to a batch of ground truth labels for
training.
The following keys are expected:
{'illuminant': the color of illuminant in RGB (an unit vector). This
is the reciprocal of RGB gains.}
mode: Indicates training, eval or inference mode: learn.ModeKeys.TRAIN,
learn.ModeKeys.EVAL, learn.ModeKeys.INFER
params: a dict with keys:
'first_bin': (float) location of the edge of the first histogram bin.
'bin_size': (float) size of each histogram bin.
'nbins': (int) number of histogram bins.
'extended_vector_length': The number of features in the extended vector.
'ellipse_w_mat', 'elllipse_b_vec': parameters describing the ellipse of
valid white points. See 'ellipse.py' for more details.
'variance': a small positive value to add to the covariance matrix to
make it definite positive.
Returns:
A model_fn_lib.ModelFnOps object that can be called by the estimator.
"""
tf.compat.v1.logging.info('Run model_fn, mode=%s', mode)
tf.compat.v1.logging.info('Params=%s', params)
tf.compat.v1.logging.info('HParams=[%s]', hparams)
# Expect input features is a dictionary
assert isinstance(features, dict)
rgb = features['rgb']
extended_feature = features['extended_feature']
weight = features['weight']
n = params['nbins']
extended_vector_length = params['extended_vector_length']
filters_extended_latent = tf.compat.v1.get_variable(
'filters_extended_latent',
shape=(extended_vector_length, n * n, NUMBER_OF_CHANNELS_FILTERS),
initializer=tf.zeros_initializer())
filters_base_latent = tf.compat.v1.get_variable(
'filters_base_latent',
shape=(n * n, NUMBER_OF_CHANNELS_FILTERS),
initializer=tf.zeros_initializer())
bias_extended_latent = tf.compat.v1.get_variable(
'bias_extended_latent',
shape=(extended_vector_length, n * n),
initializer=tf.zeros_initializer())
# Initialize the bias to a scaled+shifted 2D Hann function in [-1, 1],
# so that white point estimates are initially located at the center of the
# histogram.
hann1 = np.sin(np.pi * np.float32(np.arange(0, n)) / n)**2
bias_base_init = 2 * hann1[np.newaxis, :] * hann1[:, np.newaxis] - 1.
precond_bias = tf.cast(
fft.compute_preconditioner_vec(n, hparams['mult_bias_tv'],
hparams['mult_bias_l2']), tf.float32)
bias_base_latent_init = (fft.fft2_to_vec(
ops.r2c_fft2(
tf.convert_to_tensor(bias_base_init)[tf.newaxis, :, :,
tf.newaxis]))[0] /
precond_bias)[:, 0]
bias_base_latent = tf.compat.v1.get_variable(
'bias_base_latent', initializer=bias_base_latent_init)
with tf.control_dependencies([
tf.debugging.assert_equal(tf.math.is_nan(filters_extended_latent), False),
tf.debugging.assert_equal(tf.math.is_nan(filters_base_latent), False),
tf.debugging.assert_equal(tf.math.is_nan(bias_extended_latent), False),
tf.debugging.assert_equal(tf.math.is_nan(bias_base_latent), False)
]):
(filters_extended_fft, filters_base_fft, bias_extended_fft, bias_base_fft,
bias_extended, bias_base, precond_filters,
precond_bias) = latent_to_model(filters_extended_latent,
filters_base_latent,
bias_extended_latent, bias_base_latent,
hparams, params)
with tf.name_scope('preconditioner'):
tf.summary.histogram('precond_filters', precond_filters)
tf.summary.histogram('precond_bias', precond_bias)
mu, sigma, heatmap = evaluate_model(rgb, extended_feature,
filters_extended_fft,
filters_base_fft, bias_extended,
bias_base, params)
if mode == tf.estimator.ModeKeys.PREDICT:
predictions = dict()
predictions['uv'] = mu
if 'name' in features:
predictions['name'] = features['name']
return tf.estimator.EstimatorSpec(mode, predictions=predictions)
# Computes losses
assert isinstance(labels, dict)
true_uv = ops.rgb_to_uv(labels['illuminant'])
global_step = tf.compat.v1.train.get_global_step()
# The ground-truth white points should lie within the ellipse of valid
# white points, otherwise something has gone wrong or the ellipse needs
# to be re-fit.
true_ellipse_distance = ellipse.distance(
true_uv, np.reshape(params['ellipse_w_mat'], (2, 2)),
np.asarray(params['ellipse_b_vec']))
with tf.control_dependencies(
[tf.debugging.assert_less_equal(true_ellipse_distance, 1. + 1e-5)]):
(weighted_loss_data, data_losses) = losses.compute_data_loss(
heatmap,
mu,
sigma,
true_uv,
weight=weight,
step_size=params['bin_size'],
offset=params['first_bin'],
n=n,
rgb=rgb)
smooth_vars = [
filters_base_latent, filters_extended_latent, bias_base_latent,
bias_extended_latent
]
numer = tf.reduce_sum([tf.reduce_sum(v**2) for v in smooth_vars])
denom = tf.reduce_sum(
[tf.math.reduce_prod(tf.shape(v)) for v in smooth_vars])
loss_smooth = numer / tf.cast(denom, tf.float32)
loss = weighted_loss_data + hparams['mult_smooth'] * loss_smooth
if mode == tf.estimator.ModeKeys.EVAL:
eval_metric_ops = {
'rms_anisotropic_reproduction_error':
tf.compat.v1.metrics.root_mean_squared_error(
predictions=data_losses['anisotropic_reproduction_error'],
labels=tf.zeros(
tf.shape(
data_losses['anisotropic_reproduction_error'])),
weights=weight),
'mean_anisotropic_reproduction_error':
tf.compat.v1.metrics.mean(
values=data_losses['anisotropic_reproduction_error'],
weights=weight),
'reproduction_error':
tf.compat.v1.metrics.mean(
values=data_losses['reproduction_error'], weights=weight),
'mean_angular_error':
tf.compat.v1.metrics.mean(
values=data_losses['angular_error'], weights=weight),
'gaussian_nll':
tf.compat.v1.metrics.mean(
values=data_losses['gaussian_nll'], weights=weight)
}
return tf.estimator.EstimatorSpec(
mode, loss=loss, eval_metric_ops=eval_metric_ops)
# Smoothly decay from 1 to 0.
cosine_decay = tf.compat.v1.train.cosine_decay(1., global_step,
hparams['total_training_iterations'])
learning_rate = cosine_decay * hparams['learning_rate']
assert mode == tf.estimator.ModeKeys.TRAIN
train_op = tf.compat.v1.train.AdamOptimizer(
learning_rate=learning_rate).minimize(
loss=loss, global_step=global_step)
# Setup summaries
with tf.name_scope('summaries'):
# Save the total loss and anisotropic reproduction loss
tf.summary.scalar('loss', loss)
tf.summary.scalar('learning_rate', learning_rate)
tf.summary.scalar('reproduction_error',
tf.reduce_mean(data_losses['reproduction_error']))
tf.summary.scalar('angular_error',
tf.reduce_mean(data_losses['angular_error']))
tf.summary.scalar('loss_smooth', loss_smooth)
tf.summary.image('filters_extended',
_visualize_fft(filters_extended_fft, shift=True))
tf.summary.image('filters_base',
_visualize_fft(filters_base_fft, shift=True))
tf.summary.image('bias_extended',
_visualize_fft(bias_extended_fft, shift=False))
tf.summary.image('bias_base',
_visualize_fft(bias_base_fft, shift=False))
tf.summary.histogram('filters_extended_latent', filters_extended_latent)
tf.summary.histogram('filters_base_latent', filters_base_latent)
tf.summary.histogram('bias_extended_latent', bias_extended_latent)
tf.summary.histogram('bias_base_latent', bias_base_latent)
tf.summary.histogram(
'filters_extended_fft',
tf.stack(
[tf.math.real(filters_extended_fft),
tf.math.imag(filters_extended_fft)]))
tf.summary.histogram(
'filters_base_fft',
tf.stack(
[tf.math.real(filters_base_latent),
tf.math.imag(filters_base_latent)]))
tf.summary.histogram(
'bias_extended_fft',
tf.stack(
[tf.math.real(bias_extended_latent),
tf.math.imag(bias_extended_latent)]))
tf.summary.histogram(
'bias_base_fft',
tf.stack([tf.math.real(bias_base_fft),
tf.math.imag(bias_base_fft)]))
return tf.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op)
def main(_):
if not tf.gfile.IsDirectory(FLAGS.data_dir):
tf.logging.error('Invalid input directory: %s', FLAGS.data_dir)
return
tf.logging.info('Loading dataset')
(train_input_fn, _, _, _,
_) = (ffcc_input.input_builder_stratified(FLAGS.data_dir,
batch_size=1,
num_epochs=1,
bucket_size=1))
# Extract the ground-truth UV coordinates of the illuminants.
dataset = train_input_fn()
items = []
for (feature, label) in dataset:
items.append({
'name': feature['name'],
'uv': ops.rgb_to_uv(label['illuminant']),
})
ids = [np.asscalar(item['name'].numpy()) for item in items]
uvs = np.float32(np.concatenate([item['uv'].numpy() for item in items]))
a_mat, c_vec = fit_ellipse(uvs)
# Change the ellipse's parametrization.
w_mat, b_vec = ellipse.standard_to_general(a_mat, c_vec)
# Sanity check that all datapoints lie within the ellipse.
d = ellipse.distance(uvs, w_mat, b_vec)
np.testing.assert_array_less(d, 1. + 1e-5)
# Sanity check that projection is a no-op.
uvs_projected = ellipse.project(uvs, w_mat, b_vec)
np.testing.assert_almost_equal(uvs_projected, uvs, decimal=5)
# Generate the ellipse distance to the CSV file.
export_csv('/tmp/ellipse_distance.csv', ids, np.asarray(uvs),
np.asarray(d))
# Draw ellipse.
# Compute "tilt" of ellipse using first eigenvector.
eig_vals, eig_vecs = np.linalg.eig(np.linalg.inv(a_mat))
x, y = eig_vecs[:, 0]
theta = np.degrees(np.arctan2(y, x))
# Eigenvalues give length of ellipse along each eigenvector.
w, h = 2 * np.sqrt(eig_vals)
fig, ax = plt.subplots(1)
ax.set_xlabel('log(g/r)')
ax.set_ylabel('log(g/b)')
ax.scatter(x=uvs[:, 0], y=uvs[:, 1], marker='+', color='r')
patch = patches.Ellipse(c_vec, w, h, theta, color='g')
patch.set_clip_box(ax.bbox)
patch.set_alpha(0.2)
ax.add_artist(patch)
tf.logging.info('saving figures ')
fig.savefig('/tmp/ellipse_fit.png')
tf.logging.info('ellipse alpha = %s', 1e8)
tf.logging.info('w_mat = %s', w_mat.numpy().reshape([-1]))
tf.logging.info('b_vec = %s', b_vec.numpy())