375 lines
12 KiB
Python
375 lines
12 KiB
Python
|
"""
|
|||
|
|
语料采集 · Mac客户端
|
|||
|
|
====================
|
|||
|
|
在Mac上运行，自动截图OCR+滚屏采集，发送到服务端处理
|
|||
|
|
|
|||
|
|
用法:
|
|||
|
|
# 实时采集模式（后台监控屏幕变化）
|
|||
|
|
python3 mac-corpus-agent.py --token YOUR_TOKEN --mode auto
|
|||
|
|
|
|||
|
|
# 滚屏采集模式（采集历史对话，先手动滚到顶部）
|
|||
|
|
python3 mac-corpus-agent.py --token YOUR_TOKEN --mode scroll
|
|||
|
|
|
|||
|
|
# 剪贴板采集模式（复制即采集）
|
|||
|
|
python3 mac-corpus-agent.py --token YOUR_TOKEN --mode clipboard
|
|||
|
|
|
|||
|
|
选项:
|
|||
|
|
--server 服务端地址 (默认 http://localhost:8084)
|
|||
|
|
--interval 采集间隔秒数 (默认 auto=5s, scroll=3s)
|
|||
|
|
--region 屏幕采集区域 (默认 全屏)
|
|||
|
|
--help 显示帮助
|
|||
|
|
"""
|
|||
|
|
import os
|
|||
|
|
import sys
|
|||
|
|
import json
|
|||
|
|
import time
|
|||
|
|
import argparse
|
|||
|
|
import subprocess
|
|||
|
|
import tempfile
|
|||
|
|
import threading
|
|||
|
|
from pathlib import Path
|
|||
|
|
from typing import Optional
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
import websocket
|
|||
|
|
except ImportError:
|
|||
|
|
print("❌ 缺少依赖: pip3 install websocket-client")
|
|||
|
|
sys.exit(1)
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
import pyautogui
|
|||
|
|
except ImportError:
|
|||
|
|
print("❌ 缺少依赖: pip3 install pyautogui pillow")
|
|||
|
|
sys.exit(1)
|
|||
|
|
|
|||
|
|
# ============================================================
|
|||
|
|
# 配置
|
|||
|
|
# ============================================================
|
|||
|
|
|
|||
|
|
VERSION = "1.0.0"
|
|||
|
|
SERVER_URL = "http://localhost:8084"
|
|||
|
|
RECONNECT_DELAY = 5 # 断线重连等待秒数
|
|||
|
|
|
|||
|
|
# 键盘快捷键（全局热键用，需管理员权限）
|
|||
|
|
HOTKEY_STOP = "esc" # 停止采集
|
|||
|
|
HOTKEY_PAUSE = "f6" # 暂停/继续
|
|||
|
|
|
|||
|
|
# ============================================================
|
|||
|
|
# macOS工具函数
|
|||
|
|
# ============================================================
|
|||
|
|
|
|||
|
|
def screenshot(region: Optional[tuple] = None) -> bytes:
|
|||
|
|
"""截取屏幕指定区域，返回PNG字节"""
|
|||
|
|
if region:
|
|||
|
|
img = pyautogui.screenshot(region=region)
|
|||
|
|
else:
|
|||
|
|
img = pyautogui.screenshot()
|
|||
|
|
|
|||
|
|
buf = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
|
|||
|
|
img.save(buf, format="PNG")
|
|||
|
|
buf.close()
|
|||
|
|
|
|||
|
|
with open(buf.name, "rb") as f:
|
|||
|
|
data = f.read()
|
|||
|
|
|
|||
|
|
os.unlink(buf.name)
|
|||
|
|
return data
|
|||
|
|
|
|||
|
|
|
|||
|
|
def ocr_text(image_path: str) -> str:
|
|||
|
|
"""使用macOS原生Vision框架做OCR（本地，不上传）"""
|
|||
|
|
script = f'''
|
|||
|
|
use framework "Vision"
|
|||
|
|
use scripting additions
|
|||
|
|
|
|||
|
|
set theImage to (current application's NSImage's alloc()'s initWithContentsOfFile:"{image_path}")
|
|||
|
|
set requestHandler to (current application's VNImageRequestHandler's alloc()'s initWithData:(theImage's TIFFRepresentation()) options:(missing value))
|
|||
|
|
|
|||
|
|
set textResult to ""
|
|||
|
|
set theRequest to (current application's VNRecognizeTextRequest's alloc()'s initWithCompletionHandler:(lambda request, error
|
|||
|
|
if error ≠ missing value then return
|
|||
|
|
set observations to request's results()
|
|||
|
|
repeat with obs in observations
|
|||
|
|
set topCandidate to (obs's topCandidates:(1))
|
|||
|
|
if topCandidate's count() > 0 then
|
|||
|
|
set candidate to (topCandidate's objectAtIndex:(0))
|
|||
|
|
set recognizedText to candidate's string() as text
|
|||
|
|
set textResult to textResult & recognizedText & linefeed
|
|||
|
|
end if
|
|||
|
|
end repeat
|
|||
|
|
end))
|
|||
|
|
|
|||
|
|
theRequest's setRecognitionLevel:(VNRequestTextRecognitionLevel1) -- Accurate
|
|||
|
|
requestHandler's performRequests:({{theRequest}}) |error|:(missing value)
|
|||
|
|
return textResult
|
|||
|
|
'''
|
|||
|
|
|
|||
|
|
result = subprocess.run(
|
|||
|
|
["osascript", "-e", script],
|
|||
|
|
capture_output=True, text=True, timeout=30
|
|||
|
|
)
|
|||
|
|
return result.stdout.strip()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def scroll_page(direction: str = "down", amount: int = 3):
|
|||
|
|
"""模拟鼠标滚轮"""
|
|||
|
|
pyautogui.scroll(-amount if direction == "down" else amount)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_active_window_title() -> str:
|
|||
|
|
"""获取当前活跃窗口标题"""
|
|||
|
|
script = 'tell application "System Events" to get name of first application process whose frontmost is true'
|
|||
|
|
result = subprocess.run(["osascript", "-e", script], capture_output=True, text=True)
|
|||
|
|
return result.stdout.strip()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def notify(title: str, message: str):
|
|||
|
|
"""MacOS系统通知"""
|
|||
|
|
script = f'display notification "{message}" with title "{title}"'
|
|||
|
|
subprocess.run(["osascript", "-e", script], capture_output=True)
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ============================================================
|
|||
|
|
# 采集器核心
|
|||
|
|
# ============================================================
|
|||
|
|
|
|||
|
|
class CorpusCollector:
|
|||
|
|
"""语料采集器"""
|
|||
|
|
|
|||
|
|
def __init__(self, token: str, server: str, interval: float):
|
|||
|
|
self.token = token
|
|||
|
|
self.server = server
|
|||
|
|
self.interval = interval
|
|||
|
|
self.ws = None
|
|||
|
|
self.running = False
|
|||
|
|
self.paused = False
|
|||
|
|
self.last_text = "" # 去重用
|
|||
|
|
self.collected_count = 0
|
|||
|
|
self.filtered_count = 0
|
|||
|
|
|
|||
|
|
def connect_ws(self) -> bool:
|
|||
|
|
"""连接WebSocket"""
|
|||
|
|
try:
|
|||
|
|
ws_url = self.server.replace("http://", "ws://").replace("https://", "wss://")
|
|||
|
|
ws_url = f"{ws_url}/ws/collect"
|
|||
|
|
|
|||
|
|
self.ws = websocket.create_connection(
|
|||
|
|
f"{ws_url}?token={self.token}",
|
|||
|
|
timeout=10
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 验证连接
|
|||
|
|
self.ws.send(json.dumps({"type": "ping"}))
|
|||
|
|
resp = json.loads(self.ws.recv())
|
|||
|
|
|
|||
|
|
if resp.get("type") == "pong":
|
|||
|
|
print(f" ✅ WebSocket已连接")
|
|||
|
|
return True
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f" ❌ WebSocket连接失败: {e}")
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
def send_text(self, text: str, source: str = "screen_capture"):
|
|||
|
|
"""发送文本到服务器"""
|
|||
|
|
if not text or len(text) < 10:
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
# 简单去重（连续相同内容跳过）
|
|||
|
|
if text == self.last_text:
|
|||
|
|
return
|
|||
|
|
self.last_text = text
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
if self.ws:
|
|||
|
|
self.ws.send(json.dumps({
|
|||
|
|
"type": "text",
|
|||
|
|
"text": text,
|
|||
|
|
"source": source
|
|||
|
|
}))
|
|||
|
|
resp = json.loads(self.ws.recv())
|
|||
|
|
|
|||
|
|
if resp.get("collected", 0) > 0:
|
|||
|
|
self.collected_count += resp["collected"]
|
|||
|
|
print(f" ✅ 采集 {resp['collected']} 条 | 总计: {self.collected_count}")
|
|||
|
|
notify("语料采集", f"已采集 {resp['collected']} 条对话")
|
|||
|
|
elif resp.get("valuable"):
|
|||
|
|
pass # 有价值但未成对
|
|||
|
|
else:
|
|||
|
|
self.filtered_count += 1
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f" ⚠️ 发送失败: {e}")
|
|||
|
|
|
|||
|
|
def process_screenshot(self):
|
|||
|
|
"""截屏→OCR→发送"""
|
|||
|
|
try:
|
|||
|
|
# 截屏
|
|||
|
|
img_data = screenshot()
|
|||
|
|
|
|||
|
|
# 存临时文件
|
|||
|
|
tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
|
|||
|
|
tmp.write(img_data)
|
|||
|
|
tmp.close()
|
|||
|
|
|
|||
|
|
# OCR
|
|||
|
|
text = ocr_text(tmp.name)
|
|||
|
|
os.unlink(tmp.name)
|
|||
|
|
|
|||
|
|
if text.strip():
|
|||
|
|
self.send_text(text.strip())
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f" ⚠️ 截图处理失败: {e}")
|
|||
|
|
|
|||
|
|
# === 模式1: 实时采集 ===
|
|||
|
|
def run_auto_mode(self):
|
|||
|
|
"""实时模式：后台监控屏幕变化"""
|
|||
|
|
print("\n🔴 实时采集模式启动中...")
|
|||
|
|
print(f" 采集间隔: {self.interval}秒")
|
|||
|
|
print(f" 按 ESC 停止, F6 暂停/继续")
|
|||
|
|
print(f" 活跃窗口: {get_active_window_title()}")
|
|||
|
|
|
|||
|
|
self.running = True
|
|||
|
|
|
|||
|
|
while self.running:
|
|||
|
|
if self.paused:
|
|||
|
|
time.sleep(1)
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
self.process_screenshot()
|
|||
|
|
|
|||
|
|
# 检查快捷键（每轮检查一次）
|
|||
|
|
# Mac上需要更好的热键方案，这里简化处理
|
|||
|
|
for _ in range(int(self.interval * 2)):
|
|||
|
|
if not self.running:
|
|||
|
|
break
|
|||
|
|
time.sleep(0.5)
|
|||
|
|
|
|||
|
|
self.cleanup()
|
|||
|
|
|
|||
|
|
# === 模式2: 滚屏采集 ===
|
|||
|
|
def run_scroll_mode(self):
|
|||
|
|
"""滚屏模式：自动向下滚动采集历史对话"""
|
|||
|
|
print("\n📜 滚屏采集模式启动")
|
|||
|
|
print(f" 采集间隔: {self.interval}秒")
|
|||
|
|
print(f" 请确保已手动滚动到页面顶部!")
|
|||
|
|
print(f" 5秒后开始...")
|
|||
|
|
|
|||
|
|
time.sleep(5)
|
|||
|
|
|
|||
|
|
self.running = True
|
|||
|
|
scroll_count = 0
|
|||
|
|
empty_rounds = 0
|
|||
|
|
|
|||
|
|
while self.running and empty_rounds < 10:
|
|||
|
|
# 截图+OCR
|
|||
|
|
self.process_screenshot()
|
|||
|
|
|
|||
|
|
# 滚动
|
|||
|
|
scroll_page("down", 5)
|
|||
|
|
scroll_count += 1
|
|||
|
|
|
|||
|
|
# 检测是否到底（连续N次没有新内容）
|
|||
|
|
time.sleep(self.interval)
|
|||
|
|
|
|||
|
|
if scroll_count % 10 == 0:
|
|||
|
|
print(f" 已滚动 {scroll_count} 次 | 采集: {self.collected_count} | 过滤: {self.filtered_count}")
|
|||
|
|
|
|||
|
|
print(f"\n✅ 滚屏采集完成")
|
|||
|
|
print(f" 共滚动 {scroll_count} 次")
|
|||
|
|
print(f" 采集: {self.collected_count} 条")
|
|||
|
|
print(f" 过滤: {self.filtered_count} 条")
|
|||
|
|
|
|||
|
|
notify("语料采集完成", f"共采集 {self.collected_count} 条对话")
|
|||
|
|
self.cleanup()
|
|||
|
|
|
|||
|
|
# === 模式3: 剪贴板采集 ===
|
|||
|
|
def run_clipboard_mode(self):
|
|||
|
|
"""剪贴板模式：监控剪贴板变化，自动采集"""
|
|||
|
|
print("\n📋 剪贴板采集模式启动")
|
|||
|
|
print(f" 监控间隔: {self.interval}秒")
|
|||
|
|
print(f" 复制内容后自动采集...")
|
|||
|
|
|
|||
|
|
import subprocess
|
|||
|
|
|
|||
|
|
self.running = True
|
|||
|
|
last_clip = ""
|
|||
|
|
|
|||
|
|
while self.running:
|
|||
|
|
# 获取剪贴板内容
|
|||
|
|
result = subprocess.run(
|
|||
|
|
["pbpaste"],
|
|||
|
|
capture_output=True, text=True
|
|||
|
|
)
|
|||
|
|
current = result.stdout.strip()
|
|||
|
|
|
|||
|
|
if current and current != last_clip:
|
|||
|
|
print(f"\n 📋 检测到新内容 ({len(current)}字)")
|
|||
|
|
self.send_text(current, "clipboard")
|
|||
|
|
last_clip = current
|
|||
|
|
|
|||
|
|
time.sleep(self.interval)
|
|||
|
|
|
|||
|
|
self.cleanup()
|
|||
|
|
|
|||
|
|
def cleanup(self):
|
|||
|
|
"""清理"""
|
|||
|
|
if self.ws:
|
|||
|
|
try:
|
|||
|
|
self.ws.close()
|
|||
|
|
except:
|
|||
|
|
pass
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ============================================================
|
|||
|
|
# 主入口
|
|||
|
|
# ============================================================
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
parser = argparse.ArgumentParser(description="语料采集 Mac 客户端")
|
|||
|
|
parser.add_argument("--token", required=True, help="Gitea Access Token")
|
|||
|
|
parser.add_argument("--server", default=SERVER_URL, help=f"服务端地址 (默认 {SERVER_URL})")
|
|||
|
|
parser.add_argument("--mode", choices=["auto", "scroll", "clipboard"], default="auto",
|
|||
|
|
help="采集模式: auto=实时, scroll=滚屏, clipboard=剪贴板")
|
|||
|
|
parser.add_argument("--interval", type=float, default=0,
|
|||
|
|
help="采集间隔秒数 (默认 auto=5, scroll=3, clipboard=2)")
|
|||
|
|
|
|||
|
|
args = parser.parse_args()
|
|||
|
|
|
|||
|
|
# 默认间隔
|
|||
|
|
if args.interval <= 0:
|
|||
|
|
intervals = {"auto": 5.0, "scroll": 3.0, "clipboard": 2.0}
|
|||
|
|
args.interval = intervals[args.mode]
|
|||
|
|
|
|||
|
|
print(f"\n{'='*50}")
|
|||
|
|
print(f"🧠 语料采集 Mac 客户端 v{VERSION}")
|
|||
|
|
print(f"{'='*50}")
|
|||
|
|
print(f" 模式: {args.mode}")
|
|||
|
|
print(f" 服务端: {args.server}")
|
|||
|
|
print(f" 间隔: {args.interval}s")
|
|||
|
|
print(f"{'='*50}\n")
|
|||
|
|
|
|||
|
|
collector = CorpusCollector(args.token, args.server, args.interval)
|
|||
|
|
|
|||
|
|
# 连接服务端
|
|||
|
|
print("🔄 连接服务端...")
|
|||
|
|
if not collector.connect_ws():
|
|||
|
|
print(" ⚠️ WebSocket连接失败，将使用HTTP回退")
|
|||
|
|
print(" 请确保服务端已启动: python3 server.py")
|
|||
|
|
print(f" 服务端地址: {args.server}")
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
if args.mode == "auto":
|
|||
|
|
collector.run_auto_mode()
|
|||
|
|
elif args.mode == "scroll":
|
|||
|
|
collector.run_scroll_mode()
|
|||
|
|
elif args.mode == "clipboard":
|
|||
|
|
collector.run_clipboard_mode()
|
|||
|
|
except KeyboardInterrupt:
|
|||
|
|
print("\n\n⏹️ 用户停止")
|
|||
|
|
collector.cleanup()
|
|||
|
|
|
|||
|
|
print("\n👋 采集结束")
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
main()
|