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())