OliviaCandClaude Opus 5.5
1ad422062c
- Endpoints that wait for the camera or encode/compute (capture, rgb, depth, depth/color, pointcloud, pointcloud/roi, measure, snapshot, and the SDK calls discover/open/close) are plain `def`: FastAPI runs them in its thread pool, so one slow request no longer holds the others. Plain `def` rather than asyncio.to_thread in each route: it is FastAPI's own mechanism and keeps the route bodies unchanged. - /ws and the NATS snapshot stay async and use asyncio.to_thread for the frame wait, the encoding and the recording write. - EncodedFrame: the JPEG q85, 16-bit PNG, colormap JPEG q90, their base64 and the depth stats of a frame are computed once, on first use, and shared: parallel requests, /ws and NATS on the same frame get the same bytes. The last 4 frames' encodings are kept. - camera.cached_frame(): the whole cached frame (get_rgb/get_depth unchanged). Contract unchanged: same endpoints, same JSON fields. frame_id still counts responses (now under a lock); two parallel /api/capture served by one frame have different frame_id and the same timestamp and images. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
99 lines
3.3 KiB
Python
99 lines
3.3 KiB
Python
import asyncio
|
||
import logging
|
||
|
||
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
|
||
|
||
from src.backend.api.routes import _frame_to_dict
|
||
from src.backend.services.image_saver import ImageSaver
|
||
from src.backend.services.phase_timing import timer
|
||
|
||
try:
|
||
from src.backend.services.detector import Detector
|
||
except Exception:
|
||
Detector = None # onnxruntime is optional; YOLO features stay disabled
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
ws_router = APIRouter()
|
||
|
||
_camera = None
|
||
_saver: ImageSaver | None = None
|
||
_detector = None # Detector | None when onnxruntime is installed
|
||
|
||
|
||
def set_camera(camera):
|
||
global _camera
|
||
_camera = camera
|
||
|
||
|
||
def set_image_saver(saver: ImageSaver) -> None:
|
||
global _saver
|
||
_saver = saver
|
||
|
||
|
||
def set_detector(detector) -> None:
|
||
global _detector
|
||
_detector = detector
|
||
|
||
|
||
@ws_router.websocket("/ws")
|
||
async def websocket_stream(ws: WebSocket):
|
||
"""WebSocket stream of RGB+depth frames.
|
||
|
||
Client → server text messages:
|
||
'start' — begin streaming
|
||
'stop' — pause streaming
|
||
'close' — disconnect
|
||
'predict_on' — enable YOLO inference per frame (adds 'detections' to JSON)
|
||
'predict_off' — disable YOLO inference
|
||
'conf:0.55' — set YOLO confidence threshold (0.0–1.0)
|
||
"""
|
||
await ws.accept()
|
||
logger.info("WebSocket client connected")
|
||
streaming = False
|
||
predicting = False
|
||
|
||
try:
|
||
while True:
|
||
if streaming and _camera and _camera.is_open():
|
||
# Waiting for the frame, encoding and saving run in threads: the
|
||
# event loop keeps serving REST and the other sockets meanwhile.
|
||
frame = await asyncio.to_thread(_camera.capture, 500)
|
||
if frame:
|
||
data = await asyncio.to_thread(_frame_to_dict, frame)
|
||
if predicting and _detector is not None:
|
||
data["detections"] = await asyncio.to_thread(_detector.detect, frame.rgb)
|
||
with timer.phase("ws_send"):
|
||
await ws.send_json(data)
|
||
if _saver and _saver.is_recording:
|
||
await asyncio.to_thread(_saver.save_rec_frame, frame.rgb)
|
||
await asyncio.sleep(0.033)
|
||
|
||
try:
|
||
msg = await asyncio.wait_for(ws.receive_text(), timeout=0.01)
|
||
if msg == "start":
|
||
streaming = True
|
||
elif msg == "stop":
|
||
streaming = False
|
||
elif msg == "close":
|
||
break
|
||
elif msg == "predict_on":
|
||
predicting = _detector is not None
|
||
if not predicting:
|
||
logger.warning("predict_on requested but detector not loaded")
|
||
elif msg == "predict_off":
|
||
predicting = False
|
||
elif msg.startswith("conf:") and _detector is not None:
|
||
try:
|
||
_detector.set_conf(float(msg.split(":", 1)[1]))
|
||
except ValueError:
|
||
logger.warning(f"Invalid conf value: {msg!r}")
|
||
except asyncio.TimeoutError:
|
||
pass
|
||
|
||
if not streaming:
|
||
await asyncio.sleep(0.1)
|
||
|
||
except WebSocketDisconnect:
|
||
logger.info("WebSocket client disconnected")
|