├── File ├── TensorflowUtils.py ├── generate_tfrecord.py └── run.py └── README.md /File/TensorflowUtils.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | import scipy.misc as misc 4 | import os, sys 5 | from six.moves import urllib 6 | import tarfile 7 | import zipfile 8 | import scipy.io 9 | 10 | 11 | def get_model_data(dir_path, model_url): 12 | maybe_download_and_extract(dir_path, model_url) 13 | filename = model_url.split("/")[-1] 14 | filepath = os.path.join(dir_path, filename) 15 | if not os.path.exists(filepath): 16 | raise IOError("VGG Model not found!") 17 | data = scipy.io.loadmat(filepath) 18 | return data 19 | 20 | 21 | def maybe_download_and_extract(dir_path, url_name, is_tarfile=False, is_zipfile=False): 22 | if not os.path.exists(dir_path): 23 | os.makedirs(dir_path) 24 | filename = url_name.split('/')[-1] 25 | filepath = os.path.join(dir_path, filename) 26 | if not os.path.exists(filepath): 27 | def _progress(count, block_size, total_size): 28 | sys.stdout.write( 29 | '\r>> Downloading %s %.1f%%' % (filename, float(count * block_size) / float(total_size) * 100.0)) 30 | sys.stdout.flush() 31 | 32 | filepath, _ = urllib.request.urlretrieve(url_name, filepath, reporthook=_progress) 33 | print() 34 | statinfo = os.stat(filepath) 35 | print('Succesfully downloaded', filename, statinfo.st_size, 'bytes.') 36 | if is_tarfile: 37 | tarfile.open(filepath, 'r:gz').extractall(dir_path) 38 | elif is_zipfile: 39 | with zipfile.ZipFile(filepath) as zf: 40 | zip_dir = zf.namelist()[0] 41 | zf.extractall(dir_path) 42 | 43 | 44 | def save_image(image, save_dir, name, mean=None): 45 | """ 46 | Save image by unprocessing if mean given else just save 47 | :param mean: 48 | :param image: 49 | :param save_dir: 50 | :param name: 51 | :return: 52 | """ 53 | if mean: 54 | image = unprocess_image(image, mean) 55 | misc.imsave(os.path.join(save_dir, name + ".png"), image) 56 | 57 | 58 | def get_variable(weights, name): 59 | init = tf.constant_initializer(weights, dtype=tf.float32) 60 | if name[-1] == 'w': 61 | var = tf.get_variable(name=name, initializer=init, regularizer=tf.contrib.layers.l2_regularizer(0.001), shape=weights.shape) 62 | else: 63 | var = tf.get_variable(name=name, initializer=init, shape=weights.shape) 64 | return var 65 | 66 | 67 | def weight_variable(shape, stddev=0.02, name=None): 68 | if name is None: 69 | initial = tf.truncated_normal(shape, stddev=stddev) 70 | return tf.Variable(initial) 71 | else: 72 | return tf.get_variable(name=name, shape=shape, initializer=tf.contrib.layers.xavier_initializer(), regularizer=tf.contrib.layers.l2_regularizer(0.001)) 73 | 74 | 75 | def bias_variable(shape, name=None): 76 | initial = tf.constant(0.0, shape=shape) 77 | if name is None: 78 | return tf.Variable(initial) 79 | else: 80 | return tf.get_variable(name, initializer=initial) 81 | 82 | 83 | def get_tensor_size(tensor): 84 | from operator import mul 85 | return reduce(mul, (d.value for d in tensor.get_shape()), 1) 86 | 87 | 88 | def conv2d_basic(x, W, bias, mypadding='SAME'): 89 | conv = tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], padding=mypadding) 90 | return tf.nn.bias_add(conv, bias) 91 | 92 | 93 | def conv2d_strided(x, W, b): 94 | conv = tf.nn.conv2d(x, W, strides=[1, 2, 2, 1], padding="SAME") 95 | return tf.nn.bias_add(conv, b) 96 | 97 | 98 | def conv2d_transpose_strided(x, W, b, output_shape=None, stride = 2): 99 | if output_shape is None: 100 | output_shape = x.get_shape().as_list() 101 | output_shape[1] *= 2 102 | output_shape[2] *= 2 103 | output_shape[3] = W.get_shape().as_list()[2] 104 | 105 | conv = tf.nn.conv2d_transpose(x, W, output_shape, strides=[1, stride, stride, 1], padding="SAME") 106 | return tf.nn.bias_add(conv, b) 107 | 108 | 109 | def leaky_relu(x, alpha=0.0, name=""): 110 | return tf.maximum(alpha * x, x, name) 111 | 112 | 113 | def max_pool_2x2(x): 114 | return tf.nn.max_pool(x, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding="SAME") 115 | 116 | 117 | def avg_pool_2x2(x): 118 | return tf.nn.avg_pool(x, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding="SAME") 119 | 120 | 121 | def local_response_norm(x): 122 | return tf.nn.lrn(x, depth_radius=5, bias=2, alpha=1e-4, beta=0.75) 123 | 124 | 125 | def batch_norm(x, n_out, phase_train, scope='bn', decay=0.9, eps=1e-5): 126 | """ 127 | Code taken from http://stackoverflow.com/a/34634291/2267819 128 | """ 129 | with tf.variable_scope(scope): 130 | beta = tf.get_variable(name='beta', shape=[n_out], initializer=tf.constant_initializer(0.0) 131 | , trainable=True) 132 | gamma = tf.get_variable(name='gamma', shape=[n_out], initializer=tf.random_normal_initializer(1.0, 0.02), 133 | trainable=True) 134 | batch_mean, batch_var = tf.nn.moments(x, [0, 1, 2], name='moments') 135 | ema = tf.train.ExponentialMovingAverage(decay=decay) 136 | 137 | def mean_var_with_update(): 138 | ema_apply_op = ema.apply([batch_mean, batch_var]) 139 | with tf.control_dependencies([ema_apply_op]): 140 | return tf.identity(batch_mean), tf.identity(batch_var) 141 | 142 | mean, var = tf.cond(phase_train, 143 | mean_var_with_update, 144 | lambda: (ema.average(batch_mean), ema.average(batch_var))) 145 | normed = tf.nn.batch_normalization(x, mean, var, beta, gamma, eps) 146 | return normed 147 | 148 | 149 | def process_image(image, mean_pixel): 150 | return image - mean_pixel 151 | 152 | 153 | def unprocess_image(image, mean_pixel): 154 | return image + mean_pixel 155 | 156 | 157 | def bottleneck_unit(x, out_chan1, out_chan2, down_stride=False, up_stride=False, name=None): 158 | 159 | def conv_transpose(tensor, out_channel, shape, strides, name=None): 160 | out_shape = tensor.get_shape().as_list() 161 | in_channel = out_shape[-1] 162 | kernel = weight_variable([shape, shape, out_channel, in_channel], name=name) 163 | shape[-1] = out_channel 164 | return tf.nn.conv2d_transpose(x, kernel, output_shape=out_shape, strides=[1, strides, strides, 1], 165 | padding='SAME', name='conv_transpose') 166 | 167 | def conv(tensor, out_chans, shape, strides, name=None): 168 | in_channel = tensor.get_shape().as_list()[-1] 169 | kernel = weight_variable([shape, shape, in_channel, out_chans], name=name) 170 | return tf.nn.conv2d(x, kernel, strides=[1, strides, strides, 1], padding='SAME', name='conv') 171 | 172 | def bn(tensor, name=None): 173 | """ 174 | :param tensor: 4D tensor input 175 | :param name: name of the operation 176 | :return: local response normalized tensor - not using batch normalization :( 177 | """ 178 | return tf.nn.lrn(tensor, depth_radius=5, bias=2, alpha=1e-4, beta=0.75, name=name) 179 | 180 | in_chans = x.get_shape().as_list()[3] 181 | 182 | if down_stride or up_stride: 183 | first_stride = 2 184 | else: 185 | first_stride = 1 186 | 187 | with tf.variable_scope('res%s' % name): 188 | if in_chans == out_chan2: 189 | b1 = x 190 | else: 191 | with tf.variable_scope('branch1'): 192 | if up_stride: 193 | b1 = conv_transpose(x, out_chans=out_chan2, shape=1, strides=first_stride, 194 | name='res%s_branch1' % name) 195 | else: 196 | b1 = conv(x, out_chans=out_chan2, shape=1, strides=first_stride, name='res%s_branch1' % name) 197 | b1 = bn(b1, 'bn%s_branch1' % name, 'scale%s_branch1' % name) 198 | 199 | with tf.variable_scope('branch2a'): 200 | if up_stride: 201 | b2 = conv_transpose(x, out_chans=out_chan1, shape=1, strides=first_stride, name='res%s_branch2a' % name) 202 | else: 203 | b2 = conv(x, out_chans=out_chan1, shape=1, strides=first_stride, name='res%s_branch2a' % name) 204 | b2 = bn(b2, 'bn%s_branch2a' % name, 'scale%s_branch2a' % name) 205 | b2 = tf.nn.relu(b2, name='relu') 206 | 207 | with tf.variable_scope('branch2b'): 208 | b2 = conv(b2, out_chans=out_chan1, shape=3, strides=1, name='res%s_branch2b' % name) 209 | b2 = bn(b2, 'bn%s_branch2b' % name, 'scale%s_branch2b' % name) 210 | b2 = tf.nn.relu(b2, name='relu') 211 | 212 | with tf.variable_scope('branch2c'): 213 | b2 = conv(b2, out_chans=out_chan2, shape=1, strides=1, name='res%s_branch2c' % name) 214 | b2 = bn(b2, 'bn%s_branch2c' % name, 'scale%s_branch2c' % name) 215 | 216 | x = b1 + b2 217 | return tf.nn.relu(x, name='relu') 218 | 219 | 220 | def add_to_regularization_and_summary(var): 221 | if var is not None: 222 | tf.summary.histogram(var.op.name, var) 223 | tf.add_to_collection("reg_loss", tf.nn.l2_loss(var)) 224 | 225 | 226 | def add_activation_summary(var): 227 | if var is not None: 228 | tf.summary.histogram(var.op.name + "/activation", var) 229 | tf.summary.scalar(var.op.name + "/sparsity", tf.nn.zero_fraction(var)) 230 | 231 | 232 | def add_gradient_summary(grad, var): 233 | if grad is not None: 234 | tf.summary.histogram(var.op.name + "/gradient", grad) 235 | 236 | 237 | ################ Get one-hot label ############################################ 238 | def get_one_hot_label(batch_size, num_classes, labels): 239 | one_hot_labels = np.zeros((batch_size, num_classes)) 240 | for i in range(batch_size): 241 | one_hot_labels[i][labels[i]] = 1 242 | 243 | return one_hot_labels 244 | 245 | 246 | 247 | ######################## Read data ############################################ 248 | def read_and_decode(filename_queue, shuffle_batch=True, random_noise=False): 249 | reader = tf.TFRecordReader() 250 | _, serialized_example = reader.read(filename_queue) 251 | features = tf.parse_single_example( 252 | serialized_example, 253 | features={ 254 | 'height': tf.FixedLenFeature([], tf.int64), 255 | 'width': tf.FixedLenFeature([], tf.int64), 256 | 'name': tf.FixedLenFeature([], tf.string), 257 | 'image_raw': tf.FixedLenFeature([], tf.string), 258 | 'mask_raw': tf.FixedLenFeature([], tf.string), 259 | 'label': tf.FixedLenFeature([], tf.int64) 260 | }) 261 | 262 | image = tf.decode_raw(features['image_raw'], tf.float64) 263 | image = tf.reshape(image, [300,300,3]) 264 | 265 | if random_noise: 266 | image = tf.image.random_brightness(image, max_delta=8) 267 | image = tf.image.random_contrast(image, lower=0.85, upper=1.15) 268 | 269 | mask = tf.decode_raw(features['mask_raw'], tf.float64) 270 | mask = tf.reshape(mask, [300,300]) 271 | 272 | name = features['name'] 273 | 274 | # Convert label from a scalar uint8 tensor to an int32 scalar. 275 | label = tf.cast(features['label'], tf.int64) 276 | 277 | 278 | if shuffle_batch: 279 | images, masks, names, labels = tf.train.shuffle_batch([image, mask, name, label], 280 | batch_size=4, 281 | capacity=4000, 282 | num_threads=2, 283 | min_after_dequeue=1000) 284 | else: 285 | images, masks, names, labels = tf.train.batch([image, mask, name, label], 286 | batch_size=4, 287 | capacity=4000, 288 | num_threads=2) 289 | return images, masks, names, labels -------------------------------------------------------------------------------- /File/generate_tfrecord.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python2 2 | # -*- coding: utf-8 -*- 3 | """ 4 | Created on Thu Apr 6 09:28:58 2017 5 | 6 | @author: wzg 7 | """ 8 | 9 | import tensorflow as tf 10 | import numpy as np 11 | import matplotlib.pyplot as plt 12 | from scipy import misc 13 | import scipy.io as sio 14 | 15 | 16 | def _bytes_feature(value): 17 | return tf.train.Feature(bytes_list = tf.train.BytesList(value=[value])) 18 | 19 | def _int64_feature(value): 20 | return tf.train.Feature(int64_list = tf.train.Int64List(value=[value])) 21 | 22 | 23 | root_path = './Data/' 24 | tfrecords_filename = root_path + 'tfrecords/test.tfrecords' 25 | writer = tf.python_io.TFRecordWriter(tfrecords_filename) 26 | 27 | 28 | height = 300 29 | width = 300 30 | meanfile = sio.loadmat(root_path + 'mats/mean300.mat') 31 | meanvalue = meanfile['mean'] #mean value of images in training set 32 | 33 | txtfile = root_path + 'txt/test.txt' 34 | fr = open(txtfile) 35 | 36 | for i in fr.readlines(): 37 | item = i.split() 38 | img = np.float64(misc.imread(root_path + '/images/test_images/' + item[0])) 39 | img = img - meanvalue 40 | maskmat = sio.loadmat(root_path + '/mats/test_mats/' + item[1]) 41 | mask = np.float64(maskmat['seg_mask']) 42 | label = int(item[2]) 43 | img_raw = img.tostring() 44 | mask_raw = mask.tostring() 45 | example = tf.train.Example(features=tf.train.Features(feature={ 46 | 'height': _int64_feature(height), 47 | 'width': _int64_feature(width), 48 | 'name': _bytes_feature(item[0]), 49 | 'image_raw': _bytes_feature(img_raw), 50 | 'mask_raw': _bytes_feature(mask_raw), 51 | 'label': _int64_feature(label)})) 52 | 53 | writer.write(example.SerializeToString()) 54 | 55 | writer.close() 56 | fr.close() 57 | 58 | ################### Test Correctness ##################################### 59 | record_iterator = tf.python_io.tf_record_iterator(path=tfrecords_filename) 60 | i = 0 61 | 62 | for string_record in record_iterator: 63 | if i>0: 64 | break 65 | example = tf.train.Example() 66 | example.ParseFromString(string_record) 67 | height = int(example.features.feature['height'] 68 | .int64_list 69 | .value[0]) 70 | 71 | width = int(example.features.feature['width'] 72 | .int64_list 73 | .value[0]) 74 | 75 | name = (example.features.feature['name'] 76 | .bytes_list 77 | .value[0]) 78 | 79 | img_string = (example.features.feature['image_raw'] 80 | .bytes_list 81 | .value[0]) 82 | 83 | mask_string = (example.features.feature['mask_raw'] 84 | .bytes_list 85 | .value[0]) 86 | 87 | label = (example.features.feature['label'] 88 | .int64_list 89 | .value[0]) 90 | 91 | img = np.fromstring(img_string, dtype=np.float64) 92 | mask = np.fromstring(mask_string, dtype=np.float64) 93 | reconstructed_img = img.reshape((height,width,-1)) 94 | reconstructed_img = reconstructed_img + meanvalue 95 | reconstructed_mask = mask.reshape((height,width)) 96 | 97 | print name 98 | print 'label: ' + str(label) 99 | plt.subplot(1,2,1) 100 | plt.imshow(np.uint8(reconstructed_img)) 101 | plt.subplot(1,2,2) 102 | plt.imshow(np.uint8(reconstructed_mask)) 103 | 104 | 105 | i += 1 106 | -------------------------------------------------------------------------------- /File/run.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python2 2 | # -*- coding: utf-8 -*- 3 | """ 4 | Created on Mon May 1 17:16:14 2017 5 | 6 | @author: wzg 7 | """ 8 | 9 | from __future__ import print_function 10 | import tensorflow as tf 11 | import numpy as np 12 | import TensorflowUtils as utils 13 | from six.moves import xrange 14 | import scipy.io as sio 15 | import scipy.misc as misc 16 | 17 | 18 | FLAGS = tf.flags.FLAGS 19 | tf.flags.DEFINE_integer("batch_size", "4", "batch size for training") 20 | tf.flags.DEFINE_string("logs_dir", "logs/", "path to logs directory") 21 | tf.flags.DEFINE_string("visualize_dir", "Data/visualize/", "path to visualzie directory") 22 | tf.flags.DEFINE_float("learning_rate", "1e-6", "Learning rate for Adam Optimizer") 23 | tf.flags.DEFINE_string("model_dir", "Model/", "Path to vgg model mat") 24 | tf.flags.DEFINE_bool('debug', "False", "Debug mode: True/ False") 25 | tf.flags.DEFINE_string('mode', "test", "Mode train/ test/ visualize") 26 | 27 | root_path = './Data/' 28 | train_tfrecord_filename = root_path + 'tfrecords/train.tfrecords' 29 | test_tfrecord_filename = root_path + 'tfrecords/test.tfrecords' 30 | MODEL_URL = 'http://www.vlfeat.org/matconvnet/models/beta16/imagenet-vgg-verydeep-19.mat' 31 | ckpt_path = './Model/MTV4_model.ckpt' 32 | 33 | MAX_ITERATION = 8000 34 | NUM_OF_CLASSESS = 6 35 | WEATHER_CLASSES = 2 36 | IMAGE_SIZE = 300 37 | EPOCHS = 5 38 | RESTORE = True 39 | 40 | meanfile = sio.loadmat(root_path + 'mats/mean300.mat') 41 | meanvalue = meanfile['mean'] #mean value of images in training set 42 | 43 | 44 | 45 | def vgg_net(weights, image): 46 | layers = ( 47 | 'conv1_1', 'relu1_1', 'conv1_2', 'relu1_2', 'pool1', 48 | 49 | 'conv2_1', 'relu2_1', 'conv2_2', 'relu2_2', 'pool2', 50 | 51 | 'conv3_1', 'relu3_1', 'conv3_2', 'relu3_2', 'conv3_3', 52 | 'relu3_3', 'conv3_4', 'relu3_4', 'pool3', 53 | 54 | 'conv4_1', 'relu4_1', 'conv4_2', 'relu4_2', 'conv4_3', 55 | 'relu4_3', 'conv4_4', 'relu4_4', 'pool4', 56 | 57 | 'conv5_1', 'relu5_1', 'conv5_2', 'relu5_2', 'conv5_3', 58 | 'relu5_3', 'conv5_4', 'relu5_4' 59 | ) 60 | 61 | net = {} 62 | current = image 63 | for i, name in enumerate(layers): 64 | kind = name[:4] 65 | if kind == 'conv': 66 | kernels, bias = weights[i][0][0][0][0] 67 | # matconvnet: weights are [width, height, in_channels, out_channels] 68 | # tensorflow: weights are [height, width, in_channels, out_channels] 69 | kernels = utils.get_variable(np.transpose(kernels, (1, 0, 2, 3)), name = name + "_w") 70 | bias = utils.get_variable(bias.reshape(-1), name = name + "_b") 71 | current = utils.conv2d_basic(current, kernels, bias) 72 | elif kind == 'relu': 73 | current = tf.nn.relu(current, name = name) 74 | tf.add_to_collection('activations', current) 75 | if FLAGS.debug: 76 | utils.add_activation_summary(current) 77 | elif kind == 'pool': 78 | current = utils.avg_pool_2x2(current) 79 | net[name] = current 80 | 81 | return net 82 | 83 | 84 | def inference(image, keep_prob): 85 | 86 | print("setting up vgg initialized conv layers ...") 87 | model_data = utils.get_model_data(FLAGS.model_dir, MODEL_URL) 88 | 89 | weights = np.squeeze(model_data['layers']) 90 | 91 | with tf.variable_scope("inference"): 92 | image_net = vgg_net(weights, image) 93 | conv_final_layer = image_net["conv5_3"] 94 | 95 | pool5 = utils.max_pool_2x2(conv_final_layer) 96 | 97 | W6_1 = utils.weight_variable([7, 7, 512, 4096], name = "W6_1") 98 | b6_1 = utils.bias_variable([4096], name = "b6_1") 99 | conv6_1 = utils.conv2d_basic(pool5, W6_1, b6_1) 100 | relu6_1 = tf.nn.relu(conv6_1, name = "relu6_1") 101 | if FLAGS.debug: 102 | utils.add_activation_summary(relu6_1) 103 | relu_dropout6_1 = tf.nn.dropout(relu6_1, keep_prob = keep_prob) 104 | 105 | 106 | 107 | W7_1 = utils.weight_variable([1, 1, 4096, 4096], name = "W7_1") 108 | b7_1 = utils.bias_variable([4096], name = "b7_1") 109 | conv7_1 = utils.conv2d_basic(relu_dropout6_1, W7_1, b7_1) 110 | relu7_1 = tf.nn.relu(conv7_1, name = "relu7_1") 111 | if FLAGS.debug: 112 | utils.add_activation_summary(relu7_1) 113 | relu_dropout7_1 = tf.nn.dropout(relu7_1, keep_prob = keep_prob) 114 | 115 | 116 | W8_1 = utils.weight_variable([1, 1, 4096, NUM_OF_CLASSESS], name = "W8_1") 117 | b8_1 = utils.bias_variable([NUM_OF_CLASSESS], name = "b8_1") 118 | conv8_1 = utils.conv2d_basic(relu_dropout7_1, W8_1, b8_1) 119 | if FLAGS.debug: 120 | tf.summary.histogram("conv8_1/activation", conv8_1) 121 | tf.summary.scalar("conv8_1/sparsity", tf.nn.zero_fraction(conv8_1)) 122 | 123 | 124 | # now to upscale to actual image size 125 | deconv_shape1 = image_net["pool4"].get_shape() 126 | W_t1 = utils.weight_variable([4, 4, deconv_shape1[3].value, NUM_OF_CLASSESS], name = "W_t1") 127 | b_t1 = utils.bias_variable([deconv_shape1[3].value], name = "b_t1") 128 | conv_t1 = utils.conv2d_transpose_strided(conv8_1, W_t1, b_t1, output_shape = tf.shape(image_net["pool4"])) 129 | fuse_1 = tf.add(conv_t1, image_net["pool4"], name = "fuse_1") 130 | 131 | if FLAGS.debug: 132 | tf.summary.histogram("conv_t1/activation", conv_t1) 133 | tf.summary.scalar("conv_t1/sparsity", tf.nn.zero_fraction(conv_t1)) 134 | tf.summary.histogram("fuse_1/activation", fuse_1) 135 | tf.summary.scalar("fuse_1/sparsity", tf.nn.zero_fraction(fuse_1)) 136 | 137 | deconv_shape2 = image_net["pool3"].get_shape() 138 | W_t2 = utils.weight_variable([4, 4, deconv_shape2[3].value, deconv_shape1[3].value], name = "W_t2") 139 | b_t2 = utils.bias_variable([deconv_shape2[3].value], name = "b_t2") 140 | conv_t2 = utils.conv2d_transpose_strided(fuse_1, W_t2, b_t2, output_shape = tf.shape(image_net["pool3"])) 141 | fuse_2 = tf.add(conv_t2, image_net["pool3"], name = "fuse_2") 142 | 143 | if FLAGS.debug: 144 | tf.summary.histogram("conv_t2/activation", conv_t2) 145 | tf.summary.scalar("conv_t2/sparsity", tf.nn.zero_fraction(conv_t2)) 146 | tf.summary.histogram("fuse_2/activation", fuse_2) 147 | tf.summary.scalar("fuse_2/sparsity", tf.nn.zero_fraction(fuse_2)) 148 | 149 | shape = tf.shape(image) 150 | deconv_shape3 = tf.stack([shape[0], shape[1], shape[2], NUM_OF_CLASSESS]) 151 | W_t3 = utils.weight_variable([16, 16, NUM_OF_CLASSESS, deconv_shape2[3].value], name = "W_t3") 152 | b_t3 = utils.bias_variable([NUM_OF_CLASSESS], name = "b_t3") 153 | conv_t3 = utils.conv2d_transpose_strided(fuse_2, W_t3, b_t3, output_shape = deconv_shape3, stride = 8) 154 | 155 | annotation_pred = tf.argmax(conv_t3, axis = 3, name = "prediction") 156 | 157 | if FLAGS.debug: 158 | tf.summary.histogram("conv_t3/activation", conv_t3) 159 | tf.summary.scalar("conv_t3/sparsity", tf.nn.zero_fraction(conv_t3)) 160 | tf.summary.histogram("annotation_pred/activation", annotation_pred) 161 | tf.summary.scalar("annotation_pred/sparsity", tf.nn.zero_fraction(annotation_pred)) 162 | 163 | 164 | 165 | ############################### Another branch of multi-task architecture ################################ 166 | W6_2 = utils.weight_variable([7, 7, 512, 1024], name = "W6_2") 167 | b6_2 = utils.bias_variable([1024], name = "b6_2") 168 | conv6_2 = utils.conv2d_basic(pool5, W6_2, b6_2) 169 | relu6_2 = tf.nn.relu(conv6_2, name = "relu6_2") 170 | if FLAGS.debug: 171 | utils.add_activation_summary(relu6_2) 172 | relu_dropout6_2 = tf.nn.dropout(relu6_2, keep_prob = keep_prob) 173 | 174 | W7_2 = utils.weight_variable([1, 1, 1024, 3840], name = "W7_2") 175 | b7_2 = utils.bias_variable([3840], name = "b7_2") 176 | conv7_2 = utils.conv2d_basic(relu_dropout6_2, W7_2, b7_2) 177 | relu7_2 = tf.nn.relu(conv7_2, name = "relu7_2") 178 | if FLAGS.debug: 179 | utils.add_activation_summary(relu7_2) 180 | relu_dropout7_2 = tf.nn.dropout(relu7_2, keep_prob = keep_prob) 181 | 182 | kernel_height = conv7_2.get_shape()[1] 183 | kernel_width = conv7_2.get_shape()[2] 184 | conv7_2_gapool = tf.nn.avg_pool(relu_dropout7_2, ksize = [1, kernel_height, kernel_width, 1], 185 | strides = [1, kernel_height, kernel_width, 1], padding = "SAME") 186 | 187 | kernel_height2 = fuse_2.get_shape()[1] 188 | kernel_width2 = fuse_2.get_shape()[2] 189 | fuse_2_gapool = tf.nn.avg_pool(fuse_2, ksize=[1, kernel_height2, kernel_width2, 1], 190 | strides = [1, kernel_height2, kernel_width2, 1], padding = "SAME") 191 | 192 | concat_res = tf.concat([conv7_2_gapool, fuse_2_gapool], axis = 3) 193 | concat_res = tf.squeeze(concat_res) 194 | 195 | 196 | W8_2 = utils.weight_variable([4096, WEATHER_CLASSES], name = "W8_2") 197 | b8_2 = utils.bias_variable([WEATHER_CLASSES], name = "b8_2") 198 | logits = tf.nn.bias_add(tf.matmul(concat_res, W8_2), b8_2) 199 | 200 | 201 | return tf.expand_dims(annotation_pred, dim = 3), conv_t3, logits 202 | 203 | 204 | def train(loss_val, var_list): 205 | optimizer = tf.train.AdamOptimizer(FLAGS.learning_rate) 206 | grads = optimizer.compute_gradients(loss_val, var_list = var_list) 207 | if FLAGS.debug: 208 | for grad, var in grads: 209 | utils.add_gradient_summary(grad, var) 210 | return optimizer.apply_gradients(grads) 211 | 212 | 213 | 214 | def main(argv = None): 215 | keep_probability = tf.placeholder(tf.float32, name = "keep_probabilty") 216 | image = tf.placeholder(tf.float32, shape = [None, IMAGE_SIZE, IMAGE_SIZE, 3], name = "input_image") 217 | annotation = tf.placeholder(tf.int32, shape = [None, IMAGE_SIZE, IMAGE_SIZE, 1], name = "annotation") 218 | label = tf.placeholder(tf.float32, [None, WEATHER_CLASSES], name = 'label') 219 | 220 | pred_annotation, logits, pred_label = inference(image, keep_probability) 221 | tf.summary.image("input_image", image, max_outputs = 2) 222 | tf.summary.image("ground_truth", tf.cast(annotation, tf.uint8), max_outputs = 2) 223 | tf.summary.image("pred_annotation", tf.cast(pred_annotation, tf.uint8), max_outputs = 2) 224 | 225 | seg_loss = tf.reduce_mean((tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits, 226 | labels=tf.squeeze(annotation, squeeze_dims = [3]), 227 | name = "seg_loss"))) 228 | 229 | classification_loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits = pred_label, labels = label, 230 | name = 'classification_loss')) 231 | 232 | regular_loss = tf.reduce_mean(tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)) 233 | 234 | total_loss = seg_loss + classification_loss + regular_loss 235 | 236 | tf.summary.scalar("seg_loss", seg_loss) 237 | tf.summary.scalar("classification_loss", classification_loss) 238 | tf.summary.scalar("total_loss", total_loss) 239 | 240 | with tf.name_scope("accuracy"): 241 | correct_pred = tf.equal(tf.argmax(pred_label, 1), tf.argmax(label, 1)) 242 | accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32)) 243 | 244 | tf.summary.scalar('accuracy', accuracy) 245 | 246 | trainable_var = tf.trainable_variables() 247 | if FLAGS.debug: 248 | for var in trainable_var: 249 | utils.add_to_regularization_and_summary(var) 250 | train_op = train(total_loss, trainable_var) 251 | 252 | print("Setting up summary op...") 253 | summary_op = tf.summary.merge_all() 254 | 255 | 256 | print("Setting up Saver...") 257 | saver = tf.train.Saver() 258 | summary_writer = tf.summary.FileWriter(FLAGS.logs_dir, tf.get_default_graph()) 259 | 260 | 261 | print("Setting up dataset reader") 262 | train_filename_queue = tf.train.string_input_producer([train_tfrecord_filename], 263 | num_epochs=EPOCHS) 264 | mytrain_images, mytrain_annotations, mytrain_names, mytrain_labels = utils.read_and_decode(train_filename_queue) 265 | 266 | 267 | test_filename_queue = tf.train.string_input_producer([test_tfrecord_filename], 268 | num_epochs=EPOCHS) 269 | mytest_images, mytest_annotations, mytest_names, mytest_labels = utils.read_and_decode(test_filename_queue, shuffle_batch=False) 270 | 271 | 272 | 273 | 274 | with tf.Session(graph=tf.get_default_graph()) as sess: 275 | init_op = tf.group(tf.global_variables_initializer(), 276 | tf.local_variables_initializer()) 277 | sess.run(init_op) 278 | 279 | if RESTORE: 280 | saver.restore(sess, ckpt_path) 281 | 282 | coord = tf.train.Coordinator() 283 | threads = tf.train.start_queue_runners(coord=coord) 284 | 285 | if FLAGS.mode == "train": 286 | for itr in xrange(MAX_ITERATION): 287 | #print('Iteration: %d' %itr) 288 | train_images, train_annotations, train_labels = sess.run([mytrain_images, mytrain_annotations, mytrain_labels]) 289 | train_annotations = np.array(np.expand_dims(train_annotations, axis = 3)) 290 | one_hot_labels = utils.get_one_hot_label(FLAGS.batch_size, WEATHER_CLASSES, train_labels) 291 | feed_dict = {image: train_images, annotation: train_annotations, label: one_hot_labels, keep_probability: 0.50} 292 | 293 | sess.run(train_op, feed_dict = feed_dict) 294 | 295 | if (itr + 1) % 2000 == 0: 296 | train_acc, sum_seg_loss, sum_classification_loss, sum_total_loss = 0, 0, 0, 0 297 | 298 | ####### Train: 1 epoch == 2000 iteration ################## 299 | for j in xrange(2000): 300 | train_images, train_annotations, train_labels = sess.run([mytrain_images, mytrain_annotations, mytrain_labels]) 301 | train_annotations = np.array(np.expand_dims(train_annotations, axis = 3)) 302 | one_hot_labels = utils.get_one_hot_label(FLAGS.batch_size, WEATHER_CLASSES, train_labels) 303 | feed_dict = {image: train_images, annotation: train_annotations, label: one_hot_labels, keep_probability: 1.0} 304 | 305 | acc, train_seg_loss, train_classification_loss, train_total_loss, \ 306 | summary_str = sess.run([accuracy, seg_loss, classification_loss, total_loss, summary_op], feed_dict = feed_dict) 307 | summary_writer.add_summary(summary_str, itr + 1 - 2000 + j) 308 | 309 | train_acc += acc 310 | sum_seg_loss += train_seg_loss 311 | sum_classification_loss += train_classification_loss 312 | sum_total_loss += train_total_loss 313 | 314 | train_acc /= 2000 315 | sum_seg_loss /= 2000 316 | sum_classification_loss /= 2000 317 | sum_total_loss /= 2000 318 | print("Step: %d, Train: accuracy: %g, seg_loss: %g, classification_loss: %g, total_loss: %g" % (itr + 1, train_acc, sum_seg_loss, sum_classification_loss, sum_total_loss)) 319 | 320 | test_acc, sum_seg_loss, sum_classification_loss, sum_total_loss = 0, 0, 0, 0 321 | ######### Test: 1 epoch == 500 iteration ################## 322 | 323 | saver.save(sess, FLAGS.model_dir + "your_model_name.ckpt", itr) 324 | 325 | elif FLAGS.mode == "test": 326 | test_acc, sum_seg_loss, sum_classification_loss, sum_total_loss = 0, 0, 0, 0 327 | for itr in xrange(500): 328 | test_images, test_annotations, test_labels = sess.run([mytest_images, mytest_annotations, mytest_labels]) 329 | test_annotations = np.array(np.expand_dims(test_annotations, axis = 3)) 330 | one_hot_labels = utils.get_one_hot_label(FLAGS.batch_size, WEATHER_CLASSES, test_labels) 331 | feed_dict = {image: test_images, annotation: test_annotations, label: one_hot_labels, keep_probability: 1.0} 332 | 333 | acc, test_seg_loss, test_classification_loss, test_total_loss = \ 334 | sess.run([accuracy, seg_loss, classification_loss, total_loss], feed_dict = feed_dict) 335 | 336 | test_acc += acc 337 | sum_seg_loss += test_seg_loss 338 | sum_classification_loss += test_classification_loss 339 | sum_total_loss += test_total_loss 340 | 341 | test_acc /= 500 342 | sum_seg_loss /= 500 343 | sum_classification_loss /= 500 344 | sum_total_loss /= 500 345 | print("Test: accuracy: %g, seg_loss: %g, classification_loss: %g, total_loss: %g" \ 346 | % (test_acc, sum_seg_loss, sum_classification_loss, sum_total_loss)) 347 | 348 | elif FLAGS.mode == "visualize": 349 | for k in xrange(10): 350 | test_images, test_annotations = sess.run([mytest_images, mytest_annotations]) 351 | test_annotations = np.array(np.expand_dims(test_annotations, axis = 3)) 352 | pred = sess.run(pred_annotation, feed_dict = {image: test_images, annotation: test_annotations, keep_probability: 1.0}) 353 | test_annotations = np.squeeze(test_annotations, axis = 3) 354 | pred = np.squeeze(pred, axis = 3) 355 | 356 | for itr in range(FLAGS.batch_size): 357 | misc.imsave(FLAGS.visualize_dir + "img_" + str(k*FLAGS.batch_size+itr) + '.jpg', (test_images[itr] + meanvalue).astype(np.uint8)) 358 | sio.savemat(FLAGS.visualize_dir + 'gt_' + str(k*FLAGS.batch_size+itr) + '.mat', {'mask':test_annotations[itr].astype(np.uint8)}) 359 | sio.savemat(FLAGS.visualize_dir + 'pred_' + str(k*FLAGS.batch_size+itr) + '.mat', {'mask':pred[itr].astype(np.uint8)}) 360 | print("Saved image: %d" % (k*FLAGS.batch_size+itr)) 361 | 362 | sess.close() 363 | coord.request_stop() 364 | coord.join(threads) 365 | 366 | 367 | if __name__ == "__main__": 368 | tf.app.run() -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Multitask_Weather 2 | 3 | Environment: 4 | ------------ 5 | python 2.7
6 | tensorflow 1.4.0 7 | 8 | Usage: 9 | ----- 10 | 1. use generate_tfrecord.py to construct tfrecord files for tensorflow 11 | 2. cd to the root path 12 | 3. change the mode in run.py (train, test and visualize) 13 | 4. python File/run.py 14 | 15 | Dataset 16 | ------ 17 | The pretrained model and raw data used in our paper are available at [Baidu network disk](https://pan.baidu.com/s/1pMDE2uv) or [MEGA network disk.](https://mega.nz/#F!I7hWmLQQ!iq3XgDR3C1PkX7BmDvfPXw) 18 | 19 | 20 |
21 | 22 | `Please refer to our paper for more details.` 23 | 24 | Citation: 25 | -------- 26 | ``` 27 | @inproceedings{DBLP:conf/mm/LiWL17, 28 |   author = {Xuelong Li and Zhigang Wang and Xiaoqiang Lu}, 29 | title = {A Multi-Task Framework for Weather Recognition}, 30 | booktitle = {Proceedings of the 2017 {ACM} on Multimedia Conference, {MM} 2017, Mountain View, CA, USA, October 23-27, 2017}, 31 | pages = {1318--1326}, 32 | year = {2017} 33 | } 34 | ``` 35 | --------------------------------------------------------------------------------