用Python编写的Nets.mobilenet_v1中文标题生成器
import tensorflow as tf
from tensorflow.contrib import slim
import os
import numpy as np
class MobileNetV1:
def __init__(self):
self.width_multiplier = 1.0
self.model_dir = 'mobilenet_v1_weights'
def build_net(self):
self.inputs = tf.placeholder(dtype=tf.float32, shape=[None, 224, 224, 3], name='inputs')
self.outputs = self.mobilenet_v1(self.inputs, num_classes=1000,
is_training=False, width_multiplier=self.width_multiplier)
def load_pretrained_weights(self, sess):
var_list = slim.get_model_variables()
saver = tf.train.Saver(var_list)
saver.restore(sess, os.path.join(self.model_dir, 'mobilenet_v1_1.0_224.ckpt'))
def preprocess_image(self, image_path):
image = tf.read_file(image_path)
image = tf.image.decode_jpeg(image, channels=3)
image = tf.image.resize_images(image, [224, 224])
image = tf.image.convert_image_dtype(image, dtype=tf.float32)
image = tf.expand_dims(image, 0)
return image
def predict(self, image_path):
with tf.Session() as sess:
self.build_net()
self.load_pretrained_weights(sess)
image = self.preprocess_image(image_path)
image_np = sess.run(image)
feed_dict = {self.inputs: image_np}
predictions = sess.run(self.outputs, feed_dict=feed_dict)
return np.argmax(predictions)
def mobilenet_v1(self, inputs, num_classes=1000, is_training=True, width_multiplier=1.0):
with slim.arg_scope([slim.conv2d, slim.fully_connected],
activation_fn=self._prelu,
weights_regularizer=slim.l2_regularizer(0.00004),
biases_initializer=tf.zeros_initializer()):
with slim.arg_scope([slim.conv2d],
padding='SAME'):
network = self._conv2d(inputs, 32, stride=2, scope='conv_1')
network = self._depthwise_conv2d(network, 64, width_multiplier, stride=1, scope='conv_ds_2')
network = self._depthwise_conv2d(network, 128, width_multiplier, stride=2, scope='conv_ds_3')
network = self._depthwise_conv2d(network, 128, width_multiplier, stride=1, scope='conv_ds_4')
network = self._depthwise_conv2d(network, 256, width_multiplier, stride=2, scope='conv_ds_5')
network = self._depthwise_conv2d(network, 256, width_multiplier, stride=1, scope='conv_ds_6')
network = self._depthwise_conv2d(network, 512, width_multiplier, stride=2, scope='conv_ds_7')
num_residual_units = 5
for i in range(1, num_residual_units + 1):
network = self._depthwise_conv2d(network, 512, width_multiplier, stride=1, scope='conv_ds_' + str(7 + i))
network = self._depthwise_conv2d(network, 1024, width_multiplier, stride=2, scope='conv_ds_13')
network = self._depthwise_conv2d(network, 1024, width_multiplier, stride=1, scope='conv_ds_14')
# Add Average Pooling
network = slim.avg_pool2d(network, [7, 7], padding='VALID', scope='avgpool_15')
# Classify logits
network = slim.flatten(network, scope='flatten_16')
logits = slim.fully_connected(network, num_classes, activation_fn=None,
biases_initializer=tf.zeros_initializer(), scope='logits_17')
return logits
def _conv2d(self, inputs, filters, kernel_size=3, stride=1, scope=None):
return slim.conv2d(inputs, filters, kernel_size, stride, scope=scope)
def _depthwise_conv2d(self, inputs, filters, width_multiplier=1.0, kernel_size=3, stride=1, scope=None):
return slim.separable_conv2d(inputs, None, kernel_size, depth_multiplier=width_multiplier, stride=stride,
scope=scope)
def _prelu(self, inputs):
alpha = tf.get_variable('alpha', shape=inputs.get_shape()[-1], dtype=inputs.dtype,
initializer=tf.constant_initializer(0.25))
pos = tf.nn.relu(inputs)
neg = alpha * (inputs - abs(inputs)) * 0.5
return pos + neg
if __name__ == '__main__':
mobilenet = MobileNetV1()
image_path = 'example.jpg'
predicted_class = mobilenet.predict(image_path)
print('Predicted class:', predicted_class)
