| 72 | |
| 73 | |
| 74 | class Vgg16: |
| 75 | vgg_mean = [103.939, 116.779, 123.68] |
| 76 | |
| 77 | def __init__(self, vgg16_npy_path=None, restore_from=None): |
| 78 | # pre-trained parameters |
| 79 | try: |
| 80 | self.data_dict = np.load(vgg16_npy_path, encoding='latin1').item() |
| 81 | except FileNotFoundError: |
| 82 | print('Please download VGG16 parameters from here https://mega.nz/#!YU1FWJrA!O1ywiCS2IiOlUCtCpI6HTJOMrneN-Qdv3ywQP5poecM\nOr from my Baidu Cloud: https://pan.baidu.com/s/1Spps1Wy0bvrQHH2IMkRfpg') |
| 83 | |
| 84 | self.tfx = tf.placeholder(tf.float32, [None, 224, 224, 3]) |
| 85 | self.tfy = tf.placeholder(tf.float32, [None, 1]) |
| 86 | |
| 87 | # Convert RGB to BGR |
| 88 | red, green, blue = tf.split(axis=3, num_or_size_splits=3, value=self.tfx * 255.0) |
| 89 | bgr = tf.concat(axis=3, values=[ |
| 90 | blue - self.vgg_mean[0], |
| 91 | green - self.vgg_mean[1], |
| 92 | red - self.vgg_mean[2], |
| 93 | ]) |
| 94 | |
| 95 | # pre-trained VGG layers are fixed in fine-tune |
| 96 | conv1_1 = self.conv_layer(bgr, "conv1_1") |
| 97 | conv1_2 = self.conv_layer(conv1_1, "conv1_2") |
| 98 | pool1 = self.max_pool(conv1_2, 'pool1') |
| 99 | |
| 100 | conv2_1 = self.conv_layer(pool1, "conv2_1") |
| 101 | conv2_2 = self.conv_layer(conv2_1, "conv2_2") |
| 102 | pool2 = self.max_pool(conv2_2, 'pool2') |
| 103 | |
| 104 | conv3_1 = self.conv_layer(pool2, "conv3_1") |
| 105 | conv3_2 = self.conv_layer(conv3_1, "conv3_2") |
| 106 | conv3_3 = self.conv_layer(conv3_2, "conv3_3") |
| 107 | pool3 = self.max_pool(conv3_3, 'pool3') |
| 108 | |
| 109 | conv4_1 = self.conv_layer(pool3, "conv4_1") |
| 110 | conv4_2 = self.conv_layer(conv4_1, "conv4_2") |
| 111 | conv4_3 = self.conv_layer(conv4_2, "conv4_3") |
| 112 | pool4 = self.max_pool(conv4_3, 'pool4') |
| 113 | |
| 114 | conv5_1 = self.conv_layer(pool4, "conv5_1") |
| 115 | conv5_2 = self.conv_layer(conv5_1, "conv5_2") |
| 116 | conv5_3 = self.conv_layer(conv5_2, "conv5_3") |
| 117 | pool5 = self.max_pool(conv5_3, 'pool5') |
| 118 | |
| 119 | # detach original VGG fc layers and |
| 120 | # reconstruct your own fc layers serve for your own purpose |
| 121 | self.flatten = tf.reshape(pool5, [-1, 7*7*512]) |
| 122 | self.fc6 = tf.layers.dense(self.flatten, 256, tf.nn.relu, name='fc6') |
| 123 | self.out = tf.layers.dense(self.fc6, 1, name='out') |
| 124 | |
| 125 | self.sess = tf.Session() |
| 126 | if restore_from: |
| 127 | saver = tf.train.Saver() |
| 128 | saver.restore(self.sess, restore_from) |
| 129 | else: # training graph |
| 130 | self.loss = tf.losses.mean_squared_error(labels=self.tfy, predictions=self.out) |
| 131 | self.train_op = tf.train.RMSPropOptimizer(0.001).minimize(self.loss) |