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()
|