图像清晰化代码实战详解——基于tensorflow2.x(适合新手)
Posted Ard丶
tags:
篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了图像清晰化代码实战详解——基于tensorflow2.x(适合新手)相关的知识,希望对你有一定的参考价值。
观看前需要先大概了解卷积运算,和tensorflow2基本操作,本文主要详解代码实战,原理仅简单介绍
原理简介:
模型结构为U-net,主要通过一系列卷积和反卷积,最终将原图像大小翻倍(像素翻倍),实现清晰化。通过跨层concat(拼接),可以使模型同时具有提取抽象特征以及提取细节的能力。(注意concat是某个维度上的拼接,而不是直接相加)。模型的输入是一张图片,输出的图片长度宽度均为原图的两倍,可以达到比原图更清晰的效果。U-net结构图如下,转自知乎。
代码详解
数据集图片文件处理
数据集可以是自己在一个文件夹里放一堆图片。图片可以是风景,可以是漫画,可以是其他,尽量是同一种类型的图片(不要将现实照片与漫画放一起),这样有针对性的训练效果会比较好。
import os
#参数
path1 = 'A' #存放图片的文件夹,需要根据自己的文件夹路径改
path2 = 'B' #如果只有一个文件夹,可以把这个注释掉
#获取所有图片的路径列表
files_list1 = [os.path.join(path1, name) for name in os.listdir(path1)]
files_list2 = [os.path.join(path2, name) for name in os.listdir(path2)]
files_list = files_list1 + files_list2 #生成所有图片的路径列表
#删除不符合尺寸的图片
kk = 0
im_size = 512
for b in files_list:
img = tf.io.read_file(b)
test_pic = tf.io.decode_image(img, channels=3) #解码图片
if test_pic.shape[0] <im_size or test_pic.shape[1]<im_size:
os.remove(b)
kk += 1
print('被删除的图片有%d个'%kk)
os.listdir将列出当前路径的所有文件名,后用os.path.join把文件夹路径和文件名拼接,最后添加到一个列表里,列表包含了一个文件夹内的所有文件路径。
由于图片有RGB三个通道,所以需要用tf.io.decode_image(img, channels=3)解码图片,解码后的形状为(高度,宽度,3),用if来判断高度或宽度过小的图片,这些图片不利于训练,将被删除
以上操作将删除文件夹中图片长度或者宽度小于512的图片,这是为了在接下来的训练中,能够对数据集进行统一尺寸的裁剪。512这个数值并不是固定的,可以改成256,128之类的数值。这取决与你训练时想要用的图片大小。(由于会删除图片,可以在运行前先备份一下图片文件夹,以免删除重要图片)
处理完之后,可以进行下一步
数据预处理
导入库:
import tensorflow as tf
from tensorflow import keras
import os
数据预处理方法
#参数
train_input_size = 256 #训练时的输入图片大小,可以调整
train_target_size = train_input_size*2 #训练的目标图片大小,最大值是上文判断的大小
jpge_quality_min = 50 #模糊最小值(取值0-1,越大图片越清晰)
jpge_quality_max = 75 #模糊最大值
noise_stddev=0.015 #噪声标准差
@tf.function
def trans_func(files):
img_bin = tf.io.read_file(files)
#调用tf.cond判断图片是否为jepg(jpg),如果是则用jpeg解码方式,如果不是则用png解码方式
img_decode = tf.cond(
tf.image.is_jpeg(img_bin),
lambda: tf.image.decode_jpeg(img_bin, channels=3),
lambda: tf.image.decode_png(img_bin, channels=3)
)
target_img = tf.image.convert_image_dtype(img_decode, tf.float32) # 转为tf.float32
target_img = tf.image.random_crop(target_img, [train_target_size, train_target_size, 3]) #随机裁剪
target_img = tf.image.random_flip_left_right(target_img) # 随机左右翻转
target_img = tf.image.random_contrast(target_img, 0.25, 0.75) #随机调整图像对比度
target_img = tf.image.random_brightness(target_img, 0.2) #随机调整图像亮度
train_img = tf.image.resize(target_img, [train_input_size, train_input_size],
method=train_resize_mhd) # 缩放到固定大小,这个大小是输出的一半
train_img = tf.image.random_jpeg_quality(train_img, min_jpeg_quality=jpge_quality_min,
max_jpeg_quality=jpge_quality_max) #模糊处理
noise = tf.random.normal(shape=[train_input_size, train_input_size, 3], stddev=noise_stddev)
train_img = train_img + noise #添加噪声
return [train_img, target_img]
tf.cond()是先看第一个参数bool值,如果为True,则调用jepg(jpg)解码方法,如果为False则用png解码方式。
用tf.image.convert_image_dtype()而不是tf.cast(),虽然两者对张量操作可以达到相同的效果,但是后者在之后转unit8类型数据会出问题,导致图片编码异常,所以只能用前者。
这个方法,首先对图片进行了解码,并且转为tf.float32便于模型训练,并且进行了随机翻转和微调,增加了数据的多样性。
因为输出值target_image形状大小(高度,宽度)是输入值train_img的两倍。之后将图片resize成原来的一半大小,作为模型训练的输入值。为了让模型有到更强的效果,所以针对输入值train_img,要添加模糊化(random_jpeg_quality)以及噪声(noise),增加训练难度。
最后把[train_img, target_img]返回
tf数据集生成
path1 = 'A' #第一个存放图片的文件夹
path2 = 'B' #第二个存放图片的文件夹,如果只有一个图片文件夹可以把这个删掉
#获取所有图片路径列表
files_list1 = [os.path.join(path1, name) for name in os.listdir(path1)]
files_list2 = [os.path.join(path2, name) for name in os.listdir(path2)]
files_list = files_list1 + files_list2 #图片路径列表拼接
print('图片总数量:',len(files_list))
#生成tf的Dataset
ds = tf.data.Dataset.from_tensor_slices(files_list)
ds = ds.shuffle(10000) #随机打乱
ds = ds.map(trans_func, num_parallel_calls=-1, deterministic=False)
#调用上面的预处理函数
ds = ds.batch(8) #设置batch大小
生成dataset,之后就可以构建模型,开始训练了
U-net模型构建
class Repairer(keras.Model):
def __init__(self):
super(Repairer, self).__init__()
self.level_1_conv1 = keras.layers.Conv2D(64, 3, 1, 'same', activation='elu')
self.level_1_conv2 = keras.layers.Conv2D(64, 3, 1, 'same', activation='elu')
self.level_1_pool = keras.layers.MaxPooling2D()
self.level_2_conv1 = keras.layers.Conv2D(128, 3, 1, 'same', activation='elu')
self.level_2_conv2 = keras.layers.Conv2D(128, 3, 1, 'same', activation='elu')
self.level_2_pool = keras.layers.MaxPooling2D()
self.level_3_conv1 = keras.layers.Conv2D(256, 3, 1, 'same', activation='elu')
self.level_3_conv2 = keras.layers.Conv2D(256, 3, 1, 'same', activation='elu')
self.level_3_pool = keras.layers.MaxPooling2D()
self.level_4_conv1 = keras.layers.Conv2D(512, 3, 1, 'same', activation='elu')
self.level_4_conv2 = keras.layers.Conv2D(512, 3, 1, 'same', activation='elu')
self.level_4_pool = keras.layers.MaxPooling2D()
self.level_3_conv1t = keras.layers.Conv2DTranspose(256, 3, 2, 'same', activation='elu')
self.level_3_concat = keras.layers.Concatenate()
self.level_3_conv2t = keras.layers.Conv2DTranspose(256, 3, 1, 'same', activation='elu')
self.level_2_conv1t = keras.layers.Conv2DTranspose(128, 3, 2, 'same', activation='elu')
self.level_2_concat = keras.layers.Concatenate()
self.level_2_conv2t = keras.layers.Conv2DTranspose(128, 3, 1, 'same', activation='elu')
self.level_1_conv1t = keras.layers.Conv2DTranspose(64, 3, 2, 'same', activation='elu')
self.level_1_concat = keras.layers.Concatenate()
self.level_1_conv2t = keras.layers.Conv2DTranspose(64, 3, 1, 'same', activation='elu')
self.level_0_conv1t = keras.layers.Conv2DTranspose(64, 3, 2, 'same', activation='elu')
self.level_0_concat = keras.layers.Concatenate()
self.level_0_conv2t = keras.layers.Conv2DTranspose(64, 3, 2, 'same', activation='elu')
self.level_0_conv3t = keras.layers.Conv2D(64, 3, 1, 'same', activation='elu')
self.level_0_conv4t = keras.layers.Conv2D(64, 3, 1, 'same', activation='linear')
self.model_output = keras.layers.Conv2D(3, 3, 1, 'same', activation='elu')
def call(self,image_tffloat32):
level_1_conv1 = self.level_1_conv1(image_tffloat32)
level_1_conv2 = self.level_1_conv2(level_1_conv1)
level_1_pool = self.level_1_pool(level_1_conv2)
level_2_conv1 = self.level_2_conv1(level_1_pool)
level_2_conv2 = self.level_2_conv2(level_2_conv1)
level_2_pool = self.level_2_pool(level_2_conv2)
level_3_conv1 = self.level_3_conv1(level_2_pool)
level_3_conv2 = self.level_3_conv2(level_3_conv1)
level_3_pool = self.level_3_pool(level_3_conv2)
level_4_conv1 = self.level_4_conv1(level_3_pool)
level_4_conv2 = self.level_4_conv2(level_4_conv1)
level_4_pool = self.level_4_pool(level_4_conv2)
level_3_conv1t = self.level_3_conv1t(level_4_pool)
level_3_concat = self.level_3_concat([level_3_conv1t, level_4_conv2])
level_3_conv2t = self.level_3_conv2t(level_3_concat)
level_2_conv1t = self.level_2_conv1t(level_3_conv2t)
level_2_concat = self.level_2_concat([level_2_conv1t, level_3_conv2])
level_2_conv2t = self.level_2_conv2t(level_2_concat)
level_1_conv1t = self.level_1_conv1t(level_2_conv2t)
level_1_concat = self.level_1_concat([level_1_conv1t, level_2_conv1])
level_1_conv2t = self.level_1_conv2t(level_1_concat)
level_0_conv1t = self.level_0_conv1t(level_1_conv2t)
level_0_concat = self.level_0_concat([level_0_conv1t, level_1_conv1])
level_0_conv2t = self.level_0_conv2t(level_0_concat)
level_0_conv3t = self.level_0_conv3t(level_0_conv2t)
level_0_conv4t = self.level_0_conv4t(level_0_conv3t)
model_output = self.model_output(level_0_conv4t)
return model_output
我在这里全用了elu激活,也可以修改成别的。模型输入形状是[b,256,256,3],输出是[b,512,512,3],可以将图片高度与宽度翻倍,像素增加。b是batch大小,256是我的模型输入图高度宽度,512是模型输出图高度宽度,这个可以根据自己数据处理时的裁剪的大小来修改。(注意:由于模型对与形状的限制,高与宽必须是16的倍数)
开始训练!
epochs = 10 #迭代数据集次数
optimizer=keras.optimizers.Adam(learning_rate=0.0000003) #优化器
repairer = Repairer()
#加载曾经的权重
if os.path.exists('repairer.h5'):
print('加载权重...')
repairer.load_weights('repairer.h5')
#训练
for epoch in range(epochs):
total_loss = []
for n,(x,y) in enumerate(ds):
with tf.GradientTape() as tape:
y_pred = repairer(x)
loss = keras.backend.mean(keras.backend.abs(y - y_pred))
grads = tape.gradient(loss, repairer.trainable_variables)
optimizer.apply_gradients(zip(grads, repairer.trainable_variables))
total_loss.append(loss)
print('epoch:',epoch,' , loss:',sum(total_loss)/(n+1))
repairer.save_weights('repairer.h5')
训练并且自动保存模型权重文件,如果当前文件夹已经有模型权重文件,则加载并开始训练
生成清晰图片
#生成清晰图片保存到当前文件夹
#注意输入图片的大小,必须把图片像素裁剪成16的倍数(高宽可以不一样)
test_pic_bin = tf.io.read_file('xxx.jpg') #读取要被清晰化的图片路径
test_pic = tf.io.decode_image(test_pic_bin, channels=3) #解码图片
test_pic = tf.image.convert_image_dtype(test_pic, tf.float32) #转为float32给模型训练
tshape = test_pic.shape #获取图片形状
x_tshape = int(tshape[1]/16)*16 #输入高度,将把图像微微裁剪一点点(因为要满足16的倍数)
y_tshape = int(tshape[0]/16)*16 #输入宽度,将把图像微微裁剪一点点(因为要满足16的倍数)
print(y_tshape,x_tshape)
test_pic = tf.image.random_crop(test_pic, [y_tshape, x_tshape, 3]) #裁剪
test_pic = tf.expand_dims(test_pic, axis=0) #增加一个batch维度,便于输入模型
output = repairer(test_pic) #输入模型,生成结果
output_image = tf.clip_by_value(output,0,1) #把数值裁剪到0-1的范围内
output_image = tf.image.convert_image_dtype(output_image, dtype=tf.uint8) #转为unit8
output_image = tf.image.encode_png(output_image[0]) #编码
tf.io.write_file(r'test1.png', output_image) #保存输出的清晰图片到文件夹
需要将图片高度和宽度裁剪为16的倍数,所以先获取图片形状,将其除16,并用int化为整数,再乘16,以这个尺寸为传入图像,只会比原图稍微小一点点。
因为颜色值的范围是0-1之间,而模型有可能会预测出超出这个范围的数值,所以需要把模型输出的值全部裁剪到0-1的范围,再进行编码。
个人测试效果如下
可以看出来,在简单的任务上,该模型可以弱化锯齿效果,使线条更平滑。
将风景图放大后,可以发现物体边缘模糊的地方变得更清晰一些了。
本人水平有限,文中的不足与错误之处,恳请大家纠正。
如果我的文章对您有用,请点个赞或随意打赏。您的支持将鼓励我继续创作!
以上是关于图像清晰化代码实战详解——基于tensorflow2.x(适合新手)的主要内容,如果未能解决你的问题,请参考以下文章
实现迁徙学习-《Tensorflow 实战Google深度学习框架》代码详解
《21个项目玩转深度学习:基于TensorFlow的实践详解》高清带标签PDF版本学习下载