# _*_ coding:utf-8 _*_

import os
import cv2
import numpy as np
import matplotlib.pyplot as plt
from pylab import mpl

'''设置显示中文字体'''
mpl.rcParams["font.sans-serif"] = ["SimHei"]
'''设置正常显示符号'''
mpl.rcParams["axes.unicode_minus"] = False

def cv_show(name,img):
    cv2.imshow(name,img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

'''NMS处理函数'''
def py_nms(dets, thresh):
    '''
    Pure Python NMS baseline
    x1、y1、x2、y2以及score赋值
    (x1、y1) (x2、y2)为box的左上和右下角标
    '''
    x1 = dets[:,0]
    y1 = dets[:,1]
    x2 = dets[:,2]
    y2 = dets[:,3]
    score = dets[:,4]
    '''计算每一个候选框的面积'''
    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    '''order是按照score降序排序的'''
    order = scores.argsort()[::-1]

    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        '''计算当前概率最大矩形框与其它矩形框的相交框的坐标，会用到numpy的broadcast机制，得到的是向量'''
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.maximum(x2[i], x2[order[1:]])
        yy2 = np.maximum(y2[i], y2[order[1:]])
        '''计算相交框的面积，注意矩形框不相交时w或h算出来会是负数，用0代替'''
        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h
        '''计算重叠度IOU：重叠面积/（面积1+面积2-重叠面积）'''
        ovr = inter / (areas[i] + areas[order[1:]] - inter)
        '''找到重叠度不高于阈值的矩形框索引'''
        inds = np.where(ovr <= thresh)[0]
        '''将order序列更新，由于前面得到的矩形框索引要比矩形框在原order序列中的索引小1，所以要把这个1加回来'''
        order = order[inds + 1]
    return keep

'''获取指定目录下全部指定格式的图像'''
def get_sub_fles(ori_path):
    file_dict = {}
    for parent,dirnames,filenames in os.walk(ori_path):
        for filename in filenames:
            fullname = os.path.join(parent,filename)
            if os.path.splitext(fullname)[1] in ['.png','.PNG','.jpg','.JPG']:
                file_dict.update({fullname:filename})
    return file_dict

'''单目标匹配'''
def single_target_match(ori_pic,template_pic,result_file):
    template = cv2.imread(template_pic,1)
    h,w,channels = template.shape

    img = cv2.imread(ori_pic, 1)
    img_h,img_w,img_channels = img.shape

    '''边缘提取'''
    gray = cv2.cvtColor(template, cv2.COLOR_BGR2GRAY)
    edges = cv2.Canny(gray, threshold1 = 50, threshold2 = 150)
    '''创建一个空白画布，将原始图像复制到该画布上'''
    mask = np.zeros_like(template)
    mask_result = template.copy()

    '''在边缘区域内标记非零像素点'''
    contours, _ = cv2.findContours(edge, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    for contour in contours:
        cv2.drawContours(mask, [contour], -1, (255), thickness = cv2.FILLED)

    '''通过位运算将原始图像和掩膜相加，从而只保留边缘部分'''
    mask_result[np.where((mask == 255))] = template[np.where((mask == 255))]

    '单目标匹配方法列表'''
    methods = ['cv2.TM_CCOEFF','cv2.TM_CCOEFF_NORMED','cv2.TM_CCORR','cv2.TM_CCORR_NORMED','cv2.TM_SQDIFF','cv2.TM_SQDIFF_NORMED']

    '''设置原图图像'''
    plt.figure(figsize = ((img_w + 100) * len(methods) / 100, (img_h + 100) / 100))
    plt.subplot(1, len(methods) + 1, 1)
    plt.subplot(1, len(methods) + 1, 1).set_title("原图", fontsize = 30)
    plt.imshow(img[:,:,::-1])
    '''隐藏x、y坐标轴'''
    plt.xticks([])
    plt.yticks([])

    for meth in methods:
        img3 = img.copy()
        method = eval(meth)
        print(ori_pic,meth)

        if method in [cv2.TM_SQDIFF, cv2.TM_CCORR_NORMED]:
            match_results = cv2.matchTemplate(img, template, method, mask = mask_result)
        else:
            match_results = cv2.matchTemplate(img, template, method)

        min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(match_results)

        '''用矩形标记结果位置'''
        if method in [cv2.TM_SQDIFF, cv2.TM_SQDIFF_NORMED]:
            top_left = min_loc
            score = min_val
        else:
            top_left = max_loc
            score = max_val
        bottom_right = (top_left[0] + w, top_left[1] + h)
        cv2.rectangle(img = img3, pt1 = top_left, pt2 = bottom_right, color = (0,0,255), thickness = 4)
        cv2.putText(img3, text = str(round(score,4)), org = bottom_right, fontFace = cv2.FONT_HERSHEY_COMPLEX, color = (0,0,255), fontScale = 1, thickness = 1)

        '''设置识别结果图像'''
        plt.subplot(1, len(methods) + 1, methods.index(meth) + 2)
        plt.subplot(1, len(methods) + 1, methods.index(meth) + 2).set_title(meth, fontsize = 30)
        plt.imshow(img3[:,:,::-1])
        plt.xticks([])
        plt.xticks([])

    '''设置图像大标题'''
    '''
    plt.suptitle('图像模板匹配结果')
    '''
    '''保存图像'''
    print('Saving...')
    plt.savefig(result_file, bbox_inches = 'tight', pad_inches = 1)
    '''展示图像'''
    '''
    plt.show()
    '''

'''多目标匹配'''
def multi_target_match(ori_pic,template_pic,result_file):
    img = cv2.imread(ori_pic, 1)
    img_h,img_w,img_channels = img.shape

    template = cv2.imread(template_pic,1)
    h,w,channels = template.shape

    methods = ['cv2.TM_CCOEFF','cv2.TM_CCOEFF_NORMED','cv2.TM_CCORR','cv2.TM_CCORR_NORMED','cv2.TM_SQDIFF','cv2.TM_SQDIFF_NORMED']


    '''设置原图图像'''
    plt.figure(figsize = ((img_w + 100) * len(methods) / 100, (img_h + 100) / 100))
    plt.subplot(1, len(methods) + 1, 1)
    plt.subplot(1, len(methods) + 1, 1).set_title("原图", fontsize = 30)
    plt.imshow(img[:,:,::-1])
    '''隐藏x、y坐标轴'''
    plt.xticks([])
    plt.yticks([])

    for meth in methods:
        img3 = img.copy()
        method = eval(meth)
        print(ori_pic,meth)

        if method in [cv2.TM_SQDIFF, cv2.TM_CCORR_NORMED]:
            '''边缘提取'''
            gray = cv2.cvtColor(template, cv2.COLOR_BGR2GRAY)
            edges = cv2.Canny(gray, threshold1 = 50, threshold2 = 150) 
            '''创建一个空白画布，将原始图像复制到该画布上'''
            mask = np.zeros_like(template)
            mask_result = template.copy()

            '''在边缘区域内标记非零像素点'''
            contours, _ = cv2.findContours(edge, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            for contour in contours:
                cv2.drawContours(mask, [contour], -1, (255), thickness = cv2.FILLED)

            '''通过位运算将原始图像和掩膜相加，从而只保留边缘部分'''
            mask_result[np.where((mask == 255))] = template[np.where((mask == 255))]

            match_results = cv2.matchTemplate(img, template, method, mask = mask_result)
        else:
            match_results = cv2.matchTemplate(img, template, method)

        min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(match_results)

        if method in [cv2.TM_CCOEFF_NORMED, cv2.TM_CCORR_NORMED]:
            '''大于模板阈值的目标坐标'''
            loc = np.where(match_results >= threshold)
            '''大于模板阈值的目标置信度'''
            score = match_results[match_results >= threshold]
        elif method in [cv2.TM_CCOEFF, cv2.TM_CCORR]:
            loc = np.where(match_results >= threshold * max_val)
            score = match_results[match_results >= threshold * max_val]
        elif method in [cv2.TM_SQDIFF_NORMED]:
            loc = np.where(match_results <= (1 - threshold))
            score = match_results[match_results <= (1 - threshold)]
        elif method in [cv2.TM_SQDIFF]:
            loc = np.where(match_results <= (1 / threshold * min_val))
            score = match_results[match_results <= (1 / threshold * min_val)]
        else:
            pass

        '''将模板数据坐标进行处理成左上角、右下角的格式'''
        xmin = np.array(loc[1])
        ymin = np.array(loc[0])
        xmax = xmin + w
        ymax = ymin + h
        '''变成n行1列维度'''
        xmin = xmin.reshape(-1, 1)
        xmax = xmax.reshape(-1, 1)
        ymin = ymin.reshape(-1, 1)
        ymax = ymax.reshape(-1, 1)
        score = score.reshape(-1, 1)

        data_hlist = []
        data_hlist.append(xmin)
        data_hlist.append(ymin)
        data_hlist.append(xmax)
        data_hlist.append(ymax)
        data_hlist.append(score)
        '''将xmin、ymin、xmax、ymax、score按照列进行拼接'''
        data_hstack = np.hstack(data_hlist)

        keep_dets = py_nms(data_hstack, thresh)
        dets = data_hstack [keep_dets]

        for coord in dets:
            cv2.rectangle(img = img3, pt1 = (int(coord[0]),int(coord[1])), pt2 = (int(coord[2]),int(coord[3])), color = (0,0,255), thickness = 4)
            cv2.putText(img3, text = str(round(coord[4],4)), org = (int(coord[0]),int(coord[1])), fontFace = cv2.FONT_HERSHEY_COMPLEX, color = (0,0,255), fontScale = 1, thickness = 1)

        '''设置识别结果图像'''
        plt.subplot(1, len(methods) + 1, methods.index(meth) + 2)
        plt.subplot(1, len(methods) + 1, methods.index(meth) + 2).set_title(meth, fontsize = 30)
        plt.imshow(img3[:,:,::-1])
        plt.xticks([])
        plt.xticks([])

    '''设置图像大标题'''
    '''
    plt.suptitle('图像模板匹配结果')
    '''
    '''保存图像'''
    print('Saving...')
    plt.savefig(result_file, bbox_inches = 'tight', pad_inches = 1)
    '''展示图像'''
    '''
    plt.show()
    '''

'''入口主函数'''
if __name__ == "__main__":
    ori_pic_file = r''
    template_pic = r''
    result_file = r''

    ori_pic_dict = get_sub_files(ori_pic_file)
    for ori_pic,ori_pic_name in ori_pic_dict.items():
        save_file = os.path.join(result_file, ori_pic_name)
        '''使用单目标匹配'''
        single_target_match(ori_pic, template_pic, save_file)

        '''使用多目标匹配'''
        '''设定模板匹配阈值'''
        threshold = 0.99
        '''NMS里面的IOU交互比阈值'''
        thresh = 0.3

        multi_target_match(ori_pic, template_pic, save_file)