114 lines
4.0 KiB
Python
114 lines
4.0 KiB
Python
"""
|
||
用 CLIP 做图块相似度匹配，在大图中找模板图标位置
|
||
原理：把大图切成滑动窗口小块，用 CLIP 计算每块和模板的视觉相似度，取最高分的块
|
||
"""
|
||
import cv2
|
||
import numpy as np
|
||
from PIL import Image
|
||
import torch
|
||
from transformers import CLIPProcessor, CLIPModel
|
||
from pathlib import Path
|
||
|
||
BASE = Path(__file__).parent / "images"
|
||
MODEL_NAME = "openai/clip-vit-base-patch32" # 约600MB，小模型
|
||
|
||
|
||
def load_clip():
|
||
print("加载 CLIP 模型（约600MB，首次自动下载）...")
|
||
model = CLIPModel.from_pretrained(MODEL_NAME).to("cuda")
|
||
processor = CLIPProcessor.from_pretrained(MODEL_NAME)
|
||
print("CLIP 加载完成")
|
||
return model, processor
|
||
|
||
|
||
def find_by_clip(model, processor, main_img: np.ndarray, template_img: np.ndarray,
|
||
step=10, win_sizes=None):
|
||
"""
|
||
滑动窗口 + CLIP 相似度，找模板在大图中的最佳位置
|
||
"""
|
||
if win_sizes is None:
|
||
th, tw = template_img.shape[:2]
|
||
# 尝试原始尺寸及上下浮动
|
||
win_sizes = [(int(tw * s), int(th * s)) for s in [0.8, 0.9, 1.0, 1.1, 1.2]]
|
||
|
||
# 预处理模板
|
||
tmpl_pil = Image.fromarray(cv2.cvtColor(template_img, cv2.COLOR_BGR2RGB))
|
||
tmpl_inputs = processor(images=tmpl_pil, return_tensors="pt").to("cuda")
|
||
with torch.no_grad():
|
||
tmpl_out = model.vision_model(**tmpl_inputs)
|
||
tmpl_feat = model.visual_projection(tmpl_out.pooler_output).float()
|
||
tmpl_feat = tmpl_feat / tmpl_feat.norm(dim=-1, keepdim=True)
|
||
|
||
mh, mw = main_img.shape[:2]
|
||
best_score = -1
|
||
best_box = None
|
||
|
||
for (ww, wh) in win_sizes:
|
||
if ww > mw or wh > mh:
|
||
continue
|
||
for y in range(0, mh - wh + 1, step):
|
||
for x in range(0, mw - ww + 1, step):
|
||
crop = main_img[y:y+wh, x:x+ww]
|
||
crop_pil = Image.fromarray(cv2.cvtColor(crop, cv2.COLOR_BGR2RGB))
|
||
crop_inputs = processor(images=crop_pil, return_tensors="pt").to("cuda")
|
||
with torch.no_grad():
|
||
crop_out = model.vision_model(**crop_inputs)
|
||
crop_feat = model.visual_projection(crop_out.pooler_output).float()
|
||
crop_feat = crop_feat / crop_feat.norm(dim=-1, keepdim=True)
|
||
score = (tmpl_feat * crop_feat).sum().item()
|
||
if score > best_score:
|
||
best_score = score
|
||
best_box = (x, y, ww, wh)
|
||
|
||
return best_box, best_score
|
||
|
||
|
||
def main():
|
||
main_img = cv2.imread(str(BASE / "1.jpg"))
|
||
templates = {
|
||
"2.png": cv2.imread(str(BASE / "2.png")),
|
||
"3.png": cv2.imread(str(BASE / "3.png")),
|
||
"4.png": cv2.imread(str(BASE / "4.png")),
|
||
}
|
||
|
||
# 模板是 BGRA，转 BGR
|
||
for name in templates:
|
||
img = cv2.imread(str(BASE / name), cv2.IMREAD_UNCHANGED)
|
||
if img.shape[2] == 4:
|
||
# alpha 通道合成白底
|
||
alpha = img[:, :, 3:4] / 255.0
|
||
rgb = img[:, :, :3].astype(float)
|
||
white = np.ones_like(rgb) * 255
|
||
merged = (rgb * alpha + white * (1 - alpha)).astype(np.uint8)
|
||
templates[name] = merged
|
||
else:
|
||
templates[name] = img
|
||
|
||
model, processor = load_clip()
|
||
|
||
vis = main_img.copy()
|
||
colors = {"2.png": (0, 0, 255), "3.png": (0, 255, 0), "4.png": (255, 0, 0)}
|
||
|
||
print()
|
||
for name, tmpl in templates.items():
|
||
print(f"正在匹配 {name} ...")
|
||
box, score = find_by_clip(model, processor, main_img, tmpl, step=8)
|
||
if box:
|
||
x, y, w, h = box
|
||
cx, cy = x + w // 2, y + h // 2
|
||
color = colors[name]
|
||
cv2.rectangle(vis, (x, y), (x+w, y+h), color, 2)
|
||
cv2.circle(vis, (cx, cy), 5, color, -1)
|
||
cv2.putText(vis, name, (x, y-5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)
|
||
print(f" {name}: 中心点 ({cx}, {cy}) 相似度={score:.4f}")
|
||
else:
|
||
print(f" {name}: 未找到")
|
||
|
||
out = BASE / "result_clip.jpg"
|
||
cv2.imwrite(str(out), vis)
|
||
print(f"\n结果保存到: {out}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|