409 lines
16 KiB
Python
409 lines
16 KiB
Python
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()
|