TensorFlow杂记 - 理解读入图片，转换shape，卷积，最后再保存图片

一段有点参考意义的代码，便于新手理解HWC, CHW, NCHW, NWHC的意义。

 

# 读入图片
src_img = tf.gfile.FastGFile('0000.jpg', 'rb').read()  # 0000.jpg  500*335   [w:500, h:335]
# 解码图片，得到HWC模式的RGB图像
img_rgb_uint8 = tf.image.decode_jpeg(src_img)  # HWC
# 转换数据格式，从uint8转换为float32
img_rgb_float32 = tf.image.convert_image_dtype(img_rgb_uint8, dtype=tf.float32)
# 转换shape，有HWC转换为CHW
img_CHW_ = tf.transpose(img_rgb_float32, perm=[2, 0, 1])
# 扩展一个维度，由CHW转换为NCHW
img_NCHW_ = tf.expand_dims(img_CHW_, 0)

# 定义一个卷积核，参数2为[核宽，核高，输入通道数，输出通道数]；定义一个偏执，参数2为[等同于输出通道数]
w = tf.get_variable("weight", [3, 3, 3, 3] ,initializer=tf.truncated_normal_initializer(stddev=0.1))
# b = tf.get_variable("biase", [64], initializer=tf.truncated_normal_initializer(0.0))
# 执行卷积，添加偏执
conv_ = tf.nn.conv2d(img_NCHW_, w, strides=[1, 1, 1, 1], padding='SAME', data_format='NCHW')
# ret = tf.identity(tf.nn.bias_add(conv, b, data_format='NCHW'))

# 降维，便于保存，由NCHW降为CHW
img_to_save_CHW_ = tf.squeeze(conv_, 0)
# 转换shape，由CHW转换为HWC（原始图片读入就是HWC模式）
img_to_save_HWC_ = tf.transpose(img_to_save_CHW_, perm=[1, 2, 0])

# 转换图片数据格式，由uint8转换为float32
encode_image_u8 = tf.image.convert_image_dtype(img_to_save_HWC_, dtype=tf.uint8)
# 编码为jpeg
conv_image = tf.image.encode_jpeg(encode_image_u8)

# 有的同学运行时会出现“could not create cudnn handle: CUDNN_STATUS_ALLOC_FAILED”的错误，故加上这两句
config = tf.ConfigProto()
config.gpu_options.allow_growth = True

with tf.Session(config=config) as sess:
    sess.run(tf.global_variables_initializer())
    img_CHW = sess.run(img_CHW_)
    print(img_CHW.shape)  # CHW
    img_NCHW = sess.run(img_NCHW_)
    print(img_NCHW.shape)  # NCHW
    conv = sess.run(conv_)
    print(conv.shape)  # NCHW
    img_to_save = sess.run(img_to_save_CHW_)
    print("img_to_save shape: ", img_to_save.shape)  # CHW
    img_to_save_HWC = sess.run(img_to_save_HWC_)
    print("img_to_save_HWC shape: ", img_to_save_HWC.shape)  # HWC
    # 图像编码可以输入两种格式，一种是Tensor类型，一种是ndarray类型
    # 必须是HWC模式
    save_image = sess.run(conv_image)
    # 保存文本到本地文件
    with tf.gfile.GFile("0000_S.jpg", 'wb') as f:
        f.write(save_image)