记录用tensorflow 实现Cosin_Cross_templateMatch层

#-*- coding:utf-8 -*-
import tensorflow as tf
import cv2
import numpy as np
import matplotlib.pyplot as plt
from Integral_image import *
## read img #####
tm_path = '/home/wdh/pytorch-CycleGAN-and-pix2pix1_run/results/map_pix2pix/test_latest/real_A/2_AB_real_A.png'
ref_path = '/home/wdh/pytorch-CycleGAN-and-pix2pix1_run/results/map_pix2pix/test_latest/real_B/2_AB_real_B.png'

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 = np.copy(img_src_gray[h_tm:h_tm+h_tm,w_tm:w_tm+w_tm])

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

########build compute graph#################
tensor_tm_src = tf.convert_to_tensor(img_tm_gray,dtype=tf.float32)
tensor_ref_src = tf.convert_to_tensor(img_ref_gray,dtype=tf.float32)

Weight_tm = tf.Variable(tf.ones(img_tm_gray.shape))
Weight_ref = tf.Variable(tf.ones(img_ref_gray.shape))

tensor_tm = Weight_tm*tensor_tm_src
tensor_ref = Weight_ref*tensor_ref_src

#计算积分图，均值等
##tm的均值和平方均值
tm_avg = tf.reduce_mean(tensor_tm)
tm_avg_sq = tf.reduce_mean(tensor_tm*tensor_tm)
##ref的积分图和平方积分图
tensor_ref_sq = tensor_ref*tensor_ref
#站位，在外面计算
ref_integral_img = tf.placeholder(dtype=tf.float32,shape=(h_ref,w_ref))
ref_integral_img_sq = tf.placeholder(dtype=tf.float32,shape=(h_ref,w_ref))
ref_div_sq_mat = tf.placeholder(dtype=tf.float32,shape=(h_ref-h_tm+1,w_ref-w_tm+1))
##拼成一个大的tensor
tensor_left = tf.zeros((h_tm,w_ref-w_tm),dtype=tf.float32)
tensor_bottom = tf.zeros((h_ref-h_tm,w_ref),dtype=tf.float32)
tensor_tm_ex = tf.concat([tf.concat([tensor_tm,tensor_left],axis=1),tensor_bottom],axis=0)

###fft计算卷积
F_tm = tf.fft2d(tf.complex(tensor_tm_ex,tf.zeros(tensor_tm_ex.shape)))
F_ref = tf.fft2d(tf.complex(tensor_ref,tf.zeros(tensor_ref.shape)))
F_m = tf.multiply(F_ref,F_tm)
# F_m = F_ref * F_tm
res = tf.real(tf.ifft2d(F_m))[0:h_ref-h_tm+1,0:w_ref-w_tm+1]
res = res/(tf.sqrt(ref_div_sq_mat)*tf.sqrt(tm_avg_sq))
#训练设置
optmizer = tf.train.AdamOptimizer(0.05)
loss = -res[h_tm,w_tm]
train = optmizer.minimize(loss)

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

#######run gragh ##############
with tf.Session() as sess:
    sess.run(init)
    for i in range(20):
        tm_img_show = sess.run(tensor_tm)
        ref_img_show = sess.run(tensor_ref)
        ref_arr_sq= sess.run(tensor_ref_sq)
        ref_arr = sess.run(tensor_ref)
        in_ref_sq = calcIntag(ref_arr_sq)
        in_ref = calcIntag(ref_arr)
        ref_div_sq = calcDivmat(in_ref,h_tm,w_tm)
        sess.run(train, feed_dict={ref_integral_img: in_ref, ref_integral_img_sq: in_ref_sq, ref_div_sq_mat: ref_div_sq})
    loss_out,res_out= sess.run([loss,res],feed_dict={ref_integral_img:in_ref,ref_integral_img_sq:in_ref_sq,ref_div_sq_mat:ref_div_sq})
    print(loss_out)
    plt.figure()
    plt.subplot(221),plt.title('tm'),plt.imshow(tm_img_show,cmap='gray')
    plt.subplot(222),plt.title('ref'), plt.imshow(ref_img_show, cmap='gray')
    plt.subplot(223), plt.imshow(img_src_gray, cmap='gray')
    plt.subplot(224), plt.imshow(res_out,cmap='gray')
    plt.show()