# FC4 is by design an FCN. from __future__ import absolute_import from __future__ import division from __future__ import print_function # Things we tried but not turned out to work: (so you don't have to understand them) # FCN.build_shallow, FEED_SHALLOW # per_patch_weight, per_patch_loss # smooth_l1 # # Things to clarify: # per_pixel_est is actually per-patch estimation, each patch is the receptive field of the FCN. import cPickle as pickle import math import time import os import cv2 import sys from utils import Tee import shutil import numpy as np import tensorflow as tf from data_provider import load_data, DataProvider from utils import angular_error, print_angular_errors, LowestTrigger from config import * #from alexnet import create_alexnet as create_convnet from squeeze_net import create_convnet from summary_utils import * import random slim = tf.contrib.slim # An alternative loss function to original angular loss # not used any more def smooth_l1(x): return tf.maximum(tf.minimum(0.5 * x**2, 0.5), x - 0.5) class FCN: def __init__(self, sess=None, name=None, kwargs={}): global TRAINING_FOLDS, TEST_FOLDS self.name = name self.wd = GLOBAL_WEIGHT_DECAY TRAINING_FOLDS, TEST_FOLDS = initialize_dataset_config(**kwargs) self.training_data_provider = None self.sess = sess with slim.arg_scope( [slim.conv2d, slim.fully_connected], weights_regularizer=slim.l2_regularizer(GLOBAL_WEIGHT_DECAY)): self.build() tf.global_variables_initializer() # Store the test-time networks, for images of different resolution self.test_nets = {} self.saver = tf.train.Saver(max_to_keep=CKPTS_TO_KEEP) # There are two branches: # 1. AlexNet branch # 2. A shallower shortcut branch (*not* affected by USE_SHORTCUT), which proves to be not working # images: 0-65535 linear RGB @staticmethod def build_branches(images, dropout): images = tf.clip_by_value(images, 0.0, 65535.0) if USE_SHORTCUT: # Apply grey-world first image_mean = tf.reduce_mean(images, axis=(1, 2), keep_dims=True) images = images * (1.0 / tf.maximum(image_mean, 1e-10)) image_max = tf.reduce_max(images, axis=(1, 2, 3), keep_dims=True) images = images * (1.0 / tf.maximum(image_max, 1e-10)) else: images = images * (1.0 / 65535) #self.prob = alex_outputs['prob'] feed_to_fc = [] # Build the AlexNet branch if FEED_ALEX: with tf.variable_scope('AlexNet'): # RGB (0~1) to sRGB (0~255) with gamma # alex_images = (tf.pow(images, 1.0/config_get_input_gamma()) * 255.0) # the caffe model takes BGR, 0~255 image alex_images = (tf.pow(images, 1.0 / config_get_input_gamma()) * 255.0)[:, :, :, ::-1] #alex_images = alex_images - np.array([103.939, 116.779, 123.68])[None, None, None,:] alex_outputs = create_convnet(alex_images) alexnet_output = alex_outputs['features_out'] feed_to_fc.append(alexnet_output) # Build the shallow network branch if FEED_SHALLOW: with tf.variable_scope('Shallow'): shallow_outputs = FCN.build_shallow(images) feed_to_fc.append(shallow_outputs) feed_to_fc = tf.concat(axis=3, values=feed_to_fc) print('Feed to FC shape', feed_to_fc.get_shape()) # The FC1 here is actually a convolutional layer, since we use FCN fc1 = slim.conv2d( feed_to_fc, FC1_SIZE, [FC1_KERNEL_SIZE, FC1_KERNEL_SIZE], scope='fc1') print('FC1 shape', fc1.get_shape()) # Use dropout, training time only fc1 = slim.dropout(fc1, dropout) ############## # There are two ways to do confidence-weighted pooling # 1. Output (normalized) R, G, B and weight, average RGB by weight # 2. Just output unnormalized R, G, B, and simply take the sum and normalize. Thus # we are weighting using the length. ############## if not SEPARATE_CONFIDENCE: # Way 2 fc2 = slim.conv2d(fc1, 3, [1, 1], scope='fc2', activation_fn=None) else: # Way 1 print("Using sperate fc2") fc2_mid = slim.conv2d(fc1, 4, [1, 1], scope='fc2', activation_fn=None) fc2_mid = tf.nn.relu(fc2_mid) rgb = tf.nn.l2_normalize(fc2_mid[:, :, :, :3], 3) w, h = map(int, fc2_mid.get_shape()[1:3]) print(w, h) confidence = fc2_mid[:, :, :, 3:4] confidence = tf.reshape(confidence, shape=[-1, w * h]) print(confidence.get_shape()) confidence = tf.nn.softmax(confidence) confidence = tf.reshape(confidence, shape=[-1, w, h, 1]) fc2 = rgb * confidence print('FC2 shape', fc2.get_shape()) if USE_SHORTCUT: fc2 = fc2 * image_mean[:, :, :, ::-1] if not WEIGHTED_POOLING: # Simply average pooling fc2 = tf.nn.l2_normalize(fc2, 3) print('FC2 shape', fc2.get_shape()) if FC_POOLING: # WEIGHTED_POOLING should be true, since we should not normalize anything before feeding into FC assert WEIGHTED_POOLING assert RESIZE_TEST # Pooling using a simple FC, as noted in the paper fc2 = tf.nn.l2_normalize( slim.conv2d( tf.nn.relu(fc2), 3, [15, 15], scope='fc_pooling', activation_fn=None, padding='VALID'), 3) print('FC2 shape', fc2.get_shape()) return fc2 # Build the network def build(self): self.dropout = tf.placeholder(tf.float32, shape=(), name='dropout') # We don't use per_patch_weight any more. self.per_patch_weight = tf.placeholder( tf.float32, shape=(), name='pre_patch_weight') per_patch_weight = self.per_patch_weight self.learning_rate = tf.placeholder( tf.float32, shape=(), name='learning_rate') # ground truth, actually self.illums = tf.placeholder(tf.float32, shape=(None, 3), name='illums') # input images self.images = tf.placeholder( tf.float32, shape=(None, FCN_INPUT_SIZE, FCN_INPUT_SIZE, 3), name='images') with tf.variable_scope('FCN'): fc2 = self.build_branches(self.images, self.dropout) self.per_pixel_est = fc2 self.illum_normalized = tf.nn.l2_normalize( tf.reduce_sum(fc2, axis=(1, 2)), 1) self.train_visualization = get_visualization( self.images, self.per_pixel_est, self.illum_normalized, self.illums, (VISUALIZATION_SIZE, VISUALIZATION_SIZE)) self.global_loss = self.get_angular_loss( tf.reduce_sum(fc2, axis=(1, 2)), self.illums, LENGTH_REGULARIZATION) self.per_patch_loss = self.get_angular_loss( fc2, self.illums[:, None, None, :], LENGTH_REGULARIZATION) self.loss = ( 1 - per_patch_weight ) * self.global_loss + self.per_patch_weight * self.per_patch_loss scalar_summaries = [] scalar_summaries.append( tf.summary.scalar('per_patch_loss', self.per_patch_loss)) scalar_summaries.append( tf.summary.scalar('full_image_loss', self.global_loss)) scalar_summaries.append(tf.summary.scalar('loss', self.loss)) self.scalar_summaries = tf.summary.merge(scalar_summaries) conv_scopes = [] if FEED_ALEX: conv_scopes.append('FCN/AlexNet/conv1') if FEED_SHALLOW: conv_scopes.append('FCN/Shallow/conv1') image_summaries = [] ''' for scope in conv_scopes: with tf.variable_scope(scope, reuse=True): image_summaries.append(conv_summary(tf.get_variable("weights"), scope)) self.image_summaries = tf.summary.merge(image_summaries) ''' self.merge_summaries = tf.summary.merge_all() reg_losses = tf.add_n(slim.losses.get_regularization_losses()) self.total_loss = self.loss + reg_losses self.train_step_adam = tf.train.AdamOptimizer(self.learning_rate).minimize( self.total_loss) var_list1 = tf.get_collection( tf.GraphKeys.TRAINABLE_VARIABLES, scope='FCN/AlexNet') var_list2 = tf.get_collection( tf.GraphKeys.TRAINABLE_VARIABLES, scope='FCN/fc1') + tf.get_collection( tf.GraphKeys.TRAINABLE_VARIABLES, scope='FCN/fc2') for v in var_list1: print('list1', v.name) for v in var_list2: print('list2', v.name) opt1 = tf.train.AdamOptimizer(self.learning_rate * FINE_TUNE_LR_RATIO) opt2 = tf.train.AdamOptimizer(self.learning_rate) grads = tf.gradients(self.total_loss, var_list1 + var_list2) grads1 = grads[:len(var_list1)] grads2 = grads[len(var_list1):] train_op1 = opt1.apply_gradients(zip(grads1, var_list1)) train_op2 = opt2.apply_gradients(zip(grads2, var_list2)) self.train_step_sgd = tf.group(train_op1, train_op2) # The shallow branch - not used any more @staticmethod def build_shallow(images): net = images net = slim.conv2d( net, SHALLOW_CHANNELS[0], [11, 11], stride=4, scope='conv1', padding='SAME') net = slim.max_pool2d(net, [3, 3], padding='VALID') net = slim.conv2d( net, SHALLOW_CHANNELS[1], [5, 5], stride=1, scope='conv2', padding='SAME') net = slim.max_pool2d(net, [3, 3], padding='VALID') net = slim.max_pool2d(net, [3, 3], padding='VALID') return net def duplicate_output_to_log(self): self.tee = Tee(self.get_ckpt_folder() + 'log.txt') def get_summary_variables(self, i, j): summary_variables = [] if i >= 2 and WRITE_SUMMARY and j == 0: summary_variables.append(self.scalar_summaries) if i % IMAGE_SUMMARY_INT == 0: #summary_variables.append(self.image_summaries) pass return summary_variables def get_train_step(self, i): if i <= FORCE_ADAM or OPTIMIZER == 'ADAM': return self.train_step_adam elif OPTIMIZER == 'SGD': return self.train_step_sgd else: assert False def train(self, epochs, backup=True): trigger = LowestTrigger() if not self.try_make_ckpt_folder(): print('Warning: folder exists!!!') if backup: self.backup_scripts() print('Backup succeeded') self.train_writer = tf.summary.FileWriter( logdir=self.get_ckpt_folder() + '/training/', graph=self.sess.graph) self.validation_writer = tf.summary.FileWriter( logdir=self.get_ckpt_folder() + '/validation/') self.duplicate_output_to_log() training_batch_size = TRAINING_BATCH_SIZE print("TF", TRAINING_FOLDS) self.training_data_provider = DataProvider(True, TRAINING_FOLDS) self.training_data_provider.set_batch_size(training_batch_size) test_summary_input = tf.placeholder(tf.float32, shape=()) test_summary = tf.summary.scalar('loss', test_summary_input) batches_per_training_epoch = self.training_data_provider.data_count // training_batch_size for i in range(1, epochs + 1): # Actually we use a fixed learning rate learning_rate = BASE_LEARNING_RATE * pow(LR_DECAY, 1.0 * i / LR_DECAY_INTERVAL) epoch_starting_time = time.time() training_losses = [] for j in range(batches_per_training_epoch): batch = self.training_data_provider.get_batch() summary_variables = self.get_summary_variables(i, j) visualization = [] # visualize some training images for monitoring the process if TRAINING_VISUALIZATION and i % TRAINING_VISUALIZATION == 0: visualization.append(self.train_visualization) loss, _, per_patch_loss, global_loss, summary, ppest, vis = self.sess.run( [ self.loss, self.get_train_step(i), self.per_patch_loss, self.global_loss, summary_variables, self.per_pixel_est, visualization ], feed_dict={ self.images: batch[0], self.illums: batch[2], self.dropout: DROPOUT, self.per_patch_weight: PER_PATCH_WEIGHT, self.learning_rate: learning_rate }) for s in summary: self.train_writer.add_summary(s, i) if vis: folder = self.get_ckpt_folder() + '/training_visualization/' try: os.mkdir(folder) except: pass for k, merged in enumerate(vis[0]): summary_fn = '%s/%04d-%03d.jpg' % (folder, i, j * len(vis) + k) cv2.imwrite(summary_fn, merged[:, :, ::-1] * 255) training_losses.append(loss) training_loss = sum(training_losses) / len(training_losses) validation_loss = 10 ending = '' info = "*%s* E %4d, TL %.3f, VL %.3f, D %.3f, t %4.1f%s" % ( self.name, i, training_loss, validation_loss, validation_loss - training_loss, time.time() - epoch_starting_time, ending) print(info) saved = False if CKPT_PERIOD and i % CKPT_PERIOD == 0: self.save(i) saved = True # Do some test if TEST_PERIOD and i % TEST_PERIOD == 0: summary = i // TEST_PERIOD % 5 == 0 if MULTIRES_TEST: self.test(summary=summary, scales=[0.25], summary_key=i) self.test(summary=summary, scales=[0.5], summary_key=i) if not RESIZE_TEST: errors = self.test( summary=summary, scales=[TEST_BASE_RES], summary_key=i)[0] else: errors = self.test_resize()[0] self.validation_writer.add_summary( test_summary.eval(feed_dict={test_summary_input: np.mean(errors)}), i) if MULTIRES_TEST: self.test(summary=summary, scales=[0.25, 0.5, 1.0], summary_key=i) # Does validation loss achieve a new minimum? # It is unfortunate that the datasets have only train and test separation... # We need a validation set to prevent overfitting... if trigger.push(np.mean(errors)): error_fn = self.get_ckpt_folder() + 'error%04d.pkl' % i with open(error_fn, 'wb') as f: pickle.dump(errors, f, protocol=-1) #self.test_patch_based() if not saved: self.save(i) saved = True print('mean(errors) from fcn.py line: 330', np.mean(errors)) self.training_data_provider.stop() # Simply test the network def test_network(self, summary=False, summary_key=0): records = load_data(['m']) for r in records: scale = 1 a = 1 img = np.clip((r.img / r.img.max()), 0, 1) img = np.power(img, 2.2) img[:, :img.shape[1] // 2:, 0] *= 3 #img = np.pad(img, ((112, 112), (112, 112), (0, 0)), 'constant') if scale != 1.0: img = cv2.resize(img, (0, 0), fx=scale, fy=scale) shape = img.shape[:2] if shape not in self.test_nets: test_net = {} test_net['images'] = tf.placeholder( tf.float32, shape=(None, shape[0], shape[1], 3), name='test_images') with tf.variable_scope("FCN", reuse=True): test_net['pixels'] = FCN.build_branches(test_net['images'], 1.0) self.test_nets[shape] = test_net test_net = self.test_nets[shape] pixels = self.sess.run( test_net['pixels'], feed_dict={ test_net['images']: img[None, :, :, :], }) pixels = pixels[0].astype(np.float32) #pixels = cv2.resize(pixels, image.shape[0:2][::-1]) pixels /= np.linalg.norm(pixels, axis=2)[:, :, None] #pixels = pixels[:,:,::-1] cv2.imshow('pixels', cv2.resize(pixels, (0, 0), fx=10, fy=10)) cv2.imshow('image', cv2.resize(img, (0, 0), fx=0.5, fy=0.5)) cv2.waitKey(0) # Test each image in multiple resolutions, and then average def test_multi(self, summary=False, summary_key=0): records = load_data(['m']) summaries = [] errors = [] for r in records: all_pixels = [] scale = 1.0 a = 1 artificial = np.array((1.0, 1.0, 1.0))[None, None, :] img = r.img.copy() img = np.clip((img * artificial[:, :, ::-1] / img.max()), 0, 1) img = np.power(img, 2.2) #img[:,:img.shape[1] // 2:,0] *= 3 #img = np.pad(img, ((112, 112), (112, 112), (0, 0)), 'constant') if scale != 1.0: img = cv2.resize(img, (0, 0), fx=scale, fy=scale) shape = img.shape[:2] if shape not in self.test_nets: test_net = {} test_net['images'] = tf.placeholder( tf.float32, shape=(None, shape[0], shape[1], 3), name='test_images') with tf.variable_scope("FCN", reuse=True): test_net['pixels'] = FCN.build_branches(test_net['images'], 1.0) self.test_nets[shape] = test_net test_net = self.test_nets[shape] pixels = self.sess.run( test_net['pixels'], feed_dict={ test_net['images']: img[None, :, :, :], }) pixels = pixels[0].astype(np.float32) / artificial illum = r.illum.astype(np.float32) pixels = cv2.resize(pixels, illum.shape[0:2][::-1]) pixels /= np.linalg.norm(pixels, axis=2)[:, :, None] illum /= np.linalg.norm(illum, axis=2)[:, :, None] pixels = pixels[:, :, ::-1] cv2.imshow('pixels', cv2.resize(pixels, (0, 0), fx=2, fy=2)) cv2.imshow('illum', cv2.resize(illum, (0, 0), fx=2, fy=2)) cv2.imshow('image', img) cv2.waitKey(0) illum = illum.astype(np.float32) pixels = pixels.astype(np.float32) dot = np.sum(pixels * illum, axis=2) error = np.mean(np.arccos(np.clip(dot, -1, 1)) * 180 / math.pi) print(error) errors.append(error) errors = sorted(errors) print('Mean:', np.mean(errors), 'Median', errors[len(errors) // 2]) # Test grey world performance @staticmethod def test_naive(): t = time.time() import scipy.io std = scipy.io.loadmat('/home/yuanming/colorchecker_shi_greyworld.mat') names = map(lambda x: x[0].encode('utf8'), std['all_image_names'][0]) #print(names) records = load_data(TEST_FOLDS) errors = [] for r in records: est = np.mean(r.img, axis=(0, 1))[::-1] est /= np.linalg.norm(est) #print(r.fn, est) #est=np.array((1, 1, 1)) #est2= std['estimated_illuminants'][names.index(r.fn[:-4])] gt2 = std['groundtruth_illuminants'][names.index(r.fn[:-4])] #print(est2) error = math.degrees(angular_error(est, gt2)) errors.append(error) print("Full Image:") ret = print_angular_errors(errors) print('Test time:', time.time() - t, 'per image:', (time.time() - t) / len(records)) return errors # Just get the estimated illumination def inference(self, fn, datapacks=['g0', 'g1', 'g2']): records = load_data(datapacks) for r in records: if r.fn != fn: continue scale = 0.5 img = cv2.resize(r.img, (0, 0), fx=scale, fy=scale) shape = img.shape[:2] test_net = {} test_net['images'] = tf.placeholder( tf.float32, shape=(None, shape[0], shape[1], 3), name='test_images') with tf.variable_scope("FCN", reuse=True): test_net['pixels'] = FCN.build_branches(test_net['images'], 1.0) test_net['est'] = tf.reduce_sum(test_net['pixels'], axis=(1, 2)) self.test_nets[shape] = test_net pixels, est = self.sess.run( [test_net['pixels'], test_net['est']], feed_dict={ test_net['images']: img[None, :, :, :], }) pixels = pixels[0] est = est[0] return img, pixels, est, r.illum assert False, "Image not found" def test(self, summary=False, scales=[1.0], weights=[], summary_key=0, data=None, eval_speed=False, visualize=False): if not TEST_FOLDS: return [0] if data is None: records = load_data(TEST_FOLDS) else: records = data avg_errors = [] median_errors = [] t = time.time() summaries = [] if weights == []: weights = [1.0] * len(scales) outputs = [] ground_truth = [] avg_confidence = [] errors = [] for r in records: all_pixels = [] for scale, weight in zip(scales, weights): img = r.img if scale != 1.0: img = cv2.resize(img, (0, 0), fx=scale, fy=scale) shape = img.shape[:2] if shape not in self.test_nets: aspect_ratio = 1.0 * shape[1] / shape[0] if aspect_ratio < 1: target_shape = (MERGED_IMAGE_SIZE, MERGED_IMAGE_SIZE * aspect_ratio) else: target_shape = (MERGED_IMAGE_SIZE / aspect_ratio, MERGED_IMAGE_SIZE) target_shape = tuple(map(int, target_shape)) test_net = {} test_net['illums'] = tf.placeholder( tf.float32, shape=(None, 3), name='test_illums') test_net['images'] = tf.placeholder( tf.float32, shape=(None, shape[0], shape[1], 3), name='test_images') with tf.variable_scope("FCN", reuse=True): test_net['pixels'] = FCN.build_branches(test_net['images'], 1.0) test_net['est'] = tf.reduce_sum(test_net['pixels'], axis=(1, 2)) test_net['merged'] = get_visualization( test_net['images'], test_net['pixels'], test_net['est'], test_net['illums'], target_shape) self.test_nets[shape] = test_net test_net = self.test_nets[shape] pixels, est, merged = self.sess.run( [test_net['pixels'], test_net['est'], test_net['merged']], feed_dict={ test_net['images']: img[None, :, :, :], test_net['illums']: r.illum[None, :] }) if eval_speed: eval_batch_size = 1 eval_packed_input = img[None, :, :, :].copy() eval_packed_input = np.concatenate( [eval_packed_input for i in range(eval_batch_size)], axis=0) eval_packed_input = np.ascontiguousarray(eval_packed_input) eval_start_t = time.time() print(eval_packed_input.shape) eval_rounds = 100 images_variable = tf.Variable( tf.random_normal( eval_packed_input.shape, dtype=tf.float32, stddev=1e-1)) print(images_variable) for eval_t in range(eval_rounds): print(eval_t) pixels, est = self.sess.run( [test_net['pixels'], test_net['est']], feed_dict={ test_net['images']: #images_variable, eval_packed_input, }) eval_elapsed_t = time.time() - eval_start_t print('per image evaluation time', eval_elapsed_t / (eval_rounds * eval_batch_size)) pixels = pixels[0] #est = est[0] merged = merged[0] all_pixels.append(weight * pixels.reshape(-1, 3)) all_pixels = np.sum(np.concatenate(all_pixels, axis=0), axis=0) est = all_pixels / (np.linalg.norm(all_pixels) + 1e-7) outputs.append(est) ground_truth.append(r.illum) error = math.degrees(angular_error(est, r.illum)) errors.append(error) avg_confidence.append(np.mean(np.linalg.norm(all_pixels))) summaries.append((r.fn, error, merged)) print("Full Image:") ret = print_angular_errors(errors) ppt = (time.time() - t) / len(records) print('Test time:', time.time() - t, 'per image:', ppt) if summary: for fn, error, merged in summaries: folder = self.get_ckpt_folder() + '/test%04dsummaries_%4f/' % ( summary_key, scale) try: os.mkdir(folder) except: pass summary_fn = '%s/%5.3f-%s.png' % (folder, error, fn) cv2.imwrite(summary_fn, merged[:, :, ::-1] * 255) if visualize: for fn, error, merged in summaries: cv2.imshow('Testing', merged[:, :, ::-1]) cv2.waitKey(0) return errors, ppt, outputs, ground_truth, ret, avg_confidence # Test external images, such as sixteen or some jpegs # without ground truth # images are BGR # TODO: move network creation here def test_external(self, images, scale=1.0, fns=None, show=True, write=True): illums = [] confidence_maps = [] for img, filename in zip(images, fns): if scale != 1.0: img = cv2.resize(img, (0, 0), fx=scale, fy=scale) shape = img.shape[:2] if shape not in self.test_nets: aspect_ratio = 1.0 * shape[1] / shape[0] if aspect_ratio < 1: target_shape = (MERGED_IMAGE_SIZE, MERGED_IMAGE_SIZE * aspect_ratio) else: target_shape = (MERGED_IMAGE_SIZE / aspect_ratio, MERGED_IMAGE_SIZE) target_shape = tuple(map(int, target_shape)) test_net = {} test_net['illums'] = tf.placeholder( tf.float32, shape=(None, 3), name='test_illums') test_net['images'] = tf.placeholder( tf.float32, shape=(None, shape[0], shape[1], 3), name='test_images') with tf.variable_scope("FCN", reuse=True): test_net['pixels'] = FCN.build_branches(test_net['images'], 1.0) test_net['est'] = tf.reduce_sum(test_net['pixels'], axis=(1, 2)) test_net['merged'] = get_visualization( test_net['images'], test_net['pixels'], test_net['est'], test_net['illums'], target_shape) self.test_nets[shape] = test_net test_net = self.test_nets[shape] pixels, est, merged = self.sess.run( [test_net['pixels'], test_net['est'], test_net['merged']], feed_dict={ test_net['images']: img[None, :, :, :], test_net['illums']: [[1, 1, 1]] }) est = est[0] est /= np.linalg.norm(est) pixels = pixels[0] confidences = np.linalg.norm(pixels, axis=2) confidence_maps.append(confidences) ind = int(confidences.flatten().shape[0] * 0.95) print(confidences.mean(), confidences.max(), sorted(confidences.flatten())[ind]) merged = merged[0] illums.append(est) if show: cv2.imshow('Ret', merged[:, :, ::-1]) k = cv2.waitKey(0) % (2**20) #if k == ord('s'): filename = filename.split('/')[-1] if write: cv2.imwrite('/data/common/outputs/%s' % filename, merged[:, :, ::-1] * 255.0) try: os.makedirs('cc_outputs') except: pass corrected = np.power(img[:,:,::-1] / 65535 / est[None, None, :] * np.mean(est), 1/2.2)[:,:,::-1] cv2.imshow("corrected", corrected) cv2.waitKey(0) cv2.imwrite('cc_outputs/corrected_%s' % filename, corrected * 255.0) return illums, confidence_maps # Test how the network performs, if we just resize the image to fix the network input. def test_resize(self): records = load_data(TEST_FOLDS) t = time.time() errors = [] for r in records: img = cv2.resize(r.img, (FCN_INPUT_SIZE, FCN_INPUT_SIZE)) illum_est = self.sess.run( self.illum_normalized, feed_dict={self.images: [img], self.dropout: 1.0}) avg_error = math.degrees(angular_error(illum_est, r.illum)) errors.append(avg_error) print_angular_errors(errors) ppt = (time.time() - t) / len(records) print('Test time:', time.time() - t, 'per image:', ppt) return errors, ppt # Patch-based test def test_patch_based(self, scale, patches, pooling='median'): records = load_data(TEST_FOLDS) avg_errors = [] median_errors = [] t = time.time() def sample_patch(img): s = FCN_INPUT_SIZE x = random.randrange(0, img.shape[0] - s + 1) y = random.randrange(0, img.shape[1] - s + 1) return img[x:x + s, y:y + s] for r in records: img = cv2.resize(r.img, (0, 0), fx=scale, fy=scale) img = [sample_patch(img) for i in range(patches)] illum_est = [] batch_size = 4 for j in range((len(img) + batch_size - 1) // batch_size): illum_est.append( self.sess.run( self.illum_normalized, feed_dict={ self.images: img[j * batch_size:(j + 1) * batch_size], self.dropout: 1.0 })) illum_est = np.vstack(illum_est) med = len(illum_est) // 2 illum_est_median = np.array( [sorted(list(illum_est[:, i]))[med] for i in range(3)]) illum_est_avg = np.mean(illum_est, axis=0) avg_error = math.degrees(angular_error(illum_est_avg, r.illum)) median_error = math.degrees(angular_error(illum_est_median, r.illum)) avg_errors.append(avg_error) median_errors.append(median_error) print("Avg pooling:") print_angular_errors(avg_errors) print("Median pooling:") print_angular_errors(median_errors) ppt = (time.time() - t) / len(records) print('Test time:', time.time() - t, 'per image:', ppt) if pooling == 'median': errors = median_errors else: errors = avg_errors return errors, ppt def get_ckpt_folder(self, name=None): if name is None: name = self.name return MODEL_PATH + '/' + name + '/' def backup_scripts(self): for fn in os.listdir('.'): if fn.endswith('.py'): shutil.copy(fn, self.get_ckpt_folder()) with open(self.get_ckpt_folder() + 'config2.txt', 'w') as f: f.write(str(OVERRODE) + '\n') def try_make_ckpt_folder(self): try: os.mkdir(self.get_ckpt_folder()) except: ''' if raw_input('Warning: folder exists. Overwrite?\n') == 'yes': shutil.rmtree(self.get_ckpt_folder()) os.mkdir(self.get_ckpt_folder()) else: exit(-1) ''' return False return True def get_filename(self, key, name=None): if name is None: name = self.name return self.get_ckpt_folder(name) + str(key) + '.ckpt' def save(self, key): save_path = self.saver.save(self.sess, self.get_filename(key)) print("Model saved in file: %s" % save_path) def load(self, key, name=None): fn = self.get_filename(key, name=name) self.load_absolute(fn) def load_absolute(self, fn): self.saver.restore(self.sess, fn) print("Model %s restored." % fn) def get_angular_loss(self, vec1, vec2, length_regularization=0.0): with tf.name_scope('angular_error'): safe_v = 0.999999 if len(vec1.get_shape()) == 2: illum_normalized = tf.nn.l2_normalize(vec1, 1) _illum_normalized = tf.nn.l2_normalize(vec2, 1) dot = tf.reduce_sum(illum_normalized * _illum_normalized, 1) dot = tf.clip_by_value(dot, -safe_v, safe_v) length_loss = tf.reduce_mean( tf.maximum(tf.log(tf.reduce_sum(vec1**2, axis=1) + 1e-7), 0)) else: assert len(vec1.get_shape()) == 4 illum_normalized = tf.nn.l2_normalize(vec1, 3) _illum_normalized = tf.nn.l2_normalize(vec2, 3) dot = tf.reduce_sum(illum_normalized * _illum_normalized, 3) dot = tf.clip_by_value(dot, -safe_v, safe_v) length_loss = tf.reduce_mean( tf.maximum(tf.log(tf.reduce_sum(vec1**2, axis=3) + 1e-7), 0)) angle = tf.acos(dot) * (180 / math.pi) if SMOOTH_L1: angle = smooth_l1(angle) if ANGULAR_LOSS: return tf.reduce_mean(angle) + length_loss * length_regularization else: dot = tf.reduce_sum( (illum_normalized - _illum_normalized)**2, axis=len(illum_normalized.get_shape()) - 1) return tf.reduce_mean(dot) * 1000 + length_loss * length_regularization