diff --git a/classify_nsfw.py b/classify_nsfw.py index e5c4c7b..9b4ae18 100644 --- a/classify_nsfw.py +++ b/classify_nsfw.py @@ -38,7 +38,7 @@ def main(argv): model = OpenNsfwModel() - with tf.Session() as sess: + with tf.compat.v1.Session() as sess: input_type = InputType[args.input_type.upper()] model.build(weights_path=args.model_weights, input_type=input_type) @@ -54,7 +54,7 @@ def main(argv): import base64 fn_load_image = lambda filename: np.array([base64.urlsafe_b64encode(open(filename, "rb").read())]) - sess.run(tf.global_variables_initializer()) + sess.run(tf.compat.v1.global_variables_initializer()) image = fn_load_image(args.input_file) diff --git a/model.py b/model.py index a03d9a1..cf928f0 100644 --- a/model.py +++ b/model.py @@ -27,18 +27,18 @@ def __init__(self): def build(self, weights_path="open_nsfw-weights.npy", input_type=InputType.TENSOR): - self.weights = np.load(weights_path, encoding="latin1").item() + self.weights = np.load(weights_path, encoding="latin1", allow_pickle=True).item() self.input_tensor = None if input_type == InputType.TENSOR: - self.input = tf.placeholder(tf.float32, + self.input = tf.compat.v1.placeholder(tf.float32, shape=[None, 224, 224, 3], name="input") self.input_tensor = self.input elif input_type == InputType.BASE64_JPEG: from image_utils import load_base64_tensor - self.input = tf.placeholder(tf.string, shape=(None,), name="input") + self.input = tf.compat.v1.placeholder(tf.string, shape=(None,), name="input") self.input_tensor = load_base64_tensor(self.input) else: raise ValueError("invalid input type '{}'".format(input_type)) @@ -52,7 +52,7 @@ def build(self, weights_path="open_nsfw-weights.npy", x = self.__batch_norm("bn_1", x) x = tf.nn.relu(x) - x = tf.layers.max_pooling2d(x, pool_size=3, strides=2, padding='same') + x = tf.compat.v1.layers.max_pooling2d(x, pool_size=3, strides=2, padding='same') x = self.__conv_block(stage=0, block=0, inputs=x, filter_depths=[32, 32, 128], @@ -97,7 +97,7 @@ def build(self, weights_path="open_nsfw-weights.npy", filter_depths=[256, 256, 1024], kernel_size=3) - x = tf.layers.average_pooling2d(x, pool_size=7, strides=1, + x = tf.compat.v1.layers.average_pooling2d(x, pool_size=7, strides=1, padding="valid", name="pool") x = tf.reshape(x, shape=(-1, 1024)) @@ -123,11 +123,11 @@ def __get_weights(self, layer_name, field_name): """Layer creation and weight initialization """ def __fully_connected(self, name, inputs, num_outputs): - return tf.layers.dense( + return tf.compat.v1.layers.dense( inputs=inputs, units=num_outputs, name=name, - kernel_initializer=tf.constant_initializer( + kernel_initializer=tf.compat.v1.constant_initializer( self.__get_weights(name, "weights"), dtype=tf.float32), - bias_initializer=tf.constant_initializer( + bias_initializer=tf.compat.v1.constant_initializer( self.__get_weights(name, "biases"), dtype=tf.float32)) def __conv2d(self, name, inputs, filter_depth, kernel_size, stride=1, @@ -147,26 +147,26 @@ def __conv2d(self, name, inputs, filter_depth, kernel_size, stride=1, raise Exception('unsupported kernel size for padding: "{}"' .format(kernel_size)) - return tf.layers.conv2d( + return tf.compat.v1.layers.conv2d( inputs, filter_depth, kernel_size=(kernel_size, kernel_size), strides=(stride, stride), padding='valid', activation=None, trainable=trainable, name=name, - kernel_initializer=tf.constant_initializer( + kernel_initializer=tf.compat.v1.constant_initializer( self.__get_weights(name, "weights"), dtype=tf.float32), - bias_initializer=tf.constant_initializer( + bias_initializer=tf.compat.v1.constant_initializer( self.__get_weights(name, "biases"), dtype=tf.float32)) def __batch_norm(self, name, inputs, training=False): - return tf.layers.batch_normalization( + return tf.compat.v1.layers.batch_normalization( inputs, training=training, epsilon=self.bn_epsilon, - gamma_initializer=tf.constant_initializer( + gamma_initializer=tf.compat.v1.constant_initializer( self.__get_weights(name, "scale"), dtype=tf.float32), - beta_initializer=tf.constant_initializer( + beta_initializer=tf.compat.v1.constant_initializer( self.__get_weights(name, "offset"), dtype=tf.float32), - moving_mean_initializer=tf.constant_initializer( + moving_mean_initializer=tf.compat.v1.constant_initializer( self.__get_weights(name, "mean"), dtype=tf.float32), - moving_variance_initializer=tf.constant_initializer( + moving_variance_initializer=tf.compat.v1.constant_initializer( self.__get_weights(name, "variance"), dtype=tf.float32), name=name) diff --git a/tools/export_tflite.py b/tools/export_tflite.py index 91ac607..6c39ee3 100644 --- a/tools/export_tflite.py +++ b/tools/export_tflite.py @@ -36,13 +36,13 @@ export_path = args.target input_type = InputType[args.input_type.upper()] - with tf.Session() as sess: + with tf.compat.v1.Session() as sess: model.build(weights_path=args.model_weights, input_type=input_type) - sess.run(tf.global_variables_initializer()) + sess.run(tf.compat.v1.global_variables_initializer()) - converter = tf.contrib.lite.TFLiteConverter.from_session(sess, [model.input], [model.predictions]) + converter = tf.compat.v1.lite.TFLiteConverter.from_session(sess, [model.input], [model.predictions]) tflite_model = converter.convert() with open(export_path, "wb") as f: