template_matching.py 文件源码

def match_template_opencv(template, image, options):
    """
    Match template using OpenCV template matching implementation.
        Limited by number of channels as maximum of 3.
        Suitable for direct RGB or Gray-scale matching

    :param options: Other options:
        - distance: Distance measure to use. (euclidean | correlation | ccoeff).
            Default: 'correlation'
        - normalize: Heatmap values will be in the range of 0 to 1. Default: True
        - retain_size: Whether to retain the same size as input image. Default: True
    :return: Heatmap
    """
    # if image has more than 3 channels, use own implementation
    if len(image.shape) > 3:
        return match_template(template, image, options)

    op = _DEF_TM_OPT.copy()
    if options is not None:
        op.update(options)

    method = cv.TM_CCORR_NORMED
    if op['normalize'] and op['distance'] == 'euclidean':
        method = cv.TM_SQDIFF_NORMED
    elif op['distance'] == 'euclidean':
        method = cv.TM_SQDIFF
    elif op['normalize'] and op['distance'] == 'ccoeff':
        method = cv.TM_CCOEFF_NORMED
    elif op['distance'] == 'ccoeff':
        method = cv.TM_CCOEFF
    elif not op['normalize'] and op['distance'] == 'correlation':
        method = cv.TM_CCORR

    heatmap = cv.matchTemplate(image, template, method)

    # make minimum peak heatmap
    if method not in [cv.TM_SQDIFF, cv.TM_SQDIFF_NORMED]:
        heatmap = heatmap.max() - heatmap

    if op['normalize']:
        heatmap /= heatmap.max()

    # size
    if op['retain_size']:
        hmap = np.ones(image.shape[:2]) * heatmap.max()
        h, w = heatmap.shape
        hmap[:h, :w] = heatmap
        heatmap = hmap

    return heatmap