""" 在屏幕截图上画模板区域：支持 LightGlue 风格字典、通用匹配结果字典，或 LoFTR 命名风格的匹配对象。 LightGlue 风格：有有效四角坐标则画透视四边形，否则用大图上匹配点包一个轴对齐框。 """ from __future__ import annotations import argparse import sys from dataclasses import dataclass from pathlib import Path from typing import Any import cv2 import numpy as np @dataclass(frozen=True) class LoFTRTemplateAgainstScreenshotMatch: """匹配结果对象（原 ``loftr_template_match`` 类型；绘制轮廓用）。""" screenshot_bgr_full_size: np.ndarray homography_template_inference_to_screenshot_inference: np.ndarray | None template_width_pixels_at_inference: int template_height_pixels_at_inference: int divisor_inference_screenshot_x_to_fullsize_x: float divisor_inference_screenshot_y_to_fullsize_y: float high_confidence_match_points_on_screenshot_inference: np.ndarray template_original_width_pixels: int template_original_height_pixels: int refined_template_bbox_xywh_full_size: tuple[float, float, float, float] | None ransac_inlier_points_screenshot_inference: np.ndarray def _homography_quad_plausible_on_full_image( corners_full: np.ndarray, full_w: int, full_h: int, tw0: int, th0: int, ) -> bool: _ = tw0, th0 c = np.asarray(corners_full, dtype=np.float64).reshape(4, 2) if c.shape != (4, 2) or not np.all(np.isfinite(c)): return False margin = 4.0 if np.min(c[:, 0]) < -margin or np.min(c[:, 1]) < -margin: return False if np.max(c[:, 0]) > full_w + margin or np.max(c[:, 1]) > full_h + margin: return False xs, ys = c[:, 0], c[:, 1] area = 0.5 * abs( float(np.dot(xs, np.roll(ys, 1)) - np.dot(ys, np.roll(xs, 1))) ) img_area = float(max(1, full_w * full_h)) if area < max(16.0, 1e-4 * img_area): return False if area > img_area * 3.0: return False return True def draw_and_save_screenshot_with_template_match_outline( match_result: LoFTRTemplateAgainstScreenshotMatch, output_image_path: Path, outline_color_bgr: tuple[int, int, int] = (0, 255, 0), outline_thickness_pixels: int = 6, ) -> None: """ 在完整分辨率截图上画出模板匹配区域并写入文件。 """ output_image_path = Path(output_image_path) output_image_path.parent.mkdir(parents=True, exist_ok=True) annotated_screenshot_bgr = match_result.screenshot_bgr_full_size.copy() template_w = match_result.template_width_pixels_at_inference template_h = match_result.template_height_pixels_at_inference div_x = match_result.divisor_inference_screenshot_x_to_fullsize_x div_y = match_result.divisor_inference_screenshot_y_to_fullsize_y full_h, full_w = annotated_screenshot_bgr.shape[:2] tw0 = int(match_result.template_original_width_pixels) th0 = int(match_result.template_original_height_pixels) H = match_result.homography_template_inference_to_screenshot_inference corners_full: np.ndarray | None = None if ( H is not None and np.asarray(H).size == 9 and np.all(np.isfinite(np.asarray(H, dtype=np.float64))) and int(template_w) > 0 and int(template_h) > 0 ): tw = int(template_w) th = int(template_h) tpl_c = np.array( [ [0.0, 0.0], [float(tw - 1), 0.0], [float(tw - 1), float(th - 1)], [0.0, float(th - 1)], ], dtype=np.float32, ).reshape(1, 4, 2) ci = cv2.perspectiveTransform(tpl_c, np.asarray(H, dtype=np.float64))[0] corners_full = np.stack([ci[:, 0] / div_x, ci[:, 1] / div_y], axis=1) if not _homography_quad_plausible_on_full_image( corners_full, full_w, full_h, tw0, th0 ): corners_full = None if corners_full is not None: pts = corners_full.astype(np.int32).reshape(-1, 1, 2) cv2.polylines( annotated_screenshot_bgr, [pts], isClosed=True, color=outline_color_bgr, thickness=int(outline_thickness_pixels), lineType=cv2.LINE_AA, ) elif match_result.refined_template_bbox_xywh_full_size is not None: rx, ry, rw, rh = match_result.refined_template_bbox_xywh_full_size x0, y0 = int(round(rx)), int(round(ry)) x1, y1 = int(round(rx + rw)), int(round(ry + rh)) cv2.rectangle( annotated_screenshot_bgr, (x0, y0), (x1, y1), outline_color_bgr, int(outline_thickness_pixels), lineType=cv2.LINE_AA, ) else: inf = np.asarray( match_result.ransac_inlier_points_screenshot_inference, dtype=np.float64, ).reshape(-1, 2) if inf.shape[0] < 3: inf = np.asarray( match_result.high_confidence_match_points_on_screenshot_inference, dtype=np.float64, ).reshape(-1, 2) xy = np.stack([inf[:, 0] / div_x, inf[:, 1] / div_y], axis=1) xmin = int(np.floor(np.min(xy[:, 0]))) ymin = int(np.floor(np.min(xy[:, 1]))) xmax = int(np.ceil(np.max(xy[:, 0]))) ymax = int(np.ceil(np.max(xy[:, 1]))) cv2.rectangle( annotated_screenshot_bgr, (xmin, ymin), (xmax, ymax), outline_color_bgr, int(outline_thickness_pixels), lineType=cv2.LINE_AA, ) if not cv2.imwrite(str(output_image_path), annotated_screenshot_bgr): raise OSError(f"无法写入：{output_image_path}") def draw_and_save_screenshot_with_match_dict_outline( match_dict: dict[str, Any], output_image_path: str | Path, outline_color_bgr: tuple[int, int, int] = (0, 255, 0), outline_thickness_pixels: int = 6, ) -> None: """ 根据匹配管线输出的字典（含 ``screenshot_image_path``、四角或关键点等）在截图上画框或四边形。 """ output_image_path = Path(output_image_path) output_image_path.parent.mkdir(parents=True, exist_ok=True) scr_path = Path(str(match_dict["screenshot_image_path"])) bgr = cv2.imread(str(scr_path), cv2.IMREAD_COLOR) if bgr is None: raise FileNotFoundError(f"无法读取截图：{scr_path}") full_h, full_w = bgr.shape[:2] corners = match_dict.get("template_corners_on_screenshot_xy") if corners is not None: c = np.asarray(corners, dtype=np.float64).reshape(4, 2) if c.shape == (4, 2) and np.all(np.isfinite(c)): pts = c.astype(np.int32).reshape(-1, 1, 2) cv2.polylines( bgr, [pts], isClosed=True, color=outline_color_bgr, thickness=int(outline_thickness_pixels), lineType=cv2.LINE_AA, ) if not cv2.imwrite(str(output_image_path), bgr): raise OSError(f"无法写入：{output_image_path}") return rb = match_dict.get("refined_template_bbox_xywh_full_size") if rb is not None and len(rb) >= 4: rx, ry, rw, rh = float(rb[0]), float(rb[1]), float(rb[2]), float(rb[3]) x0, y0 = int(round(rx)), int(round(ry)) x1, y1 = int(round(rx + rw)), int(round(ry + rh)) cv2.rectangle( bgr, (x0, y0), (x1, y1), outline_color_bgr, int(outline_thickness_pixels), lineType=cv2.LINE_AA, ) if not cv2.imwrite(str(output_image_path), bgr): raise OSError(f"无法写入：{output_image_path}") return mk = match_dict.get("matched_keypoints_original_xy") if mk is None: raise ValueError("字典中缺少可用的 template_corners、bbox 或 matched_keypoints") xy = np.asarray(mk, dtype=np.float64).reshape(-1, 2) if xy.shape[0] < 1: raise ValueError("matched_keypoints_original_xy 为空") xmin = max(0, int(np.floor(np.min(xy[:, 0])))) ymin = max(0, int(np.floor(np.min(xy[:, 1])))) xmax = min(full_w - 1, int(np.ceil(np.max(xy[:, 0])))) ymax = min(full_h - 1, int(np.ceil(np.max(xy[:, 1])))) cv2.rectangle( bgr, (xmin, ymin), (xmax, ymax), outline_color_bgr, int(outline_thickness_pixels), lineType=cv2.LINE_AA, ) if not cv2.imwrite(str(output_image_path), bgr): raise OSError(f"无法写入：{output_image_path}") def main() -> None: parser = argparse.ArgumentParser( description="根据匹配 JSON 字典在截图上画模板区域（测试用）。" ) parser.add_argument( "json_path", type=Path, nargs="?", help="含 screenshot_image_path、template_corners_on_screenshot_xy 等的 JSON", ) parser.add_argument( "-o", "--out", type=Path, default=Path("output") / "match_outline.png", help="输出图片路径", ) args = parser.parse_args() if args.json_path is None: parser.print_help() raise SystemExit(2) import json data = json.loads(Path(args.json_path).read_text(encoding="utf-8")) draw_and_save_screenshot_with_match_dict_outline(data, args.out) print(f"已写入：{args.out.resolve()}", flush=True) if __name__ == "__main__": if sys.platform == "win32": for stream in (sys.stdout, sys.stderr): try: stream.reconfigure(encoding="utf-8") except Exception: pass main()