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 """ WeChat Window Live
WeChat window live stream
""" 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()