import argparse import ctypes import ctypes.wintypes import json import os import platform 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 try: from PIL import ImageGrab except ImportError: ImageGrab = None try: import Quartz except ImportError: Quartz = None IS_WINDOWS = platform.system() == "Windows" IS_DARWIN = platform.system() == "Darwin" DWMWA_EXTENDED_FRAME_BOUNDS = 9 PROCESS_QUERY_LIMITED_INFORMATION = 0x1000 if IS_WINDOWS: ctypes.windll.kernel32.OpenProcess.restype = ctypes.wintypes.HANDLE ctypes.windll.kernel32.QueryFullProcessImageNameW.argtypes = [ ctypes.wintypes.HANDLE, ctypes.wintypes.DWORD, ctypes.wintypes.LPWSTR, ctypes.POINTER(ctypes.wintypes.DWORD), ] ctypes.windll.user32.EnumWindows.argtypes = [ctypes.c_void_p, ctypes.wintypes.LPARAM] 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 _window_matches_wechat(owner, title): owner = owner or "" title = title or "" owner_lower = owner.lower() title_lower = title.lower() return ( owner in {"微信", "WeChat"} or owner_lower in {"wechat", "wechat.exe", "weixin", "weixin.exe"} or "wechat" in owner_lower or title == "微信" or "wechat" in title_lower ) def _windows_text(hwnd): user32 = ctypes.windll.user32 length = user32.GetWindowTextLengthW(hwnd) if length <= 0: return "" buffer = ctypes.create_unicode_buffer(length + 1) copied = user32.GetWindowTextW(hwnd, buffer, length + 1) if copied <= 0: return "" return buffer.value.strip() def _windows_process_name(pid): if not pid: return "" kernel32 = ctypes.windll.kernel32 handle = kernel32.OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, False, pid) if not handle: return "" try: size = ctypes.wintypes.DWORD(32768) buffer = ctypes.create_unicode_buffer(size.value) ok = kernel32.QueryFullProcessImageNameW(handle, 0, buffer, ctypes.byref(size)) if not ok: return "" return os.path.basename(buffer.value) finally: kernel32.CloseHandle(handle) def _windows_rect(hwnd): rect = ctypes.wintypes.RECT() try: result = ctypes.windll.dwmapi.DwmGetWindowAttribute( hwnd, DWMWA_EXTENDED_FRAME_BOUNDS, ctypes.byref(rect), ctypes.sizeof(rect), ) except AttributeError: result = -1 if result != 0: ok = ctypes.windll.user32.GetWindowRect(hwnd, ctypes.byref(rect)) if not ok: return None if rect.right <= rect.left or rect.bottom <= rect.top: return None return rect def _find_windows_wechat_window(): user32 = ctypes.windll.user32 windows = [] enum_proc = ctypes.WINFUNCTYPE(ctypes.c_bool, ctypes.wintypes.HWND, ctypes.wintypes.LPARAM) @enum_proc def callback(hwnd, _lparam): if not user32.IsWindowVisible(hwnd) or user32.IsIconic(hwnd): return True rect = _windows_rect(hwnd) if rect is None: return True width = rect.right - rect.left height = rect.bottom - rect.top if width < 80 or height < 80: return True title = _windows_text(hwnd) pid = ctypes.wintypes.DWORD() user32.GetWindowThreadProcessId(hwnd, ctypes.byref(pid)) owner = _windows_process_name(pid.value) if not _window_matches_wechat(owner, title): return True windows.append( { "id": int(hwnd), "owner": owner or "Windows", "title": title, "x": float(rect.left), "y": float(rect.top), "width": float(width), "height": float(height), "platform": "windows", "pid": int(pid.value) if pid.value else None, } ) return True user32.EnumWindows(callback, 0) if not windows: return None windows.sort(key=lambda item: (item["owner"].lower() != "wechat.exe", -item["width"] * item["height"])) return windows[0] def _capture_windows_window_image(window_id): if ImageGrab is None: raise RuntimeError("Pillow ImageGrab is required for Windows window capture") rect = _windows_rect(ctypes.wintypes.HWND(int(window_id))) if rect is None: return None bbox = (rect.left, rect.top, rect.right, rect.bottom) try: image = ImageGrab.grab(bbox=bbox, all_screens=True) except TypeError: image = ImageGrab.grab(bbox=bbox) return image.convert("RGB") def find_wechat_window(): if IS_WINDOWS: return _find_windows_wechat_window() if not IS_DARWIN or Quartz is None: return None 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 _window_matches_wechat(owner, 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)), "platform": "darwin", } return None def capture_window_image(window_id): if IS_WINDOWS: return _capture_windows_window_image(window_id) if not IS_DARWIN or Quartz is None: return None 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()