用tensorflow的conv2d卷积实现均值blur

import numpy as np
import cv2
import tensorflow as tf
import time


blur_size = 31


s_time = time.time()

# (height, width, channel)BGR  -- > HWC,RGB
# img = cv2.imread('img_src.jpg')[:,:,::-1]
img = cv2.imread('img_src.jpg')
(img_b, img_g, img_r) = cv2.split(img)

#print("img   shape:{}".format(img.shape))
#print("img_b shape:{}".format(img_b.shape))
read_time = time.time()
print("\n read time：{} s \n".format(read_time - s_time))

# NHWC
input_b = np.expand_dims(img_b, axis=(0, 3)).astype(np.float64)
input_g = np.expand_dims(img_g, axis=(0, 3)).astype(np.float64)
input_r = np.expand_dims(img_r, axis=(0, 3)).astype(np.float64)

#print("input_b shape:{}".format(input_b.shape))
expand_time = time.time()
print("\n expand time：{} s \n".format(expand_time - read_time))


# filter_height, filter_width, in_channels, out_channels
mean_filter = tf.ones((blur_size, blur_size, 1, 1), dtype=tf.float64) / (blur_size * blur_size)


mean_filter_time = time.time()
print("\n mean_filter time：{} s \n".format(mean_filter_time - expand_time))

img_mean_tf_b = tf.nn.conv2d(input_b, filter=mean_filter, strides=[1,1,1,1], padding='SAME', name='conv_b')
img_mean_tf_g = tf.nn.conv2d(input_g, filter=mean_filter, strides=[1,1,1,1], padding='SAME', name='conv_g')
img_mean_tf_r = tf.nn.conv2d(input_r, filter=mean_filter, strides=[1,1,1,1], padding='SAME', name='conv_r')

img_mean_time = time.time()
print("\n img_mean_time time：{} s \n".format(img_mean_time - mean_filter_time))

with tf.compat.v1.Session() as sess:
	img_mean_tf_b = sess.run(img_mean_tf_b)[0, :, :, :]
	img_mean_tf_g = sess.run(img_mean_tf_g)[0, :, :, :]
	img_mean_tf_r = sess.run(img_mean_tf_r)[0, :, :, :]

run_time = time.time()
print("\n run time：{} s \n".format(run_time - img_mean_time))

e_time = time.time()
print("\n [conv2d]Total time：{} s, blur_size: {}, img: {} \n".format(e_time - s_time, blur_size, img.shape))


img_mean_tf_b = img_mean_tf_b.astype(np.uint8)
img_mean_tf_g = img_mean_tf_g.astype(np.uint8)
img_mean_tf_r = img_mean_tf_r.astype(np.uint8)
img_merged = cv2.merge([img_mean_tf_b, img_mean_tf_g, img_mean_tf_r])  #合并R、G、B分量

cv2.imwrite("./image_blur.png", img_merged)