import argparse import json import signal import subprocess import sys 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 BASE_DIR = Path(__file__).resolve().parent MODEL_PATH = BASE_DIR / "best.onnx" DEBUG_DIR = BASE_DIR / "ouptsw" / "monitor" VISUALIZE_DIR = BASE_DIR / "ouptsw" / "realtime_view" CLASS_NAMES = { 0: "input_box", } IOU_THRESHOLD = 0.5 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="Realtime FFmpeg screen capture + ONNX detection demo." ) parser.add_argument("--model", default=str(MODEL_PATH), help="ONNX model path") parser.add_argument( "--input", default="3:none", help="FFmpeg avfoundation input, e.g. '3:none'. Use --list-devices to inspect.", ) parser.add_argument("--fps", type=float, default=2.0, help="capture FPS") parser.add_argument("--confidence", type=float, default=0.1, help="minimum score") parser.add_argument("--target-class", type=int, default=0, help="target class id") parser.add_argument( "--max-frames", type=int, default=0, help="stop after N frames; 0 means run until interrupted", ) parser.add_argument( "--emit-empty", action="store_true", help="emit no_detection events when target is not found", ) parser.add_argument( "--debug-dir", default="", help="save annotated frames into this directory", ) parser.add_argument( "--debug-every", type=int, default=0, help="save one debug image every N frames; 0 disables debug images", ) parser.add_argument( "--visualize-dir", default="", help="write latest.jpg/latest.json/index.html for browser visualization", ) parser.add_argument( "--visualize-every", type=int, default=1, help="update visualization every N frames when --visualize-dir is enabled", ) parser.add_argument( "--open-visualizer", action="store_true", help="open the visualization HTML in the default browser", ) parser.add_argument( "--serve", action="store_true", help="serve a realtime MJPEG stream and dashboard over HTTP", ) parser.add_argument("--host", default="127.0.0.1", help="HTTP server host") parser.add_argument("--port", type=int, default=8765, help="HTTP server port") parser.add_argument( "--display", action="store_true", help="open an ffplay window with the annotated realtime video", ) parser.add_argument( "--display-title", default="WeChat Vision ONNX Preview", help="ffplay window title when --display is enabled", ) parser.add_argument( "--display-width", type=int, default=720, help="ffplay window width when --display is enabled", ) parser.add_argument( "--display-height", type=int, default=450, help="ffplay window height when --display is enabled", ) parser.add_argument( "--display-left", type=int, default=0, help="ffplay window left position when --display is enabled", ) parser.add_argument( "--display-top", type=int, default=0, help="ffplay window top position when --display is enabled", ) parser.add_argument( "--list-devices", action="store_true", help="list FFmpeg avfoundation devices and exit", ) return parser.parse_args() def log_error(message): print(message, file=sys.stderr, flush=True) def utc_now(): return datetime.now(timezone.utc).isoformat() def input_shape(input_meta): return [dim if isinstance(dim, int) else 1 for dim in input_meta.shape] def preprocess(image, shape): _, _, height, width = shape resized = image.convert("RGB").resize((width, height)) array = np.asarray(resized, dtype=np.float32) / 255.0 return np.transpose(array, (2, 0, 1))[None] def iou(box, boxes): x1 = np.maximum(box[0], boxes[:, 0]) y1 = np.maximum(box[1], boxes[:, 1]) x2 = np.minimum(box[2], boxes[:, 2]) y2 = np.minimum(box[3], boxes[:, 3]) intersection = np.maximum(0, x2 - x1) * np.maximum(0, y2 - y1) box_area = np.maximum(0, box[2] - box[0]) * np.maximum(0, box[3] - box[1]) boxes_area = np.maximum(0, boxes[:, 2] - boxes[:, 0]) * np.maximum( 0, boxes[:, 3] - boxes[:, 1], ) return intersection / np.maximum(box_area + boxes_area - intersection, 1e-6) def nms(detections, confidence): detections = detections[detections[:, 4] >= confidence] if len(detections) == 0: return detections detections = detections[np.argsort(detections[:, 4])[::-1]] kept = [] while len(detections) > 0: best = detections[0] kept.append(best) if len(detections) == 1: break detections = detections[1:][iou(best[:4], detections[1:, :4]) < IOU_THRESHOLD] return np.array(kept) def scale_detections(detections, original_size, shape): if len(detections) == 0: return detections _, _, input_height, input_width = shape original_width, original_height = original_size scaled = detections.copy() scaled[:, [0, 2]] *= original_width / input_width scaled[:, [1, 3]] *= original_height / input_height scaled[:, [0, 2]] = np.clip(scaled[:, [0, 2]], 0, original_width) scaled[:, [1, 3]] = np.clip(scaled[:, [1, 3]], 0, original_height) return scaled def detect(session, input_meta, image, target_class, confidence): shape = input_shape(input_meta) tensor = preprocess(image, shape) raw = session.run(None, {input_meta.name: tensor})[0][0] detections = raw[raw[:, 5].astype(int) == target_class] detections = nms(detections, confidence) return scale_detections(detections, image.size, shape) def emit(event): print(json.dumps(event, ensure_ascii=False, separators=(",", ":")), flush=True) def write_json_atomic(path, data): path.parent.mkdir(parents=True, exist_ok=True) temp_path = path.with_suffix(path.suffix + ".tmp") temp_path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8") temp_path.replace(path) def save_image_atomic(image, path): path.parent.mkdir(parents=True, exist_ok=True) temp_path = path.with_suffix(path.suffix + ".tmp") image.save(temp_path, format="JPEG", quality=85) temp_path.replace(path) def image_to_jpeg_bytes(image): buffer = BytesIO() image.save(buffer, format="JPEG", quality=85) return buffer.getvalue() def detection_event(detection, frame_index, frame_size, class_name): x1, y1, x2, y2, score, class_id = detection cx = (x1 + x2) / 2 cy = (y1 + y2) / 2 return { "type": "detection", "timestamp": utc_now(), "frame_index": frame_index, "class_id": int(class_id), "class_name": class_name, "confidence": round(float(score), 6), "bbox_screen": [round(float(x1), 2), round(float(y1), 2), round(float(x2), 2), round(float(y2), 2)], "center_screen": [round(float(cx), 2), round(float(cy), 2)], "frame_size": [frame_size[0], frame_size[1]], } def no_detection_event(frame_index, frame_size, target_class, class_name): return { "type": "no_detection", "timestamp": utc_now(), "frame_index": frame_index, "class_id": target_class, "class_name": class_name, "frame_size": [frame_size[0], frame_size[1]], } def save_debug_image(image, detections, path, class_name): debug = draw_detections(image, detections, class_name) path.parent.mkdir(parents=True, exist_ok=True) debug.save(path) def draw_detections(image, detections, class_name): debug = image.convert("RGB") draw = ImageDraw.Draw(debug) for detection in detections: x1, y1, x2, y2, score, _ = detection cx = (x1 + x2) / 2 cy = (y1 + y2) / 2 draw.rectangle((x1, y1, x2, y2), outline="red", width=3) draw.ellipse((cx - 5, cy - 5, cx + 5, cy + 5), fill="yellow", outline="red") draw.text( (x1, max(0, y1 - 14)), f"{class_name} ({cx:.1f}, {cy:.1f}) {score:.2f}", fill="yellow", ) return debug def draw_frame_info(image, event): draw = ImageDraw.Draw(image) frame_index = event.get("frame_index", "-") timestamp = event.get("timestamp", "") text = f"frame={frame_index} {timestamp}" draw.rectangle((12, 12, 620, 52), fill=(0, 0, 0)) draw.text((22, 22), text, fill=(80, 255, 120)) return image def write_visualizer_html(path): path.parent.mkdir(parents=True, exist_ok=True) html = """ WeChat Vision Realtime
latest detection frame
""" path.write_text(html, encoding="utf-8") def dashboard_html(frame_url="frame.jpg", json_url="latest.json"): return f""" WeChat Vision Live Stream
waiting for frames
""" def make_stream_handler(state): class Handler(BaseHTTPRequestHandler): protocol_version = "HTTP/1.1" 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("/frame.jpg"): self.send_response(200) self.send_header("Content-Type", "image/jpeg") self.send_header("Cache-Control", "no-store, no-cache, must-revalidate, max-age=0") 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"): self.send_response(200) self.send_header("Content-Type", "text/html; charset=utf-8") self.send_header("Cache-Control", "no-store") self.end_headers() self.wfile.write(dashboard_html().encode("utf-8")) return if self.path.startswith("/latest.json"): _, event, sequence = state.snapshot() payload = { "sequence": sequence, "updated_at": utc_now(), "event": event or {"type": "waiting", "timestamp": utc_now()}, } 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("/frame.jpg"): jpeg, _, sequence = state.snapshot() if not jpeg: self.send_response(503) self.send_header("Cache-Control", "no-store") self.end_headers() return self.send_response(200) self.send_header("Content-Type", "image/jpeg") self.send_header("Cache-Control", "no-store, no-cache, must-revalidate, max-age=0") self.send_header("Pragma", "no-cache") self.send_header("X-Frame-Sequence", str(sequence)) self.send_header("Content-Length", str(len(jpeg))) self.end_headers() self.wfile.write(jpeg) return if self.path.startswith("/stream.mjpg"): self.send_response(200) self.send_header("Age", "0") 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 = state.jpeg last_sequence = 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_stream_server(host, port, state): server = ThreadingHTTPServer((host, port), make_stream_handler(state)) thread = threading.Thread(target=server.serve_forever, daemon=True) thread.start() return server def update_visualizer(image, detections, event, output_dir, class_name): annotated = draw_detections(image, detections, class_name) annotated = draw_frame_info(annotated, event) save_image_atomic(annotated, output_dir / "latest.jpg") payload = { "updated_at": utc_now(), "event": event, "frame_index": event.get("frame_index"), "class_id": event.get("class_id"), "class_name": event.get("class_name"), "confidence": event.get("confidence"), "bbox_screen": event.get("bbox_screen"), "center_screen": event.get("center_screen"), "frame_size": event.get("frame_size"), "timestamp": event.get("timestamp"), } write_json_atomic(output_dir / "latest.json", payload) def ffmpeg_command(args): return [ "ffmpeg", "-hide_banner", "-loglevel", "error", "-f", "avfoundation", "-framerate", str(args.fps), "-i", args.input, "-f", "image2pipe", "-vcodec", "mjpeg", "-", ] def ffplay_command(args): return [ "ffplay", "-hide_banner", "-loglevel", "error", "-fflags", "nobuffer", "-flags", "low_delay", "-framedrop", "-sync", "ext", "-window_title", args.display_title, "-left", str(args.display_left), "-top", str(args.display_top), "-x", str(args.display_width), "-y", str(args.display_height), "-f", "mjpeg", "-i", "pipe:0", ] def start_ffplay(args): return subprocess.Popen( ffplay_command(args), stdin=subprocess.PIPE, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, bufsize=0, ) def write_display_frame(display_process, jpeg): if display_process is None or display_process.stdin is None: return False if display_process.poll() is not None: return False try: display_process.stdin.write(jpeg) display_process.stdin.flush() return True except (BrokenPipeError, OSError): return False def list_devices(): command = [ "ffmpeg", "-hide_banner", "-f", "avfoundation", "-list_devices", "true", "-i", "", ] subprocess.run(command, check=False) def read_jpeg_frames(stream): buffer = bytearray() while True: chunk = stream.read(65536) if not chunk: break buffer.extend(chunk) while True: start = buffer.find(b"\xff\xd8") if start < 0: if len(buffer) > 1024 * 1024: del buffer[:-2] break end = buffer.find(b"\xff\xd9", start + 2) if end < 0: if start > 0: del buffer[:start] break frame = bytes(buffer[start : end + 2]) del buffer[: end + 2] yield frame def run_monitor(args): model_path = Path(args.model) class_name = CLASS_NAMES.get(args.target_class, f"class_{args.target_class}") visualize_dir = Path(args.visualize_dir) if args.visualize_dir else None stream_state = StreamState() if args.serve else None stream_server = None display_process = None if stream_state: stream_server = start_stream_server(args.host, args.port, stream_state) url = f"http://{args.host}:{args.port}/" emit({"type": "stream_server_ready", "url": url}) if args.open_visualizer: webbrowser.open(url) if visualize_dir: write_visualizer_html(visualize_dir / "index.html") emit({"type": "visualizer_ready", "path": str(visualize_dir / "index.html")}) if args.open_visualizer and not args.serve: webbrowser.open((visualize_dir / "index.html").resolve().as_uri()) if args.display: display_process = start_ffplay(args) emit({"type": "display_started", "title": args.display_title}) session = ort.InferenceSession(str(model_path)) input_meta = session.get_inputs()[0] emit( { "type": "monitor_started", "timestamp": utc_now(), "model": str(model_path), "providers": session.get_providers(), "input": args.input, "fps": args.fps, "target_class": args.target_class, "target_name": class_name, "confidence": args.confidence, } ) process = subprocess.Popen( ffmpeg_command(args), stdout=subprocess.PIPE, stderr=subprocess.PIPE, bufsize=0, ) if process.stdout is None: raise RuntimeError("ffmpeg stdout pipe is not available") frame_index = 0 try: for frame_bytes in read_jpeg_frames(process.stdout): frame_index += 1 image = Image.open(BytesIO(frame_bytes)).convert("RGB") detections = detect( session, input_meta, image, args.target_class, args.confidence, ) if len(detections) > 0: best = detections[np.argmax(detections[:, 4])] event = detection_event(best, frame_index, image.size, class_name) emit(event) elif args.emit_empty: event = no_detection_event(frame_index, image.size, args.target_class, class_name) emit(event) else: event = no_detection_event(frame_index, image.size, args.target_class, class_name) if args.debug_dir and args.debug_every > 0 and frame_index % args.debug_every == 0: save_debug_image( image, detections, Path(args.debug_dir) / f"frame_{frame_index:06d}.jpg", class_name, ) if ( visualize_dir and args.visualize_every > 0 and frame_index % args.visualize_every == 0 ): update_visualizer(image, detections, event, visualize_dir, class_name) annotated = None if stream_state: annotated = draw_detections(image, detections, class_name) stream_state.update(image_to_jpeg_bytes(annotated), event) if display_process: if annotated is None: annotated = draw_detections(image, detections, class_name) annotated = draw_frame_info(annotated, event) if not write_display_frame(display_process, image_to_jpeg_bytes(annotated)): emit({"type": "display_closed", "timestamp": utc_now()}) break if args.max_frames > 0 and frame_index >= args.max_frames: break finally: process.terminate() try: process.wait(timeout=2) except subprocess.TimeoutExpired: process.kill() stderr = b"" if process.stderr is not None: try: stderr = process.stderr.read() except Exception: stderr = b"" if frame_index == 0 and stderr: log_error(stderr.decode(errors="replace").strip()) if display_process: if display_process.stdin: try: display_process.stdin.close() except Exception: pass display_process.terminate() try: display_process.wait(timeout=2) except subprocess.TimeoutExpired: display_process.kill() emit( { "type": "monitor_stopped", "timestamp": utc_now(), "frames": frame_index, } ) if stream_server: stream_server.shutdown() def main(): def stop_handler(signum, frame): raise KeyboardInterrupt signal.signal(signal.SIGTERM, stop_handler) signal.signal(signal.SIGINT, stop_handler) args = parse_args() if args.list_devices: list_devices() return if args.debug_every > 0 and not args.debug_dir: args.debug_dir = str(DEBUG_DIR) if args.open_visualizer and not args.visualize_dir and not args.serve: args.visualize_dir = str(VISUALIZE_DIR) try: run_monitor(args) except KeyboardInterrupt: emit({"type": "monitor_interrupted", "timestamp": utc_now()}) if __name__ == "__main__": main()