import cv2
import numpy as np

base = 'c:/Users/27942/Desktop/codes/codex_jxs_code/images/'
main_img = cv2.imread(base + '1.jpg')
main_gray = cv2.cvtColor(main_img, cv2.COLOR_BGR2GRAY)

def match_template_alpha(main_gray, alpha, name):
    h, w = alpha.shape[:2]
    best = None

    # 策略1: 直接用alpha通道匹配灰度图
    # 策略2: 用alpha做mask，只比较非透明区域
    # 策略3: 对alpha二值化后匹配

    _, alpha_bin = cv2.threshold(alpha, 10, 255, cv2.THRESH_BINARY)

    scales = [1.0] + list(np.linspace(0.65, 1.35, 29))
    for scale in scales:
        rw, rh = int(w * scale), int(h * scale)
        if rh < 10 or rw < 10:
            continue
        if rh > main_gray.shape[0] or rw > main_gray.shape[1]:
            continue

        resized_alpha = cv2.resize(alpha, (rw, rh))
        resized_bin = cv2.resize(alpha_bin, (rw, rh))
        _, resized_mask = cv2.threshold(resized_bin, 10, 255, cv2.THRESH_BINARY)

        # 用mask匹配（只看非透明区域）
        res = cv2.matchTemplate(main_gray, resized_alpha, cv2.TM_CCOEFF_NORMED, mask=resized_mask)
        _, max_val, _, max_loc = cv2.minMaxLoc(res)

        if best is None or max_val > best['conf']:
            best = {
                'x': max_loc[0], 'y': max_loc[1],
                'w': rw, 'h': rh,
                'conf': max_val, 'scale': scale,
                'cx': max_loc[0] + rw // 2,
                'cy': max_loc[1] + rh // 2,
            }

    return best

results = {}
for name in ['2.png', '3.png', '4.png']:
    template = cv2.imread(base + name, cv2.IMREAD_UNCHANGED)
    alpha = template[:, :, 3]
    best = match_template_alpha(main_gray, alpha, name)
    results[name] = best
    x, y, cx, cy = best['x'], best['y'], best['cx'], best['cy']
    conf, scale = best['conf'], best['scale']
    print(name + ': pos=(' + str(x) + ',' + str(y) + ') center=(' + str(cx) + ',' + str(cy) + ') conf=' + str(round(conf, 4)) + ' scale=' + str(round(scale, 2)))

# 可视化
colors = [(0, 0, 255), (0, 255, 0), (255, 0, 0)]
vis = main_img.copy()
for i, name in enumerate(['2.png', '3.png', '4.png']):
    m = results[name]
    cv2.rectangle(vis, (m['x'], m['y']), (m['x'] + m['w'], m['y'] + m['h']), colors[i], 2)
    cv2.putText(vis, name, (m['x'], m['y'] - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, colors[i], 1)
    cv2.circle(vis, (m['cx'], m['cy']), 5, colors[i], -1)

cv2.imwrite(base + 'result.jpg', vis)
print('result saved to images/result.jpg')