记录使用tensorflow实现大卷积核卷积的代码

#-*- coding:utf-8 -*-
import tensorflow as tf
import cv2
import numpy as np
import matplotlib.pyplot as plt
## read img #####
tm_path = '/home/wdh/pytorch-CycleGAN-and-pix2pix1_run/datasets/maps_0/testA/1_A.jpg'
ref_path = '/home/wdh/pytorch-CycleGAN-and-pix2pix1_run/datasets/maps_0/testB/1_B.jpg'

img_ref_bgr = cv2.imread(ref_path)
img_tm_bgr = cv2.imread(tm_path)
img_ref_gray = (cv2.cvtColor(img_ref_bgr,cv2.COLOR_BGR2GRAY))/255.0
img_src_gray = (cv2.cvtColor(img_tm_bgr,cv2.COLOR_BGR2GRAY))/255.0

[h_ref,w_ref] = img_ref_gray.shape
[h_tm,w_tm] = img_src_gray.shape
h_tm = int(h_tm/3)
w_tm = int(w_tm/3)

img_tm_gray = np.zeros([h_ref,w_ref])
img_tm_gray[0:h_tm,0:w_tm] = np.copy(img_src_gray[h_tm:h_tm+h_tm,w_tm:w_tm+w_tm])

###########################

########build compute graph#################
Var_tm = tf.Variable(img_tm_gray,dtype=tf.float32)
Var_ref = tf.Variable(img_ref_gray,dtype=tf.float32)
F_tm = tf.fft2d(tf.complex(Var_tm,tf.zeros(Var_tm.shape)))
F_ref = tf.fft2d(tf.complex(Var_ref,tf.zeros(Var_ref.shape)))

F_m = tf.multiply(F_ref,F_tm)
# F_m = F_ref * F_tm
res = tf.real(tf.ifft2d(F_m))

optmizer = tf.train.AdamOptimizer(0.05)
loss = -tf.reduce_max(res)

train = optmizer.minimize(loss)

init = tf.global_variables_initializer()
############################################

#######run gragh ##############
with tf.Session() as sess:
    sess.run(init)
    for _ in range(500):
        sess.run(train)
    d = sess.run(Var_tm)
    c = sess.run(Var_ref)
    b = sess.run(loss)
    a = np.array(sess.run(res))
    print(b)
    print(np.max(a),np.min(a))
    plt.figure()
    plt.subplot(221),plt.imshow(d,cmap='gray')
    plt.subplot(222), plt.imshow(c, cmap='gray')
    plt.subplot(223), plt.imshow(img_src_gray, cmap='gray')
    plt.subplot(224), plt.imshow(a, cmap='gray')
    plt.show()