899 lines
32 KiB
Python
899 lines
32 KiB
Python
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 = """<!doctype html>
|
|
<html lang="zh-CN">
|
|
<head>
|
|
<meta charset="utf-8" />
|
|
<meta name="viewport" content="width=device-width, initial-scale=1" />
|
|
<title>WeChat Vision Realtime</title>
|
|
<style>
|
|
:root { color-scheme: dark; font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; }
|
|
body { margin: 0; background: #09090b; color: #f4f4f5; }
|
|
.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; min-height: 60vh; }
|
|
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 id="frame" alt="latest detection frame" /></main>
|
|
<aside>
|
|
<h1>WeChat Vision Realtime</h1>
|
|
<div class="item"><div class="label">Status</div><div id="status" class="value">loading</div></div>
|
|
<div class="item"><div class="label">Frame</div><div id="frameIndex" class="value">-</div></div>
|
|
<div class="item"><div class="label">Class</div><div id="className" 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">Frame Size</div><div id="size" class="value">-</div></div>
|
|
<div class="item"><div class="label">Timestamp</div><div id="time" class="value">-</div></div>
|
|
</aside>
|
|
</div>
|
|
<script>
|
|
async function refresh() {
|
|
const cacheBust = Date.now();
|
|
document.getElementById('frame').src = 'latest.jpg?t=' + cacheBust;
|
|
try {
|
|
const response = await fetch('latest.json?t=' + cacheBust);
|
|
const data = await response.json();
|
|
const hit = data.event && data.event.type === 'detection';
|
|
const status = document.getElementById('status');
|
|
status.textContent = hit ? 'detection' : 'no detection';
|
|
status.className = 'value ' + (hit ? 'ok' : 'miss');
|
|
document.getElementById('frameIndex').textContent = data.frame_index ?? '-';
|
|
document.getElementById('className').textContent = data.class_name ?? '-';
|
|
document.getElementById('confidence').textContent = data.confidence ?? '-';
|
|
document.getElementById('center').textContent = JSON.stringify(data.center_screen ?? '-');
|
|
document.getElementById('bbox').textContent = JSON.stringify(data.bbox_screen ?? '-');
|
|
document.getElementById('size').textContent = JSON.stringify(data.frame_size ?? '-');
|
|
document.getElementById('time').textContent = data.timestamp ?? '-';
|
|
} catch (error) {
|
|
document.getElementById('status').textContent = 'waiting for latest.json';
|
|
}
|
|
}
|
|
refresh();
|
|
setInterval(refresh, 500);
|
|
</script>
|
|
</body>
|
|
</html>
|
|
"""
|
|
path.write_text(html, encoding="utf-8")
|
|
|
|
|
|
def dashboard_html(frame_url="frame.jpg", json_url="latest.json"):
|
|
return f"""<!doctype html>
|
|
<html lang="zh-CN">
|
|
<head>
|
|
<meta charset="utf-8" />
|
|
<meta name="viewport" content="width=device-width, initial-scale=1" />
|
|
<title>WeChat Vision Live Stream</title>
|
|
<style>
|
|
:root {{ color-scheme: dark; font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; }}
|
|
body {{ margin: 0; background: #09090b; color: #f4f4f5; }}
|
|
.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; min-height: 60vh; position: relative; }}
|
|
canvas {{ display: block; width: 100%; height: auto; }}
|
|
.overlay {{ position: absolute; left: 12px; top: 12px; padding: 6px 10px; border-radius: 999px; background: rgba(9, 9, 11, 0.78); color: #86efac; font: 12px ui-monospace, SFMono-Regular, Menlo, monospace; }}
|
|
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; }}
|
|
.pill {{ display: inline-flex; align-items: center; gap: 8px; padding: 6px 10px; border-radius: 999px; background: #052e16; color: #86efac; font-size: 12px; margin-bottom: 10px; }}
|
|
.dot {{ width: 8px; height: 8px; border-radius: 50%; background: #22c55e; box-shadow: 0 0 14px #22c55e; }}
|
|
.ok {{ color: #86efac; }}
|
|
.miss {{ color: #fca5a5; }}
|
|
@media (max-width: 900px) {{ .wrap {{ grid-template-columns: 1fr; }} }}
|
|
</style>
|
|
</head>
|
|
<body>
|
|
<div class="wrap">
|
|
<main class="stage">
|
|
<canvas id="preview"></canvas>
|
|
<div id="overlay" class="overlay">waiting for frames</div>
|
|
</main>
|
|
<aside>
|
|
<h1>WeChat Vision Live</h1>
|
|
<div class="pill"><span class="dot"></span><span>Canvas realtime preview</span></div>
|
|
<div class="item"><div class="label">Frame URL</div><div class="value">/{frame_url}</div></div>
|
|
<div class="item"><div class="label">Sequence</div><div id="sequence" class="value">-</div></div>
|
|
<div class="item"><div class="label">Metadata FPS</div><div id="metaFps" class="value">-</div></div>
|
|
<div class="item"><div class="label">Status</div><div id="status" class="value">loading</div></div>
|
|
<div class="item"><div class="label">Frame</div><div id="frameIndex" class="value">-</div></div>
|
|
<div class="item"><div class="label">Class</div><div id="className" 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">Frame Size</div><div id="size" class="value">-</div></div>
|
|
<div class="item"><div class="label">Timestamp</div><div id="time" class="value">-</div></div>
|
|
</aside>
|
|
</div>
|
|
<script>
|
|
const canvas = document.getElementById('preview');
|
|
const context = canvas.getContext('2d');
|
|
let previousSequence = null;
|
|
let previousTime = null;
|
|
let drawnSequence = null;
|
|
|
|
async function drawFrame(sequence) {{
|
|
const response = await fetch('{frame_url}?seq=' + sequence + '&t=' + Date.now(), {{ cache: 'no-store' }});
|
|
if (!response.ok) return;
|
|
const blob = await response.blob();
|
|
const bitmap = await createImageBitmap(blob);
|
|
if (canvas.width !== bitmap.width || canvas.height !== bitmap.height) {{
|
|
canvas.width = bitmap.width;
|
|
canvas.height = bitmap.height;
|
|
}}
|
|
context.drawImage(bitmap, 0, 0);
|
|
bitmap.close();
|
|
drawnSequence = sequence;
|
|
}}
|
|
|
|
async function refreshMeta() {{
|
|
try {{
|
|
const response = await fetch('{json_url}?t=' + Date.now(), {{ cache: 'no-store' }});
|
|
const data = await response.json();
|
|
const event = data.event || data;
|
|
const now = performance.now();
|
|
if (previousSequence !== null && previousTime !== null && data.sequence !== undefined) {{
|
|
const elapsed = Math.max((now - previousTime) / 1000, 0.001);
|
|
const fps = (data.sequence - previousSequence) / elapsed;
|
|
document.getElementById('metaFps').textContent = fps.toFixed(1);
|
|
}}
|
|
previousSequence = data.sequence;
|
|
previousTime = now;
|
|
const hit = event.type === 'detection';
|
|
const status = document.getElementById('status');
|
|
status.textContent = hit ? 'detection' : event.type || 'waiting';
|
|
status.className = 'value ' + (hit ? 'ok' : 'miss');
|
|
document.getElementById('sequence').textContent = data.sequence ?? '-';
|
|
document.getElementById('frameIndex').textContent = event.frame_index ?? '-';
|
|
document.getElementById('className').textContent = event.class_name ?? '-';
|
|
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('size').textContent = JSON.stringify(event.frame_size ?? '-');
|
|
document.getElementById('time').textContent = event.timestamp ?? '-';
|
|
document.getElementById('overlay').textContent = `seq=${{data.sequence ?? '-'}} frame=${{event.frame_index ?? '-'}} ${{hit ? 'detection' : event.type || 'waiting'}}`;
|
|
if (data.sequence !== undefined && data.sequence > 0 && data.sequence !== drawnSequence) {{
|
|
await drawFrame(data.sequence);
|
|
}}
|
|
}} catch (error) {{
|
|
document.getElementById('status').textContent = 'waiting for preview';
|
|
}}
|
|
}}
|
|
refreshMeta();
|
|
setInterval(refreshMeta, 200);
|
|
</script>
|
|
</body>
|
|
</html>
|
|
"""
|
|
|
|
|
|
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()
|