Initial commit: RealSense D435 API server

Drop-in replacement for the Nuwa-HP60C camera container, exposing RGB +
aligned depth + intrinsics via REST/WebSocket so the existing
pestello-pose-detector pipeline can consume frames without changes.

Key adaptations for D4xx on USB 2.0 (this robot has no USB 3 port exposed
externally — see STATUS.md):
- pyrealsense2 built from source; the librealsense v2.55.1 install rule
  doesn't ship the main .so, so we copy it from the build dir and add an
  __init__.py that re-exports the submodule.
- Color stream uses YUYV format on USB 2.x (BGR8/RGB8 time out on this
  bus); converted to BGR via cv2.cvtColor in capture().
- rs.align is skipped on USB 2 (incompatible with YUYV); color is
  cv2-resized to depth shape for approximate per-pixel alignment.
- Driver auto-detects usb_type_descriptor and falls back to the
  full-resolution BGR8 + rs.align pipeline when on USB 3.

Tested: /api/{health,status,discover,rgb,depth,depth/color,capture,
intrinsics} all return 200 with valid payloads at 640x480.
This commit is contained in:
robot committed 2026-05-19 00:58:23 +02:00
commit 07d95792e3
58 files changed
+6987

No files matched your search

View File
Whitespace-only changes.
View File
Whitespace-only changes.
+9
View File
@@ -0,0 +1,9 @@
from fastapi import Header, HTTPException
from src.backend.config import settings
async def verify_api_key(x_api_key: str | None = Header(None)):
if not settings.api_key:
return
if x_api_key != settings.api_key:
raise HTTPException(status_code=403, detail="API key non valida")
+537
View File
@@ -0,0 +1,537 @@
"""REST API routes for the RealSense D435 camera server."""
import base64
import logging
import cv2
import numpy as np
from fastapi import APIRouter, Depends, HTTPException, Query
from fastapi.responses import Response, JSONResponse
from src.backend.api.dependencies import verify_api_key
from src.backend.models.api import ok, fail
from src.backend.services.camera import RealSenseD435, CameraFrame
from src.backend.services.image_saver import ImageSaver
logger = logging.getLogger(__name__)
router = APIRouter(dependencies=[Depends(verify_api_key)])
_camera: RealSenseD435 | None = None
_saver: ImageSaver | None = None
_frame_count = 0
def set_camera(camera: RealSenseD435) -> None:
global _camera
_camera = camera
def set_image_saver(saver: ImageSaver) -> None:
global _saver
_saver = saver
# ---------------------------------------------------------------------------
# Image encoding helpers
# ---------------------------------------------------------------------------
def _encode_rgb_jpeg(rgb: np.ndarray, quality: int = 85) -> bytes:
_, buf = cv2.imencode(".jpg", rgb, [cv2.IMWRITE_JPEG_QUALITY, quality])
return buf.tobytes()
def _encode_depth_png(depth: np.ndarray) -> bytes:
_, buf = cv2.imencode(".png", depth)
return buf.tobytes()
def _depth_colormap(depth: np.ndarray) -> np.ndarray:
valid = depth[depth > 0]
if len(valid) == 0:
return np.zeros((*depth.shape, 3), dtype=np.uint8)
p5, p95 = np.percentile(valid, [5, 95])
vis = np.zeros_like(depth, dtype=np.uint8)
mask = depth > 0
vis[mask] = np.clip(
(depth[mask].astype(float) - p5) / max(p95 - p5, 1) * 255, 0, 255
).astype(np.uint8)
return cv2.applyColorMap(vis, cv2.COLORMAP_JET)
def _frame_to_dict(frame: CameraFrame, include_images: bool = True) -> dict:
global _frame_count
_frame_count += 1
result = {
"frame_id": _frame_count,
"timestamp": frame.timestamp,
"rgb_shape": list(frame.rgb.shape),
"depth_shape": list(frame.depth.shape),
"depth_min_mm": int(frame.depth[frame.depth > 0].min()) if (frame.depth > 0).any() else 0,
"depth_max_mm": int(frame.depth.max()),
"depth_valid_pixels": int((frame.depth > 0).sum()),
}
if include_images:
result["rgb_jpeg_b64"] = base64.b64encode(_encode_rgb_jpeg(frame.rgb)).decode()
result["depth_png_b64"] = base64.b64encode(_encode_depth_png(frame.depth)).decode()
colormap = _depth_colormap(frame.depth)
result["depth_color_b64"] = base64.b64encode(_encode_rgb_jpeg(colormap, 90)).decode()
return result
# ---------------------------------------------------------------------------
# Endpoints
# ---------------------------------------------------------------------------
# /health is mounted as public (no API key) in main.py
def _calibration_payload() -> dict | None:
"""Full calibration (intrinsics + distortion + extrinsics) or None."""
if not _camera or not _camera.intrinsics:
return None
i = _camera.intrinsics
e = _camera.extrinsics
dir_ = _camera.distortion_ir
drgb = _camera.distortion_rgb
def _d(d):
return None if d is None else {"k1": d.k1, "k2": d.k2, "k3": d.k3, "p1": d.p1, "p2": d.p2}
return {
"ir": {
"fx": i.fx_ir, "fy": i.fy_ir, "cx": i.cx_ir, "cy": i.cy_ir,
"distortion": _d(dir_),
},
"rgb": {
"fx": i.fx_rgb, "fy": i.fy_rgb, "cx": i.cx_rgb, "cy": i.cy_rgb,
"distortion": _d(drgb),
},
"extrinsics_ir_to_rgb": None if e is None else {
"rotation": list(e.rotation),
"translation": list(e.translation),
},
}
@router.get("/status")
async def status():
"""Camera status and intrinsics."""
return ok({
"sdk_initialized": _camera.is_initialized() if _camera else False,
"camera_open": _camera.is_open() if _camera else False,
"color_available": _camera.is_color_available() if (_camera and _camera.is_open()) else False,
"usb_type": _camera.usb_type() if (_camera and _camera.is_open()) else None,
"frame_count": _frame_count,
"resolution": {"width": 640, "height": 480},
"intrinsics": _calibration_payload(),
})
@router.get("/discover")
async def discover(timeout_ms: int = Query(default=5000, ge=1000, le=30000)):
"""Discover available RealSense cameras on USB bus."""
if not _camera or not _camera.is_initialized():
return fail("SDK_NOT_INITIALIZED", "SDK not initialized")
devices = _camera.discover(timeout_ms=timeout_ms)
return ok({
"cameras": [
{
"index": d.index,
"vid": f"{d.vid:04x}",
"pid": f"{d.pid:04x}",
"serial": d.serial,
"bus": d.bus,
"device": d.device,
}
for d in devices
],
"count": len(devices),
})
@router.post("/open")
async def open_camera(device_index: int = Query(default=0, ge=0)):
"""Open a discovered camera by index and start streaming.
Idempotent: if a camera is already open, returns the current open state
instead of failing — the server only supports one open camera at a time.
"""
if not _camera or not _camera.is_initialized():
return fail("SDK_NOT_INITIALIZED", "SDK not initialized")
if _camera.is_open():
return ok({
"status": "open",
"device_index": _camera.device_index,
"intrinsics": _calibration_payload(),
"already_open": True,
})
if _camera.open(device_index):
return ok({
"status": "open",
"device_index": device_index,
"intrinsics": _calibration_payload(),
"already_open": False,
})
else:
return fail("OPEN_FAILED", "Failed to open camera")
@router.post("/close")
async def close_camera():
"""Close the active camera and stop streaming."""
if not _camera or not _camera.is_open():
return fail("CAMERA_NOT_OPEN", "No camera is open")
_camera.close()
return ok({"status": "closed"})
@router.get("/capture")
async def capture():
"""Capture a single RGB+depth frame. Returns JSON with base64-encoded images."""
if not _camera or not _camera.is_open():
return fail("CAMERA_NOT_OPEN", "Camera not open")
frame = _camera.capture()
if frame is None:
return fail("CAPTURE_TIMEOUT", "Frame capture timeout")
return JSONResponse(ok(_frame_to_dict(frame)))
@router.get("/rgb")
async def get_rgb():
"""Get RGB frame as JPEG image."""
if not _camera or not _camera.is_open():
raise HTTPException(503, "Camera not open")
rgb = _camera.get_rgb()
if rgb is None:
raise HTTPException(504, "No RGB frame")
return Response(content=_encode_rgb_jpeg(rgb), media_type="image/jpeg")
@router.get("/depth")
async def get_depth():
"""Get depth frame as 16-bit PNG (values in mm)."""
if not _camera or not _camera.is_open():
raise HTTPException(503, "Camera not open")
depth = _camera.get_depth()
if depth is None:
raise HTTPException(504, "No depth frame")
return Response(content=_encode_depth_png(depth), media_type="image/png")
@router.get("/depth/color")
async def get_depth_colormap():
"""Get depth frame as JET colormap JPEG."""
if not _camera or not _camera.is_open():
raise HTTPException(503, "Camera not open")
depth = _camera.get_depth()
if depth is None:
raise HTTPException(504, "No depth frame")
colormap = _depth_colormap(depth)
return Response(content=_encode_rgb_jpeg(colormap, 90), media_type="image/jpeg")
@router.get("/pointcloud", summary="3D Point cloud", tags=["capture"])
async def get_pointcloud(
step: int = Query(default=4, ge=1, le=16),
max_depth_mm: int = Query(default=4000, ge=100, le=10000),
):
"""Get 3D point cloud as JSON array of [x,y,z,r,g,b] points.
Points are computed server-side using depth + RGB + intrinsics.
step: subsampling factor (4 = every 4th pixel, faster)
max_depth_mm: discard points beyond this depth
"""
if not _camera or not _camera.is_open():
return fail("NO_CAMERA", "Camera not open")
frame = _camera.capture()
if frame is None:
return fail("TIMEOUT", "Frame capture timeout")
depth = frame.depth
rgb = frame.rgb # BGR format
# Use intrinsics if available, otherwise use typical HP60C defaults
if _camera.intrinsics:
fx = _camera.intrinsics.fx_ir
fy = _camera.intrinsics.fy_ir
cx = _camera.intrinsics.cx_ir
cy = _camera.intrinsics.cy_ir
else:
# Default intrinsics for 640x480 HP60C
fx, fy = 445.0, 445.0
cx, cy = 320.0, 240.0
h, w = depth.shape
points = []
for y in range(0, h, step):
for x in range(0, w, step):
z_val = float(depth[y, x])
if z_val <= 0 or z_val > max_depth_mm:
continue
z_m = z_val / 1000.0
x_m = (x - cx) * z_m / fx
y_m = (y - cy) * z_m / fy
b_val, g_val, r_val = rgb[y, x] # BGR to RGB
points.append([
round(x_m, 4), round(-y_m, 4), round(-z_m, 4),
int(r_val), int(g_val), int(b_val)
])
return ok({
"points": points,
"count": len(points),
"step": step,
"intrinsics": {
"fx": fx, "fy": fy, "cx": cx, "cy": cy
}
})
@router.get("/pointcloud/roi", summary="3D Point cloud within ROI", tags=["capture"])
async def get_pointcloud_roi(
u_min: int = Query(..., ge=0, description="ROI left pixel (inclusive)"),
v_min: int = Query(..., ge=0, description="ROI top pixel (inclusive)"),
u_max: int = Query(..., ge=1, description="ROI right pixel (exclusive)"),
v_max: int = Query(..., ge=1, description="ROI bottom pixel (exclusive)"),
step: int = Query(default=1, ge=1, le=16),
max_depth_mm: int = Query(default=4000, ge=100, le=10000),
):
"""Get 3D point cloud restricted to a pixel-space ROI.
Coordinates are RAW image pixels (no flip applied) in the depth/RGB
frame (typically 640x480). Points are returned as [x, -y, -z, r, g, b]
— same convention as /pointcloud so the frontend renderer is compatible.
"""
if not _camera or not _camera.is_open():
return fail("NO_CAMERA", "Camera not open")
frame = _camera.capture()
if frame is None:
return fail("TIMEOUT", "Frame capture timeout")
depth = frame.depth
rgb = frame.rgb # BGR format
h, w = depth.shape
u0 = max(0, min(w, u_min))
v0 = max(0, min(h, v_min))
u1 = max(0, min(w, u_max))
v1 = max(0, min(h, v_max))
if u1 <= u0 or v1 <= v0:
return fail("INVALID_ROI", f"Empty ROI ({u0},{v0})-({u1},{v1}) in {w}x{h}")
if _camera.intrinsics:
fx = _camera.intrinsics.fx_ir
fy = _camera.intrinsics.fy_ir
cx = _camera.intrinsics.cx_ir
cy = _camera.intrinsics.cy_ir
else:
fx, fy = 445.0, 445.0
cx, cy = 320.0, 240.0
# Vectorized: restrict both arrays to the ROI and filter by depth mask.
roi_depth = depth[v0:v1:step, u0:u1:step].astype(np.float32)
roi_rgb = rgb[v0:v1:step, u0:u1:step]
mask = (roi_depth > 0) & (roi_depth <= max_depth_mm)
if not mask.any():
return ok({
"points": [],
"count": 0,
"step": step,
"roi": {"u_min": u0, "v_min": v0, "u_max": u1, "v_max": v1},
"intrinsics": {"fx": fx, "fy": fy, "cx": cx, "cy": cy},
})
ys_local, xs_local = np.where(mask)
zs_mm = roi_depth[ys_local, xs_local]
# Absolute pixel coordinates within the full image (needed for fx/cx back-projection).
xs_abs = u0 + xs_local * step
ys_abs = v0 + ys_local * step
zs_m = zs_mm / 1000.0
xs_m = (xs_abs - cx) * zs_m / fx
ys_m = (ys_abs - cy) * zs_m / fy
bgr = roi_rgb[ys_local, xs_local] # (N, 3) BGR
# Match /pointcloud convention: [x, -y, -z, r, g, b]
points = np.stack([
np.round(xs_m, 4),
np.round(-ys_m, 4),
np.round(-zs_m, 4),
bgr[:, 2].astype(np.float32),
bgr[:, 1].astype(np.float32),
bgr[:, 0].astype(np.float32),
], axis=1).tolist()
return ok({
"points": points,
"count": len(points),
"step": step,
"roi": {"u_min": u0, "v_min": v0, "u_max": u1, "v_max": v1},
"intrinsics": {"fx": fx, "fy": fy, "cx": cx, "cy": cy},
})
@router.get("/measure", summary="3D measurement at pixel", tags=["capture"])
async def measure_point(
u: int = Query(..., ge=0, description="Pixel X in raw image coords (0..W-1)"),
v: int = Query(..., ge=0, description="Pixel Y in raw image coords (0..H-1)"),
radius: int = Query(default=6, ge=0, le=50, description="Search radius (px) for nearest valid depth"),
):
"""Find nearest valid-depth pixel around (u, v) and back-project to 3D.
Useful for testing RGB/depth alignment: click on RGB at a pixel, get the
closest point-cloud point within `radius`. Coordinates are RAW image pixels
(no flip applied) in the depth/RGB frame (typically 640x480).
"""
if not _camera or not _camera.is_open():
return fail("CAMERA_NOT_OPEN", "Camera not open")
frame = _camera.capture()
if frame is None:
return fail("CAPTURE_TIMEOUT", "Frame capture timeout")
depth = frame.depth
h, w = depth.shape
if u >= w or v >= h:
return fail("OUT_OF_RANGE", f"Pixel ({u},{v}) outside image ({w}x{h})")
u0 = max(0, u - radius)
u1 = min(w, u + radius + 1)
v0 = max(0, v - radius)
v1 = min(h, v + radius + 1)
window = depth[v0:v1, u0:u1]
ys, xs = np.where(window > 0)
if len(xs) == 0:
return fail("NO_DEPTH", f"No valid depth within radius {radius}")
abs_xs = xs + u0
abs_ys = ys + v0
d2 = (abs_xs - u) ** 2 + (abs_ys - v) ** 2
idx = int(np.argmin(d2))
u_found = int(abs_xs[idx])
v_found = int(abs_ys[idx])
z_mm = int(depth[v_found, u_found])
if _camera.intrinsics:
fx = _camera.intrinsics.fx_ir
fy = _camera.intrinsics.fy_ir
cx = _camera.intrinsics.cx_ir
cy = _camera.intrinsics.cy_ir
else:
fx, fy = 445.0, 445.0
cx, cy = 320.0, 240.0
z_m = z_mm / 1000.0
x_m = (u_found - cx) * z_m / fx
y_m = (v_found - cy) * z_m / fy
dist_mm = float(np.sqrt(x_m * x_m + y_m * y_m + z_m * z_m) * 1000.0)
return ok({
"u": u,
"v": v,
"u_found": u_found,
"v_found": v_found,
"pixel_offset": round(float(np.sqrt(float(d2[idx]))), 2),
"radius": radius,
"depth_mm": z_mm,
"distance_mm": round(dist_mm, 1),
"point_m": {
"x": round(x_m, 4),
"y": round(y_m, 4),
"z": round(z_m, 4),
},
})
@router.post("/snapshot")
async def snapshot(
flip_h: bool = Query(default=False),
flip_v: bool = Query(default=False),
):
"""Capture a single RGB frame and save it as JPG in images_dir."""
if not _camera or not _camera.is_open():
return fail("CAMERA_NOT_OPEN", "Camera not open")
if _saver is None:
return fail("SAVER_NOT_READY", "Image saver not initialized")
frame = _camera.capture()
if frame is None:
return fail("CAPTURE_TIMEOUT", "Frame capture timeout")
filename = _saver.save_snapshot(frame.rgb, flip_h, flip_v)
return ok({"filename": filename})
@router.post("/recording/start")
async def recording_start(
flip_h: bool = Query(default=False),
flip_v: bool = Query(default=False),
interval: int = Query(default=1, ge=1, le=600),
):
"""Start continuous frame recording (applies to WebSocket stream loop).
interval: save 1 frame every N streamed frames (1 = every frame).
"""
if _saver is None:
return fail("SAVER_NOT_READY", "Image saver not initialized")
if _saver.is_recording:
return fail("ALREADY_RECORDING", "Recording already in progress")
status = _saver.start_recording(flip_h, flip_v, interval)
return ok(status)
@router.post("/recording/stop")
async def recording_stop():
"""Stop continuous frame recording."""
if _saver is None:
return fail("SAVER_NOT_READY", "Image saver not initialized")
if not _saver.is_recording:
return fail("NOT_RECORDING", "No recording in progress")
status = _saver.stop_recording()
return ok(status)
@router.get("/recording/status")
async def recording_status():
"""Current recording state."""
if _saver is None:
return fail("SAVER_NOT_READY", "Image saver not initialized")
return ok(_saver.status())
@router.get("/intrinsics")
async def get_intrinsics():
"""Get camera intrinsic parameters."""
if not _camera:
return fail("INTRINSICS_NOT_AVAILABLE", "Camera driver not ready")
if not _camera.is_open():
return fail("INTRINSICS_NOT_AVAILABLE", "Camera not open")
if not _camera.intrinsics:
return fail(
"INTRINSICS_NOT_AVAILABLE",
"Camera open but SDK did not return calibration (AS_SDK_GetCamParameter). "
"Check docker logs for 'GetCamParameter ... failed'.",
)
return ok(_calibration_payload())
+94
View File
@@ -0,0 +1,94 @@
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
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():
frame = _camera.capture(timeout_ms=500)
if frame:
data = _frame_to_dict(frame)
if predicting and _detector is not None:
data["detections"] = await asyncio.to_thread(_detector.detect, frame.rgb)
await ws.send_json(data)
if _saver and _saver.is_recording:
_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")
+51
View File
@@ -0,0 +1,51 @@
from pydantic_settings import BaseSettings
class Settings(BaseSettings):
# Server
api_host: str = "0.0.0.0"
api_port: int = 8878
# NATS
nats_url: str = "nats://nats:4222"
nats_token: str = ""
# Service identity
service_name: str = "realsense-d435"
service_domain: str = "vision"
# Security
api_key: str = ""
# Storage
images_dir: str = "/app/images"
# RealSense stream config
camera_width: int = 640
camera_height: int = 480
camera_fps: int = 30
camera_auto_open: bool = True
# YOLO detector (optional — disabled if model file or onnxruntime missing)
yolo_onnx_path: str = "/app/models/best.onnx"
yolo_classes: str = ""
yolo_conf_default: float = 0.5
@property
def yolo_class_list(self) -> list[str]:
return [s.strip() for s in self.yolo_classes.split(",") if s.strip()]
# CORS (comma-separated string, parsed into list)
cors_origins: str = "http://localhost:5173,http://localhost,http://localhost:80"
@property
def cors_origins_list(self) -> list[str]:
return [s.strip() for s in self.cors_origins.split(",")]
model_config = {
"env_file": ".env",
"env_file_encoding": "utf-8",
}
settings = Settings()
+147
View File
@@ -0,0 +1,147 @@
import logging
from contextlib import asynccontextmanager
from pathlib import Path
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from src.backend.config import settings
from src.backend.api.routes import (
router,
set_camera as set_routes_camera,
set_image_saver as set_routes_saver,
)
from src.backend.api.websocket import (
ws_router,
set_camera as set_ws_camera,
set_image_saver as set_ws_saver,
set_detector as set_ws_detector,
)
from src.backend.nats_client.connection import nats_manager
from src.backend.nats_client.handlers import set_camera as set_nats_camera
from src.backend.services.camera import RealSenseD435
from src.backend.services.image_saver import ImageSaver
try:
from src.backend.services.detector import Detector
except Exception:
Detector = None # YOLO/onnxruntime optional
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(levelname)s %(message)s")
logger = logging.getLogger(__name__)
@asynccontextmanager
async def lifespan(app: FastAPI):
# Startup
camera = RealSenseD435(
width=settings.camera_width,
height=settings.camera_height,
fps=settings.camera_fps,
)
sdk_ok = camera.init()
if sdk_ok:
logger.info("SDK initialized")
else:
logger.error(
"SDK init failed — API will stay up so /api/health can be queried, "
"but camera endpoints will return SDK_NOT_INITIALIZED"
)
if sdk_ok and settings.camera_auto_open:
try:
devices = camera.discover()
if devices:
if camera.open(0):
logger.info(
f"Auto-opened camera 0 (serial={devices[0].serial})"
)
else:
logger.warning("Auto-open requested but open(0) failed")
else:
logger.warning("Auto-open requested but no RealSense devices found")
except Exception as exc:
logger.warning(f"Auto-open failed: {exc}")
# Inject camera into all layers
set_routes_camera(camera)
set_ws_camera(camera)
set_nats_camera(camera)
# Image saver (snapshot + continuous recording to images_dir)
saver = ImageSaver(Path(settings.images_dir))
set_routes_saver(saver)
set_ws_saver(saver)
logger.info(f"Image saver ready: {settings.images_dir}")
# YOLO detector (optional — disabled if onnxruntime or model file is missing)
if Detector is not None:
try:
detector = Detector(
settings.yolo_onnx_path,
settings.yolo_class_list,
settings.yolo_conf_default,
)
set_ws_detector(detector)
except FileNotFoundError as exc:
logger.warning(f"YOLO detector disabled (model not found): {exc}")
except Exception as exc:
logger.warning(f"YOLO detector disabled: {exc}")
else:
logger.info("YOLO detector disabled (onnxruntime not installed)")
# Connect NATS (non-blocking, graceful if unavailable)
await nats_manager.connect()
yield
# Shutdown
await nats_manager.disconnect()
camera.close()
camera.deinit()
logger.info("SDK deinitialized")
app = FastAPI(
title=settings.service_name,
description="REST API + WebSocket for Intel RealSense D4xx depth camera (pyrealsense2)",
version="1.0.0",
docs_url="/apidocs",
redoc_url="/redoc",
openapi_url="/openapi.json",
lifespan=lifespan,
)
app.add_middleware(
CORSMiddleware,
allow_origins=settings.cors_origins_list,
allow_methods=["*"],
allow_headers=["*"],
)
app.include_router(router, prefix="/api")
app.include_router(ws_router)
# Serve React frontend at /gui (built files in src/frontend/react/dist)
_GUI_DIR = Path(__file__).parent.parent / "frontend" / "react" / "dist"
if _GUI_DIR.is_dir():
app.mount("/gui", StaticFiles(directory=str(_GUI_DIR), html=True), name="gui")
logger.info(f"React GUI mounted at /gui from {_GUI_DIR}")
# Public health endpoint (no API key required) for Docker healthchecks
@app.get("/api/health", summary="Health check", tags=["system"])
async def health():
"""Service health check — no API key required."""
from src.backend.models.api import ok
from src.backend.nats_client.connection import nats_manager
from src.backend.api.routes import _camera
return ok({
"status": "ok",
"sdk_initialized": _camera.is_initialized() if _camera else False,
"camera_open": _camera.is_open() if _camera else False,
"nats_connected": nats_manager.is_connected,
})
View File
Whitespace-only changes.
+54
View File
@@ -0,0 +1,54 @@
from pydantic import BaseModel
from typing import Any
class ApiError(BaseModel):
code: str
message: str
class ApiResponse(BaseModel):
success: bool
data: Any | None = None
error: ApiError | None = None
def ok(data: Any) -> dict:
return {"success": True, "data": data, "error": None}
def fail(code: str, message: str) -> dict:
return {"success": False, "data": None, "error": {"code": code, "message": message}}
class DeviceInfo(BaseModel):
index: int
vid: str
pid: str
serial: str
bus: int
device: int
class IntrinsicsInfo(BaseModel):
fx: float
fy: float
cx: float
cy: float
class CameraIntrinsicsInfo(BaseModel):
ir: IntrinsicsInfo
rgb: IntrinsicsInfo
class FrameInfo(BaseModel):
frame_id: int
timestamp: float
rgb_shape: list[int]
depth_shape: list[int]
depth_min_mm: int
depth_max_mm: int
depth_valid_pixels: int
rgb_jpeg_b64: str | None = None
depth_png_b64: str | None = None
+26
View File
@@ -0,0 +1,26 @@
from pydantic import BaseModel
class NatsResponse(BaseModel):
success: bool
data: dict | None = None
error: dict | None = None
class SnapshotRequest(BaseModel):
camera_id: str = "default"
include_depth: bool = True
class SnapshotResponse(BaseModel):
rgb_jpeg_b64: str
depth_png_b64: str | None = None
timestamp: float
width: int
height: int
class CameraStatusMessage(BaseModel):
camera_open: bool
sdk_initialized: bool
frame_count: int
View File
Whitespace-only changes.
+58
View File
@@ -0,0 +1,58 @@
import logging
import nats
from nats.aio.client import Client
from src.backend.config import settings
logger = logging.getLogger(__name__)
class NatsManager:
def __init__(self):
self._nc: Client | None = None
@property
def is_connected(self) -> bool:
return self._nc is not None and self._nc.is_connected
@property
def client(self) -> Client:
if not self.is_connected:
raise RuntimeError("NATS non connesso")
return self._nc
async def connect(self):
if not settings.nats_url:
logger.info("NATS URL not configured, skipping connection")
return
try:
self._nc = await nats.connect(
settings.nats_url,
token=settings.nats_token if settings.nats_token else None,
name=settings.service_name,
reconnect_time_wait=2,
max_reconnect_attempts=-1,
)
logger.info(f"Connected to NATS at {settings.nats_url}")
await self._register_handlers()
await self._publish_discovery()
except Exception as e:
logger.warning(f"NATS connection failed: {e}. Running without NATS.")
self._nc = None
async def disconnect(self):
if self.is_connected:
await self._nc.drain()
self._nc = None
logger.info("Disconnected from NATS")
async def _register_handlers(self):
from src.backend.nats_client.handlers import register_all
await register_all(self._nc)
async def _publish_discovery(self):
from src.backend.nats_client.discovery import publish_discovery
await publish_discovery(self._nc)
nats_manager = NatsManager()
+52
View File
@@ -0,0 +1,52 @@
import json
from nats.aio.client import Client
from src.backend.config import settings
from src.backend.nats_client.subjects import Subjects
CAPABILITIES = {
"service": settings.service_name,
"domain": settings.service_domain,
"version": "1.0.0",
"description": "Intel RealSense D4xx stereo depth camera server",
"endpoints": {
"rest": {
"health": {"method": "GET", "path": "/api/health"},
"discover": {"method": "GET", "path": "/api/discover"},
"open": {"method": "POST", "path": "/api/open"},
"close": {"method": "POST", "path": "/api/close"},
"capture": {"method": "GET", "path": "/api/capture"},
"rgb": {"method": "GET", "path": "/api/rgb"},
"depth": {"method": "GET", "path": "/api/depth"},
},
},
"nats": {
"actions": [
{
"subject": Subjects.action("snapshot"),
"pattern": "request/reply",
"description": "Capture RGB+depth frame",
"request_schema": {"include_depth": "boolean"},
"response_schema": {"rgb_jpeg_b64": "string", "depth_png_b64": "string", "timestamp": "number"},
},
{
"subject": Subjects.action("status"),
"pattern": "request/reply",
"description": "Get camera status",
"response_schema": {"camera_open": "boolean", "sdk_initialized": "boolean"},
},
],
},
}
async def publish_discovery(nc: Client):
await nc.publish(
Subjects.DISCOVERY_ANNOUNCE,
json.dumps(CAPABILITIES).encode(),
)
async def handle_discovery_request(msg):
await msg.respond(json.dumps(CAPABILITIES).encode())
await nc.subscribe(Subjects.DISCOVERY_REQUEST, cb=handle_discovery_request)
await nc.subscribe(Subjects.SERVICE_DISCOVERY, cb=handle_discovery_request)
+66
View File
@@ -0,0 +1,66 @@
import base64
import json
import logging
from nats.aio.client import Client
from src.backend.nats_client.subjects import Subjects
logger = logging.getLogger(__name__)
_camera = None
def set_camera(camera):
global _camera
_camera = camera
def _nats_response(success: bool, data=None, error=None) -> bytes:
return json.dumps({
"success": success,
"data": data,
"error": error,
}).encode()
async def register_all(nc: Client):
async def handle_snapshot(msg):
try:
if not _camera or not _camera.is_open():
await msg.respond(_nats_response(
False, error={"code": "NO_CAMERA", "message": "Camera not open"}
))
return
import cv2
frame = _camera.capture()
if frame is None:
await msg.respond(_nats_response(
False, error={"code": "TIMEOUT", "message": "Frame capture timeout"}
))
return
_, rgb_buf = cv2.imencode(".jpg", frame.rgb, [cv2.IMWRITE_JPEG_QUALITY, 85])
_, depth_buf = cv2.imencode(".png", frame.depth)
await msg.respond(_nats_response(True, data={
"rgb_jpeg_b64": base64.b64encode(rgb_buf.tobytes()).decode(),
"depth_png_b64": base64.b64encode(depth_buf.tobytes()).decode(),
"timestamp": frame.timestamp,
"width": frame.rgb.shape[1],
"height": frame.rgb.shape[0],
}))
except Exception as e:
logger.error(f"Snapshot handler error: {e}")
await msg.respond(_nats_response(
False, error={"code": "PROCESSING_ERROR", "message": str(e)}
))
async def handle_status(msg):
await msg.respond(_nats_response(True, data={
"camera_open": _camera.is_open() if _camera else False,
"sdk_initialized": _camera.is_initialized() if _camera else False,
}))
await nc.subscribe(Subjects.action("snapshot"), cb=handle_snapshot)
await nc.subscribe(Subjects.action("status"), cb=handle_status)
+18
View File
@@ -0,0 +1,18 @@
from src.backend.config import settings
DOMAIN = settings.service_domain
SERVICE = settings.service_name
class Subjects:
DISCOVERY_REQUEST = "_discovery.request"
DISCOVERY_ANNOUNCE = "_discovery.announce"
SERVICE_DISCOVERY = f"_discovery.{SERVICE}"
@staticmethod
def action(action: str) -> str:
return f"{DOMAIN}.{SERVICE}.{action}"
@staticmethod
def action_reply(action: str) -> str:
return f"{DOMAIN}.{SERVICE}.{action}.reply"
View File
Whitespace-only changes.
+413
View File
@@ -0,0 +1,413 @@
"""Intel RealSense D4xx depth camera driver via pyrealsense2.
Provides RGB (BGR8) and aligned Depth (uint16 mm) frames plus
intrinsics/extrinsics through the same public surface as the
Nuwa HP60C driver, so the routes/websocket/frontend stay unchanged.
In this driver the "ir_*" intrinsic fields hold the **depth** sensor
intrinsics (the D4xx depth stream is rectified from the IR stereo pair),
and "rgb_*" hold the **color** sensor intrinsics. After alignment both
streams share the color intrinsics, but we still expose the raw depth
intrinsics so client code that wants to back-project the unaligned depth
can do so.
Usage:
cam = RealSenseD435()
cam.init()
devs = cam.discover()
cam.open(0)
frame = cam.capture()
cam.close()
cam.deinit()
"""
import logging
import time
from dataclasses import dataclass
from typing import Optional
import cv2
import numpy as np
import pyrealsense2 as rs
logger = logging.getLogger(__name__)
@dataclass
class CameraFrame:
"""Container for a single capture: aligned color + depth + timestamp."""
rgb: np.ndarray # (H, W, 3) uint8 BGR
depth: np.ndarray # (H, W) uint16 millimeters
timestamp: float
@dataclass
class CameraIntrinsics:
"""Per-sensor intrinsics. `*_ir` = depth, `*_rgb` = color."""
fx_ir: float
fy_ir: float
cx_ir: float
cy_ir: float
fx_rgb: float
fy_rgb: float
cx_rgb: float
cy_rgb: float
@dataclass
class CameraExtrinsics:
"""Rigid transform depth → color. Row-major 3x3 rotation, translation in mm."""
rotation: tuple
translation: tuple
@dataclass
class Distortion:
"""Brown-Conrady distortion (k1, k2, k3, p1, p2)."""
k1: float
k2: float
k3: float
p1: float
p2: float
@dataclass
class DiscoveredDevice:
index: int
vid: int
pid: int
serial: str
bus: int
device: int
class RealSenseD435:
"""Driver for Intel RealSense D4xx (tested on D435) via pyrealsense2.
Lifecycle: init() → discover() → open(index) → capture/get_rgb/get_depth → close() → deinit()
"""
def __init__(
self,
width: int = 640,
height: int = 480,
fps: int = 30,
frame_cache_ms: float = 100.0,
):
self._width = width
self._height = height
self._fps = fps
self._frame_cache_ms = frame_cache_ms
self._ctx: Optional[rs.context] = None
self._pipeline: Optional[rs.pipeline] = None
self._profile = None
self._align: Optional[rs.align] = None
self._depth_scale: float = 0.001
self._initialized = False
self._open = False
self._discovered: list[DiscoveredDevice] = []
self._last_frame: Optional[CameraFrame] = None
self._color_available: bool = True
self._color_format = rs.format.bgr8
self._usb_type: str = "unknown"
self.device_index: Optional[int] = None
self.intrinsics: Optional[CameraIntrinsics] = None
self.extrinsics: Optional[CameraExtrinsics] = None
self.distortion_ir: Optional[Distortion] = None
self.distortion_rgb: Optional[Distortion] = None
# ------------------------------------------------------------------
# Lifecycle
# ------------------------------------------------------------------
def is_initialized(self) -> bool:
return self._initialized
def is_open(self) -> bool:
return self._open
def is_color_available(self) -> bool:
"""False when the pipeline runs depth-only (USB 2.x fallback)."""
return self._color_available
def usb_type(self) -> str:
return self._usb_type
def init(self) -> bool:
try:
self._ctx = rs.context()
self._initialized = True
logger.info("RealSense SDK initialized")
return True
except Exception as e:
logger.error(f"RealSense init failed: {e}")
return False
def deinit(self) -> None:
if self._open:
self.close()
self._ctx = None
self._initialized = False
logger.info("RealSense SDK deinitialized")
# ------------------------------------------------------------------
# Discovery
# ------------------------------------------------------------------
def discover(self, timeout_ms: int = 5000) -> list[DiscoveredDevice]:
"""Enumerate RealSense devices on USB. `timeout_ms` kept for API parity."""
if not self._initialized:
raise RuntimeError("SDK not initialized, call init() first")
devices = self._ctx.query_devices()
self._discovered = []
for i, dev in enumerate(devices):
try:
serial = dev.get_info(rs.camera_info.serial_number)
try:
pid_hex = dev.get_info(rs.camera_info.product_id)
pid = int(pid_hex, 16) if pid_hex else 0
except Exception:
pid = 0
bus = 0
addr = 0
try:
phys = dev.get_info(rs.camera_info.physical_port)
parts = phys.rsplit("/", 1)[-1].split("-")
if len(parts) >= 2:
try:
bus = int(parts[0])
except ValueError:
pass
except Exception:
pass
self._discovered.append(DiscoveredDevice(
index=i, vid=0x8086, pid=pid, serial=serial, bus=bus, device=addr,
))
except Exception as e:
logger.warning(f"Failed to read device {i} info: {e}")
logger.info(f"Discovered {len(self._discovered)} RealSense device(s)")
return self._discovered
# ------------------------------------------------------------------
# Stream open/close
# ------------------------------------------------------------------
def open(self, device_index: int = 0) -> bool:
if not self._initialized:
logger.error("Cannot open: SDK not initialized")
return False
if self._open:
logger.warning("Camera already open")
return True
if not self._discovered:
self.discover()
if device_index < 0 or device_index >= len(self._discovered):
logger.error(
f"Invalid device_index {device_index} ({len(self._discovered)} devices found)"
)
return False
serial = self._discovered[device_index].serial
# Probe the USB type via the SDK. D4xx firmware can't sustain both
# depth and color streams simultaneously on USB 2.x — wait_for_frames
# times out indefinitely. We fall back to depth-only there.
usb_type = "unknown"
for dev in self._ctx.query_devices():
try:
if dev.get_info(rs.camera_info.serial_number) == serial:
if dev.supports(rs.camera_info.usb_type_descriptor):
usb_type = dev.get_info(rs.camera_info.usb_type_descriptor)
break
except Exception:
continue
self._usb_type = usb_type
is_usb2 = usb_type.startswith("2.")
if is_usb2:
# On USB 2.x the D4xx firmware only sustains the *native* YUYV
# color format — BGR8/RGB8 trigger a wait_for_frames timeout. We
# convert YUYV → BGR in capture() so callers stay format-agnostic.
# 640x480@15 is the highest tested combination that holds without
# dropping frames on this Jetson; actual delivery rate is ~4 fps
# due to USB 2 bandwidth, but wait_for_frames stays under 300 ms.
color_w, color_h, color_fps = 640, 480, 15
depth_w, depth_h, depth_fps = 640, 480, 15
color_fmt = rs.format.yuyv
logger.warning(
f"D4xx on USB {usb_type} — using USB2 profile "
f"(color YUYV {color_w}x{color_h}@{color_fps}, "
f"depth {depth_w}x{depth_h}@{depth_fps}, ~4 fps actual). "
f"Move to a USB 3 port for full framerate."
)
else:
color_w, color_h, color_fps = self._width, self._height, self._fps
depth_w, depth_h, depth_fps = self._width, self._height, self._fps
color_fmt = rs.format.bgr8
try:
self._pipeline = rs.pipeline(self._ctx)
config = rs.config()
config.enable_device(serial)
config.enable_stream(
rs.stream.depth, depth_w, depth_h, rs.format.z16, depth_fps
)
config.enable_stream(
rs.stream.color, color_w, color_h, color_fmt, color_fps
)
self._profile = self._pipeline.start(config)
self._color_available = True
self._color_format = color_fmt
# rs.align fails when the color stream is YUYV (the only color
# format the D4xx can sustain on USB 2). Skip hardware alignment
# there; we fall back to a numpy resize in capture() so the rgb
# and depth arrays share the same shape for downstream consumers.
self._align = None if is_usb2 else rs.align(rs.stream.color)
depth_sensor = self._profile.get_device().first_depth_sensor()
self._depth_scale = depth_sensor.get_depth_scale()
self._read_calibration()
self.device_index = device_index
self._open = True
logger.info(
f"RealSense D4xx opened: serial={serial} usb={usb_type} "
f"color={color_w}x{color_h}@{color_fps}({color_fmt}) "
f"depth={depth_w}x{depth_h}@{depth_fps} depth_scale={self._depth_scale}"
)
return True
except Exception as e:
logger.error(f"Open failed: {e}")
self._teardown_pipeline()
return False
def close(self) -> None:
if not self._open:
return
self._teardown_pipeline()
self.device_index = None
self._last_frame = None
logger.info("RealSense camera closed")
def _teardown_pipeline(self) -> None:
if self._pipeline is not None:
try:
self._pipeline.stop()
except Exception as e:
logger.warning(f"pipeline.stop failed: {e}")
self._pipeline = None
self._profile = None
self._align = None
self._open = False
# ------------------------------------------------------------------
# Calibration
# ------------------------------------------------------------------
def _read_calibration(self) -> None:
color_stream = self._profile.get_stream(rs.stream.color).as_video_stream_profile()
depth_stream = self._profile.get_stream(rs.stream.depth).as_video_stream_profile()
ci = color_stream.get_intrinsics()
di = depth_stream.get_intrinsics()
self.intrinsics = CameraIntrinsics(
fx_ir=di.fx, fy_ir=di.fy, cx_ir=di.ppx, cy_ir=di.ppy,
fx_rgb=ci.fx, fy_rgb=ci.fy, cx_rgb=ci.ppx, cy_rgb=ci.ppy,
)
try:
ext = depth_stream.get_extrinsics_to(color_stream)
r = list(ext.rotation)
rotation = (
r[0], r[3], r[6],
r[1], r[4], r[7],
r[2], r[5], r[8],
)
t = list(ext.translation)
translation = (t[0] * 1000.0, t[1] * 1000.0, t[2] * 1000.0)
self.extrinsics = CameraExtrinsics(rotation=rotation, translation=translation)
except Exception as e:
logger.warning(f"Extrinsics read failed: {e}")
self.extrinsics = None
# rs intrinsics.coeffs order is [k1, k2, p1, p2, k3].
def _to_distortion(coeffs) -> Distortion:
if not coeffs or len(coeffs) < 5:
return Distortion(0.0, 0.0, 0.0, 0.0, 0.0)
return Distortion(
k1=coeffs[0], k2=coeffs[1], k3=coeffs[4], p1=coeffs[2], p2=coeffs[3],
)
self.distortion_ir = _to_distortion(list(di.coeffs))
self.distortion_rgb = _to_distortion(list(ci.coeffs))
# ------------------------------------------------------------------
# Frame capture
# ------------------------------------------------------------------
def capture(self, timeout_ms: int = 5000) -> Optional[CameraFrame]:
if not self._open:
return None
try:
frames = self._pipeline.wait_for_frames(timeout_ms)
except Exception as e:
logger.warning(f"wait_for_frames failed: {e}")
return None
if self._align is not None:
frames = self._align.process(frames)
color_frame = frames.get_color_frame()
depth_frame = frames.get_depth_frame()
if not color_frame or not depth_frame:
return None
if self._color_format == rs.format.yuyv:
# SDK exposes YUYV as a 1-channel view of the packed buffer;
# rebuild it as (H, W, 2) uint8 so cv2 can decode it to BGR.
h = color_frame.get_height()
w = color_frame.get_width()
buf = np.frombuffer(color_frame.get_data(), dtype=np.uint8)
rgb = cv2.cvtColor(buf.reshape(h, w, 2), cv2.COLOR_YUV2BGR_YUYV)
else:
rgb = np.asanyarray(color_frame.get_data())
depth_raw = np.asanyarray(depth_frame.get_data())
# On USB 2 (no hardware alignment) the color and depth streams have
# different native resolutions. Resize color to depth shape so
# downstream endpoints (/pointcloud, /measure) that index rgb[y, x]
# and depth[y, x] in lockstep continue to work. This is an
# *approximate* alignment — depth↔color extrinsics are ignored.
if rgb.shape[:2] != depth_raw.shape[:2]:
rgb = cv2.resize(
rgb, (depth_raw.shape[1], depth_raw.shape[0]),
interpolation=cv2.INTER_LINEAR,
)
# depth_scale ~ 0.001 m/unit on D4xx → raw uint16 is already millimeters.
if abs(self._depth_scale - 0.001) < 1e-6:
depth_mm = depth_raw
else:
depth_mm = (depth_raw.astype(np.float32) * self._depth_scale * 1000.0).astype(np.uint16)
frame = CameraFrame(rgb=rgb, depth=depth_mm, timestamp=time.time())
self._last_frame = frame
return frame
def _cached_or_capture(self) -> Optional[CameraFrame]:
if (
self._last_frame is not None
and (time.time() - self._last_frame.timestamp) * 1000.0 < self._frame_cache_ms
):
return self._last_frame
return self.capture()
def get_rgb(self) -> Optional[np.ndarray]:
f = self._cached_or_capture()
return f.rgb if f else None
def get_depth(self) -> Optional[np.ndarray]:
f = self._cached_or_capture()
return f.depth if f else None
+92
View File
@@ -0,0 +1,92 @@
"""ONNX YOLO detector with letterbox preprocessing.
Assumes the exported ONNX model has end-to-end NMS embedded:
output shape (1, max_det, 6) where the 6 columns are
[x1, y1, x2, y2, conf, cls] in letterboxed input space.
"""
import logging
from pathlib import Path
import cv2
import numpy as np
import onnxruntime as ort
logger = logging.getLogger(__name__)
def _letterbox(img: np.ndarray, new_size: int = 640, color: tuple = (114, 114, 114)):
"""Resize keeping aspect ratio and pad to (new_size, new_size)."""
h, w = img.shape[:2]
r = min(new_size / h, new_size / w)
new_w, new_h = int(round(w * r)), int(round(h * r))
pad_w = (new_size - new_w) / 2
pad_h = (new_size - new_h) / 2
if (w, h) != (new_w, new_h):
img = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
top, bottom = int(round(pad_h - 0.1)), int(round(pad_h + 0.1))
left, right = int(round(pad_w - 0.1)), int(round(pad_w + 0.1))
img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)
return img, r, left, top
class Detector:
def __init__(self, onnx_path: str | Path, class_names: list[str], conf_threshold: float = 0.5):
path = Path(onnx_path)
if not path.is_file():
raise FileNotFoundError(f"ONNX model not found: {path}")
opts = ort.SessionOptions()
opts.intra_op_num_threads = 4
opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
self.session = ort.InferenceSession(str(path), opts, providers=["CPUExecutionProvider"])
inp = self.session.get_inputs()[0]
self.input_name = inp.name
self.imgsz = int(inp.shape[-1])
self.class_names = class_names
self.conf_threshold = conf_threshold
# Warmup
dummy = np.zeros((self.imgsz, self.imgsz, 3), dtype=np.uint8)
self.detect(dummy)
logger.info(f"Detector ready: model={path.name} imgsz={self.imgsz} classes={class_names}")
def set_conf(self, value: float) -> None:
self.conf_threshold = float(np.clip(value, 0.0, 1.0))
def detect(self, bgr: np.ndarray) -> list[dict]:
h, w = bgr.shape[:2]
padded, scale, pad_x, pad_y = _letterbox(bgr, self.imgsz)
rgb = cv2.cvtColor(padded, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
tensor = np.expand_dims(rgb.transpose(2, 0, 1), 0)
outputs = self.session.run(None, {self.input_name: tensor})
dets = outputs[0][0] # (max_det, 6)
mask = dets[:, 4] >= self.conf_threshold
dets = dets[mask]
if len(dets) == 0:
return []
best = int(np.argmax(dets[:, 4]))
dets = dets[best:best + 1]
boxes = dets[:, :4].copy()
boxes[:, [0, 2]] = (boxes[:, [0, 2]] - pad_x) / scale
boxes[:, [1, 3]] = (boxes[:, [1, 3]] - pad_y) / scale
boxes[:, [0, 2]] = boxes[:, [0, 2]].clip(0, w - 1)
boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, h - 1)
results = []
for box, det in zip(boxes, dets):
cls_idx = int(det[5])
label = self.class_names[cls_idx] if 0 <= cls_idx < len(self.class_names) else str(cls_idx)
results.append({
"bbox": [round(float(box[0]), 1), round(float(box[1]), 1),
round(float(box[2]), 1), round(float(box[3]), 1)],
"conf": round(float(det[4]), 3),
"cls": cls_idx,
"label": label,
})
return results
+105
View File
@@ -0,0 +1,105 @@
import logging
from datetime import datetime
from pathlib import Path
from threading import Lock
import cv2
import numpy as np
logger = logging.getLogger(__name__)
def _flip_code(flip_h: bool, flip_v: bool) -> int | None:
if flip_h and flip_v:
return -1
if flip_h:
return 1
if flip_v:
return 0
return None
def _timestamp() -> str:
now = datetime.now()
return now.strftime("%Y%m%d_%H%M%S_") + f"{now.microsecond // 1000:03d}"
class ImageSaver:
def __init__(self, images_dir: Path, jpeg_quality: int = 85):
self._dir = images_dir
self._dir.mkdir(parents=True, exist_ok=True)
self._quality = jpeg_quality
self._lock = Lock()
self._recording = False
self._rec_flip_h = False
self._rec_flip_v = False
self._rec_started_at: str | None = None
self._rec_count = 0
self._rec_interval = 1
self._rec_seen = 0
@property
def is_recording(self) -> bool:
return self._recording
def _encode_jpeg(self, rgb: np.ndarray, flip_h: bool, flip_v: bool) -> bytes:
img = rgb
code = _flip_code(flip_h, flip_v)
if code is not None:
img = cv2.flip(img, code)
_, buf = cv2.imencode(".jpg", img, [cv2.IMWRITE_JPEG_QUALITY, self._quality])
return buf.tobytes()
def save_snapshot(self, rgb: np.ndarray, flip_h: bool, flip_v: bool) -> str:
data = self._encode_jpeg(rgb, flip_h, flip_v)
filename = f"snap_{_timestamp()}.jpg"
(self._dir / filename).write_bytes(data)
logger.info(f"Snapshot saved: {filename}")
return filename
def start_recording(self, flip_h: bool, flip_v: bool, interval: int = 1) -> dict:
if interval < 1:
interval = 1
with self._lock:
if self._recording:
raise RuntimeError("already recording")
self._recording = True
self._rec_flip_h = flip_h
self._rec_flip_v = flip_v
self._rec_started_at = datetime.now().isoformat(timespec="seconds")
self._rec_count = 0
self._rec_seen = 0
self._rec_interval = interval
logger.info(
f"Recording started (flip_h={flip_h}, flip_v={flip_v}, interval={interval})"
)
return self.status()
def stop_recording(self) -> dict:
with self._lock:
snapshot = self.status()
self._recording = False
logger.info(f"Recording stopped: {snapshot['frames_written']} frames")
return snapshot
def save_rec_frame(self, rgb: np.ndarray) -> None:
if not self._recording:
return
self._rec_seen += 1
if (self._rec_seen - 1) % self._rec_interval != 0:
return
data = self._encode_jpeg(rgb, self._rec_flip_h, self._rec_flip_v)
filename = f"rec_{_timestamp()}.jpg"
(self._dir / filename).write_bytes(data)
self._rec_count += 1
def status(self) -> dict:
return {
"recording": self._recording,
"frames_written": self._rec_count,
"frames_seen": self._rec_seen,
"interval": self._rec_interval,
"started_at": self._rec_started_at,
"flip_h": self._rec_flip_h,
"flip_v": self._rec_flip_v,
}