tfe keras 任意大小的灰度转换

灰度转换"/>

tfe keras 任意大小的灰度转换

原文链接: tfe keras 任意大小的灰度转换

上一篇: tfe 在静态图中的使用

下一篇: Keras 多输入多输出模型搭建

效果

Keras的model输入时必须是tensor类型,并且对类型很敏感,需要手动设置float32,否则报错

由于不规则形状无法使用批次计算,所以每次都是一张一张的计算loss,然后求和

完整代码

import tensorflow as tf
import tensorflow.contrib.slim as slim
from tensorflow import keras
import tensorflow.contrib.eager as tfe
import os
import os
import numpy as np
import scipy.misc as sm
import matplotlib.pyplot as plttf.enable_eager_execution()
root = r'D:\data\coco\train2014'
paths = np.array([os.path.join(root, name)for name in os.listdir(root)
])batch_size = 4def get_data():images = []grays = []for path in np.random.choice(paths, batch_size):image = sm.imread(path, mode='RGB').astype(np.float32) / 127.5 - 1.images.append(image)gray = sm.imread(path, mode='I').astype(np.float32) / 255.gray = np.expand_dims(gray, axis=2)grays.append(gray)# print(image.shape, gray.shape) # (510, 640, 3) (510, 640, 1)# 输入 -1,1  输出 0,1return images, graysclass Model(tf.keras.Model):def __init__(self):super(Model, self).__init__()in_x = keras.Input(batch_shape=(None, None, None, 3), dtype=tf.float32)net1 = keras.layers.Conv2D(32, 1, padding="SAME")(in_x)net3 = keras.layers.Conv2D(32, 3, padding="SAME")(in_x)net5 = keras.layers.Conv2D(32, 5, padding="SAME")(in_x)net_sum = keras.layers.add([net1, net3, net5])net = keras.layers.concatenate([net1, net3, net5, net_sum])net = keras.layers.Conv2D(1, 1, padding="SAME")(net)self.model = keras.Model(inputs=[in_x], outputs=[net])def call(self, in_x):return self.model(in_x)def loss(model, images, targets):ls = [tf.reduce_mean((model(tf.expand_dims(image, 0)) - tf.expand_dims(gray, 0)) ** 2)for image, gray in zip(images, targets)]return tf.reduce_sum(ls)def main():lr = .01optimizer = tf.train.AdamOptimizer(lr)model = Model()show_step = 10train_step = 200for i in range(1, 1 + train_step):images, grays = get_data()optimizer.minimize(lambda: loss(model, images, grays))if not i % show_step:print(loss(model, images, grays).numpy())#  输入model的必须是tf的tensorimage = images[0]image = tf.expand_dims(image, 0)gray = model(image)image = ((images[0] + 1) * 127.5).astype(np.uint8)gray = np.squeeze(gray, axis=2)plt.imshow(image)plt.show()plt.imshow(gray, cmap='gray')plt.show()# 保存训练参数checkpoint_dir = './save/'checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")root = tfe.Checkpoint(optimizer=optimizer,model=model,optimizer_step=tf.train.get_or_create_global_step())root.save(file_prefix=checkpoint_prefix)def test():lr = .01optimizer = tf.train.AdamOptimizer(lr)model = Model()# 载入训练参数checkpoint_dir = './save/'root = tfe.Checkpoint(optimizer=optimizer,model=model,optimizer_step=tf.train.get_or_create_global_step())root.restore(tf.train.latest_checkpoint(checkpoint_dir))images, grays = get_data()#  输入model的必须是tf的tensorimage = images[0]image = tf.expand_dims(image, 0)gray = model(image)image = ((images[0] + 1) * 127.5).astype(np.uint8)gray = np.squeeze(gray, axis=2)plt.imshow(image)plt.show()plt.imshow(gray, cmap='gray')plt.show()if __name__ == '__main__':# main()test()

loss

1.3053823
0.029689737
0.0134448
0.02779741
0.01262876
0.010925731
0.0044341944
0.0021283496