feat: add realtime wechat vision monitor
This commit is contained in:
parent
486f2af11f
commit
f9661f5109
9
.gitignore
vendored
9
.gitignore
vendored
@ -7,6 +7,15 @@ src-tauri/gen/
|
||||
|
||||
agent/agent
|
||||
agent/agent-*
|
||||
agent/config.toml
|
||||
|
||||
wechat_vision/venv/
|
||||
wechat_vision/__pycache__/
|
||||
wechat_vision/ouptsw/
|
||||
wechat_vision/test/
|
||||
|
||||
.codegraph/daemon.pid
|
||||
.vscode/
|
||||
|
||||
.DS_Store
|
||||
*.log
|
||||
|
||||
252
agent/ax_probe.go
Normal file
252
agent/ax_probe.go
Normal file
@ -0,0 +1,252 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"sort"
|
||||
"strings"
|
||||
)
|
||||
|
||||
type AXProbeResult struct {
|
||||
AppName string `json:"app_name"`
|
||||
PID int `json:"pid"`
|
||||
Nodes []AXNode `json:"nodes"`
|
||||
}
|
||||
|
||||
type AXNode struct {
|
||||
Role string `json:"role"`
|
||||
Title string `json:"title"`
|
||||
Value string `json:"value"`
|
||||
Description string `json:"description"`
|
||||
Frame [4]float64 `json:"frame"`
|
||||
Depth int `json:"depth"`
|
||||
}
|
||||
|
||||
func runAXProbe() error {
|
||||
logInfo("macOS Accessibility 微信 UI 树探测启动")
|
||||
|
||||
config, err := loadConfig("config.toml")
|
||||
if err != nil {
|
||||
return fmt.Errorf("读取 config.toml 失败：%w", err)
|
||||
}
|
||||
|
||||
regionsPath := appDataPath(config, config.Agent.RegionsPath)
|
||||
annotation, err := loadAnnotation(regionsPath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("读取标注文件失败：%w", err)
|
||||
}
|
||||
summaries, warnings := inferMissingRegionTypes(summarizeRegions(annotation.Regions))
|
||||
for _, warning := range warnings {
|
||||
logWarning(warning)
|
||||
}
|
||||
|
||||
probe, err := probeWeChatAXTree()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
logInfo(fmt.Sprintf("AX 节点总数=%d pid=%d app=%s", len(probe.Nodes), probe.PID, probe.AppName))
|
||||
|
||||
textNodes := usefulAXTextNodes(probe.Nodes)
|
||||
logInfo(fmt.Sprintf("可读文本/按钮节点=%d", len(textNodes)))
|
||||
for index, node := range firstNodes(textNodes, 80) {
|
||||
logJSON("ax", fmt.Sprintf("#%02d depth=%d role=%s frame=[%.0f,%.0f,%.0f,%.0f] text=%s", index+1, node.Depth, node.Role, node.Frame[0], node.Frame[1], node.Frame[2], node.Frame[3], axNodeText(node)))
|
||||
}
|
||||
|
||||
if chat, ok := regionByType(summaries)["chat_content"]; ok {
|
||||
chatNodes := nodesInsideBox(textNodes, chat.BBoxScreen)
|
||||
sort.SliceStable(chatNodes, func(i, j int) bool {
|
||||
if chatNodes[i].Frame[1] == chatNodes[j].Frame[1] {
|
||||
return chatNodes[i].Frame[0] < chatNodes[j].Frame[0]
|
||||
}
|
||||
return chatNodes[i].Frame[1] < chatNodes[j].Frame[1]
|
||||
})
|
||||
logInfo(fmt.Sprintf("chat_content 区域内可读节点=%d", len(chatNodes)))
|
||||
for index, node := range firstNodes(chatNodes, 80) {
|
||||
logJSON("ax-chat", fmt.Sprintf("#%02d role=%s frame=[%.0f,%.0f,%.0f,%.0f] text=%s", index+1, node.Role, node.Frame[0], node.Frame[1], node.Frame[2], node.Frame[3], axNodeText(node)))
|
||||
}
|
||||
} else {
|
||||
logWarning("未找到 chat_content 标注区域，跳过区域过滤")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func probeWeChatAXTree() (AXProbeResult, error) {
|
||||
script := `
|
||||
import AppKit
|
||||
import ApplicationServices
|
||||
import Foundation
|
||||
|
||||
struct Node: Encodable {
|
||||
let role: String
|
||||
let title: String
|
||||
let value: String
|
||||
let description: String
|
||||
let frame: [Double]
|
||||
let depth: Int
|
||||
}
|
||||
|
||||
struct Result: Encodable {
|
||||
let app_name: String
|
||||
let pid: Int
|
||||
let nodes: [Node]
|
||||
}
|
||||
|
||||
func stringAttr(_ element: AXUIElement, _ key: String) -> String {
|
||||
var value: CFTypeRef?
|
||||
let error = AXUIElementCopyAttributeValue(element, key as CFString, &value)
|
||||
if error != .success || value == nil { return "" }
|
||||
if let str = value as? String { return clean(str) }
|
||||
if let attr = value as? NSAttributedString { return clean(attr.string) }
|
||||
return ""
|
||||
}
|
||||
|
||||
func clean(_ value: String) -> String {
|
||||
return value
|
||||
.replacingOccurrences(of: "\n", with: " ")
|
||||
.replacingOccurrences(of: "\r", with: " ")
|
||||
.replacingOccurrences(of: "\t", with: " ")
|
||||
}
|
||||
|
||||
func frameOf(_ element: AXUIElement) -> [Double] {
|
||||
var positionValue: CFTypeRef?
|
||||
var sizeValue: CFTypeRef?
|
||||
var point = CGPoint.zero
|
||||
var size = CGSize.zero
|
||||
if AXUIElementCopyAttributeValue(element, kAXPositionAttribute as CFString, &positionValue) == .success,
|
||||
let positionValue,
|
||||
CFGetTypeID(positionValue) == AXValueGetTypeID() {
|
||||
AXValueGetValue(positionValue as! AXValue, .cgPoint, &point)
|
||||
}
|
||||
if AXUIElementCopyAttributeValue(element, kAXSizeAttribute as CFString, &sizeValue) == .success,
|
||||
let sizeValue,
|
||||
CFGetTypeID(sizeValue) == AXValueGetTypeID() {
|
||||
AXValueGetValue(sizeValue as! AXValue, .cgSize, &size)
|
||||
}
|
||||
return [Double(point.x), Double(point.y), Double(size.width), Double(size.height)]
|
||||
}
|
||||
|
||||
func childrenOf(_ element: AXUIElement) -> [AXUIElement] {
|
||||
var value: CFTypeRef?
|
||||
let error = AXUIElementCopyAttributeValue(element, kAXChildrenAttribute as CFString, &value)
|
||||
if error != .success || value == nil { return [] }
|
||||
return value as? [AXUIElement] ?? []
|
||||
}
|
||||
|
||||
func walk(_ element: AXUIElement, depth: Int, nodes: inout [Node], visited: inout Int) {
|
||||
if visited > 2500 || depth > 20 { return }
|
||||
visited += 1
|
||||
let node = Node(
|
||||
role: stringAttr(element, kAXRoleAttribute),
|
||||
title: stringAttr(element, kAXTitleAttribute),
|
||||
value: stringAttr(element, kAXValueAttribute),
|
||||
description: stringAttr(element, kAXDescriptionAttribute),
|
||||
frame: frameOf(element),
|
||||
depth: depth
|
||||
)
|
||||
nodes.append(node)
|
||||
for child in childrenOf(element) {
|
||||
walk(child, depth: depth + 1, nodes: &nodes, visited: &visited)
|
||||
}
|
||||
}
|
||||
|
||||
let trusted = AXIsProcessTrustedWithOptions([kAXTrustedCheckOptionPrompt.takeUnretainedValue() as String: true] as CFDictionary)
|
||||
if !trusted {
|
||||
fputs("Accessibility permission is not granted\n", stderr)
|
||||
exit(10)
|
||||
}
|
||||
|
||||
let apps = NSWorkspace.shared.runningApplications.filter { app in
|
||||
let name = app.localizedName ?? ""
|
||||
return name == "微信" || name == "WeChat" || name.lowercased().contains("wechat")
|
||||
}
|
||||
|
||||
guard let app = apps.first else {
|
||||
fputs("WeChat app not found\n", stderr)
|
||||
exit(2)
|
||||
}
|
||||
|
||||
app.activate(options: [.activateIgnoringOtherApps])
|
||||
usleep(300000)
|
||||
|
||||
let root = AXUIElementCreateApplication(app.processIdentifier)
|
||||
var nodes: [Node] = []
|
||||
var visited = 0
|
||||
walk(root, depth: 0, nodes: &nodes, visited: &visited)
|
||||
let result = Result(app_name: app.localizedName ?? "", pid: Int(app.processIdentifier), nodes: nodes)
|
||||
let data = try JSONEncoder().encode(result)
|
||||
print(String(data: data, encoding: .utf8)!)
|
||||
`
|
||||
|
||||
output, err := runSwift(script)
|
||||
if err != nil {
|
||||
return AXProbeResult{}, fmt.Errorf("AX probe failed: %w", err)
|
||||
}
|
||||
|
||||
var result AXProbeResult
|
||||
if err := json.Unmarshal([]byte(output), &result); err != nil {
|
||||
preview := output
|
||||
if len(preview) > 300 {
|
||||
preview = preview[:300]
|
||||
}
|
||||
return AXProbeResult{}, fmt.Errorf("parse AX JSON failed: %w output_prefix=%q", err, preview)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func usefulAXTextNodes(nodes []AXNode) []AXNode {
|
||||
result := make([]AXNode, 0)
|
||||
seen := map[string]bool{}
|
||||
for _, node := range nodes {
|
||||
text := strings.TrimSpace(axNodeText(node))
|
||||
if text == "" {
|
||||
continue
|
||||
}
|
||||
if node.Frame[2] <= 0 || node.Frame[3] <= 0 {
|
||||
continue
|
||||
}
|
||||
key := fmt.Sprintf("%s|%s|%.0f|%.0f|%.0f|%.0f", node.Role, text, node.Frame[0], node.Frame[1], node.Frame[2], node.Frame[3])
|
||||
if seen[key] {
|
||||
continue
|
||||
}
|
||||
seen[key] = true
|
||||
result = append(result, node)
|
||||
}
|
||||
sort.SliceStable(result, func(i, j int) bool {
|
||||
if result[i].Frame[1] == result[j].Frame[1] {
|
||||
return result[i].Frame[0] < result[j].Frame[0]
|
||||
}
|
||||
return result[i].Frame[1] < result[j].Frame[1]
|
||||
})
|
||||
return result
|
||||
}
|
||||
|
||||
func axNodeText(node AXNode) string {
|
||||
parts := []string{node.Value, node.Title, node.Description}
|
||||
for _, part := range parts {
|
||||
part = strings.TrimSpace(part)
|
||||
if part != "" {
|
||||
return part
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func nodesInsideBox(nodes []AXNode, box [4]float64) []AXNode {
|
||||
result := make([]AXNode, 0)
|
||||
for _, node := range nodes {
|
||||
cx := node.Frame[0] + node.Frame[2]/2
|
||||
cy := node.Frame[1] + node.Frame[3]/2
|
||||
if cx >= box[0] && cx <= box[2] && cy >= box[1] && cy <= box[3] {
|
||||
result = append(result, node)
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func firstNodes(nodes []AXNode, limit int) []AXNode {
|
||||
if len(nodes) <= limit {
|
||||
return nodes
|
||||
}
|
||||
return nodes[:limit]
|
||||
}
|
||||
274
agent/chat_sync.go
Normal file
274
agent/chat_sync.go
Normal file
@ -0,0 +1,274 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"crypto/sha1"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
func runChatSync() error {
|
||||
logInfo("聊天记录同步 demo 启动")
|
||||
|
||||
config, err := loadConfig("config.toml")
|
||||
if err != nil {
|
||||
return fmt.Errorf("读取 config.toml 失败：%w", err)
|
||||
}
|
||||
|
||||
regionsPath := appDataPath(config, config.Agent.RegionsPath)
|
||||
screenshotPath := appDataPath(config, config.Agent.ScreenshotPath)
|
||||
archivePath := appDataPath(config, config.Agent.ChatRecordsPath)
|
||||
|
||||
annotation, err := loadAnnotation(regionsPath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("读取标注文件失败：%w", err)
|
||||
}
|
||||
summaries, warnings := inferMissingRegionTypes(summarizeRegions(annotation.Regions))
|
||||
for _, warning := range warnings {
|
||||
logWarning(warning)
|
||||
}
|
||||
if err := validateChatSyncRegions(summaries); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
archive, existed, err := loadChatArchive(archivePath)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !existed || len(archive.Messages) == 0 {
|
||||
logInfo("本地没有聊天记录，先滚动到聊天记录顶部")
|
||||
if err := scrollChatToTop(config, summaries, screenshotPath); err != nil {
|
||||
return err
|
||||
}
|
||||
} else {
|
||||
logInfo(fmt.Sprintf("已读取本地聊天记录 messages=%d，将从当前视图继续向下读取", len(archive.Messages)))
|
||||
}
|
||||
|
||||
if archive.App == "" {
|
||||
archive.App = "wechat"
|
||||
archive.CreatedAt = time.Now().Format(time.RFC3339)
|
||||
}
|
||||
archive.Source = TaskSource{RegionsPath: regionsPath, ScreenshotPath: screenshotPath}
|
||||
|
||||
seenNoNewPages := 0
|
||||
for pageIndex := 0; pageIndex < config.Agent.ChatSyncMaxPages; pageIndex++ {
|
||||
window, err := findWeChatWindow()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := captureWindow(window, screenshotPath); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
result, err := extractChatPage(config, summaries, screenshotPath)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if result.ContactName != "" && archive.Contact == "" {
|
||||
archive.Contact = result.ContactName
|
||||
}
|
||||
|
||||
newCount := appendChatMessages(&archive, result.Messages, pageIndex)
|
||||
archive.Pages = append(archive.Pages, ChatPage{
|
||||
PageIndex: pageIndex,
|
||||
Direction: "down",
|
||||
NewCount: newCount,
|
||||
SeenAt: time.Now().Format(time.RFC3339),
|
||||
})
|
||||
logTask(fmt.Sprintf("聊天页 page=%d 提取 messages=%d new=%d summary=%s", pageIndex+1, len(result.Messages), newCount, result.PageSummary))
|
||||
|
||||
if newCount == 0 {
|
||||
seenNoNewPages++
|
||||
} else {
|
||||
seenNoNewPages = 0
|
||||
}
|
||||
if seenNoNewPages >= 2 {
|
||||
logInfo("连续页面没有新增聊天记录，停止同步")
|
||||
break
|
||||
}
|
||||
|
||||
if pageIndex < config.Agent.ChatSyncMaxPages-1 {
|
||||
if err := scrollChat(config, summaries, -absInt(config.Agent.ScrollDeltaY)); err != nil {
|
||||
return err
|
||||
}
|
||||
time.Sleep(800 * time.Millisecond)
|
||||
}
|
||||
}
|
||||
|
||||
archive.UpdatedAt = time.Now().Format(time.RFC3339)
|
||||
if err := saveChatArchive(archivePath, archive); err != nil {
|
||||
return err
|
||||
}
|
||||
logInfo(fmt.Sprintf("聊天记录同步完成，messages=%d path=%s", len(archive.Messages), archivePath))
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func validateChatSyncRegions(summaries []RegionSummary) error {
|
||||
index := regionByType(summaries)
|
||||
if _, exists := index["chat_content"]; !exists {
|
||||
return fmt.Errorf("missing required region: chat_content")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func scrollChatToTop(config Config, summaries []RegionSummary, screenshotPath string) error {
|
||||
for index := 0; index < config.Agent.ChatSyncTopScrolls; index++ {
|
||||
if err := scrollChat(config, summaries, absInt(config.Agent.ScrollDeltaY)); err != nil {
|
||||
return err
|
||||
}
|
||||
logTask(fmt.Sprintf("向上滚动到历史顶部 round=%d/%d", index+1, config.Agent.ChatSyncTopScrolls))
|
||||
time.Sleep(350 * time.Millisecond)
|
||||
}
|
||||
|
||||
window, err := findWeChatWindow()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return captureWindow(window, screenshotPath)
|
||||
}
|
||||
|
||||
func scrollChat(config Config, summaries []RegionSummary, deltaY int) error {
|
||||
chat, ok := regionByType(summaries)["chat_content"]
|
||||
if !ok {
|
||||
return fmt.Errorf("chat_content region not found")
|
||||
}
|
||||
window, err := findWeChatWindow()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
pointX := window.X + chat.CenterSource[0]
|
||||
pointY := window.Y + chat.CenterSource[1]
|
||||
logTask(fmt.Sprintf("滚动聊天区域 point=[%.1f, %.1f] delta_y=%d", pointX, pointY, deltaY))
|
||||
return scrollAt(pointX, pointY, deltaY)
|
||||
}
|
||||
|
||||
func extractChatPage(config Config, summaries []RegionSummary, screenshotPath string) (ChatExtractResult, error) {
|
||||
imageDataURL, err := loadImageDataURL(screenshotPath)
|
||||
if err != nil {
|
||||
return ChatExtractResult{}, err
|
||||
}
|
||||
prompt, err := buildChatExtractPrompt(summaries)
|
||||
if err != nil {
|
||||
return ChatExtractResult{}, err
|
||||
}
|
||||
logThink("正在用豆包多模态抽取当前聊天区域消息")
|
||||
raw, err := requestLLM(config, prompt, imageDataURL)
|
||||
if err != nil {
|
||||
return ChatExtractResult{}, err
|
||||
}
|
||||
jsonText, err := extractJSON(raw)
|
||||
if err != nil {
|
||||
_ = saveRawResponse(config, raw)
|
||||
return ChatExtractResult{}, err
|
||||
}
|
||||
var result ChatExtractResult
|
||||
if err := json.Unmarshal([]byte(jsonText), &result); err != nil {
|
||||
_ = saveRawResponse(config, raw)
|
||||
return ChatExtractResult{}, err
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func buildChatExtractPrompt(summaries []RegionSummary) (string, error) {
|
||||
regionsJSON, err := json.MarshalIndent(summaries, "", " ")
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return `你是微信聊天记录抽取器。
|
||||
|
||||
请只识别截图中 chat_content 标注区域里的聊天消息，不要识别联系人列表、输入框或其它区域。
|
||||
|
||||
要求：
|
||||
- 按消息在页面中的从上到下顺序输出。
|
||||
- sender 可填 "me"、"other" 或能识别到的昵称；无法判断填 "unknown"。
|
||||
- role 可填 "me"、"other"、"system"。
|
||||
- time 如果画面中没有明确时间则留空字符串。
|
||||
- content 必须是消息文本，不能编造。
|
||||
- 只能返回合法 JSON，不要 Markdown，不要解释。
|
||||
|
||||
标注区域：
|
||||
` + string(regionsJSON) + `
|
||||
|
||||
返回结构：
|
||||
{
|
||||
"contact_name": "unknown",
|
||||
"page_summary": "当前页聊天内容摘要",
|
||||
"messages": [
|
||||
{"sender":"other","role":"other","time":"","content":"消息文本"}
|
||||
]
|
||||
}`, nil
|
||||
}
|
||||
|
||||
func loadChatArchive(path string) (ChatArchive, bool, error) {
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
return ChatArchive{}, false, nil
|
||||
}
|
||||
return ChatArchive{}, false, err
|
||||
}
|
||||
var archive ChatArchive
|
||||
if err := json.Unmarshal(data, &archive); err != nil {
|
||||
return ChatArchive{}, false, err
|
||||
}
|
||||
return archive, true, nil
|
||||
}
|
||||
|
||||
func saveChatArchive(path string, archive ChatArchive) error {
|
||||
if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil {
|
||||
return err
|
||||
}
|
||||
data, err := json.MarshalIndent(archive, "", " ")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return os.WriteFile(path, data, 0644)
|
||||
}
|
||||
|
||||
func appendChatMessages(archive *ChatArchive, messages []ChatMessage, pageIndex int) int {
|
||||
existing := make(map[string]bool, len(archive.Messages))
|
||||
for _, message := range archive.Messages {
|
||||
existing[message.ID] = true
|
||||
}
|
||||
|
||||
newCount := 0
|
||||
for _, message := range messages {
|
||||
message.Content = strings.TrimSpace(message.Content)
|
||||
if message.Content == "" {
|
||||
continue
|
||||
}
|
||||
message.PageIndex = pageIndex
|
||||
message.SeenAt = time.Now().Format(time.RFC3339)
|
||||
message.ID = chatMessageID(message)
|
||||
if existing[message.ID] {
|
||||
continue
|
||||
}
|
||||
existing[message.ID] = true
|
||||
archive.Messages = append(archive.Messages, message)
|
||||
newCount++
|
||||
}
|
||||
return newCount
|
||||
}
|
||||
|
||||
func chatMessageID(message ChatMessage) string {
|
||||
key := strings.Join([]string{
|
||||
strings.TrimSpace(message.Sender),
|
||||
strings.TrimSpace(message.Role),
|
||||
strings.TrimSpace(message.Time),
|
||||
strings.TrimSpace(message.Content),
|
||||
}, "|")
|
||||
sum := sha1.Sum([]byte(key))
|
||||
return hex.EncodeToString(sum[:])
|
||||
}
|
||||
|
||||
func absInt(value int) int {
|
||||
if value < 0 {
|
||||
return -value
|
||||
}
|
||||
return value
|
||||
}
|
||||
@ -11,3 +11,9 @@ screenshot_path = "data/screenshots/WeChat.jpg"
|
||||
task_output_path = "data/tasks/latest_task.json"
|
||||
task_history_dir = "data/tasks/history"
|
||||
dry_run = true
|
||||
observe_mode = "normal" # normal | scroll_first
|
||||
max_scrolls = 1
|
||||
scroll_delta_y = 6
|
||||
chat_records_path = "data/chats/wechat_chat_records.json"
|
||||
chat_sync_max_pages = 3
|
||||
chat_sync_top_scrolls = 8
|
||||
|
||||
@ -50,6 +50,24 @@ func applyConfigDefaults(config *Config) {
|
||||
if config.Agent.TaskHistoryDir == "" {
|
||||
config.Agent.TaskHistoryDir = "data/tasks/history"
|
||||
}
|
||||
if config.Agent.ObserveMode == "" {
|
||||
config.Agent.ObserveMode = "normal"
|
||||
}
|
||||
if config.Agent.MaxScrolls <= 0 {
|
||||
config.Agent.MaxScrolls = 1
|
||||
}
|
||||
if config.Agent.ScrollDeltaY == 0 {
|
||||
config.Agent.ScrollDeltaY = 6
|
||||
}
|
||||
if config.Agent.ChatRecordsPath == "" {
|
||||
config.Agent.ChatRecordsPath = "data/chats/wechat_chat_records.json"
|
||||
}
|
||||
if config.Agent.ChatSyncMaxPages <= 0 {
|
||||
config.Agent.ChatSyncMaxPages = 3
|
||||
}
|
||||
if config.Agent.ChatSyncTopScrolls <= 0 {
|
||||
config.Agent.ChatSyncTopScrolls = 8
|
||||
}
|
||||
config.Agent.DryRun = true
|
||||
}
|
||||
|
||||
|
||||
178
agent/executor.go
Normal file
178
agent/executor.go
Normal file
@ -0,0 +1,178 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"math"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
func maybeScrollBeforeObserve(config Config, summaries []RegionSummary, screenshotPath string) error {
|
||||
if config.Agent.ObserveMode != "scroll_first" {
|
||||
return nil
|
||||
}
|
||||
if runtime.GOOS != "darwin" {
|
||||
return fmt.Errorf("observe_mode=scroll_first currently supports macOS only")
|
||||
}
|
||||
|
||||
chat, ok := regionByType(summaries)["chat_content"]
|
||||
if !ok {
|
||||
return fmt.Errorf("cannot scroll before observe: chat_content region not found")
|
||||
}
|
||||
|
||||
maxScrolls := config.Agent.MaxScrolls
|
||||
if maxScrolls <= 0 {
|
||||
maxScrolls = 1
|
||||
}
|
||||
|
||||
for index := 0; index < maxScrolls; index++ {
|
||||
window, err := findWeChatWindow()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
pointX := window.X + chat.CenterSource[0]
|
||||
pointY := window.Y + chat.CenterSource[1]
|
||||
logTask(fmt.Sprintf("观察前滚动聊天区域：round=%d point=[%.1f, %.1f] delta_y=%d window=[%.0f, %.0f, %.0f, %.0f]", index+1, pointX, pointY, config.Agent.ScrollDeltaY, window.X, window.Y, window.Width, window.Height))
|
||||
if err := scrollAt(pointX, pointY, config.Agent.ScrollDeltaY); err != nil {
|
||||
return err
|
||||
}
|
||||
time.Sleep(700 * time.Millisecond)
|
||||
|
||||
if err := captureWindow(window, screenshotPath); err != nil {
|
||||
return err
|
||||
}
|
||||
logInfo("已滚动并覆盖更新截图 WeChat.jpg")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func findWeChatWindow() (WindowBounds, error) {
|
||||
script := `
|
||||
import CoreGraphics
|
||||
import Foundation
|
||||
|
||||
let options = CGWindowListOption(arrayLiteral: [.optionOnScreenOnly, .excludeDesktopElements])
|
||||
guard let infos = CGWindowListCopyWindowInfo(options, kCGNullWindowID) as? [[String: Any]] else {
|
||||
exit(1)
|
||||
}
|
||||
|
||||
for info in infos {
|
||||
let owner = info[kCGWindowOwnerName as String] as? String ?? ""
|
||||
let title = info[kCGWindowName as String] as? String ?? ""
|
||||
let layer = info[kCGWindowLayer as String] as? Int ?? -1
|
||||
if layer == 0 && (owner.contains("微信") || owner.lowercased().contains("wechat") || title == "微信") {
|
||||
if let bounds = info[kCGWindowBounds as String] as? [String: Any],
|
||||
let x = bounds["X"] as? Double,
|
||||
let y = bounds["Y"] as? Double,
|
||||
let w = bounds["Width"] as? Double,
|
||||
let h = bounds["Height"] as? Double {
|
||||
let id = info[kCGWindowNumber as String] as? Int ?? 0
|
||||
let result: [String: Any] = ["x": x, "y": y, "width": w, "height": h, "owner": owner, "title": title, "id": id]
|
||||
let data = try! JSONSerialization.data(withJSONObject: result)
|
||||
print(String(data: data, encoding: .utf8)!)
|
||||
exit(0)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
exit(2)
|
||||
`
|
||||
|
||||
output, err := runSwift(script)
|
||||
if err != nil {
|
||||
return WindowBounds{}, fmt.Errorf("find WeChat window failed: %w", err)
|
||||
}
|
||||
|
||||
var raw struct {
|
||||
X float64 `json:"x"`
|
||||
Y float64 `json:"y"`
|
||||
Width float64 `json:"width"`
|
||||
Height float64 `json:"height"`
|
||||
Owner string `json:"owner"`
|
||||
Title string `json:"title"`
|
||||
ID int `json:"id"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(output), &raw); err != nil {
|
||||
return WindowBounds{}, err
|
||||
}
|
||||
|
||||
return WindowBounds(raw), nil
|
||||
}
|
||||
|
||||
func scrollAt(x float64, y float64, deltaY int) error {
|
||||
if deltaY == 0 {
|
||||
deltaY = 6
|
||||
}
|
||||
|
||||
script := fmt.Sprintf(`
|
||||
import CoreGraphics
|
||||
import Foundation
|
||||
import AppKit
|
||||
|
||||
NSWorkspace.shared.runningApplications
|
||||
.first { $0.localizedName == "微信" || $0.localizedName == "WeChat" }?
|
||||
.activate(options: [.activateIgnoringOtherApps])
|
||||
usleep(300000)
|
||||
|
||||
let point = CGPoint(x: %.3f, y: %.3f)
|
||||
let move = CGEvent(mouseEventSource: nil, mouseType: .mouseMoved, mouseCursorPosition: point, mouseButton: .left)
|
||||
move?.post(tap: .cghidEventTap)
|
||||
usleep(80000)
|
||||
let down = CGEvent(mouseEventSource: nil, mouseType: .leftMouseDown, mouseCursorPosition: point, mouseButton: .left)
|
||||
let up = CGEvent(mouseEventSource: nil, mouseType: .leftMouseUp, mouseCursorPosition: point, mouseButton: .left)
|
||||
down?.post(tap: .cghidEventTap)
|
||||
usleep(60000)
|
||||
up?.post(tap: .cghidEventTap)
|
||||
usleep(120000)
|
||||
|
||||
let delta = Int32(%d)
|
||||
for _ in 0..<8 {
|
||||
let lineScroll = CGEvent(scrollWheelEvent2Source: nil, units: .line, wheelCount: 1, wheel1: delta, wheel2: 0, wheel3: 0)
|
||||
lineScroll?.post(tap: .cghidEventTap)
|
||||
usleep(45000)
|
||||
}
|
||||
|
||||
for _ in 0..<4 {
|
||||
let pixelScroll = CGEvent(scrollWheelEvent2Source: nil, units: .pixel, wheelCount: 1, wheel1: delta * 80, wheel2: 0, wheel3: 0)
|
||||
pixelScroll?.post(tap: .cghidEventTap)
|
||||
usleep(45000)
|
||||
}
|
||||
`, x, y, deltaY)
|
||||
|
||||
_, err := runSwift(script)
|
||||
return err
|
||||
}
|
||||
|
||||
func captureWindow(window WindowBounds, screenshotPath string) error {
|
||||
if err := os.MkdirAll(filepath.Dir(screenshotPath), 0755); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
rect := fmt.Sprintf("%d,%d,%d,%d", int(math.Round(window.X)), int(math.Round(window.Y)), int(math.Round(window.Width)), int(math.Round(window.Height)))
|
||||
command := exec.Command("screencapture", "-x", "-R", rect, screenshotPath)
|
||||
if output, err := command.CombinedOutput(); err != nil {
|
||||
return fmt.Errorf("screencapture failed: %w output=%s", err, string(output))
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func runSwift(script string) (string, error) {
|
||||
command := exec.Command("swift", "-")
|
||||
command.Stdin = strings.NewReader(script)
|
||||
var stderr bytes.Buffer
|
||||
command.Stderr = &stderr
|
||||
output, err := command.Output()
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("%w stderr=%s", err, stderr.String())
|
||||
}
|
||||
return strings.TrimSpace(string(output)), nil
|
||||
}
|
||||
@ -6,6 +6,22 @@ import (
|
||||
)
|
||||
|
||||
func main() {
|
||||
if len(os.Args) > 1 && os.Args[1] == "ax-probe" {
|
||||
if err := runAXProbe(); err != nil {
|
||||
logError(err.Error())
|
||||
os.Exit(1)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
if len(os.Args) > 1 && os.Args[1] == "sync-chat" {
|
||||
if err := runChatSync(); err != nil {
|
||||
logError(err.Error())
|
||||
os.Exit(1)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
if err := runOnce(); err != nil {
|
||||
logError(err.Error())
|
||||
os.Exit(1)
|
||||
@ -41,6 +57,10 @@ func runOnce() error {
|
||||
logThink(fmt.Sprintf("区域识别：%s center=[%.1f, %.1f]", summary.Type, summary.CenterScreen[0], summary.CenterScreen[1]))
|
||||
}
|
||||
|
||||
if err := maybeScrollBeforeObserve(config, summaries, screenshotPath); err != nil {
|
||||
return fmt.Errorf("滚动获取更多聊天记录失败：%w", err)
|
||||
}
|
||||
|
||||
imageDataURL, err := loadImageDataURL(screenshotPath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("读取截图失败：%w", err)
|
||||
|
||||
@ -19,7 +19,8 @@ func buildPrompt(summaries []RegionSummary) (string, error) {
|
||||
2. 识别联系人或会话对象，如果无法识别则填 unknown。
|
||||
3. 判断是否需要回复。
|
||||
4. 如果需要回复，生成 reply_text。
|
||||
5. 生成 dry-run 动作计划。
|
||||
5. 如果聊天上下文不足以判断是否回复，可以规划滚动聊天区域查看更多记录。
|
||||
6. 生成 dry-run 动作计划。
|
||||
|
||||
限制：
|
||||
- 只能返回合法 JSON，不要 Markdown，不要解释。
|
||||
@ -28,7 +29,13 @@ func buildPrompt(summaries []RegionSummary) (string, error) {
|
||||
- 如果区域 type 是 custom，请重点参考 description 判断该区域的用途和可执行动作。
|
||||
- description 是用户标注时填写的用途说明，优先级高于模型自行猜测。
|
||||
- 微信发送消息使用键入回车键，不要规划点击发送按钮。
|
||||
- action.type 只能是 click、type_text、press_enter、scroll、wait、noop。
|
||||
- 如果当前聊天内容不足以判断是否回复，可以返回 scroll 动作。
|
||||
- scroll 只能作用于 chat_content 区域。
|
||||
- delta_y < 0 表示向上滚动查看更多历史消息；delta_y > 0 表示向下滚动回到最新消息。
|
||||
- 每次最多返回一个 scroll 动作，abs(delta_y) 不要超过 5。
|
||||
- 如果滚动后需要重新识别，请追加 observe_again 动作。
|
||||
- 不要在同一个任务里同时包含 scroll 和 type_text/press_enter；上下文不足时先滚动观察，不要直接回复。
|
||||
- action.type 只能是 click、type_text、press_enter、scroll、observe_again、wait、noop。
|
||||
- 当前是 dry_run，只生成计划，不要假设动作已执行。
|
||||
- 如果无需回复，task_type 使用 wechat_observe 或 noop，并给出 observations。
|
||||
|
||||
|
||||
@ -16,12 +16,13 @@ var allowedTaskTypes = map[string]bool{
|
||||
}
|
||||
|
||||
var allowedActionTypes = map[string]bool{
|
||||
"click": true,
|
||||
"type_text": true,
|
||||
"press_enter": true,
|
||||
"scroll": true,
|
||||
"wait": true,
|
||||
"noop": true,
|
||||
"click": true,
|
||||
"type_text": true,
|
||||
"press_enter": true,
|
||||
"observe_again": true,
|
||||
"scroll": true,
|
||||
"wait": true,
|
||||
"noop": true,
|
||||
}
|
||||
|
||||
func finalizeTask(task Task, summaries []RegionSummary, config Config, regionsPath string, screenshotPath string) Task {
|
||||
@ -68,11 +69,26 @@ func finalizeTask(task Task, summaries []RegionSummary, config Config, regionsPa
|
||||
if summary, exists := regionIndex[action.TargetRegion]; exists {
|
||||
point := summary.CenterScreen
|
||||
action.Point = &point
|
||||
} else if action.Type == "click" {
|
||||
} else if action.Type == "click" || action.Type == "scroll" {
|
||||
action.Type = "noop"
|
||||
action.Reason = "target_region not found"
|
||||
}
|
||||
}
|
||||
if action.Type == "scroll" {
|
||||
if action.TargetRegion != "chat_content" {
|
||||
action.Type = "noop"
|
||||
action.Reason = "scroll only allowed on chat_content"
|
||||
}
|
||||
if action.DeltaY > 5 {
|
||||
action.DeltaY = 5
|
||||
}
|
||||
if action.DeltaY < -5 {
|
||||
action.DeltaY = -5
|
||||
}
|
||||
if action.DeltaY == 0 {
|
||||
action.DeltaY = -3
|
||||
}
|
||||
}
|
||||
if action.Type == "type_text" && len([]rune(action.Text)) > 500 {
|
||||
runes := []rune(action.Text)
|
||||
action.Text = string(runes[:500])
|
||||
@ -119,6 +135,8 @@ func executionMessage(action Action) string {
|
||||
return "dry_run: 跳过按回车发送"
|
||||
case "scroll":
|
||||
return fmt.Sprintf("dry_run: 跳过滑动 %s delta_y=%d", action.TargetRegion, action.DeltaY)
|
||||
case "observe_again":
|
||||
return "dry_run: 跳过重新截图并识别"
|
||||
case "wait":
|
||||
return fmt.Sprintf("dry_run: 跳过等待 %dms", action.DurationMs)
|
||||
default:
|
||||
|
||||
@ -12,12 +12,28 @@ type VolcengineConfig struct {
|
||||
}
|
||||
|
||||
type AgentConfig struct {
|
||||
AppDataDir string `toml:"app_data_dir"`
|
||||
RegionsPath string `toml:"regions_path"`
|
||||
ScreenshotPath string `toml:"screenshot_path"`
|
||||
TaskOutputPath string `toml:"task_output_path"`
|
||||
TaskHistoryDir string `toml:"task_history_dir"`
|
||||
DryRun bool `toml:"dry_run"`
|
||||
AppDataDir string `toml:"app_data_dir"`
|
||||
RegionsPath string `toml:"regions_path"`
|
||||
ScreenshotPath string `toml:"screenshot_path"`
|
||||
TaskOutputPath string `toml:"task_output_path"`
|
||||
TaskHistoryDir string `toml:"task_history_dir"`
|
||||
DryRun bool `toml:"dry_run"`
|
||||
ObserveMode string `toml:"observe_mode"`
|
||||
MaxScrolls int `toml:"max_scrolls"`
|
||||
ScrollDeltaY int `toml:"scroll_delta_y"`
|
||||
ChatRecordsPath string `toml:"chat_records_path"`
|
||||
ChatSyncMaxPages int `toml:"chat_sync_max_pages"`
|
||||
ChatSyncTopScrolls int `toml:"chat_sync_top_scrolls"`
|
||||
}
|
||||
|
||||
type WindowBounds struct {
|
||||
X float64
|
||||
Y float64
|
||||
Width float64
|
||||
Height float64
|
||||
Owner string
|
||||
Title string
|
||||
ID int
|
||||
}
|
||||
|
||||
type LogEntry struct {
|
||||
@ -138,3 +154,36 @@ type openAIResponse struct {
|
||||
Message string `json:"message"`
|
||||
} `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
type ChatArchive struct {
|
||||
App string `json:"app"`
|
||||
Contact string `json:"contact"`
|
||||
Source TaskSource `json:"source"`
|
||||
Messages []ChatMessage `json:"messages"`
|
||||
Pages []ChatPage `json:"pages"`
|
||||
CreatedAt string `json:"created_at"`
|
||||
UpdatedAt string `json:"updated_at"`
|
||||
}
|
||||
|
||||
type ChatMessage struct {
|
||||
ID string `json:"id"`
|
||||
Sender string `json:"sender"`
|
||||
Role string `json:"role"`
|
||||
Time string `json:"time,omitempty"`
|
||||
Content string `json:"content"`
|
||||
PageIndex int `json:"page_index"`
|
||||
SeenAt string `json:"seen_at"`
|
||||
}
|
||||
|
||||
type ChatPage struct {
|
||||
PageIndex int `json:"page_index"`
|
||||
Direction string `json:"direction"`
|
||||
NewCount int `json:"new_count"`
|
||||
SeenAt string `json:"seen_at"`
|
||||
}
|
||||
|
||||
type ChatExtractResult struct {
|
||||
ContactName string `json:"contact_name"`
|
||||
PageSummary string `json:"page_summary"`
|
||||
Messages []ChatMessage `json:"messages"`
|
||||
}
|
||||
|
||||
163
wechat_vision/app.py
Normal file
163
wechat_vision/app.py
Normal file
@ -0,0 +1,163 @@
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import onnxruntime as ort
|
||||
from PIL import Image, ImageDraw
|
||||
|
||||
|
||||
MODEL_PATH = Path(__file__).with_name("best.onnx")
|
||||
TEST_DIR = Path(__file__).with_name("test")
|
||||
OUTPUT_DIR = Path(__file__).with_name("ouptsw")
|
||||
IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".bmp", ".webp"}
|
||||
CONFIDENCE_THRESHOLD = 0.1
|
||||
IOU_THRESHOLD = 0.5
|
||||
|
||||
|
||||
def _dtype(onnx_type):
|
||||
if onnx_type == "tensor(float)":
|
||||
return np.float32
|
||||
if onnx_type == "tensor(float16)":
|
||||
return np.float16
|
||||
if onnx_type == "tensor(double)":
|
||||
return np.float64
|
||||
if onnx_type in {"tensor(int64)", "tensor(uint64)"}:
|
||||
return np.int64
|
||||
if onnx_type in {"tensor(int32)", "tensor(uint32)"}:
|
||||
return np.int32
|
||||
raise ValueError(f"Unsupported input type: {onnx_type}")
|
||||
|
||||
|
||||
def _shape(input_meta):
|
||||
return [dim if isinstance(dim, int) else 1 for dim in input_meta.shape]
|
||||
|
||||
|
||||
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):
|
||||
detections = detections[detections[:, 4] >= CONFIDENCE_THRESHOLD]
|
||||
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 _preprocess(image, input_shape):
|
||||
_, _, height, width = input_shape
|
||||
resized = image.resize((width, height))
|
||||
array = np.asarray(resized, dtype=np.float32) / 255.0
|
||||
return np.transpose(array, (2, 0, 1))[None]
|
||||
|
||||
|
||||
def _scale_detections(detections, original_size, input_shape):
|
||||
_, _, input_height, input_width = input_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 _draw_detections(image, detections):
|
||||
draw = ImageDraw.Draw(image)
|
||||
for detection in detections:
|
||||
x1, y1, x2, y2, score, class_id = detection
|
||||
cx = (x1 + x2) / 2
|
||||
cy = (y1 + y2) / 2
|
||||
label = f"({cx:.1f}, {cy:.1f}) {score:.2f}"
|
||||
|
||||
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.line((cx - 10, cy, cx + 10, cy), fill="red", width=2)
|
||||
draw.line((cx, cy - 10, cx, cy + 10), fill="red", width=2)
|
||||
draw.text((x1, max(0, y1 - 14)), label, fill="yellow")
|
||||
|
||||
|
||||
def main():
|
||||
session = ort.InferenceSession(str(MODEL_PATH))
|
||||
input_meta = session.get_inputs()[0]
|
||||
input_shape = _shape(input_meta)
|
||||
|
||||
print(f"model: {MODEL_PATH}")
|
||||
print(f"providers: {session.get_providers()}")
|
||||
|
||||
print("inputs:")
|
||||
feeds = {}
|
||||
for input_meta in session.get_inputs():
|
||||
print(
|
||||
f"- name={input_meta.name}, type={input_meta.type}, "
|
||||
f"shape={input_meta.shape}"
|
||||
)
|
||||
feeds[input_meta.name] = np.zeros(
|
||||
_shape(input_meta),
|
||||
dtype=_dtype(input_meta.type),
|
||||
)
|
||||
|
||||
print("outputs:")
|
||||
for output_meta in session.get_outputs():
|
||||
print(
|
||||
f"- name={output_meta.name}, type={output_meta.type}, "
|
||||
f"shape={output_meta.shape}"
|
||||
)
|
||||
|
||||
outputs = session.run(None, feeds)
|
||||
print("dummy inference: ok")
|
||||
for index, output in enumerate(outputs):
|
||||
print(f"- output[{index}]: shape={output.shape}, dtype={output.dtype}")
|
||||
|
||||
OUTPUT_DIR.mkdir(exist_ok=True)
|
||||
result_lines = ["image,class_id,score,center_x,center_y,x1,y1,x2,y2"]
|
||||
|
||||
for image_path in sorted(TEST_DIR.iterdir()):
|
||||
if image_path.suffix.lower() not in IMAGE_SUFFIXES:
|
||||
continue
|
||||
|
||||
image = Image.open(image_path).convert("RGB")
|
||||
tensor = _preprocess(image, input_shape)
|
||||
detections = session.run(None, {input_meta.name: tensor})[0][0]
|
||||
detections = _scale_detections(_nms(detections), image.size, input_shape)
|
||||
|
||||
annotated = image.copy()
|
||||
_draw_detections(annotated, detections)
|
||||
output_path = OUTPUT_DIR / image_path.name
|
||||
annotated.save(output_path)
|
||||
|
||||
print(f"{image_path.name}: {len(detections)} target(s) -> {output_path}")
|
||||
for detection in detections:
|
||||
x1, y1, x2, y2, score, class_id = detection
|
||||
cx = (x1 + x2) / 2
|
||||
cy = (y1 + y2) / 2
|
||||
print(f" center=({cx:.1f}, {cy:.1f}), score={score:.3f}")
|
||||
result_lines.append(
|
||||
f"{image_path.name},{int(class_id)},{score:.6f},{cx:.2f},{cy:.2f},"
|
||||
f"{x1:.2f},{y1:.2f},{x2:.2f},{y2:.2f}"
|
||||
)
|
||||
|
||||
(OUTPUT_DIR / "coordinates.csv").write_text("\n".join(result_lines) + "\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
BIN
wechat_vision/best.onnx
Normal file
BIN
wechat_vision/best.onnx
Normal file
Binary file not shown.
898
wechat_vision/ffmpeg_realtime_detect.py
Normal file
898
wechat_vision/ffmpeg_realtime_detect.py
Normal file
@ -0,0 +1,898 @@
|
||||
import argparse
|
||||
import json
|
||||
import signal
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
import webbrowser
|
||||
from datetime import datetime, timezone
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import onnxruntime as ort
|
||||
from PIL import Image, ImageDraw
|
||||
|
||||
|
||||
BASE_DIR = Path(__file__).resolve().parent
|
||||
MODEL_PATH = BASE_DIR / "best.onnx"
|
||||
DEBUG_DIR = BASE_DIR / "ouptsw" / "monitor"
|
||||
VISUALIZE_DIR = BASE_DIR / "ouptsw" / "realtime_view"
|
||||
CLASS_NAMES = {
|
||||
0: "input_box",
|
||||
}
|
||||
IOU_THRESHOLD = 0.5
|
||||
|
||||
|
||||
class StreamState:
|
||||
def __init__(self):
|
||||
self.condition = threading.Condition()
|
||||
self.jpeg = None
|
||||
self.event = None
|
||||
self.sequence = 0
|
||||
|
||||
def update(self, jpeg, event):
|
||||
with self.condition:
|
||||
self.jpeg = jpeg
|
||||
self.event = event
|
||||
self.sequence += 1
|
||||
self.condition.notify_all()
|
||||
|
||||
def snapshot(self):
|
||||
with self.condition:
|
||||
return self.jpeg, self.event, self.sequence
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Realtime FFmpeg screen capture + ONNX detection demo."
|
||||
)
|
||||
parser.add_argument("--model", default=str(MODEL_PATH), help="ONNX model path")
|
||||
parser.add_argument(
|
||||
"--input",
|
||||
default="3:none",
|
||||
help="FFmpeg avfoundation input, e.g. '3:none'. Use --list-devices to inspect.",
|
||||
)
|
||||
parser.add_argument("--fps", type=float, default=2.0, help="capture FPS")
|
||||
parser.add_argument("--confidence", type=float, default=0.1, help="minimum score")
|
||||
parser.add_argument("--target-class", type=int, default=0, help="target class id")
|
||||
parser.add_argument(
|
||||
"--max-frames",
|
||||
type=int,
|
||||
default=0,
|
||||
help="stop after N frames; 0 means run until interrupted",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--emit-empty",
|
||||
action="store_true",
|
||||
help="emit no_detection events when target is not found",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--debug-dir",
|
||||
default="",
|
||||
help="save annotated frames into this directory",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--debug-every",
|
||||
type=int,
|
||||
default=0,
|
||||
help="save one debug image every N frames; 0 disables debug images",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--visualize-dir",
|
||||
default="",
|
||||
help="write latest.jpg/latest.json/index.html for browser visualization",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--visualize-every",
|
||||
type=int,
|
||||
default=1,
|
||||
help="update visualization every N frames when --visualize-dir is enabled",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--open-visualizer",
|
||||
action="store_true",
|
||||
help="open the visualization HTML in the default browser",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--serve",
|
||||
action="store_true",
|
||||
help="serve a realtime MJPEG stream and dashboard over HTTP",
|
||||
)
|
||||
parser.add_argument("--host", default="127.0.0.1", help="HTTP server host")
|
||||
parser.add_argument("--port", type=int, default=8765, help="HTTP server port")
|
||||
parser.add_argument(
|
||||
"--display",
|
||||
action="store_true",
|
||||
help="open an ffplay window with the annotated realtime video",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--display-title",
|
||||
default="WeChat Vision ONNX Preview",
|
||||
help="ffplay window title when --display is enabled",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--display-width",
|
||||
type=int,
|
||||
default=720,
|
||||
help="ffplay window width when --display is enabled",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--display-height",
|
||||
type=int,
|
||||
default=450,
|
||||
help="ffplay window height when --display is enabled",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--display-left",
|
||||
type=int,
|
||||
default=0,
|
||||
help="ffplay window left position when --display is enabled",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--display-top",
|
||||
type=int,
|
||||
default=0,
|
||||
help="ffplay window top position when --display is enabled",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--list-devices",
|
||||
action="store_true",
|
||||
help="list FFmpeg avfoundation devices and exit",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def log_error(message):
|
||||
print(message, file=sys.stderr, flush=True)
|
||||
|
||||
|
||||
def utc_now():
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
|
||||
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):
|
||||
if len(detections) == 0:
|
||||
return detections
|
||||
|
||||
_, _, 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(session, input_meta, image, target_class, confidence):
|
||||
shape = input_shape(input_meta)
|
||||
tensor = preprocess(image, shape)
|
||||
raw = session.run(None, {input_meta.name: tensor})[0][0]
|
||||
detections = raw[raw[:, 5].astype(int) == target_class]
|
||||
detections = nms(detections, confidence)
|
||||
return scale_detections(detections, image.size, shape)
|
||||
|
||||
|
||||
def emit(event):
|
||||
print(json.dumps(event, ensure_ascii=False, separators=(",", ":")), flush=True)
|
||||
|
||||
|
||||
def write_json_atomic(path, data):
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
temp_path = path.with_suffix(path.suffix + ".tmp")
|
||||
temp_path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
temp_path.replace(path)
|
||||
|
||||
|
||||
def save_image_atomic(image, path):
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
temp_path = path.with_suffix(path.suffix + ".tmp")
|
||||
image.save(temp_path, format="JPEG", quality=85)
|
||||
temp_path.replace(path)
|
||||
|
||||
|
||||
def image_to_jpeg_bytes(image):
|
||||
buffer = BytesIO()
|
||||
image.save(buffer, format="JPEG", quality=85)
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
def detection_event(detection, frame_index, frame_size, class_name):
|
||||
x1, y1, x2, y2, score, class_id = detection
|
||||
cx = (x1 + x2) / 2
|
||||
cy = (y1 + y2) / 2
|
||||
return {
|
||||
"type": "detection",
|
||||
"timestamp": utc_now(),
|
||||
"frame_index": frame_index,
|
||||
"class_id": int(class_id),
|
||||
"class_name": class_name,
|
||||
"confidence": round(float(score), 6),
|
||||
"bbox_screen": [round(float(x1), 2), round(float(y1), 2), round(float(x2), 2), round(float(y2), 2)],
|
||||
"center_screen": [round(float(cx), 2), round(float(cy), 2)],
|
||||
"frame_size": [frame_size[0], frame_size[1]],
|
||||
}
|
||||
|
||||
|
||||
def no_detection_event(frame_index, frame_size, target_class, class_name):
|
||||
return {
|
||||
"type": "no_detection",
|
||||
"timestamp": utc_now(),
|
||||
"frame_index": frame_index,
|
||||
"class_id": target_class,
|
||||
"class_name": class_name,
|
||||
"frame_size": [frame_size[0], frame_size[1]],
|
||||
}
|
||||
|
||||
|
||||
def save_debug_image(image, detections, path, class_name):
|
||||
debug = draw_detections(image, detections, class_name)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
debug.save(path)
|
||||
|
||||
|
||||
def draw_detections(image, detections, class_name):
|
||||
debug = image.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"{class_name} ({cx:.1f}, {cy:.1f}) {score:.2f}",
|
||||
fill="yellow",
|
||||
)
|
||||
return debug
|
||||
|
||||
|
||||
def draw_frame_info(image, event):
|
||||
draw = ImageDraw.Draw(image)
|
||||
frame_index = event.get("frame_index", "-")
|
||||
timestamp = event.get("timestamp", "")
|
||||
text = f"frame={frame_index} {timestamp}"
|
||||
draw.rectangle((12, 12, 620, 52), fill=(0, 0, 0))
|
||||
draw.text((22, 22), text, fill=(80, 255, 120))
|
||||
return image
|
||||
|
||||
|
||||
def write_visualizer_html(path):
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
html = """<!doctype html>
|
||||
<html lang="zh-CN">
|
||||
<head>
|
||||
<meta charset="utf-8" />
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1" />
|
||||
<title>WeChat Vision Realtime</title>
|
||||
<style>
|
||||
:root { color-scheme: dark; font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; }
|
||||
body { margin: 0; background: #09090b; color: #f4f4f5; }
|
||||
.wrap { display: grid; grid-template-columns: minmax(0, 1fr) 360px; gap: 16px; padding: 16px; }
|
||||
.stage { background: #18181b; border: 1px solid #27272a; border-radius: 14px; overflow: hidden; min-height: 60vh; }
|
||||
img { display: block; width: 100%; height: auto; }
|
||||
aside { background: #18181b; border: 1px solid #27272a; border-radius: 14px; padding: 16px; }
|
||||
h1 { margin: 0 0 12px; font-size: 18px; }
|
||||
.item { padding: 10px 0; border-top: 1px solid #27272a; }
|
||||
.label { color: #a1a1aa; font-size: 12px; margin-bottom: 4px; }
|
||||
.value { font-family: ui-monospace, SFMono-Regular, Menlo, monospace; word-break: break-all; }
|
||||
.ok { color: #86efac; }
|
||||
.miss { color: #fca5a5; }
|
||||
@media (max-width: 900px) { .wrap { grid-template-columns: 1fr; } }
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="wrap">
|
||||
<main class="stage"><img id="frame" alt="latest detection frame" /></main>
|
||||
<aside>
|
||||
<h1>WeChat Vision Realtime</h1>
|
||||
<div class="item"><div class="label">Status</div><div id="status" class="value">loading</div></div>
|
||||
<div class="item"><div class="label">Frame</div><div id="frameIndex" class="value">-</div></div>
|
||||
<div class="item"><div class="label">Class</div><div id="className" class="value">-</div></div>
|
||||
<div class="item"><div class="label">Confidence</div><div id="confidence" class="value">-</div></div>
|
||||
<div class="item"><div class="label">Center Screen</div><div id="center" class="value">-</div></div>
|
||||
<div class="item"><div class="label">BBox Screen</div><div id="bbox" class="value">-</div></div>
|
||||
<div class="item"><div class="label">Frame Size</div><div id="size" class="value">-</div></div>
|
||||
<div class="item"><div class="label">Timestamp</div><div id="time" class="value">-</div></div>
|
||||
</aside>
|
||||
</div>
|
||||
<script>
|
||||
async function refresh() {
|
||||
const cacheBust = Date.now();
|
||||
document.getElementById('frame').src = 'latest.jpg?t=' + cacheBust;
|
||||
try {
|
||||
const response = await fetch('latest.json?t=' + cacheBust);
|
||||
const data = await response.json();
|
||||
const hit = data.event && data.event.type === 'detection';
|
||||
const status = document.getElementById('status');
|
||||
status.textContent = hit ? 'detection' : 'no detection';
|
||||
status.className = 'value ' + (hit ? 'ok' : 'miss');
|
||||
document.getElementById('frameIndex').textContent = data.frame_index ?? '-';
|
||||
document.getElementById('className').textContent = data.class_name ?? '-';
|
||||
document.getElementById('confidence').textContent = data.confidence ?? '-';
|
||||
document.getElementById('center').textContent = JSON.stringify(data.center_screen ?? '-');
|
||||
document.getElementById('bbox').textContent = JSON.stringify(data.bbox_screen ?? '-');
|
||||
document.getElementById('size').textContent = JSON.stringify(data.frame_size ?? '-');
|
||||
document.getElementById('time').textContent = data.timestamp ?? '-';
|
||||
} catch (error) {
|
||||
document.getElementById('status').textContent = 'waiting for latest.json';
|
||||
}
|
||||
}
|
||||
refresh();
|
||||
setInterval(refresh, 500);
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
"""
|
||||
path.write_text(html, encoding="utf-8")
|
||||
|
||||
|
||||
def dashboard_html(frame_url="frame.jpg", json_url="latest.json"):
|
||||
return f"""<!doctype html>
|
||||
<html lang="zh-CN">
|
||||
<head>
|
||||
<meta charset="utf-8" />
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1" />
|
||||
<title>WeChat Vision Live Stream</title>
|
||||
<style>
|
||||
:root {{ color-scheme: dark; font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; }}
|
||||
body {{ margin: 0; background: #09090b; color: #f4f4f5; }}
|
||||
.wrap {{ display: grid; grid-template-columns: minmax(0, 1fr) 360px; gap: 16px; padding: 16px; }}
|
||||
.stage {{ background: #18181b; border: 1px solid #27272a; border-radius: 14px; overflow: hidden; min-height: 60vh; position: relative; }}
|
||||
canvas {{ display: block; width: 100%; height: auto; }}
|
||||
.overlay {{ position: absolute; left: 12px; top: 12px; padding: 6px 10px; border-radius: 999px; background: rgba(9, 9, 11, 0.78); color: #86efac; font: 12px ui-monospace, SFMono-Regular, Menlo, monospace; }}
|
||||
aside {{ background: #18181b; border: 1px solid #27272a; border-radius: 14px; padding: 16px; }}
|
||||
h1 {{ margin: 0 0 12px; font-size: 18px; }}
|
||||
.item {{ padding: 10px 0; border-top: 1px solid #27272a; }}
|
||||
.label {{ color: #a1a1aa; font-size: 12px; margin-bottom: 4px; }}
|
||||
.value {{ font-family: ui-monospace, SFMono-Regular, Menlo, monospace; word-break: break-all; }}
|
||||
.pill {{ display: inline-flex; align-items: center; gap: 8px; padding: 6px 10px; border-radius: 999px; background: #052e16; color: #86efac; font-size: 12px; margin-bottom: 10px; }}
|
||||
.dot {{ width: 8px; height: 8px; border-radius: 50%; background: #22c55e; box-shadow: 0 0 14px #22c55e; }}
|
||||
.ok {{ color: #86efac; }}
|
||||
.miss {{ color: #fca5a5; }}
|
||||
@media (max-width: 900px) {{ .wrap {{ grid-template-columns: 1fr; }} }}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="wrap">
|
||||
<main class="stage">
|
||||
<canvas id="preview"></canvas>
|
||||
<div id="overlay" class="overlay">waiting for frames</div>
|
||||
</main>
|
||||
<aside>
|
||||
<h1>WeChat Vision Live</h1>
|
||||
<div class="pill"><span class="dot"></span><span>Canvas realtime preview</span></div>
|
||||
<div class="item"><div class="label">Frame URL</div><div class="value">/{frame_url}</div></div>
|
||||
<div class="item"><div class="label">Sequence</div><div id="sequence" class="value">-</div></div>
|
||||
<div class="item"><div class="label">Metadata FPS</div><div id="metaFps" class="value">-</div></div>
|
||||
<div class="item"><div class="label">Status</div><div id="status" class="value">loading</div></div>
|
||||
<div class="item"><div class="label">Frame</div><div id="frameIndex" class="value">-</div></div>
|
||||
<div class="item"><div class="label">Class</div><div id="className" class="value">-</div></div>
|
||||
<div class="item"><div class="label">Confidence</div><div id="confidence" class="value">-</div></div>
|
||||
<div class="item"><div class="label">Center Screen</div><div id="center" class="value">-</div></div>
|
||||
<div class="item"><div class="label">BBox Screen</div><div id="bbox" class="value">-</div></div>
|
||||
<div class="item"><div class="label">Frame Size</div><div id="size" class="value">-</div></div>
|
||||
<div class="item"><div class="label">Timestamp</div><div id="time" class="value">-</div></div>
|
||||
</aside>
|
||||
</div>
|
||||
<script>
|
||||
const canvas = document.getElementById('preview');
|
||||
const context = canvas.getContext('2d');
|
||||
let previousSequence = null;
|
||||
let previousTime = null;
|
||||
let drawnSequence = null;
|
||||
|
||||
async function drawFrame(sequence) {{
|
||||
const response = await fetch('{frame_url}?seq=' + sequence + '&t=' + Date.now(), {{ cache: 'no-store' }});
|
||||
if (!response.ok) return;
|
||||
const blob = await response.blob();
|
||||
const bitmap = await createImageBitmap(blob);
|
||||
if (canvas.width !== bitmap.width || canvas.height !== bitmap.height) {{
|
||||
canvas.width = bitmap.width;
|
||||
canvas.height = bitmap.height;
|
||||
}}
|
||||
context.drawImage(bitmap, 0, 0);
|
||||
bitmap.close();
|
||||
drawnSequence = sequence;
|
||||
}}
|
||||
|
||||
async function refreshMeta() {{
|
||||
try {{
|
||||
const response = await fetch('{json_url}?t=' + Date.now(), {{ cache: 'no-store' }});
|
||||
const data = await response.json();
|
||||
const event = data.event || data;
|
||||
const now = performance.now();
|
||||
if (previousSequence !== null && previousTime !== null && data.sequence !== undefined) {{
|
||||
const elapsed = Math.max((now - previousTime) / 1000, 0.001);
|
||||
const fps = (data.sequence - previousSequence) / elapsed;
|
||||
document.getElementById('metaFps').textContent = fps.toFixed(1);
|
||||
}}
|
||||
previousSequence = data.sequence;
|
||||
previousTime = now;
|
||||
const hit = event.type === 'detection';
|
||||
const status = document.getElementById('status');
|
||||
status.textContent = hit ? 'detection' : event.type || 'waiting';
|
||||
status.className = 'value ' + (hit ? 'ok' : 'miss');
|
||||
document.getElementById('sequence').textContent = data.sequence ?? '-';
|
||||
document.getElementById('frameIndex').textContent = event.frame_index ?? '-';
|
||||
document.getElementById('className').textContent = event.class_name ?? '-';
|
||||
document.getElementById('confidence').textContent = event.confidence ?? '-';
|
||||
document.getElementById('center').textContent = JSON.stringify(event.center_screen ?? '-');
|
||||
document.getElementById('bbox').textContent = JSON.stringify(event.bbox_screen ?? '-');
|
||||
document.getElementById('size').textContent = JSON.stringify(event.frame_size ?? '-');
|
||||
document.getElementById('time').textContent = event.timestamp ?? '-';
|
||||
document.getElementById('overlay').textContent = `seq=${{data.sequence ?? '-'}} frame=${{event.frame_index ?? '-'}} ${{hit ? 'detection' : event.type || 'waiting'}}`;
|
||||
if (data.sequence !== undefined && data.sequence > 0 && data.sequence !== drawnSequence) {{
|
||||
await drawFrame(data.sequence);
|
||||
}}
|
||||
}} catch (error) {{
|
||||
document.getElementById('status').textContent = 'waiting for preview';
|
||||
}}
|
||||
}}
|
||||
refreshMeta();
|
||||
setInterval(refreshMeta, 200);
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
"""
|
||||
|
||||
|
||||
def make_stream_handler(state):
|
||||
class Handler(BaseHTTPRequestHandler):
|
||||
protocol_version = "HTTP/1.1"
|
||||
|
||||
def log_message(self, format, *args):
|
||||
return
|
||||
|
||||
def do_HEAD(self):
|
||||
if self.path == "/" or self.path.startswith("/index.html"):
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "text/html; charset=utf-8")
|
||||
self.send_header("Cache-Control", "no-store")
|
||||
self.end_headers()
|
||||
return
|
||||
if self.path.startswith("/latest.json"):
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "application/json; charset=utf-8")
|
||||
self.send_header("Cache-Control", "no-store")
|
||||
self.end_headers()
|
||||
return
|
||||
if self.path.startswith("/frame.jpg"):
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "image/jpeg")
|
||||
self.send_header("Cache-Control", "no-store, no-cache, must-revalidate, max-age=0")
|
||||
self.end_headers()
|
||||
return
|
||||
if self.path.startswith("/stream.mjpg"):
|
||||
self.send_response(200)
|
||||
self.send_header("Cache-Control", "no-cache, private")
|
||||
self.send_header("Pragma", "no-cache")
|
||||
self.send_header("Content-Type", "multipart/x-mixed-replace; boundary=frame")
|
||||
self.end_headers()
|
||||
return
|
||||
self.send_error(404)
|
||||
|
||||
def do_GET(self):
|
||||
if self.path == "/" or self.path.startswith("/index.html"):
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "text/html; charset=utf-8")
|
||||
self.send_header("Cache-Control", "no-store")
|
||||
self.end_headers()
|
||||
self.wfile.write(dashboard_html().encode("utf-8"))
|
||||
return
|
||||
|
||||
if self.path.startswith("/latest.json"):
|
||||
_, event, sequence = state.snapshot()
|
||||
payload = {
|
||||
"sequence": sequence,
|
||||
"updated_at": utc_now(),
|
||||
"event": event or {"type": "waiting", "timestamp": utc_now()},
|
||||
}
|
||||
data = json.dumps(payload, ensure_ascii=False).encode("utf-8")
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "application/json; charset=utf-8")
|
||||
self.send_header("Cache-Control", "no-store")
|
||||
self.send_header("Content-Length", str(len(data)))
|
||||
self.end_headers()
|
||||
self.wfile.write(data)
|
||||
return
|
||||
|
||||
if self.path.startswith("/frame.jpg"):
|
||||
jpeg, _, sequence = state.snapshot()
|
||||
if not jpeg:
|
||||
self.send_response(503)
|
||||
self.send_header("Cache-Control", "no-store")
|
||||
self.end_headers()
|
||||
return
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "image/jpeg")
|
||||
self.send_header("Cache-Control", "no-store, no-cache, must-revalidate, max-age=0")
|
||||
self.send_header("Pragma", "no-cache")
|
||||
self.send_header("X-Frame-Sequence", str(sequence))
|
||||
self.send_header("Content-Length", str(len(jpeg)))
|
||||
self.end_headers()
|
||||
self.wfile.write(jpeg)
|
||||
return
|
||||
|
||||
if self.path.startswith("/stream.mjpg"):
|
||||
self.send_response(200)
|
||||
self.send_header("Age", "0")
|
||||
self.send_header("Cache-Control", "no-cache, private")
|
||||
self.send_header("Pragma", "no-cache")
|
||||
self.send_header("Content-Type", "multipart/x-mixed-replace; boundary=frame")
|
||||
self.end_headers()
|
||||
|
||||
last_sequence = -1
|
||||
while True:
|
||||
with state.condition:
|
||||
state.condition.wait_for(lambda: state.sequence != last_sequence, timeout=5)
|
||||
jpeg = state.jpeg
|
||||
last_sequence = state.sequence
|
||||
if not jpeg:
|
||||
continue
|
||||
try:
|
||||
self.wfile.write(b"--frame\r\n")
|
||||
self.wfile.write(b"Content-Type: image/jpeg\r\n")
|
||||
self.wfile.write(f"Content-Length: {len(jpeg)}\r\n\r\n".encode("ascii"))
|
||||
self.wfile.write(jpeg)
|
||||
self.wfile.write(b"\r\n")
|
||||
except (BrokenPipeError, ConnectionResetError):
|
||||
break
|
||||
return
|
||||
|
||||
self.send_error(404)
|
||||
|
||||
return Handler
|
||||
|
||||
|
||||
def start_stream_server(host, port, state):
|
||||
server = ThreadingHTTPServer((host, port), make_stream_handler(state))
|
||||
thread = threading.Thread(target=server.serve_forever, daemon=True)
|
||||
thread.start()
|
||||
return server
|
||||
|
||||
|
||||
def update_visualizer(image, detections, event, output_dir, class_name):
|
||||
annotated = draw_detections(image, detections, class_name)
|
||||
annotated = draw_frame_info(annotated, event)
|
||||
save_image_atomic(annotated, output_dir / "latest.jpg")
|
||||
payload = {
|
||||
"updated_at": utc_now(),
|
||||
"event": event,
|
||||
"frame_index": event.get("frame_index"),
|
||||
"class_id": event.get("class_id"),
|
||||
"class_name": event.get("class_name"),
|
||||
"confidence": event.get("confidence"),
|
||||
"bbox_screen": event.get("bbox_screen"),
|
||||
"center_screen": event.get("center_screen"),
|
||||
"frame_size": event.get("frame_size"),
|
||||
"timestamp": event.get("timestamp"),
|
||||
}
|
||||
write_json_atomic(output_dir / "latest.json", payload)
|
||||
|
||||
|
||||
def ffmpeg_command(args):
|
||||
return [
|
||||
"ffmpeg",
|
||||
"-hide_banner",
|
||||
"-loglevel",
|
||||
"error",
|
||||
"-f",
|
||||
"avfoundation",
|
||||
"-framerate",
|
||||
str(args.fps),
|
||||
"-i",
|
||||
args.input,
|
||||
"-f",
|
||||
"image2pipe",
|
||||
"-vcodec",
|
||||
"mjpeg",
|
||||
"-",
|
||||
]
|
||||
|
||||
|
||||
def ffplay_command(args):
|
||||
return [
|
||||
"ffplay",
|
||||
"-hide_banner",
|
||||
"-loglevel",
|
||||
"error",
|
||||
"-fflags",
|
||||
"nobuffer",
|
||||
"-flags",
|
||||
"low_delay",
|
||||
"-framedrop",
|
||||
"-sync",
|
||||
"ext",
|
||||
"-window_title",
|
||||
args.display_title,
|
||||
"-left",
|
||||
str(args.display_left),
|
||||
"-top",
|
||||
str(args.display_top),
|
||||
"-x",
|
||||
str(args.display_width),
|
||||
"-y",
|
||||
str(args.display_height),
|
||||
"-f",
|
||||
"mjpeg",
|
||||
"-i",
|
||||
"pipe:0",
|
||||
]
|
||||
|
||||
|
||||
def start_ffplay(args):
|
||||
return subprocess.Popen(
|
||||
ffplay_command(args),
|
||||
stdin=subprocess.PIPE,
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.PIPE,
|
||||
bufsize=0,
|
||||
)
|
||||
|
||||
|
||||
def write_display_frame(display_process, jpeg):
|
||||
if display_process is None or display_process.stdin is None:
|
||||
return False
|
||||
if display_process.poll() is not None:
|
||||
return False
|
||||
try:
|
||||
display_process.stdin.write(jpeg)
|
||||
display_process.stdin.flush()
|
||||
return True
|
||||
except (BrokenPipeError, OSError):
|
||||
return False
|
||||
|
||||
|
||||
def list_devices():
|
||||
command = [
|
||||
"ffmpeg",
|
||||
"-hide_banner",
|
||||
"-f",
|
||||
"avfoundation",
|
||||
"-list_devices",
|
||||
"true",
|
||||
"-i",
|
||||
"",
|
||||
]
|
||||
subprocess.run(command, check=False)
|
||||
|
||||
|
||||
def read_jpeg_frames(stream):
|
||||
buffer = bytearray()
|
||||
while True:
|
||||
chunk = stream.read(65536)
|
||||
if not chunk:
|
||||
break
|
||||
buffer.extend(chunk)
|
||||
|
||||
while True:
|
||||
start = buffer.find(b"\xff\xd8")
|
||||
if start < 0:
|
||||
if len(buffer) > 1024 * 1024:
|
||||
del buffer[:-2]
|
||||
break
|
||||
|
||||
end = buffer.find(b"\xff\xd9", start + 2)
|
||||
if end < 0:
|
||||
if start > 0:
|
||||
del buffer[:start]
|
||||
break
|
||||
|
||||
frame = bytes(buffer[start : end + 2])
|
||||
del buffer[: end + 2]
|
||||
yield frame
|
||||
|
||||
|
||||
def run_monitor(args):
|
||||
model_path = Path(args.model)
|
||||
class_name = CLASS_NAMES.get(args.target_class, f"class_{args.target_class}")
|
||||
visualize_dir = Path(args.visualize_dir) if args.visualize_dir else None
|
||||
stream_state = StreamState() if args.serve else None
|
||||
stream_server = None
|
||||
display_process = None
|
||||
if stream_state:
|
||||
stream_server = start_stream_server(args.host, args.port, stream_state)
|
||||
url = f"http://{args.host}:{args.port}/"
|
||||
emit({"type": "stream_server_ready", "url": url})
|
||||
if args.open_visualizer:
|
||||
webbrowser.open(url)
|
||||
if visualize_dir:
|
||||
write_visualizer_html(visualize_dir / "index.html")
|
||||
emit({"type": "visualizer_ready", "path": str(visualize_dir / "index.html")})
|
||||
if args.open_visualizer and not args.serve:
|
||||
webbrowser.open((visualize_dir / "index.html").resolve().as_uri())
|
||||
|
||||
if args.display:
|
||||
display_process = start_ffplay(args)
|
||||
emit({"type": "display_started", "title": args.display_title})
|
||||
|
||||
session = ort.InferenceSession(str(model_path))
|
||||
input_meta = session.get_inputs()[0]
|
||||
emit(
|
||||
{
|
||||
"type": "monitor_started",
|
||||
"timestamp": utc_now(),
|
||||
"model": str(model_path),
|
||||
"providers": session.get_providers(),
|
||||
"input": args.input,
|
||||
"fps": args.fps,
|
||||
"target_class": args.target_class,
|
||||
"target_name": class_name,
|
||||
"confidence": args.confidence,
|
||||
}
|
||||
)
|
||||
|
||||
process = subprocess.Popen(
|
||||
ffmpeg_command(args),
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
bufsize=0,
|
||||
)
|
||||
|
||||
if process.stdout is None:
|
||||
raise RuntimeError("ffmpeg stdout pipe is not available")
|
||||
|
||||
frame_index = 0
|
||||
try:
|
||||
for frame_bytes in read_jpeg_frames(process.stdout):
|
||||
frame_index += 1
|
||||
image = Image.open(BytesIO(frame_bytes)).convert("RGB")
|
||||
detections = detect(
|
||||
session,
|
||||
input_meta,
|
||||
image,
|
||||
args.target_class,
|
||||
args.confidence,
|
||||
)
|
||||
|
||||
if len(detections) > 0:
|
||||
best = detections[np.argmax(detections[:, 4])]
|
||||
event = detection_event(best, frame_index, image.size, class_name)
|
||||
emit(event)
|
||||
elif args.emit_empty:
|
||||
event = no_detection_event(frame_index, image.size, args.target_class, class_name)
|
||||
emit(event)
|
||||
else:
|
||||
event = no_detection_event(frame_index, image.size, args.target_class, class_name)
|
||||
|
||||
if args.debug_dir and args.debug_every > 0 and frame_index % args.debug_every == 0:
|
||||
save_debug_image(
|
||||
image,
|
||||
detections,
|
||||
Path(args.debug_dir) / f"frame_{frame_index:06d}.jpg",
|
||||
class_name,
|
||||
)
|
||||
|
||||
if (
|
||||
visualize_dir
|
||||
and args.visualize_every > 0
|
||||
and frame_index % args.visualize_every == 0
|
||||
):
|
||||
update_visualizer(image, detections, event, visualize_dir, class_name)
|
||||
|
||||
annotated = None
|
||||
if stream_state:
|
||||
annotated = draw_detections(image, detections, class_name)
|
||||
stream_state.update(image_to_jpeg_bytes(annotated), event)
|
||||
|
||||
if display_process:
|
||||
if annotated is None:
|
||||
annotated = draw_detections(image, detections, class_name)
|
||||
annotated = draw_frame_info(annotated, event)
|
||||
if not write_display_frame(display_process, image_to_jpeg_bytes(annotated)):
|
||||
emit({"type": "display_closed", "timestamp": utc_now()})
|
||||
break
|
||||
|
||||
if args.max_frames > 0 and frame_index >= args.max_frames:
|
||||
break
|
||||
finally:
|
||||
process.terminate()
|
||||
try:
|
||||
process.wait(timeout=2)
|
||||
except subprocess.TimeoutExpired:
|
||||
process.kill()
|
||||
|
||||
stderr = b""
|
||||
if process.stderr is not None:
|
||||
try:
|
||||
stderr = process.stderr.read()
|
||||
except Exception:
|
||||
stderr = b""
|
||||
if frame_index == 0 and stderr:
|
||||
log_error(stderr.decode(errors="replace").strip())
|
||||
|
||||
if display_process:
|
||||
if display_process.stdin:
|
||||
try:
|
||||
display_process.stdin.close()
|
||||
except Exception:
|
||||
pass
|
||||
display_process.terminate()
|
||||
try:
|
||||
display_process.wait(timeout=2)
|
||||
except subprocess.TimeoutExpired:
|
||||
display_process.kill()
|
||||
|
||||
emit(
|
||||
{
|
||||
"type": "monitor_stopped",
|
||||
"timestamp": utc_now(),
|
||||
"frames": frame_index,
|
||||
}
|
||||
)
|
||||
|
||||
if stream_server:
|
||||
stream_server.shutdown()
|
||||
|
||||
|
||||
def main():
|
||||
def stop_handler(signum, frame):
|
||||
raise KeyboardInterrupt
|
||||
|
||||
signal.signal(signal.SIGTERM, stop_handler)
|
||||
signal.signal(signal.SIGINT, stop_handler)
|
||||
|
||||
args = parse_args()
|
||||
if args.list_devices:
|
||||
list_devices()
|
||||
return
|
||||
|
||||
if args.debug_every > 0 and not args.debug_dir:
|
||||
args.debug_dir = str(DEBUG_DIR)
|
||||
if args.open_visualizer and not args.visualize_dir and not args.serve:
|
||||
args.visualize_dir = str(VISUALIZE_DIR)
|
||||
|
||||
try:
|
||||
run_monitor(args)
|
||||
except KeyboardInterrupt:
|
||||
emit({"type": "monitor_interrupted", "timestamp": utc_now()})
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
180
wechat_vision/screen_srk_click_send.py
Normal file
180
wechat_vision/screen_srk_click_send.py
Normal file
@ -0,0 +1,180 @@
|
||||
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()
|
||||
36
wechat_vision/start_realtime_stream.sh
Executable file
36
wechat_vision/start_realtime_stream.sh
Executable file
@ -0,0 +1,36 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
cd "$(dirname "$0")"
|
||||
mkdir -p ouptsw
|
||||
|
||||
PID_FILE="ouptsw/realtime_stream.pid"
|
||||
LOG_FILE="ouptsw/realtime_stream.log"
|
||||
|
||||
# Clean up stale FFmpeg capture children from a previously killed monitor.
|
||||
pkill -f "avfoundation -framerate 2.0 -i 3:none" 2>/dev/null || true
|
||||
pkill -f "ffplay .*WeChat Vision ONNX Preview" 2>/dev/null || true
|
||||
|
||||
if [[ -f "$PID_FILE" ]]; then
|
||||
OLD_PID="$(cat "$PID_FILE")"
|
||||
if [[ -n "$OLD_PID" ]] && kill -0 "$OLD_PID" 2>/dev/null; then
|
||||
printf 'Realtime stream is already running: pid=%s\n' "$OLD_PID"
|
||||
printf 'Open http://127.0.0.1:8765/\n'
|
||||
exit 0
|
||||
fi
|
||||
fi
|
||||
|
||||
nohup ./venv/bin/python ffmpeg_realtime_detect.py \
|
||||
--input 3:none \
|
||||
--fps 2 \
|
||||
--confidence 0.05 \
|
||||
--target-class 0 \
|
||||
--display \
|
||||
> "$LOG_FILE" 2>&1 &
|
||||
|
||||
PID="$!"
|
||||
printf '%s' "$PID" > "$PID_FILE"
|
||||
|
||||
printf 'Realtime stream started: pid=%s\n' "$PID"
|
||||
printf 'An ffplay preview window should open automatically.\n'
|
||||
printf 'Log: %s\n' "$LOG_FILE"
|
||||
38
wechat_vision/start_wechat_algorithm_live.sh
Executable file
38
wechat_vision/start_wechat_algorithm_live.sh
Executable file
@ -0,0 +1,38 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
cd "$(dirname "$0")"
|
||||
mkdir -p ouptsw
|
||||
|
||||
PID_FILE="ouptsw/wechat_algorithm_live.pid"
|
||||
LOG_FILE="ouptsw/wechat_algorithm_live.log"
|
||||
|
||||
if [[ -f "$PID_FILE" ]]; then
|
||||
OLD_PID="$(cat "$PID_FILE")"
|
||||
if [[ -n "$OLD_PID" ]] && kill -0 "$OLD_PID" 2>/dev/null; then
|
||||
printf 'WeChat algorithm live monitor is already running: pid=%s\n' "$OLD_PID"
|
||||
printf 'Open http://127.0.0.1:8765/\n'
|
||||
exit 0
|
||||
fi
|
||||
fi
|
||||
|
||||
nohup ./venv/bin/python wechat_algorithm_live.py \
|
||||
--fps 30 \
|
||||
--host 127.0.0.1 \
|
||||
--port 8765 \
|
||||
--click-on-badge \
|
||||
--llm-on-click \
|
||||
--llm-on-chat-change \
|
||||
--reply-on-chat-change \
|
||||
--send-reply \
|
||||
--typing-chunk-size 2 \
|
||||
--typing-delay 0.08 \
|
||||
--open-browser \
|
||||
> "$LOG_FILE" 2>&1 &
|
||||
|
||||
PID="$!"
|
||||
printf '%s' "$PID" > "$PID_FILE"
|
||||
|
||||
printf 'WeChat algorithm live monitor started: pid=%s\n' "$PID"
|
||||
printf 'Open http://127.0.0.1:8765/\n'
|
||||
printf 'Log: %s\n' "$LOG_FILE"
|
||||
35
wechat_vision/start_wechat_window_live.sh
Executable file
35
wechat_vision/start_wechat_window_live.sh
Executable file
@ -0,0 +1,35 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
cd "$(dirname "$0")"
|
||||
mkdir -p ouptsw
|
||||
|
||||
PID_FILE="ouptsw/wechat_window_live.pid"
|
||||
LOG_FILE="ouptsw/wechat_window_live.log"
|
||||
|
||||
if [[ -f "$PID_FILE" ]]; then
|
||||
OLD_PID="$(cat "$PID_FILE")"
|
||||
if [[ -n "$OLD_PID" ]] && kill -0 "$OLD_PID" 2>/dev/null; then
|
||||
printf 'WeChat window live stream is already running: pid=%s\n' "$OLD_PID"
|
||||
printf 'Open http://127.0.0.1:8765/\n'
|
||||
exit 0
|
||||
fi
|
||||
fi
|
||||
|
||||
nohup ./venv/bin/python wechat_window_live.py \
|
||||
--fps 60 \
|
||||
--infer-interval 3 \
|
||||
--confidence 0.05 \
|
||||
--target-class 0 \
|
||||
--host 127.0.0.1 \
|
||||
--port 8765 \
|
||||
--open-browser \
|
||||
> "$LOG_FILE" 2>&1 &
|
||||
|
||||
PID="$!"
|
||||
printf '%s' "$PID" > "$PID_FILE"
|
||||
|
||||
printf 'WeChat window live stream started: pid=%s\n' "$PID"
|
||||
printf 'Open http://127.0.0.1:8765/\n'
|
||||
printf 'Stream http://127.0.0.1:8765/stream.mjpg\n'
|
||||
printf 'Log: %s\n' "$LOG_FILE"
|
||||
24
wechat_vision/stop_realtime_stream.sh
Executable file
24
wechat_vision/stop_realtime_stream.sh
Executable file
@ -0,0 +1,24 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
cd "$(dirname "$0")"
|
||||
|
||||
PID_FILE="ouptsw/realtime_stream.pid"
|
||||
if [[ ! -f "$PID_FILE" ]]; then
|
||||
printf 'Realtime stream is not running: missing %s\n' "$PID_FILE"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
PID="$(cat "$PID_FILE")"
|
||||
if [[ -z "$PID" ]] || ! kill -0 "$PID" 2>/dev/null; then
|
||||
rm -f "$PID_FILE"
|
||||
printf 'Realtime stream is not running. Removed stale pid file.\n'
|
||||
exit 0
|
||||
fi
|
||||
|
||||
kill "$PID"
|
||||
sleep 1
|
||||
pkill -f "avfoundation -framerate 2.0 -i 3:none" 2>/dev/null || true
|
||||
pkill -f "ffplay .*WeChat Vision ONNX Preview" 2>/dev/null || true
|
||||
rm -f "$PID_FILE"
|
||||
printf 'Realtime stream stopped: pid=%s\n' "$PID"
|
||||
22
wechat_vision/stop_wechat_algorithm_live.sh
Executable file
22
wechat_vision/stop_wechat_algorithm_live.sh
Executable file
@ -0,0 +1,22 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
cd "$(dirname "$0")"
|
||||
|
||||
PID_FILE="ouptsw/wechat_algorithm_live.pid"
|
||||
if [[ ! -f "$PID_FILE" ]]; then
|
||||
printf 'WeChat algorithm live monitor is not running: missing %s\n' "$PID_FILE"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
PID="$(cat "$PID_FILE")"
|
||||
if [[ -z "$PID" ]] || ! kill -0 "$PID" 2>/dev/null; then
|
||||
rm -f "$PID_FILE"
|
||||
printf 'WeChat algorithm live monitor is not running. Removed stale pid file.\n'
|
||||
exit 0
|
||||
fi
|
||||
|
||||
kill "$PID"
|
||||
sleep 1
|
||||
rm -f "$PID_FILE"
|
||||
printf 'WeChat algorithm live monitor stopped: pid=%s\n' "$PID"
|
||||
22
wechat_vision/stop_wechat_window_live.sh
Executable file
22
wechat_vision/stop_wechat_window_live.sh
Executable file
@ -0,0 +1,22 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
cd "$(dirname "$0")"
|
||||
|
||||
PID_FILE="ouptsw/wechat_window_live.pid"
|
||||
if [[ ! -f "$PID_FILE" ]]; then
|
||||
printf 'WeChat window live stream is not running: missing %s\n' "$PID_FILE"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
PID="$(cat "$PID_FILE")"
|
||||
if [[ -z "$PID" ]] || ! kill -0 "$PID" 2>/dev/null; then
|
||||
rm -f "$PID_FILE"
|
||||
printf 'WeChat window live stream is not running. Removed stale pid file.\n'
|
||||
exit 0
|
||||
fi
|
||||
|
||||
kill "$PID"
|
||||
sleep 1
|
||||
rm -f "$PID_FILE"
|
||||
printf 'WeChat window live stream stopped: pid=%s\n' "$PID"
|
||||
773
wechat_vision/wechat_algorithm_live.py
Normal file
773
wechat_vision/wechat_algorithm_live.py
Normal file
@ -0,0 +1,773 @@
|
||||
import argparse
|
||||
import base64
|
||||
import json
|
||||
import signal
|
||||
import threading
|
||||
import time
|
||||
import webbrowser
|
||||
from datetime import datetime, timezone
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
from urllib import request as urlrequest
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image, ImageDraw
|
||||
import pyautogui
|
||||
import pyperclip
|
||||
import tomllib
|
||||
|
||||
from wechat_window_live import capture_window_image, find_wechat_window, image_to_jpeg_bytes
|
||||
|
||||
|
||||
DEFAULT_REGIONS = Path.home() / "Library/Application Support/com.tauri.dev/data/regions/wechat.json"
|
||||
DEFAULT_AGENT_CONFIG = Path(__file__).resolve().parent.parent / "agent" / "config.toml"
|
||||
|
||||
|
||||
class StreamState:
|
||||
def __init__(self):
|
||||
self.condition = threading.Condition()
|
||||
self.jpeg = None
|
||||
self.event = None
|
||||
self.sequence = 0
|
||||
|
||||
def update(self, jpeg, event):
|
||||
with self.condition:
|
||||
self.jpeg = jpeg
|
||||
self.event = event
|
||||
self.sequence += 1
|
||||
self.condition.notify_all()
|
||||
|
||||
def snapshot(self):
|
||||
with self.condition:
|
||||
return self.jpeg, self.event, self.sequence
|
||||
|
||||
|
||||
class SharedFrame:
|
||||
def __init__(self):
|
||||
self.lock = threading.Lock()
|
||||
self.image = None
|
||||
self.window = None
|
||||
self.boxes = {}
|
||||
self.frame_index = 0
|
||||
self.latest_detection = {"events": [], "badge_events": [], "change_events": [], "timestamp": utc_now()}
|
||||
self.latest_llm = None
|
||||
self.stop = threading.Event()
|
||||
|
||||
def update_frame(self, image, window, boxes, frame_index):
|
||||
with self.lock:
|
||||
self.image = image
|
||||
self.window = window
|
||||
self.boxes = boxes
|
||||
self.frame_index = frame_index
|
||||
|
||||
def snapshot(self):
|
||||
with self.lock:
|
||||
return self.image, self.window, dict(self.boxes), self.frame_index, dict(self.latest_detection), self.latest_llm
|
||||
|
||||
def update_detection(self, detection):
|
||||
with self.lock:
|
||||
self.latest_detection = detection
|
||||
|
||||
def update_llm(self, reading):
|
||||
with self.lock:
|
||||
self.latest_llm = reading
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Pure algorithm WeChat new-message monitor.")
|
||||
parser.add_argument("--regions", default=str(DEFAULT_REGIONS), help="wechat.json annotation path")
|
||||
parser.add_argument("--fps", type=float, default=30.0, help="preview FPS target")
|
||||
parser.add_argument("--host", default="127.0.0.1", help="HTTP host")
|
||||
parser.add_argument("--port", type=int, default=8765, help="HTTP port")
|
||||
parser.add_argument("--max-frames", type=int, default=0, help="stop after N frames; 0 runs forever")
|
||||
parser.add_argument("--open-browser", action="store_true", help="open browser preview")
|
||||
parser.add_argument("--chat-diff-threshold", type=int, default=34, help="chat content pixel diff threshold")
|
||||
parser.add_argument("--dry-run", action="store_true", default=True, help="draw only; never click")
|
||||
parser.add_argument("--click-on-badge", action="store_true", help="click the first detected unread contact row")
|
||||
parser.add_argument("--llm-on-click", action="store_true", help="after clicking a contact, ask LLM to read current chat content")
|
||||
parser.add_argument("--llm-on-chat-change", action="store_true", help="ask LLM to read current chat when chat_content diff indicates a new message")
|
||||
parser.add_argument("--reply-on-chat-change", action="store_true", help="send reply_text for current-chat diff events")
|
||||
parser.add_argument("--send-reply", action="store_true", help="paste LLM reply_text into input_box and press Enter")
|
||||
parser.add_argument("--typing-chunk-size", type=int, default=2, help="characters pasted per simulated typing chunk")
|
||||
parser.add_argument("--typing-delay", type=float, default=0.08, help="seconds between simulated typing chunks")
|
||||
parser.add_argument("--agent-config", default=str(DEFAULT_AGENT_CONFIG), help="agent/config.toml path for LLM settings")
|
||||
parser.add_argument("--click-cooldown", type=float, default=30.0, help="seconds before clicking the same contact row again")
|
||||
parser.add_argument("--chat-change-cooldown", type=float, default=30.0, help="seconds before handling current chat diff again")
|
||||
parser.add_argument("--after-click-wait", type=float, default=1.0, help="seconds to wait before capturing chat after click")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def utc_now():
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
|
||||
def emit(event):
|
||||
print(json.dumps(json_safe(event), ensure_ascii=False, separators=(",", ":")), flush=True)
|
||||
|
||||
|
||||
def emit_llm_reading(reading):
|
||||
payload = {
|
||||
"type": "llm_chat_reading_print",
|
||||
"timestamp": utc_now(),
|
||||
"contact_name": reading.get("contact_name", "unknown") if isinstance(reading, dict) else "unknown",
|
||||
"latest_user_message": reading.get("latest_user_message", "") if isinstance(reading, dict) else "",
|
||||
"visible_messages": reading.get("visible_messages", []) if isinstance(reading, dict) else [],
|
||||
"summary": reading.get("summary", "") if isinstance(reading, dict) else "",
|
||||
"reply_text": reading.get("reply_text", "") if isinstance(reading, dict) else "",
|
||||
"send_status": reading.get("send_status", "not_sent") if isinstance(reading, dict) else "not_sent",
|
||||
"send_error": reading.get("send_error", "") if isinstance(reading, dict) else "",
|
||||
}
|
||||
emit(payload)
|
||||
|
||||
|
||||
def json_safe(value):
|
||||
if isinstance(value, dict):
|
||||
return {key: json_safe(item) for key, item in value.items()}
|
||||
if isinstance(value, list):
|
||||
return [json_safe(item) for item in value]
|
||||
if isinstance(value, tuple):
|
||||
return [json_safe(item) for item in value]
|
||||
if isinstance(value, np.integer):
|
||||
return int(value)
|
||||
if isinstance(value, np.floating):
|
||||
return float(value)
|
||||
return value
|
||||
|
||||
|
||||
def load_regions(path):
|
||||
data = json.loads(Path(path).read_text(encoding="utf-8"))
|
||||
regions = {}
|
||||
for region in data.get("regions", []):
|
||||
regions.setdefault(region.get("type"), region)
|
||||
if region.get("type") == "custom" and "未读消息" in region.get("description", ""):
|
||||
regions["message_button"] = region
|
||||
return data, regions
|
||||
|
||||
|
||||
def load_agent_config(path):
|
||||
with Path(path).open("rb") as file:
|
||||
return tomllib.load(file)
|
||||
|
||||
|
||||
def image_data_url(image):
|
||||
buffer = BytesIO()
|
||||
image.save(buffer, format="JPEG", quality=88)
|
||||
encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
|
||||
return "data:image/jpeg;base64," + encoded
|
||||
|
||||
|
||||
def request_chat_reading(config, chat_image):
|
||||
volcengine = config.get("volcengine", {})
|
||||
base_url = (volcengine.get("base_url") or "").rstrip("/")
|
||||
api_key = volcengine.get("api_key") or ""
|
||||
model = volcengine.get("model") or ""
|
||||
if not base_url or not api_key or not model:
|
||||
raise RuntimeError("missing volcengine base_url/api_key/model in agent config")
|
||||
|
||||
prompt = """你是微信聊天截图读取和回复草稿助手。
|
||||
请读取图片中当前聊天区域的可见消息，不要编造看不见的内容。
|
||||
请基于可见上下文生成一条自然、简短、适合直接发送的中文回复草稿。
|
||||
这里只生成草稿，不要假设已经发送，也不要执行任何操作。
|
||||
返回合法 JSON，格式如下：
|
||||
{
|
||||
"contact_name": "unknown",
|
||||
"latest_user_message": "",
|
||||
"visible_messages": [
|
||||
{"sender":"unknown", "role":"user|me|system|unknown", "content":""}
|
||||
],
|
||||
"summary": "",
|
||||
"reply_text": ""
|
||||
}
|
||||
"""
|
||||
payload = {
|
||||
"model": model,
|
||||
"temperature": 0.2,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": prompt},
|
||||
{"type": "image_url", "image_url": {"url": image_data_url(chat_image)}},
|
||||
],
|
||||
}
|
||||
],
|
||||
}
|
||||
data = json.dumps(payload).encode("utf-8")
|
||||
req = urlrequest.Request(
|
||||
base_url + "/chat/completions",
|
||||
data=data,
|
||||
headers={"Authorization": "Bearer " + api_key, "Content-Type": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
with urlrequest.urlopen(req, timeout=90) as response:
|
||||
body = json.loads(response.read().decode("utf-8"))
|
||||
choices = body.get("choices") or []
|
||||
if not choices:
|
||||
raise RuntimeError("LLM returned no choices")
|
||||
content = choices[0].get("message", {}).get("content", "")
|
||||
return extract_json_object(content)
|
||||
|
||||
|
||||
def extract_json_object(text):
|
||||
stripped = text.strip()
|
||||
if stripped.startswith("```json"):
|
||||
stripped = stripped[len("```json"):].strip()
|
||||
if stripped.startswith("```"):
|
||||
stripped = stripped[3:].strip()
|
||||
if stripped.endswith("```"):
|
||||
stripped = stripped[:-3].strip()
|
||||
start = stripped.find("{")
|
||||
end = stripped.rfind("}")
|
||||
if start >= 0 and end > start:
|
||||
stripped = stripped[start:end + 1]
|
||||
try:
|
||||
return json.loads(stripped)
|
||||
except json.JSONDecodeError:
|
||||
return {"raw_text": text[:1200]}
|
||||
|
||||
|
||||
def scaled_box(region, annotation, image_size, window):
|
||||
# Annotation bbox_image is stored in physical pixels at annotation-time
|
||||
# scaleFactor. Convert back to logical window coords, then map to the
|
||||
# current capture's physical pixels. This stays accurate when the WeChat
|
||||
# window is moved or resized horizontally.
|
||||
annotation_scale = region.get("scaleFactor") or annotation.get("scaleFactor") or 1
|
||||
current_scale_x = image_size[0] / max(window["width"], 1)
|
||||
current_scale_y = image_size[1] / max(window["height"], 1)
|
||||
x1, y1, x2, y2 = region["bbox_image"]
|
||||
return [
|
||||
int((x1 / annotation_scale) * current_scale_x),
|
||||
int((y1 / annotation_scale) * current_scale_y),
|
||||
int((x2 / annotation_scale) * current_scale_x),
|
||||
int((y2 / annotation_scale) * current_scale_y),
|
||||
]
|
||||
|
||||
|
||||
def clamp_box(box, image_size):
|
||||
width, height = image_size
|
||||
x1, y1, x2, y2 = box
|
||||
x1, x2 = sorted((x1, x2))
|
||||
y1, y2 = sorted((y1, y2))
|
||||
return [max(0, x1), max(0, y1), min(width, x2), min(height, y2)]
|
||||
|
||||
|
||||
def valid_box(box):
|
||||
return box and len(box) == 4 and box[2] > box[0] and box[3] > box[1]
|
||||
|
||||
|
||||
def components(mask, min_area=12):
|
||||
h, w = mask.shape
|
||||
visited = np.zeros_like(mask, dtype=bool)
|
||||
found = []
|
||||
ys, xs = np.where(mask)
|
||||
for start_x, start_y in zip(xs, ys):
|
||||
if visited[start_y, start_x]:
|
||||
continue
|
||||
stack = [(start_x, start_y)]
|
||||
visited[start_y, start_x] = True
|
||||
area = 0
|
||||
min_x = max_x = start_x
|
||||
min_y = max_y = start_y
|
||||
while stack:
|
||||
x, y = stack.pop()
|
||||
area += 1
|
||||
min_x, max_x = min(min_x, x), max(max_x, x)
|
||||
min_y, max_y = min(min_y, y), max(max_y, y)
|
||||
for nx, ny in ((x - 1, y), (x + 1, y), (x, y - 1), (x, y + 1)):
|
||||
if 0 <= nx < w and 0 <= ny < h and mask[ny, nx] and not visited[ny, nx]:
|
||||
visited[ny, nx] = True
|
||||
stack.append((nx, ny))
|
||||
if area >= min_area:
|
||||
found.append({"bbox": [min_x, min_y, max_x + 1, max_y + 1], "area": area})
|
||||
return found
|
||||
|
||||
|
||||
def detect_red_badges(image, contact_box):
|
||||
x1, y1, x2, y2 = contact_box
|
||||
crop = np.asarray(image.crop((x1, y1, x2, y2)).convert("RGB"))
|
||||
r, g, b = crop[:, :, 0], crop[:, :, 1], crop[:, :, 2]
|
||||
mask = (r > 175) & (g < 105) & (b < 105) & ((r.astype(int) - g.astype(int)) > 65) & ((r.astype(int) - b.astype(int)) > 65)
|
||||
digit_mask = (r > 215) & (g > 215) & (b > 215) & (np.maximum.reduce([r, g, b]) - np.minimum.reduce([r, g, b]) < 42)
|
||||
results = []
|
||||
for comp in components(mask, min_area=18):
|
||||
bx1, by1, bx2, by2 = comp["bbox"]
|
||||
bw, bh = bx2 - bx1, by2 - by1
|
||||
if not (6 <= bw <= 52 and 6 <= bh <= 52):
|
||||
continue
|
||||
if comp["area"] / max(bw * bh, 1) < 0.22:
|
||||
continue
|
||||
pad_x = max(2, bw // 5)
|
||||
pad_y = max(2, bh // 5)
|
||||
inner_x1 = min(bx2, bx1 + pad_x)
|
||||
inner_y1 = min(by2, by1 + pad_y)
|
||||
inner_x2 = max(inner_x1, bx2 - pad_x)
|
||||
inner_y2 = max(inner_y1, by2 - pad_y)
|
||||
digit_pixels = int(digit_mask[inner_y1:inner_y2, inner_x1:inner_x2].sum())
|
||||
# Require white digit strokes inside the red badge. This intentionally
|
||||
# ignores plain red dots and unrelated red UI fragments.
|
||||
if digit_pixels < 3:
|
||||
continue
|
||||
gx1, gy1, gx2, gy2 = x1 + bx1, y1 + by1, x1 + bx2, y1 + by2
|
||||
center_y = (gy1 + gy2) // 2
|
||||
# WeChat conversation rows are roughly 64-76 logical px tall. The
|
||||
# captured image is usually Retina scale 2, so 132 image px is a good
|
||||
# first approximation and less noisy than deriving height from badge.
|
||||
row_h = 132
|
||||
row = [x1, max(y1, center_y - row_h // 2), x2, min(y2, center_y + row_h // 2)]
|
||||
click = [(row[0] + row[2]) // 2, (row[1] + row[2] * 0 + row[3]) // 2]
|
||||
results.append({"type": "new_message_badge", "badge_box": [gx1, gy1, gx2, gy2], "contact_row": row, "click_image": click, "score": round(float(comp["area"] / max(bw * bh, 1)), 3), "digit_pixels": digit_pixels})
|
||||
return results
|
||||
|
||||
|
||||
def detect_chat_change(image, previous_gray, chat_box, threshold):
|
||||
x1, y1, x2, y2 = chat_box
|
||||
gray = np.asarray(image.crop((x1, y1, x2, y2)).convert("L"), dtype=np.int16)
|
||||
if previous_gray is None or previous_gray.shape != gray.shape:
|
||||
return gray, []
|
||||
diff = np.abs(gray - previous_gray)
|
||||
mask = diff > threshold
|
||||
crop_h, crop_w = gray.shape
|
||||
changed_ratio = float(mask.sum()) / max(crop_w * crop_h, 1)
|
||||
if changed_ratio > 0.28:
|
||||
# Window moves, scrolls, or large reflows are not new-message events.
|
||||
return gray, []
|
||||
results = []
|
||||
for comp in components(mask, min_area=120):
|
||||
bx1, by1, bx2, by2 = comp["bbox"]
|
||||
bw, bh = bx2 - bx1, by2 - by1
|
||||
if bw < 24 or bh < 10:
|
||||
continue
|
||||
if by1 < crop_h * 0.38:
|
||||
continue
|
||||
if comp["area"] > crop_w * crop_h * 0.18:
|
||||
continue
|
||||
# Focus on localized lower-area changes where new messages appear.
|
||||
gx1, gy1, gx2, gy2 = x1 + bx1, y1 + by1, x1 + bx2, y1 + by2
|
||||
results.append({"type": "current_chat_new_message", "change_box": [gx1, gy1, gx2, gy2], "area": comp["area"], "changed_ratio": round(changed_ratio, 4)})
|
||||
return gray, results[:5]
|
||||
|
||||
|
||||
def image_to_screen_box(box, window, image_size):
|
||||
scale_x = image_size[0] / max(window["width"], 1)
|
||||
scale_y = image_size[1] / max(window["height"], 1)
|
||||
x1, y1, x2, y2 = box
|
||||
return [round(window["x"] + x1 / scale_x, 2), round(window["y"] + y1 / scale_y, 2), round(window["x"] + x2 / scale_x, 2), round(window["y"] + y2 / scale_y, 2)]
|
||||
|
||||
|
||||
def point_to_screen(point, window, image_size):
|
||||
scale_x = image_size[0] / max(window["width"], 1)
|
||||
scale_y = image_size[1] / max(window["height"], 1)
|
||||
return [round(window["x"] + point[0] / scale_x, 2), round(window["y"] + point[1] / scale_y, 2)]
|
||||
|
||||
|
||||
def box_center_screen(box, window, image_size):
|
||||
cx = (box[0] + box[2]) / 2
|
||||
cy = (box[1] + box[3]) / 2
|
||||
return point_to_screen([cx, cy], window, image_size)
|
||||
|
||||
|
||||
def text_chunks(text, chunk_size):
|
||||
runes = list(text)
|
||||
size = max(1, chunk_size)
|
||||
for index in range(0, len(runes), size):
|
||||
yield "".join(runes[index:index + size])
|
||||
|
||||
|
||||
def send_reply_text(reply_text, input_box, window, image_size, chunk_size=2, typing_delay=0.08):
|
||||
text = (reply_text or "").strip()
|
||||
if not text:
|
||||
return {"send_status": "skipped", "send_error": "empty reply_text"}
|
||||
if not valid_box(input_box):
|
||||
return {"send_status": "failed", "send_error": "invalid input_box"}
|
||||
point = box_center_screen(input_box, window, image_size)
|
||||
pyautogui.click(point[0], point[1])
|
||||
time.sleep(0.15)
|
||||
chunk_count = 0
|
||||
for chunk in text_chunks(text, chunk_size):
|
||||
pyperclip.copy(chunk)
|
||||
pyautogui.hotkey("command", "v")
|
||||
chunk_count += 1
|
||||
time.sleep(max(0.01, typing_delay))
|
||||
pyautogui.press("enter")
|
||||
return {"send_status": "sent", "send_error": "", "input_click_screen": point, "typing_mode": "chunked_paste", "typing_chunk_size": max(1, chunk_size), "typing_chunks": chunk_count}
|
||||
|
||||
|
||||
def decorate_events(events, window, image_size):
|
||||
decorated = []
|
||||
for event in events:
|
||||
item = dict(event)
|
||||
for key in ("badge_box", "contact_row", "change_box"):
|
||||
if key in item:
|
||||
item[key + "_screen"] = image_to_screen_box(item[key], window, image_size)
|
||||
if "click_image" in item:
|
||||
item["click_screen"] = point_to_screen(item["click_image"], window, image_size)
|
||||
item["dry_run"] = True
|
||||
decorated.append(item)
|
||||
return decorated
|
||||
|
||||
|
||||
def draw_events(image, events, boxes):
|
||||
draw = ImageDraw.Draw(image)
|
||||
for name, box in boxes.items():
|
||||
if not valid_box(box):
|
||||
continue
|
||||
color = {"contact_list": "#666666", "chat_content": "#5555ff", "input_box": "#22c55e", "message_button": "#888888"}.get(name, "#666666")
|
||||
draw.rectangle(box, outline=color, width=2)
|
||||
draw.text((box[0] + 4, box[1] + 4), name, fill=color)
|
||||
for event in events:
|
||||
if event["type"] == "new_message_badge":
|
||||
if valid_box(event.get("badge_box")):
|
||||
draw.rectangle(event["badge_box"], outline="red", width=4)
|
||||
if valid_box(event.get("contact_row")):
|
||||
draw.rectangle(event["contact_row"], outline="yellow", width=4)
|
||||
draw.text((event["contact_row"][0] + 8, event["contact_row"][1] + 8), "new message", fill="yellow")
|
||||
elif event["type"] in {"chat_content_changed", "current_chat_new_message"}:
|
||||
if valid_box(event.get("change_box")):
|
||||
draw.rectangle(event["change_box"], outline="#00aaff", width=3)
|
||||
draw.text((event["change_box"][0] + 8, event["change_box"][1] + 8), "chat diff", fill="#00aaff")
|
||||
return image
|
||||
|
||||
|
||||
def draw_overlay(image, event):
|
||||
draw = ImageDraw.Draw(image)
|
||||
text = f"frame={event.get('frame_index')} badges={event.get('badge_count')} changes={event.get('change_count')} click={event.get('click_enabled')} llm={event.get('llm_enabled')}"
|
||||
draw.rectangle((12, 12, 980, 56), fill=(0, 0, 0))
|
||||
draw.text((22, 24), text, fill=(80, 255, 120))
|
||||
return image
|
||||
|
||||
|
||||
def html_page():
|
||||
return """<!doctype html>
|
||||
<html lang="zh-CN">
|
||||
<head>
|
||||
<meta charset="utf-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1">
|
||||
<title>WeChat Monitor</title>
|
||||
<style>
|
||||
body{margin:0;background:#09090b;color:#f4f4f5;font-family:-apple-system,BlinkMacSystemFont,'Segoe UI',sans-serif}
|
||||
.wrap{display:grid;grid-template-columns:minmax(0,1fr)460px;gap:16px;padding:16px}
|
||||
.stage{background:#18181b;border:1px solid #27272a;border-radius:14px;overflow:hidden;align-self:start}
|
||||
img{display:block;width:100%;height:auto}
|
||||
aside{background:#18181b;border:1px solid #27272a;border-radius:14px;padding:16px;max-height:calc(100vh - 32px);overflow:auto}
|
||||
h1{margin:0 0 12px;font-size:18px}.item{padding:12px 0;border-top:1px solid #27272a}.label{color:#a1a1aa;font-size:12px;margin-bottom:6px}
|
||||
.value{font-family:ui-monospace,SFMono-Regular,Menlo,monospace;white-space:pre-wrap;word-break:break-all;font-size:12px}
|
||||
.msg{padding:8px 10px;border-radius:10px;margin:8px 0;background:#27272a}.msg.user{background:#1e3a8a}.msg.me{background:#14532d}.role{font-size:11px;color:#d4d4d8;margin-bottom:4px}.content{font-size:14px;line-height:1.45}
|
||||
.reply{padding:12px;border-radius:12px;background:#052e16;color:#bbf7d0;font-size:15px;line-height:1.5}.empty{color:#71717a}.ok{color:#86efac}.warn{color:#facc15}
|
||||
@media(max-width:1000px){.wrap{grid-template-columns:1fr}aside{max-height:none}}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="wrap">
|
||||
<main class="stage"><img src="/stream.mjpg" alt="live stream"></main>
|
||||
<aside>
|
||||
<h1>WeChat Monitor</h1>
|
||||
<div class="item"><div class="label">Sequence / Status</div><div id="status" class="value">-</div></div>
|
||||
<div class="item"><div class="label">监测事件</div><div id="eventSummary" class="value">-</div></div>
|
||||
<div class="item"><div class="label">最新用户消息</div><div id="latest" class="value empty">等待 LLM 读取...</div></div>
|
||||
<div class="item"><div class="label">识别到的聊天记录</div><div id="messages"></div></div>
|
||||
<div class="item"><div class="label">建议回复（未发送）</div><div id="reply" class="reply empty">等待生成...</div></div>
|
||||
<div class="item"><div class="label">摘要</div><div id="summary" class="value empty">-</div></div>
|
||||
</aside>
|
||||
</div>
|
||||
<script>
|
||||
function escapeHtml(text){return String(text ?? '').replace(/[&<>'"]/g,c=>({'&':'&','<':'<','>':'>',"'":''','"':'"'}[c]))}
|
||||
function messageHtml(message){
|
||||
const role = message.role || 'unknown';
|
||||
const cls = role === 'user' ? 'user' : (role === 'me' ? 'me' : '');
|
||||
const sender = message.sender || role;
|
||||
return `<div class="msg ${cls}"><div class="role">${escapeHtml(sender)} · ${escapeHtml(role)}</div><div class="content">${escapeHtml(message.content || '')}</div></div>`;
|
||||
}
|
||||
async function refresh(){
|
||||
try{
|
||||
const res = await fetch('/latest.json?t=' + Date.now(), {cache:'no-store'});
|
||||
const data = await res.json();
|
||||
const event = data.event || {};
|
||||
const llm = event.latest_llm || event.llm_chat_reading || null;
|
||||
document.getElementById('status').textContent = `seq=${data.sequence ?? '-'} frame=${event.frame_index ?? '-'} badges=${event.badge_count ?? 0} changes=${event.change_count ?? 0}`;
|
||||
document.getElementById('eventSummary').textContent = JSON.stringify({type:event.type, click_enabled:event.click_enabled, llm_enabled:event.llm_enabled, events:event.events || []}, null, 2);
|
||||
if(llm){
|
||||
document.getElementById('latest').textContent = llm.latest_user_message || '';
|
||||
document.getElementById('latest').className = 'value';
|
||||
const messages = Array.isArray(llm.visible_messages) ? llm.visible_messages : [];
|
||||
document.getElementById('messages').innerHTML = messages.length ? messages.map(messageHtml).join('') : '<div class="empty">未识别到可见消息</div>';
|
||||
const reply = llm.reply_text || '';
|
||||
document.getElementById('reply').textContent = reply || '未生成建议回复';
|
||||
document.getElementById('reply').className = 'reply' + (reply ? '' : ' empty');
|
||||
document.getElementById('summary').textContent = llm.summary || '-';
|
||||
document.getElementById('summary').className = 'value';
|
||||
}
|
||||
}catch(error){
|
||||
document.getElementById('status').textContent = '连接中断或服务未启动';
|
||||
}
|
||||
}
|
||||
refresh();
|
||||
setInterval(refresh, 500);
|
||||
</script>
|
||||
</body>
|
||||
</html>"""
|
||||
|
||||
|
||||
def make_handler(state):
|
||||
class Handler(BaseHTTPRequestHandler):
|
||||
def log_message(self, format, *args):
|
||||
return
|
||||
|
||||
def do_GET(self):
|
||||
if self.path == "/" or self.path.startswith("/index.html"):
|
||||
data = html_page().encode("utf-8")
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "text/html; charset=utf-8")
|
||||
self.send_header("Content-Length", str(len(data)))
|
||||
self.end_headers()
|
||||
self.wfile.write(data)
|
||||
return
|
||||
if self.path.startswith("/latest.json"):
|
||||
_, event, sequence = state.snapshot()
|
||||
data = json.dumps(json_safe({"sequence": sequence, "updated_at": utc_now(), "event": event or {"type": "waiting"}}), ensure_ascii=False).encode("utf-8")
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "application/json; charset=utf-8")
|
||||
self.send_header("Cache-Control", "no-store")
|
||||
self.send_header("Content-Length", str(len(data)))
|
||||
self.end_headers()
|
||||
self.wfile.write(data)
|
||||
return
|
||||
if self.path.startswith("/stream.mjpg"):
|
||||
self.send_response(200)
|
||||
self.send_header("Cache-Control", "no-cache, private")
|
||||
self.send_header("Pragma", "no-cache")
|
||||
self.send_header("Content-Type", "multipart/x-mixed-replace; boundary=frame")
|
||||
self.end_headers()
|
||||
last_sequence = -1
|
||||
while True:
|
||||
with state.condition:
|
||||
state.condition.wait_for(lambda: state.sequence != last_sequence, timeout=5)
|
||||
jpeg, last_sequence = state.jpeg, state.sequence
|
||||
if not jpeg:
|
||||
continue
|
||||
try:
|
||||
self.wfile.write(b"--frame\r\nContent-Type: image/jpeg\r\n")
|
||||
self.wfile.write(f"Content-Length: {len(jpeg)}\r\n\r\n".encode("ascii"))
|
||||
self.wfile.write(jpeg)
|
||||
self.wfile.write(b"\r\n")
|
||||
except (BrokenPipeError, ConnectionResetError):
|
||||
break
|
||||
return
|
||||
self.send_error(404)
|
||||
return Handler
|
||||
|
||||
|
||||
def start_server(host, port, state):
|
||||
server = ThreadingHTTPServer((host, port), make_handler(state))
|
||||
threading.Thread(target=server.serve_forever, daemon=True).start()
|
||||
return server
|
||||
|
||||
|
||||
def boxes_for_frame(annotation, regions, image, window):
|
||||
boxes = {name: clamp_box(scaled_box(regions[name], annotation, image.size, window), image.size) for name in ("contact_list", "chat_content")}
|
||||
if "input_box" in regions:
|
||||
boxes["input_box"] = clamp_box(scaled_box(regions["input_box"], annotation, image.size, window), image.size)
|
||||
if "message_button" in regions:
|
||||
boxes["message_button"] = clamp_box(scaled_box(regions["message_button"], annotation, image.size, window), image.size)
|
||||
return boxes
|
||||
|
||||
|
||||
def preview_loop(args, annotation, regions, stream_state, shared):
|
||||
interval = 1 / max(args.fps, 1)
|
||||
frame_index = 0
|
||||
while not shared.stop.is_set() and (args.max_frames <= 0 or frame_index < args.max_frames):
|
||||
start = time.time()
|
||||
window = find_wechat_window()
|
||||
if not window:
|
||||
event = {"type": "wechat_window_not_found", "timestamp": utc_now()}
|
||||
stream_state.update(image_to_jpeg_bytes(Image.new("RGB", (960, 540), "black")), event)
|
||||
time.sleep(interval)
|
||||
continue
|
||||
|
||||
image = capture_window_image(window["id"])
|
||||
if image is None:
|
||||
time.sleep(interval)
|
||||
continue
|
||||
|
||||
frame_index += 1
|
||||
boxes = boxes_for_frame(annotation, regions, image, window)
|
||||
shared.update_frame(image, window, boxes, frame_index)
|
||||
_, _, _, _, detection, latest_llm = shared.snapshot()
|
||||
badge_events = detection.get("badge_events", [])
|
||||
change_events = detection.get("change_events", [])
|
||||
event = {
|
||||
"type": "algorithm_monitor",
|
||||
"timestamp": utc_now(),
|
||||
"frame_index": frame_index,
|
||||
"badge_count": len(badge_events),
|
||||
"change_count": len(change_events),
|
||||
"events": detection.get("events", []),
|
||||
"window": window,
|
||||
"dry_run": not args.click_on_badge,
|
||||
"click_enabled": args.click_on_badge,
|
||||
"llm_enabled": args.llm_on_click,
|
||||
"preview_only": True,
|
||||
}
|
||||
if latest_llm:
|
||||
event["latest_llm"] = latest_llm
|
||||
annotated = draw_events(image.copy(), badge_events + change_events, boxes)
|
||||
annotated = draw_overlay(annotated, event)
|
||||
stream_state.update(image_to_jpeg_bytes(annotated), event)
|
||||
time.sleep(max(0, interval - (time.time() - start)))
|
||||
shared.stop.set()
|
||||
|
||||
|
||||
def llm_worker(args, annotation, regions, agent_config, shared, job_event):
|
||||
last_click_at = {}
|
||||
last_chat_change_at = 0.0
|
||||
while not shared.stop.is_set():
|
||||
if not job_event.wait(timeout=0.5):
|
||||
continue
|
||||
job_event.clear()
|
||||
image, window, boxes, frame_index, detection, _ = shared.snapshot()
|
||||
if image is None or window is None:
|
||||
continue
|
||||
target = next((item for item in detection.get("events", []) if item.get("type") == "new_message_badge"), None)
|
||||
chat_change = next((item for item in detection.get("events", []) if item.get("type") == "current_chat_new_message"), None)
|
||||
trigger = "unread_badge" if target else ("current_chat_change" if chat_change else "")
|
||||
if not trigger:
|
||||
continue
|
||||
|
||||
if trigger == "unread_badge":
|
||||
click_key = str(round(target["click_screen"][1] / 10) * 10)
|
||||
if time.time() - last_click_at.get(click_key, 0) < args.click_cooldown:
|
||||
continue
|
||||
if args.click_on_badge:
|
||||
pyautogui.click(target["click_screen"][0], target["click_screen"][1])
|
||||
last_click_at[click_key] = time.time()
|
||||
else:
|
||||
continue
|
||||
if not args.llm_on_click:
|
||||
continue
|
||||
elif trigger == "current_chat_change":
|
||||
if not args.llm_on_chat_change:
|
||||
continue
|
||||
if time.time() - last_chat_change_at < args.chat_change_cooldown:
|
||||
continue
|
||||
last_chat_change_at = time.time()
|
||||
|
||||
time.sleep(args.after_click_wait)
|
||||
refreshed_window = find_wechat_window() or window
|
||||
refreshed_image = capture_window_image(refreshed_window["id"])
|
||||
if refreshed_image is None:
|
||||
continue
|
||||
refreshed_chat_box = clamp_box(scaled_box(regions["chat_content"], annotation, refreshed_image.size, refreshed_window), refreshed_image.size)
|
||||
if not valid_box(refreshed_chat_box):
|
||||
continue
|
||||
x1, y1, x2, y2 = refreshed_chat_box
|
||||
chat_crop = refreshed_image.crop((x1, y1, x2, y2))
|
||||
try:
|
||||
reading = request_chat_reading(agent_config, chat_crop)
|
||||
reading["send_status"] = "not_sent"
|
||||
reading["trigger"] = trigger
|
||||
should_send = args.send_reply and (trigger == "unread_badge" or args.reply_on_chat_change)
|
||||
if should_send:
|
||||
input_region = regions.get("input_box")
|
||||
if input_region:
|
||||
input_box = clamp_box(scaled_box(input_region, annotation, refreshed_image.size, refreshed_window), refreshed_image.size)
|
||||
send_result = send_reply_text(
|
||||
reading.get("reply_text", ""),
|
||||
input_box,
|
||||
refreshed_window,
|
||||
refreshed_image.size,
|
||||
args.typing_chunk_size,
|
||||
args.typing_delay,
|
||||
)
|
||||
reading.update(send_result)
|
||||
else:
|
||||
reading["send_status"] = "failed"
|
||||
reading["send_error"] = "missing input_box region"
|
||||
shared.update_llm(reading)
|
||||
emit_llm_reading(reading)
|
||||
except Exception as error:
|
||||
emit({"type": "llm_error", "timestamp": utc_now(), "error": str(error)})
|
||||
|
||||
|
||||
def detection_loop(args, shared, job_event):
|
||||
previous_chat_gray = None
|
||||
interval = 1.0
|
||||
while not shared.stop.is_set():
|
||||
start = time.time()
|
||||
image, window, boxes, frame_index, _, _ = shared.snapshot()
|
||||
if image is not None and window is not None and "contact_list" in boxes and "chat_content" in boxes:
|
||||
badge_events = detect_red_badges(image, boxes["contact_list"])
|
||||
previous_chat_gray, change_events = detect_chat_change(image, previous_chat_gray, boxes["chat_content"], args.chat_diff_threshold)
|
||||
events = decorate_events(badge_events + change_events, window, image.size)
|
||||
detection = {
|
||||
"timestamp": utc_now(),
|
||||
"frame_index": frame_index,
|
||||
"events": events,
|
||||
"badge_events": badge_events,
|
||||
"change_events": change_events,
|
||||
}
|
||||
shared.update_detection(detection)
|
||||
event = {
|
||||
"type": "algorithm_detection",
|
||||
"timestamp": utc_now(),
|
||||
"frame_index": frame_index,
|
||||
"badge_count": len(badge_events),
|
||||
"change_count": len(change_events),
|
||||
"events": events,
|
||||
"window": window,
|
||||
"detect_interval_sec": interval,
|
||||
}
|
||||
if badge_events or change_events:
|
||||
emit(event)
|
||||
if badge_events and args.click_on_badge:
|
||||
job_event.set()
|
||||
elif change_events and args.llm_on_chat_change:
|
||||
job_event.set()
|
||||
time.sleep(max(0, interval - (time.time() - start)))
|
||||
|
||||
|
||||
def run(args):
|
||||
annotation, regions = load_regions(args.regions)
|
||||
agent_config = load_agent_config(args.agent_config) if args.llm_on_click else None
|
||||
required = ["contact_list", "chat_content"]
|
||||
missing = [name for name in required if name not in regions]
|
||||
if missing:
|
||||
raise RuntimeError(f"missing required regions: {missing}")
|
||||
state = StreamState()
|
||||
server = start_server(args.host, args.port, state)
|
||||
url = f"http://{args.host}:{args.port}/"
|
||||
emit({"type": "server_started", "url": url, "optimized": True})
|
||||
if args.open_browser:
|
||||
webbrowser.open(url)
|
||||
|
||||
shared = SharedFrame()
|
||||
job_event = threading.Event()
|
||||
threads = [
|
||||
threading.Thread(target=preview_loop, args=(args, annotation, regions, state, shared), daemon=True),
|
||||
threading.Thread(target=detection_loop, args=(args, shared, job_event), daemon=True),
|
||||
]
|
||||
if args.click_on_badge and args.llm_on_click:
|
||||
threads.append(threading.Thread(target=llm_worker, args=(args, annotation, regions, agent_config, shared, job_event), daemon=True))
|
||||
for thread in threads:
|
||||
thread.start()
|
||||
try:
|
||||
while not shared.stop.is_set():
|
||||
time.sleep(0.5)
|
||||
finally:
|
||||
shared.stop.set()
|
||||
server.shutdown()
|
||||
|
||||
|
||||
def main():
|
||||
def stop_handler(signum, frame):
|
||||
raise KeyboardInterrupt
|
||||
signal.signal(signal.SIGTERM, stop_handler)
|
||||
signal.signal(signal.SIGINT, stop_handler)
|
||||
args = parse_args()
|
||||
try:
|
||||
run(args)
|
||||
except KeyboardInterrupt:
|
||||
emit({"type": "stopped", "timestamp": utc_now()})
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
408
wechat_vision/wechat_window_live.py
Normal file
408
wechat_vision/wechat_window_live.py
Normal file
@ -0,0 +1,408 @@
|
||||
import argparse
|
||||
import json
|
||||
import signal
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
import webbrowser
|
||||
from datetime import datetime, timezone
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import onnxruntime as ort
|
||||
from PIL import Image, ImageDraw
|
||||
|
||||
import Quartz
|
||||
|
||||
from ffmpeg_realtime_detect import (
|
||||
CLASS_NAMES,
|
||||
MODEL_PATH,
|
||||
detect,
|
||||
detection_event,
|
||||
draw_detections,
|
||||
image_to_jpeg_bytes,
|
||||
no_detection_event,
|
||||
)
|
||||
|
||||
|
||||
class StreamState:
|
||||
def __init__(self):
|
||||
self.condition = threading.Condition()
|
||||
self.jpeg = None
|
||||
self.event = None
|
||||
self.sequence = 0
|
||||
|
||||
def update(self, jpeg, event):
|
||||
with self.condition:
|
||||
self.jpeg = jpeg
|
||||
self.event = event
|
||||
self.sequence += 1
|
||||
self.condition.notify_all()
|
||||
|
||||
def snapshot(self):
|
||||
with self.condition:
|
||||
return self.jpeg, self.event, self.sequence
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Capture only the WeChat window, run ONNX, and expose a live stream."
|
||||
)
|
||||
parser.add_argument("--model", default=str(MODEL_PATH), help="ONNX model path")
|
||||
parser.add_argument("--fps", type=float, default=2.0, help="capture FPS")
|
||||
parser.add_argument(
|
||||
"--infer-interval",
|
||||
type=float,
|
||||
default=3.0,
|
||||
help="seconds between ONNX inference runs; preview frames reuse the latest detections",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log-every",
|
||||
type=int,
|
||||
default=60,
|
||||
help="log every N preview frames in addition to inference frames; 0 disables periodic logs",
|
||||
)
|
||||
parser.add_argument("--confidence", type=float, default=0.05, help="minimum score")
|
||||
parser.add_argument("--target-class", type=int, default=0, help="target class id")
|
||||
parser.add_argument("--host", default="127.0.0.1", help="HTTP host")
|
||||
parser.add_argument("--port", type=int, default=8765, help="HTTP port")
|
||||
parser.add_argument("--max-frames", type=int, default=0, help="stop after N frames; 0 runs forever")
|
||||
parser.add_argument("--open-browser", action="store_true", help="open the live preview in browser")
|
||||
parser.add_argument("--open-ffplay", action="store_true", help="open ffplay for the MJPEG stream")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def utc_now():
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
|
||||
def emit(event):
|
||||
print(json.dumps(event, ensure_ascii=False, separators=(",", ":")), flush=True)
|
||||
|
||||
|
||||
def find_wechat_window():
|
||||
options = Quartz.kCGWindowListOptionOnScreenOnly | Quartz.kCGWindowListExcludeDesktopElements
|
||||
windows = Quartz.CGWindowListCopyWindowInfo(options, Quartz.kCGNullWindowID) or []
|
||||
for window in windows:
|
||||
owner = window.get(Quartz.kCGWindowOwnerName, "") or ""
|
||||
title = window.get(Quartz.kCGWindowName, "") or ""
|
||||
layer = window.get(Quartz.kCGWindowLayer, -1)
|
||||
if layer != 0:
|
||||
continue
|
||||
if owner == "微信" or owner == "WeChat" or "wechat" in owner.lower() or title == "微信":
|
||||
bounds = window.get(Quartz.kCGWindowBounds, {}) or {}
|
||||
return {
|
||||
"id": int(window.get(Quartz.kCGWindowNumber)),
|
||||
"owner": owner,
|
||||
"title": title,
|
||||
"x": float(bounds.get("X", 0)),
|
||||
"y": float(bounds.get("Y", 0)),
|
||||
"width": float(bounds.get("Width", 0)),
|
||||
"height": float(bounds.get("Height", 0)),
|
||||
}
|
||||
return None
|
||||
|
||||
|
||||
def capture_window_image(window_id):
|
||||
image_ref = Quartz.CGWindowListCreateImage(
|
||||
Quartz.CGRectNull,
|
||||
Quartz.kCGWindowListOptionIncludingWindow,
|
||||
window_id,
|
||||
Quartz.kCGWindowImageBoundsIgnoreFraming,
|
||||
)
|
||||
if image_ref is None:
|
||||
return None
|
||||
|
||||
width = Quartz.CGImageGetWidth(image_ref)
|
||||
height = Quartz.CGImageGetHeight(image_ref)
|
||||
bytes_per_row = Quartz.CGImageGetBytesPerRow(image_ref)
|
||||
provider = Quartz.CGImageGetDataProvider(image_ref)
|
||||
data = Quartz.CGDataProviderCopyData(provider)
|
||||
array = np.frombuffer(data, dtype=np.uint8)
|
||||
array = array.reshape((height, bytes_per_row))[:, : width * 4]
|
||||
bgra = array.reshape((height, width, 4))
|
||||
rgba = bgra[:, :, [2, 1, 0, 3]]
|
||||
return Image.fromarray(rgba, "RGBA").convert("RGB")
|
||||
|
||||
|
||||
def screen_event(event, window):
|
||||
event = dict(event)
|
||||
event["window"] = window
|
||||
frame_width, frame_height = event.get("frame_size") or [window["width"], window["height"]]
|
||||
scale_x = frame_width / max(window["width"], 1)
|
||||
scale_y = frame_height / max(window["height"], 1)
|
||||
if event.get("bbox_screen"):
|
||||
x1, y1, x2, y2 = event["bbox_screen"]
|
||||
event["bbox_image"] = [x1, y1, x2, y2]
|
||||
wx1, wy1, wx2, wy2 = x1 / scale_x, y1 / scale_y, x2 / scale_x, y2 / scale_y
|
||||
event["bbox_window"] = [round(wx1, 2), round(wy1, 2), round(wx2, 2), round(wy2, 2)]
|
||||
event["bbox_screen"] = [
|
||||
round(window["x"] + wx1, 2),
|
||||
round(window["y"] + wy1, 2),
|
||||
round(window["x"] + wx2, 2),
|
||||
round(window["y"] + wy2, 2),
|
||||
]
|
||||
if event.get("center_screen"):
|
||||
cx, cy = event["center_screen"]
|
||||
event["center_image"] = [cx, cy]
|
||||
wcx, wcy = cx / scale_x, cy / scale_y
|
||||
event["center_window"] = [round(wcx, 2), round(wcy, 2)]
|
||||
event["center_screen"] = [round(window["x"] + wcx, 2), round(window["y"] + wcy, 2)]
|
||||
return event
|
||||
|
||||
|
||||
def event_from_detections(detections, frame_index, frame_size, class_name, target_class):
|
||||
if len(detections) > 0:
|
||||
best = detections[np.argmax(detections[:, 4])]
|
||||
return detection_event(best, frame_index, frame_size, class_name)
|
||||
return no_detection_event(frame_index, frame_size, target_class, class_name)
|
||||
|
||||
|
||||
def draw_overlay(image, event):
|
||||
draw = ImageDraw.Draw(image)
|
||||
text = f"frame={event.get('frame_index', '-')} {event.get('timestamp', '')}"
|
||||
draw.rectangle((12, 12, 720, 56), fill=(0, 0, 0))
|
||||
draw.text((22, 24), text, fill=(80, 255, 120))
|
||||
return image
|
||||
|
||||
|
||||
def html_page():
|
||||
return """<!doctype html>
|
||||
<html lang="zh-CN">
|
||||
<head>
|
||||
<meta charset="utf-8" />
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1" />
|
||||
<title>WeChat Window Live</title>
|
||||
<style>
|
||||
body { margin: 0; background: #09090b; color: #f4f4f5; font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; }
|
||||
.wrap { display: grid; grid-template-columns: minmax(0, 1fr) 360px; gap: 16px; padding: 16px; }
|
||||
.stage { background: #18181b; border: 1px solid #27272a; border-radius: 14px; overflow: hidden; }
|
||||
img { display: block; width: 100%; height: auto; }
|
||||
aside { background: #18181b; border: 1px solid #27272a; border-radius: 14px; padding: 16px; }
|
||||
h1 { margin: 0 0 12px; font-size: 18px; }
|
||||
.item { padding: 10px 0; border-top: 1px solid #27272a; }
|
||||
.label { color: #a1a1aa; font-size: 12px; margin-bottom: 4px; }
|
||||
.value { font-family: ui-monospace, SFMono-Regular, Menlo, monospace; word-break: break-all; }
|
||||
.ok { color: #86efac; }
|
||||
.miss { color: #fca5a5; }
|
||||
@media (max-width: 900px) { .wrap { grid-template-columns: 1fr; } }
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="wrap">
|
||||
<main class="stage"><img src="/stream.mjpg" alt="WeChat window live stream" /></main>
|
||||
<aside>
|
||||
<h1>WeChat Window Live</h1>
|
||||
<div class="item"><div class="label">Stream</div><div class="value">/stream.mjpg</div></div>
|
||||
<div class="item"><div class="label">Sequence</div><div id="sequence" class="value">-</div></div>
|
||||
<div class="item"><div class="label">Status</div><div id="status" class="value">-</div></div>
|
||||
<div class="item"><div class="label">Confidence</div><div id="confidence" class="value">-</div></div>
|
||||
<div class="item"><div class="label">Center Screen</div><div id="center" class="value">-</div></div>
|
||||
<div class="item"><div class="label">BBox Screen</div><div id="bbox" class="value">-</div></div>
|
||||
<div class="item"><div class="label">Window</div><div id="window" class="value">-</div></div>
|
||||
</aside>
|
||||
</div>
|
||||
<script>
|
||||
async function refresh() {
|
||||
try {
|
||||
const res = await fetch('/latest.json?t=' + Date.now(), { cache: 'no-store' });
|
||||
const data = await res.json();
|
||||
const event = data.event || {};
|
||||
const hit = event.type === 'detection';
|
||||
document.getElementById('sequence').textContent = data.sequence ?? '-';
|
||||
const status = document.getElementById('status');
|
||||
status.textContent = event.type ?? 'waiting';
|
||||
status.className = 'value ' + (hit ? 'ok' : 'miss');
|
||||
document.getElementById('confidence').textContent = event.confidence ?? '-';
|
||||
document.getElementById('center').textContent = JSON.stringify(event.center_screen ?? '-');
|
||||
document.getElementById('bbox').textContent = JSON.stringify(event.bbox_screen ?? '-');
|
||||
document.getElementById('window').textContent = JSON.stringify(event.window ?? '-');
|
||||
} catch (e) {}
|
||||
}
|
||||
refresh();
|
||||
setInterval(refresh, 300);
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
"""
|
||||
|
||||
|
||||
def make_handler(state):
|
||||
class Handler(BaseHTTPRequestHandler):
|
||||
def log_message(self, format, *args):
|
||||
return
|
||||
|
||||
def do_HEAD(self):
|
||||
if self.path == "/" or self.path.startswith("/index.html"):
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "text/html; charset=utf-8")
|
||||
self.send_header("Cache-Control", "no-store")
|
||||
self.end_headers()
|
||||
return
|
||||
if self.path.startswith("/latest.json"):
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "application/json; charset=utf-8")
|
||||
self.send_header("Cache-Control", "no-store")
|
||||
self.end_headers()
|
||||
return
|
||||
if self.path.startswith("/stream.mjpg"):
|
||||
self.send_response(200)
|
||||
self.send_header("Cache-Control", "no-cache, private")
|
||||
self.send_header("Pragma", "no-cache")
|
||||
self.send_header("Content-Type", "multipart/x-mixed-replace; boundary=frame")
|
||||
self.end_headers()
|
||||
return
|
||||
self.send_error(404)
|
||||
|
||||
def do_GET(self):
|
||||
if self.path == "/" or self.path.startswith("/index.html"):
|
||||
data = html_page().encode("utf-8")
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "text/html; charset=utf-8")
|
||||
self.send_header("Cache-Control", "no-store")
|
||||
self.send_header("Content-Length", str(len(data)))
|
||||
self.end_headers()
|
||||
self.wfile.write(data)
|
||||
return
|
||||
if self.path.startswith("/latest.json"):
|
||||
_, event, sequence = state.snapshot()
|
||||
payload = {"sequence": sequence, "updated_at": utc_now(), "event": event or {"type": "waiting"}}
|
||||
data = json.dumps(payload, ensure_ascii=False).encode("utf-8")
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "application/json; charset=utf-8")
|
||||
self.send_header("Cache-Control", "no-store")
|
||||
self.send_header("Content-Length", str(len(data)))
|
||||
self.end_headers()
|
||||
self.wfile.write(data)
|
||||
return
|
||||
if self.path.startswith("/stream.mjpg"):
|
||||
self.send_response(200)
|
||||
self.send_header("Cache-Control", "no-cache, private")
|
||||
self.send_header("Pragma", "no-cache")
|
||||
self.send_header("Content-Type", "multipart/x-mixed-replace; boundary=frame")
|
||||
self.end_headers()
|
||||
last_sequence = -1
|
||||
while True:
|
||||
with state.condition:
|
||||
state.condition.wait_for(lambda: state.sequence != last_sequence, timeout=5)
|
||||
jpeg, _, last_sequence = state.jpeg, state.event, state.sequence
|
||||
if not jpeg:
|
||||
continue
|
||||
try:
|
||||
self.wfile.write(b"--frame\r\n")
|
||||
self.wfile.write(b"Content-Type: image/jpeg\r\n")
|
||||
self.wfile.write(f"Content-Length: {len(jpeg)}\r\n\r\n".encode("ascii"))
|
||||
self.wfile.write(jpeg)
|
||||
self.wfile.write(b"\r\n")
|
||||
except (BrokenPipeError, ConnectionResetError):
|
||||
break
|
||||
return
|
||||
self.send_error(404)
|
||||
|
||||
return Handler
|
||||
|
||||
|
||||
def start_server(host, port, state):
|
||||
server = ThreadingHTTPServer((host, port), make_handler(state))
|
||||
thread = threading.Thread(target=server.serve_forever, daemon=True)
|
||||
thread.start()
|
||||
return server
|
||||
|
||||
|
||||
def open_ffplay(url):
|
||||
return subprocess.Popen(
|
||||
["ffplay", "-hide_banner", "-loglevel", "error", "-fflags", "nobuffer", "-flags", "low_delay", "-framedrop", url],
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL,
|
||||
)
|
||||
|
||||
|
||||
def run(args):
|
||||
state = StreamState()
|
||||
server = start_server(args.host, args.port, state)
|
||||
url = f"http://{args.host}:{args.port}/"
|
||||
stream_url = f"http://{args.host}:{args.port}/stream.mjpg"
|
||||
emit({"type": "server_started", "url": url, "stream_url": stream_url})
|
||||
if args.open_browser:
|
||||
webbrowser.open(url)
|
||||
ffplay_process = open_ffplay(stream_url) if args.open_ffplay else None
|
||||
|
||||
session = ort.InferenceSession(str(Path(args.model)))
|
||||
input_meta = session.get_inputs()[0]
|
||||
class_name = CLASS_NAMES.get(args.target_class, f"class_{args.target_class}")
|
||||
interval = 1 / max(args.fps, 0.1)
|
||||
frame_index = 0
|
||||
last_infer_at = 0.0
|
||||
last_infer_size = None
|
||||
last_detections = np.empty((0, 6))
|
||||
try:
|
||||
while args.max_frames <= 0 or frame_index < args.max_frames:
|
||||
start = time.time()
|
||||
window = find_wechat_window()
|
||||
if not window:
|
||||
event = {"type": "wechat_window_not_found", "timestamp": utc_now()}
|
||||
state.update(image_to_jpeg_bytes(Image.new("RGB", (960, 540), "black")), event)
|
||||
emit(event)
|
||||
time.sleep(interval)
|
||||
continue
|
||||
|
||||
image = capture_window_image(window["id"])
|
||||
if image is None:
|
||||
event = {"type": "capture_failed", "timestamp": utc_now(), "window": window}
|
||||
emit(event)
|
||||
time.sleep(interval)
|
||||
continue
|
||||
|
||||
frame_index += 1
|
||||
now = time.time()
|
||||
should_infer = (
|
||||
now - last_infer_at >= args.infer_interval
|
||||
or last_infer_size != image.size
|
||||
or frame_index == 1
|
||||
)
|
||||
if should_infer:
|
||||
last_detections = detect(session, input_meta, image, args.target_class, args.confidence)
|
||||
last_infer_at = now
|
||||
last_infer_size = image.size
|
||||
|
||||
event = event_from_detections(
|
||||
last_detections,
|
||||
frame_index,
|
||||
image.size,
|
||||
class_name,
|
||||
args.target_class,
|
||||
)
|
||||
event["inferred"] = should_infer
|
||||
event["inference_age_ms"] = round((now - last_infer_at) * 1000, 1)
|
||||
event = screen_event(event, window)
|
||||
annotated = draw_detections(image, last_detections, class_name)
|
||||
annotated = draw_overlay(annotated, event)
|
||||
state.update(image_to_jpeg_bytes(annotated), event)
|
||||
if should_infer or (args.log_every > 0 and frame_index % args.log_every == 0):
|
||||
emit(event)
|
||||
|
||||
elapsed = time.time() - start
|
||||
time.sleep(max(0, interval - elapsed))
|
||||
finally:
|
||||
server.shutdown()
|
||||
if ffplay_process:
|
||||
ffplay_process.terminate()
|
||||
|
||||
|
||||
def main():
|
||||
def stop_handler(signum, frame):
|
||||
raise KeyboardInterrupt
|
||||
|
||||
signal.signal(signal.SIGTERM, stop_handler)
|
||||
signal.signal(signal.SIGINT, stop_handler)
|
||||
args = parse_args()
|
||||
try:
|
||||
run(args)
|
||||
except KeyboardInterrupt:
|
||||
emit({"type": "stopped", "timestamp": utc_now()})
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Loading…
x
Reference in New Issue
Block a user