wechat_ai/wechat_vision/screen_srk_click_send.py

import argparse
import time
from pathlib import Path

import numpy as np
import onnxruntime as ort
import pyautogui
import pyperclip
from PIL import ImageDraw


BASE_DIR = Path(__file__).resolve().parent
MODEL_PATH = BASE_DIR / "best.onnx"
OUTPUT_DIR = BASE_DIR / "ouptsw"
CLASS_MAP = {"srk": 0}
IOU_THRESHOLD = 0.5


def parse_args():
    parser = argparse.ArgumentParser(
        description="Screenshot screen, detect srk target, click center, type text, and send."
    )
    parser.add_argument("--class-name", default="srk", help="target class name")
    parser.add_argument("--confidence", type=float, default=0.1, help="minimum score")
    parser.add_argument("--text", default="你好", help="text to paste after clicking")
    parser.add_argument("--dry-run", action="store_true", help="detect only; do not click/type")
    parser.add_argument(
        "--debug-image",
        default=str(OUTPUT_DIR / "screen_srk_debug.png"),
        help="path to save screenshot with detection annotation",
    )
    parser.add_argument(
        "--delay",
        type=float,
        default=1.0,
        help="seconds to wait before taking screenshot",
    )
    return parser.parse_args()


def input_shape(input_meta):
    return [dim if isinstance(dim, int) else 1 for dim in input_meta.shape]


def preprocess(image, shape):
    _, _, height, width = shape
    resized = image.convert("RGB").resize((width, height))
    array = np.asarray(resized, dtype=np.float32) / 255.0
    return np.transpose(array, (2, 0, 1))[None]


def iou(box, boxes):
    x1 = np.maximum(box[0], boxes[:, 0])
    y1 = np.maximum(box[1], boxes[:, 1])
    x2 = np.minimum(box[2], boxes[:, 2])
    y2 = np.minimum(box[3], boxes[:, 3])
    intersection = np.maximum(0, x2 - x1) * np.maximum(0, y2 - y1)
    box_area = np.maximum(0, box[2] - box[0]) * np.maximum(0, box[3] - box[1])
    boxes_area = np.maximum(0, boxes[:, 2] - boxes[:, 0]) * np.maximum(
        0,
        boxes[:, 3] - boxes[:, 1],
    )
    return intersection / np.maximum(box_area + boxes_area - intersection, 1e-6)


def nms(detections, confidence):
    detections = detections[detections[:, 4] >= confidence]
    if len(detections) == 0:
        return detections

    detections = detections[np.argsort(detections[:, 4])[::-1]]
    kept = []
    while len(detections) > 0:
        best = detections[0]
        kept.append(best)
        if len(detections) == 1:
            break
        detections = detections[1:][iou(best[:4], detections[1:, :4]) < IOU_THRESHOLD]
    return np.array(kept)


def scale_detections(detections, original_size, shape):
    _, _, input_height, input_width = shape
    original_width, original_height = original_size
    scaled = detections.copy()
    scaled[:, [0, 2]] *= original_width / input_width
    scaled[:, [1, 3]] *= original_height / input_height
    scaled[:, [0, 2]] = np.clip(scaled[:, [0, 2]], 0, original_width)
    scaled[:, [1, 3]] = np.clip(scaled[:, [1, 3]], 0, original_height)
    return scaled


def detect_srk(session, input_meta, screenshot, class_id, confidence):
    shape = input_shape(input_meta)
    tensor = preprocess(screenshot, shape)
    detections = session.run(None, {input_meta.name: tensor})[0][0]
    detections = detections[detections[:, 5].astype(int) == class_id]
    detections = nms(detections, confidence)
    return scale_detections(detections, screenshot.size, shape)


def save_debug_image(screenshot, detections, path):
    debug = screenshot.convert("RGB")
    draw = ImageDraw.Draw(debug)
    for detection in detections:
        x1, y1, x2, y2, score, _ = detection
        cx = (x1 + x2) / 2
        cy = (y1 + y2) / 2
        draw.rectangle((x1, y1, x2, y2), outline="red", width=3)
        draw.ellipse((cx - 5, cy - 5, cx + 5, cy + 5), fill="yellow", outline="red")
        draw.text((x1, max(0, y1 - 14)), f"srk ({cx:.1f}, {cy:.1f}) {score:.2f}", fill="yellow")
    path = Path(path)
    path.parent.mkdir(exist_ok=True)
    debug.save(path)


def screen_click_position(x, y, screenshot_size):
    screen_width, screen_height = pyautogui.size()
    screenshot_width, screenshot_height = screenshot_size
    return (
        x * screen_width / screenshot_width,
        y * screen_height / screenshot_height,
    )


def click_type_send(x, y, text):
    pyautogui.click(x, y)
    time.sleep(0.2)
    pyperclip.copy(text)
    pyautogui.hotkey("command", "v")
    time.sleep(0.1)
    pyautogui.press("enter")


def main():
    args = parse_args()
    if args.class_name not in CLASS_MAP:
        raise ValueError(f"Unknown class name: {args.class_name}. CLASS_MAP={CLASS_MAP}")

    pyautogui.FAILSAFE = True
    time.sleep(args.delay)

    session = ort.InferenceSession(str(MODEL_PATH))
    input_meta = session.get_inputs()[0]
    screenshot = pyautogui.screenshot()
    detections = detect_srk(
        session,
        input_meta,
        screenshot,
        CLASS_MAP[args.class_name],
        args.confidence,
    )

    save_debug_image(screenshot, detections, args.debug_image)
    if len(detections) == 0:
        print(f"No {args.class_name} target found. debug_image={args.debug_image}")
        return

    target = detections[np.argmax(detections[:, 4])]
    x1, y1, x2, y2, score, class_id = target
    cx = (x1 + x2) / 2
    cy = (y1 + y2) / 2
    click_x, click_y = screen_click_position(cx, cy, screenshot.size)
    print(
        f"target={args.class_name}, class_id={int(class_id)}, score={score:.3f}, "
        f"image_center=({cx:.1f}, {cy:.1f}), click_center=({click_x:.1f}, {click_y:.1f}), "
        f"screenshot_size={screenshot.size}, screen_size={pyautogui.size()}, "
        f"debug_image={args.debug_image}"
    )

    if args.dry_run:
        print("dry-run enabled; skip click/type/send")
        return

    click_type_send(click_x, click_y, args.text)
    print("clicked target, pasted text, and pressed enter")


if __name__ == "__main__":
    main()