feat: Complete Rust port of WiFi-DensePose with modular crates

Major changes:
- Organized Python v1 implementation into v1/ subdirectory
- Created Rust workspace with 9 modular crates:
  - wifi-densepose-core: Core types, traits, errors
  - wifi-densepose-signal: CSI processing, phase sanitization, FFT
  - wifi-densepose-nn: Neural network inference (ONNX/Candle/tch)
  - wifi-densepose-api: Axum-based REST/WebSocket API
  - wifi-densepose-db: SQLx database layer
  - wifi-densepose-config: Configuration management
  - wifi-densepose-hardware: Hardware abstraction
  - wifi-densepose-wasm: WebAssembly bindings
  - wifi-densepose-cli: Command-line interface

Documentation:
- ADR-001: Workspace structure
- ADR-002: Signal processing library selection
- ADR-003: Neural network inference strategy
- DDD domain model with bounded contexts

Testing:
- 69 tests passing across all crates
- Signal processing: 45 tests
- Neural networks: 21 tests
- Core: 3 doc tests

Performance targets:
- 10x faster CSI processing (~0.5ms vs ~5ms)
- 5x lower memory usage (~100MB vs ~500MB)
- WASM support for browser deployment
This commit is contained in:
Claude
2026-01-13 03:11:16 +00:00
parent 5101504b72
commit 6ed69a3d48
427 changed files with 90993 additions and 0 deletions
+7
View File
@@ -0,0 +1,7 @@
"""
API routers package
"""
from . import pose, stream, health
__all__ = ["pose", "stream", "health"]
+419
View File
@@ -0,0 +1,419 @@
"""
Health check API endpoints
"""
import logging
import psutil
from typing import Dict, Any, Optional
from datetime import datetime, timedelta
from fastapi import APIRouter, Depends, HTTPException, Request
from pydantic import BaseModel, Field
from src.api.dependencies import get_current_user
from src.config.settings import get_settings
logger = logging.getLogger(__name__)
router = APIRouter()
# Response models
class ComponentHealth(BaseModel):
"""Health status for a system component."""
name: str = Field(..., description="Component name")
status: str = Field(..., description="Health status (healthy, degraded, unhealthy)")
message: Optional[str] = Field(default=None, description="Status message")
last_check: datetime = Field(..., description="Last health check timestamp")
uptime_seconds: Optional[float] = Field(default=None, description="Component uptime")
metrics: Optional[Dict[str, Any]] = Field(default=None, description="Component metrics")
class SystemHealth(BaseModel):
"""Overall system health status."""
status: str = Field(..., description="Overall system status")
timestamp: datetime = Field(..., description="Health check timestamp")
uptime_seconds: float = Field(..., description="System uptime")
components: Dict[str, ComponentHealth] = Field(..., description="Component health status")
system_metrics: Dict[str, Any] = Field(..., description="System-level metrics")
class ReadinessCheck(BaseModel):
"""System readiness check result."""
ready: bool = Field(..., description="Whether system is ready to serve requests")
timestamp: datetime = Field(..., description="Readiness check timestamp")
checks: Dict[str, bool] = Field(..., description="Individual readiness checks")
message: str = Field(..., description="Readiness status message")
# Health check endpoints
@router.get("/health", response_model=SystemHealth)
async def health_check(request: Request):
"""Comprehensive system health check."""
try:
# Get services from app state
hardware_service = getattr(request.app.state, 'hardware_service', None)
pose_service = getattr(request.app.state, 'pose_service', None)
stream_service = getattr(request.app.state, 'stream_service', None)
timestamp = datetime.utcnow()
components = {}
overall_status = "healthy"
# Check hardware service
if hardware_service:
try:
hw_health = await hardware_service.health_check()
components["hardware"] = ComponentHealth(
name="Hardware Service",
status=hw_health["status"],
message=hw_health.get("message"),
last_check=timestamp,
uptime_seconds=hw_health.get("uptime_seconds"),
metrics=hw_health.get("metrics")
)
if hw_health["status"] != "healthy":
overall_status = "degraded" if overall_status == "healthy" else "unhealthy"
except Exception as e:
logger.error(f"Hardware service health check failed: {e}")
components["hardware"] = ComponentHealth(
name="Hardware Service",
status="unhealthy",
message=f"Health check failed: {str(e)}",
last_check=timestamp
)
overall_status = "unhealthy"
else:
components["hardware"] = ComponentHealth(
name="Hardware Service",
status="unavailable",
message="Service not initialized",
last_check=timestamp
)
overall_status = "degraded"
# Check pose service
if pose_service:
try:
pose_health = await pose_service.health_check()
components["pose"] = ComponentHealth(
name="Pose Service",
status=pose_health["status"],
message=pose_health.get("message"),
last_check=timestamp,
uptime_seconds=pose_health.get("uptime_seconds"),
metrics=pose_health.get("metrics")
)
if pose_health["status"] != "healthy":
overall_status = "degraded" if overall_status == "healthy" else "unhealthy"
except Exception as e:
logger.error(f"Pose service health check failed: {e}")
components["pose"] = ComponentHealth(
name="Pose Service",
status="unhealthy",
message=f"Health check failed: {str(e)}",
last_check=timestamp
)
overall_status = "unhealthy"
else:
components["pose"] = ComponentHealth(
name="Pose Service",
status="unavailable",
message="Service not initialized",
last_check=timestamp
)
overall_status = "degraded"
# Check stream service
if stream_service:
try:
stream_health = await stream_service.health_check()
components["stream"] = ComponentHealth(
name="Stream Service",
status=stream_health["status"],
message=stream_health.get("message"),
last_check=timestamp,
uptime_seconds=stream_health.get("uptime_seconds"),
metrics=stream_health.get("metrics")
)
if stream_health["status"] != "healthy":
overall_status = "degraded" if overall_status == "healthy" else "unhealthy"
except Exception as e:
logger.error(f"Stream service health check failed: {e}")
components["stream"] = ComponentHealth(
name="Stream Service",
status="unhealthy",
message=f"Health check failed: {str(e)}",
last_check=timestamp
)
overall_status = "unhealthy"
else:
components["stream"] = ComponentHealth(
name="Stream Service",
status="unavailable",
message="Service not initialized",
last_check=timestamp
)
overall_status = "degraded"
# Get system metrics
system_metrics = get_system_metrics()
# Calculate system uptime (placeholder - would need actual startup time)
uptime_seconds = 0.0 # TODO: Implement actual uptime tracking
return SystemHealth(
status=overall_status,
timestamp=timestamp,
uptime_seconds=uptime_seconds,
components=components,
system_metrics=system_metrics
)
except Exception as e:
logger.error(f"Health check failed: {e}")
raise HTTPException(
status_code=500,
detail=f"Health check failed: {str(e)}"
)
@router.get("/ready", response_model=ReadinessCheck)
async def readiness_check(request: Request):
"""Check if system is ready to serve requests."""
try:
timestamp = datetime.utcnow()
checks = {}
# Check if services are available in app state
if hasattr(request.app.state, 'pose_service') and request.app.state.pose_service:
try:
checks["pose_ready"] = await request.app.state.pose_service.is_ready()
except Exception as e:
logger.warning(f"Pose service readiness check failed: {e}")
checks["pose_ready"] = False
else:
checks["pose_ready"] = False
if hasattr(request.app.state, 'stream_service') and request.app.state.stream_service:
try:
checks["stream_ready"] = await request.app.state.stream_service.is_ready()
except Exception as e:
logger.warning(f"Stream service readiness check failed: {e}")
checks["stream_ready"] = False
else:
checks["stream_ready"] = False
# Hardware service check (basic availability)
checks["hardware_ready"] = True # Basic readiness - API is responding
# Check system resources
checks["memory_available"] = check_memory_availability()
checks["disk_space_available"] = check_disk_space()
# Application is ready if at least the basic services are available
# For now, we'll consider it ready if the API is responding
ready = True # Basic readiness
message = "System is ready" if ready else "System is not ready"
if not ready:
failed_checks = [name for name, status in checks.items() if not status]
message += f". Failed checks: {', '.join(failed_checks)}"
return ReadinessCheck(
ready=ready,
timestamp=timestamp,
checks=checks,
message=message
)
except Exception as e:
logger.error(f"Readiness check failed: {e}")
return ReadinessCheck(
ready=False,
timestamp=datetime.utcnow(),
checks={},
message=f"Readiness check failed: {str(e)}"
)
@router.get("/live")
async def liveness_check():
"""Simple liveness check for load balancers."""
return {
"status": "alive",
"timestamp": datetime.utcnow().isoformat()
}
@router.get("/metrics")
async def get_health_metrics(
request: Request,
current_user: Optional[Dict] = Depends(get_current_user)
):
"""Get detailed system metrics."""
try:
metrics = get_system_metrics()
# Add additional metrics if authenticated
if current_user:
metrics.update(get_detailed_metrics())
return {
"timestamp": datetime.utcnow().isoformat(),
"metrics": metrics
}
except Exception as e:
logger.error(f"Error getting system metrics: {e}")
raise HTTPException(
status_code=500,
detail=f"Failed to get system metrics: {str(e)}"
)
@router.get("/version")
async def get_version_info():
"""Get application version information."""
settings = get_settings()
return {
"name": settings.app_name,
"version": settings.version,
"environment": settings.environment,
"debug": settings.debug,
"timestamp": datetime.utcnow().isoformat()
}
def get_system_metrics() -> Dict[str, Any]:
"""Get basic system metrics."""
try:
# CPU metrics
cpu_percent = psutil.cpu_percent(interval=1)
cpu_count = psutil.cpu_count()
# Memory metrics
memory = psutil.virtual_memory()
memory_metrics = {
"total_gb": round(memory.total / (1024**3), 2),
"available_gb": round(memory.available / (1024**3), 2),
"used_gb": round(memory.used / (1024**3), 2),
"percent": memory.percent
}
# Disk metrics
disk = psutil.disk_usage('/')
disk_metrics = {
"total_gb": round(disk.total / (1024**3), 2),
"free_gb": round(disk.free / (1024**3), 2),
"used_gb": round(disk.used / (1024**3), 2),
"percent": round((disk.used / disk.total) * 100, 2)
}
# Network metrics (basic)
network = psutil.net_io_counters()
network_metrics = {
"bytes_sent": network.bytes_sent,
"bytes_recv": network.bytes_recv,
"packets_sent": network.packets_sent,
"packets_recv": network.packets_recv
}
return {
"cpu": {
"percent": cpu_percent,
"count": cpu_count
},
"memory": memory_metrics,
"disk": disk_metrics,
"network": network_metrics
}
except Exception as e:
logger.error(f"Error getting system metrics: {e}")
return {}
def get_detailed_metrics() -> Dict[str, Any]:
"""Get detailed system metrics (requires authentication)."""
try:
# Process metrics
process = psutil.Process()
process_metrics = {
"pid": process.pid,
"cpu_percent": process.cpu_percent(),
"memory_mb": round(process.memory_info().rss / (1024**2), 2),
"num_threads": process.num_threads(),
"create_time": datetime.fromtimestamp(process.create_time()).isoformat()
}
# Load average (Unix-like systems)
load_avg = None
try:
load_avg = psutil.getloadavg()
except AttributeError:
# Windows doesn't have load average
pass
# Temperature sensors (if available)
temperatures = {}
try:
temps = psutil.sensors_temperatures()
for name, entries in temps.items():
temperatures[name] = [
{"label": entry.label, "current": entry.current}
for entry in entries
]
except AttributeError:
# Not available on all systems
pass
detailed = {
"process": process_metrics
}
if load_avg:
detailed["load_average"] = {
"1min": load_avg[0],
"5min": load_avg[1],
"15min": load_avg[2]
}
if temperatures:
detailed["temperatures"] = temperatures
return detailed
except Exception as e:
logger.error(f"Error getting detailed metrics: {e}")
return {}
def check_memory_availability() -> bool:
"""Check if sufficient memory is available."""
try:
memory = psutil.virtual_memory()
# Consider system ready if less than 90% memory is used
return memory.percent < 90.0
except Exception:
return False
def check_disk_space() -> bool:
"""Check if sufficient disk space is available."""
try:
disk = psutil.disk_usage('/')
# Consider system ready if more than 1GB free space
free_gb = disk.free / (1024**3)
return free_gb > 1.0
except Exception:
return False
+420
View File
@@ -0,0 +1,420 @@
"""
Pose estimation API endpoints
"""
import logging
from typing import List, Optional, Dict, Any
from datetime import datetime, timedelta
from fastapi import APIRouter, Depends, HTTPException, Query, BackgroundTasks
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field
from src.api.dependencies import (
get_pose_service,
get_hardware_service,
get_current_user,
require_auth
)
from src.services.pose_service import PoseService
from src.services.hardware_service import HardwareService
from src.config.settings import get_settings
logger = logging.getLogger(__name__)
router = APIRouter()
# Request/Response models
class PoseEstimationRequest(BaseModel):
"""Request model for pose estimation."""
zone_ids: Optional[List[str]] = Field(
default=None,
description="Specific zones to analyze (all zones if not specified)"
)
confidence_threshold: Optional[float] = Field(
default=None,
ge=0.0,
le=1.0,
description="Minimum confidence threshold for detections"
)
max_persons: Optional[int] = Field(
default=None,
ge=1,
le=50,
description="Maximum number of persons to detect"
)
include_keypoints: bool = Field(
default=True,
description="Include detailed keypoint data"
)
include_segmentation: bool = Field(
default=False,
description="Include DensePose segmentation masks"
)
class PersonPose(BaseModel):
"""Person pose data model."""
person_id: str = Field(..., description="Unique person identifier")
confidence: float = Field(..., description="Detection confidence score")
bounding_box: Dict[str, float] = Field(..., description="Person bounding box")
keypoints: Optional[List[Dict[str, Any]]] = Field(
default=None,
description="Body keypoints with coordinates and confidence"
)
segmentation: Optional[Dict[str, Any]] = Field(
default=None,
description="DensePose segmentation data"
)
zone_id: Optional[str] = Field(
default=None,
description="Zone where person is detected"
)
activity: Optional[str] = Field(
default=None,
description="Detected activity"
)
timestamp: datetime = Field(..., description="Detection timestamp")
class PoseEstimationResponse(BaseModel):
"""Response model for pose estimation."""
timestamp: datetime = Field(..., description="Analysis timestamp")
frame_id: str = Field(..., description="Unique frame identifier")
persons: List[PersonPose] = Field(..., description="Detected persons")
zone_summary: Dict[str, int] = Field(..., description="Person count per zone")
processing_time_ms: float = Field(..., description="Processing time in milliseconds")
metadata: Dict[str, Any] = Field(default_factory=dict, description="Additional metadata")
class HistoricalDataRequest(BaseModel):
"""Request model for historical pose data."""
start_time: datetime = Field(..., description="Start time for data query")
end_time: datetime = Field(..., description="End time for data query")
zone_ids: Optional[List[str]] = Field(
default=None,
description="Filter by specific zones"
)
aggregation_interval: Optional[int] = Field(
default=300,
ge=60,
le=3600,
description="Aggregation interval in seconds"
)
include_raw_data: bool = Field(
default=False,
description="Include raw detection data"
)
# Endpoints
@router.get("/current", response_model=PoseEstimationResponse)
async def get_current_pose_estimation(
request: PoseEstimationRequest = Depends(),
pose_service: PoseService = Depends(get_pose_service),
current_user: Optional[Dict] = Depends(get_current_user)
):
"""Get current pose estimation from WiFi signals."""
try:
logger.info(f"Processing pose estimation request from user: {current_user.get('id') if current_user else 'anonymous'}")
# Get current pose estimation
result = await pose_service.estimate_poses(
zone_ids=request.zone_ids,
confidence_threshold=request.confidence_threshold,
max_persons=request.max_persons,
include_keypoints=request.include_keypoints,
include_segmentation=request.include_segmentation
)
return PoseEstimationResponse(**result)
except Exception as e:
logger.error(f"Error in pose estimation: {e}")
raise HTTPException(
status_code=500,
detail=f"Pose estimation failed: {str(e)}"
)
@router.post("/analyze", response_model=PoseEstimationResponse)
async def analyze_pose_data(
request: PoseEstimationRequest,
background_tasks: BackgroundTasks,
pose_service: PoseService = Depends(get_pose_service),
current_user: Dict = Depends(require_auth)
):
"""Trigger pose analysis with custom parameters."""
try:
logger.info(f"Custom pose analysis requested by user: {current_user['id']}")
# Trigger analysis
result = await pose_service.analyze_with_params(
zone_ids=request.zone_ids,
confidence_threshold=request.confidence_threshold,
max_persons=request.max_persons,
include_keypoints=request.include_keypoints,
include_segmentation=request.include_segmentation
)
# Schedule background processing if needed
if request.include_segmentation:
background_tasks.add_task(
pose_service.process_segmentation_data,
result["frame_id"]
)
return PoseEstimationResponse(**result)
except Exception as e:
logger.error(f"Error in pose analysis: {e}")
raise HTTPException(
status_code=500,
detail=f"Pose analysis failed: {str(e)}"
)
@router.get("/zones/{zone_id}/occupancy")
async def get_zone_occupancy(
zone_id: str,
pose_service: PoseService = Depends(get_pose_service),
current_user: Optional[Dict] = Depends(get_current_user)
):
"""Get current occupancy for a specific zone."""
try:
occupancy = await pose_service.get_zone_occupancy(zone_id)
if occupancy is None:
raise HTTPException(
status_code=404,
detail=f"Zone '{zone_id}' not found"
)
return {
"zone_id": zone_id,
"current_occupancy": occupancy["count"],
"max_occupancy": occupancy.get("max_occupancy"),
"persons": occupancy["persons"],
"timestamp": occupancy["timestamp"]
}
except HTTPException:
raise
except Exception as e:
logger.error(f"Error getting zone occupancy: {e}")
raise HTTPException(
status_code=500,
detail=f"Failed to get zone occupancy: {str(e)}"
)
@router.get("/zones/summary")
async def get_zones_summary(
pose_service: PoseService = Depends(get_pose_service),
current_user: Optional[Dict] = Depends(get_current_user)
):
"""Get occupancy summary for all zones."""
try:
summary = await pose_service.get_zones_summary()
return {
"timestamp": datetime.utcnow(),
"total_persons": summary["total_persons"],
"zones": summary["zones"],
"active_zones": summary["active_zones"]
}
except Exception as e:
logger.error(f"Error getting zones summary: {e}")
raise HTTPException(
status_code=500,
detail=f"Failed to get zones summary: {str(e)}"
)
@router.post("/historical")
async def get_historical_data(
request: HistoricalDataRequest,
pose_service: PoseService = Depends(get_pose_service),
current_user: Dict = Depends(require_auth)
):
"""Get historical pose estimation data."""
try:
# Validate time range
if request.end_time <= request.start_time:
raise HTTPException(
status_code=400,
detail="End time must be after start time"
)
# Limit query range to prevent excessive data
max_range = timedelta(days=7)
if request.end_time - request.start_time > max_range:
raise HTTPException(
status_code=400,
detail="Query range cannot exceed 7 days"
)
data = await pose_service.get_historical_data(
start_time=request.start_time,
end_time=request.end_time,
zone_ids=request.zone_ids,
aggregation_interval=request.aggregation_interval,
include_raw_data=request.include_raw_data
)
return {
"query": {
"start_time": request.start_time,
"end_time": request.end_time,
"zone_ids": request.zone_ids,
"aggregation_interval": request.aggregation_interval
},
"data": data["aggregated_data"],
"raw_data": data.get("raw_data") if request.include_raw_data else None,
"total_records": data["total_records"]
}
except HTTPException:
raise
except Exception as e:
logger.error(f"Error getting historical data: {e}")
raise HTTPException(
status_code=500,
detail=f"Failed to get historical data: {str(e)}"
)
@router.get("/activities")
async def get_detected_activities(
zone_id: Optional[str] = Query(None, description="Filter by zone ID"),
limit: int = Query(10, ge=1, le=100, description="Maximum number of activities"),
pose_service: PoseService = Depends(get_pose_service),
current_user: Optional[Dict] = Depends(get_current_user)
):
"""Get recently detected activities."""
try:
activities = await pose_service.get_recent_activities(
zone_id=zone_id,
limit=limit
)
return {
"activities": activities,
"total_count": len(activities),
"zone_id": zone_id
}
except Exception as e:
logger.error(f"Error getting activities: {e}")
raise HTTPException(
status_code=500,
detail=f"Failed to get activities: {str(e)}"
)
@router.post("/calibrate")
async def calibrate_pose_system(
background_tasks: BackgroundTasks,
pose_service: PoseService = Depends(get_pose_service),
hardware_service: HardwareService = Depends(get_hardware_service),
current_user: Dict = Depends(require_auth)
):
"""Calibrate the pose estimation system."""
try:
logger.info(f"Pose system calibration initiated by user: {current_user['id']}")
# Check if calibration is already in progress
if await pose_service.is_calibrating():
raise HTTPException(
status_code=409,
detail="Calibration already in progress"
)
# Start calibration process
calibration_id = await pose_service.start_calibration()
# Schedule background calibration task
background_tasks.add_task(
pose_service.run_calibration,
calibration_id
)
return {
"calibration_id": calibration_id,
"status": "started",
"estimated_duration_minutes": 5,
"message": "Calibration process started"
}
except HTTPException:
raise
except Exception as e:
logger.error(f"Error starting calibration: {e}")
raise HTTPException(
status_code=500,
detail=f"Failed to start calibration: {str(e)}"
)
@router.get("/calibration/status")
async def get_calibration_status(
pose_service: PoseService = Depends(get_pose_service),
current_user: Dict = Depends(require_auth)
):
"""Get current calibration status."""
try:
status = await pose_service.get_calibration_status()
return {
"is_calibrating": status["is_calibrating"],
"calibration_id": status.get("calibration_id"),
"progress_percent": status.get("progress_percent", 0),
"current_step": status.get("current_step"),
"estimated_remaining_minutes": status.get("estimated_remaining_minutes"),
"last_calibration": status.get("last_calibration")
}
except Exception as e:
logger.error(f"Error getting calibration status: {e}")
raise HTTPException(
status_code=500,
detail=f"Failed to get calibration status: {str(e)}"
)
@router.get("/stats")
async def get_pose_statistics(
hours: int = Query(24, ge=1, le=168, description="Hours of data to analyze"),
pose_service: PoseService = Depends(get_pose_service),
current_user: Optional[Dict] = Depends(get_current_user)
):
"""Get pose estimation statistics."""
try:
end_time = datetime.utcnow()
start_time = end_time - timedelta(hours=hours)
stats = await pose_service.get_statistics(
start_time=start_time,
end_time=end_time
)
return {
"period": {
"start_time": start_time,
"end_time": end_time,
"hours": hours
},
"statistics": stats
}
except Exception as e:
logger.error(f"Error getting statistics: {e}")
raise HTTPException(
status_code=500,
detail=f"Failed to get statistics: {str(e)}"
)
+468
View File
@@ -0,0 +1,468 @@
"""
WebSocket streaming API endpoints
"""
import json
import logging
from typing import Dict, List, Optional, Any
from datetime import datetime
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends, HTTPException, Query
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field
from src.api.dependencies import (
get_stream_service,
get_pose_service,
get_current_user_ws,
require_auth
)
from src.api.websocket.connection_manager import ConnectionManager
from src.services.stream_service import StreamService
from src.services.pose_service import PoseService
logger = logging.getLogger(__name__)
router = APIRouter()
# Initialize connection manager
connection_manager = ConnectionManager()
# Request/Response models
class StreamSubscriptionRequest(BaseModel):
"""Request model for stream subscription."""
zone_ids: Optional[List[str]] = Field(
default=None,
description="Zones to subscribe to (all zones if not specified)"
)
stream_types: List[str] = Field(
default=["pose_data"],
description="Types of data to stream"
)
min_confidence: float = Field(
default=0.5,
ge=0.0,
le=1.0,
description="Minimum confidence threshold for streaming"
)
max_fps: int = Field(
default=30,
ge=1,
le=60,
description="Maximum frames per second"
)
include_metadata: bool = Field(
default=True,
description="Include metadata in stream"
)
class StreamStatus(BaseModel):
"""Stream status model."""
is_active: bool = Field(..., description="Whether streaming is active")
connected_clients: int = Field(..., description="Number of connected clients")
streams: List[Dict[str, Any]] = Field(..., description="Active streams")
uptime_seconds: float = Field(..., description="Stream uptime in seconds")
# WebSocket endpoints
@router.websocket("/pose")
async def websocket_pose_stream(
websocket: WebSocket,
zone_ids: Optional[str] = Query(None, description="Comma-separated zone IDs"),
min_confidence: float = Query(0.5, ge=0.0, le=1.0),
max_fps: int = Query(30, ge=1, le=60),
token: Optional[str] = Query(None, description="Authentication token")
):
"""WebSocket endpoint for real-time pose data streaming."""
client_id = None
try:
# Accept WebSocket connection
await websocket.accept()
# Check authentication if enabled
from src.config.settings import get_settings
settings = get_settings()
if settings.enable_authentication and not token:
await websocket.send_json({
"type": "error",
"message": "Authentication token required"
})
await websocket.close(code=1008)
return
# Parse zone IDs
zone_list = None
if zone_ids:
zone_list = [zone.strip() for zone in zone_ids.split(",") if zone.strip()]
# Register client with connection manager
client_id = await connection_manager.connect(
websocket=websocket,
stream_type="pose",
zone_ids=zone_list,
min_confidence=min_confidence,
max_fps=max_fps
)
logger.info(f"WebSocket client {client_id} connected for pose streaming")
# Send initial connection confirmation
await websocket.send_json({
"type": "connection_established",
"client_id": client_id,
"timestamp": datetime.utcnow().isoformat(),
"config": {
"zone_ids": zone_list,
"min_confidence": min_confidence,
"max_fps": max_fps
}
})
# Keep connection alive and handle incoming messages
while True:
try:
# Wait for client messages (ping, config updates, etc.)
message = await websocket.receive_text()
data = json.loads(message)
await handle_websocket_message(client_id, data, websocket)
except WebSocketDisconnect:
break
except json.JSONDecodeError:
await websocket.send_json({
"type": "error",
"message": "Invalid JSON format"
})
except Exception as e:
logger.error(f"Error handling WebSocket message: {e}")
await websocket.send_json({
"type": "error",
"message": "Internal server error"
})
except WebSocketDisconnect:
logger.info(f"WebSocket client {client_id} disconnected")
except Exception as e:
logger.error(f"WebSocket error: {e}")
finally:
if client_id:
await connection_manager.disconnect(client_id)
@router.websocket("/events")
async def websocket_events_stream(
websocket: WebSocket,
event_types: Optional[str] = Query(None, description="Comma-separated event types"),
zone_ids: Optional[str] = Query(None, description="Comma-separated zone IDs"),
token: Optional[str] = Query(None, description="Authentication token")
):
"""WebSocket endpoint for real-time event streaming."""
client_id = None
try:
await websocket.accept()
# Check authentication if enabled
from src.config.settings import get_settings
settings = get_settings()
if settings.enable_authentication and not token:
await websocket.send_json({
"type": "error",
"message": "Authentication token required"
})
await websocket.close(code=1008)
return
# Parse parameters
event_list = None
if event_types:
event_list = [event.strip() for event in event_types.split(",") if event.strip()]
zone_list = None
if zone_ids:
zone_list = [zone.strip() for zone in zone_ids.split(",") if zone.strip()]
# Register client
client_id = await connection_manager.connect(
websocket=websocket,
stream_type="events",
zone_ids=zone_list,
event_types=event_list
)
logger.info(f"WebSocket client {client_id} connected for event streaming")
# Send confirmation
await websocket.send_json({
"type": "connection_established",
"client_id": client_id,
"timestamp": datetime.utcnow().isoformat(),
"config": {
"event_types": event_list,
"zone_ids": zone_list
}
})
# Handle messages
while True:
try:
message = await websocket.receive_text()
data = json.loads(message)
await handle_websocket_message(client_id, data, websocket)
except WebSocketDisconnect:
break
except Exception as e:
logger.error(f"Error in events WebSocket: {e}")
except WebSocketDisconnect:
logger.info(f"Events WebSocket client {client_id} disconnected")
except Exception as e:
logger.error(f"Events WebSocket error: {e}")
finally:
if client_id:
await connection_manager.disconnect(client_id)
async def handle_websocket_message(client_id: str, data: Dict[str, Any], websocket: WebSocket):
"""Handle incoming WebSocket messages."""
message_type = data.get("type")
if message_type == "ping":
await websocket.send_json({
"type": "pong",
"timestamp": datetime.utcnow().isoformat()
})
elif message_type == "update_config":
# Update client configuration
config = data.get("config", {})
await connection_manager.update_client_config(client_id, config)
await websocket.send_json({
"type": "config_updated",
"timestamp": datetime.utcnow().isoformat(),
"config": config
})
elif message_type == "get_status":
# Send current status
status = await connection_manager.get_client_status(client_id)
await websocket.send_json({
"type": "status",
"timestamp": datetime.utcnow().isoformat(),
"status": status
})
else:
await websocket.send_json({
"type": "error",
"message": f"Unknown message type: {message_type}"
})
# HTTP endpoints for stream management
@router.get("/status", response_model=StreamStatus)
async def get_stream_status(
stream_service: StreamService = Depends(get_stream_service)
):
"""Get current streaming status."""
try:
status = await stream_service.get_status()
connections = await connection_manager.get_connection_stats()
# Calculate uptime (simplified for now)
uptime_seconds = 0.0
if status.get("running", False):
uptime_seconds = 3600.0 # Default 1 hour for demo
return StreamStatus(
is_active=status.get("running", False),
connected_clients=connections.get("total_clients", status["connections"]["active"]),
streams=[{
"type": "pose_stream",
"active": status.get("running", False),
"buffer_size": status["buffers"]["pose_buffer_size"]
}],
uptime_seconds=uptime_seconds
)
except Exception as e:
logger.error(f"Error getting stream status: {e}")
raise HTTPException(
status_code=500,
detail=f"Failed to get stream status: {str(e)}"
)
@router.post("/start")
async def start_streaming(
stream_service: StreamService = Depends(get_stream_service),
current_user: Dict = Depends(require_auth)
):
"""Start the streaming service."""
try:
logger.info(f"Starting streaming service by user: {current_user['id']}")
if await stream_service.is_active():
return JSONResponse(
status_code=200,
content={"message": "Streaming service is already active"}
)
await stream_service.start()
return {
"message": "Streaming service started successfully",
"timestamp": datetime.utcnow().isoformat()
}
except Exception as e:
logger.error(f"Error starting streaming: {e}")
raise HTTPException(
status_code=500,
detail=f"Failed to start streaming: {str(e)}"
)
@router.post("/stop")
async def stop_streaming(
stream_service: StreamService = Depends(get_stream_service),
current_user: Dict = Depends(require_auth)
):
"""Stop the streaming service."""
try:
logger.info(f"Stopping streaming service by user: {current_user['id']}")
await stream_service.stop()
await connection_manager.disconnect_all()
return {
"message": "Streaming service stopped successfully",
"timestamp": datetime.utcnow().isoformat()
}
except Exception as e:
logger.error(f"Error stopping streaming: {e}")
raise HTTPException(
status_code=500,
detail=f"Failed to stop streaming: {str(e)}"
)
@router.get("/clients")
async def get_connected_clients(
current_user: Dict = Depends(require_auth)
):
"""Get list of connected WebSocket clients."""
try:
clients = await connection_manager.get_connected_clients()
return {
"total_clients": len(clients),
"clients": clients,
"timestamp": datetime.utcnow().isoformat()
}
except Exception as e:
logger.error(f"Error getting connected clients: {e}")
raise HTTPException(
status_code=500,
detail=f"Failed to get connected clients: {str(e)}"
)
@router.delete("/clients/{client_id}")
async def disconnect_client(
client_id: str,
current_user: Dict = Depends(require_auth)
):
"""Disconnect a specific WebSocket client."""
try:
logger.info(f"Disconnecting client {client_id} by user: {current_user['id']}")
success = await connection_manager.disconnect(client_id)
if not success:
raise HTTPException(
status_code=404,
detail=f"Client {client_id} not found"
)
return {
"message": f"Client {client_id} disconnected successfully",
"timestamp": datetime.utcnow().isoformat()
}
except HTTPException:
raise
except Exception as e:
logger.error(f"Error disconnecting client: {e}")
raise HTTPException(
status_code=500,
detail=f"Failed to disconnect client: {str(e)}"
)
@router.post("/broadcast")
async def broadcast_message(
message: Dict[str, Any],
stream_type: Optional[str] = Query(None, description="Target stream type"),
zone_ids: Optional[List[str]] = Query(None, description="Target zone IDs"),
current_user: Dict = Depends(require_auth)
):
"""Broadcast a message to connected WebSocket clients."""
try:
logger.info(f"Broadcasting message by user: {current_user['id']}")
# Add metadata to message
broadcast_data = {
**message,
"broadcast_timestamp": datetime.utcnow().isoformat(),
"sender": current_user["id"]
}
# Broadcast to matching clients
sent_count = await connection_manager.broadcast(
data=broadcast_data,
stream_type=stream_type,
zone_ids=zone_ids
)
return {
"message": "Broadcast sent successfully",
"recipients": sent_count,
"timestamp": datetime.utcnow().isoformat()
}
except Exception as e:
logger.error(f"Error broadcasting message: {e}")
raise HTTPException(
status_code=500,
detail=f"Failed to broadcast message: {str(e)}"
)
@router.get("/metrics")
async def get_streaming_metrics():
"""Get streaming performance metrics."""
try:
metrics = await connection_manager.get_metrics()
return {
"metrics": metrics,
"timestamp": datetime.utcnow().isoformat()
}
except Exception as e:
logger.error(f"Error getting streaming metrics: {e}")
raise HTTPException(
status_code=500,
detail=f"Failed to get streaming metrics: {str(e)}"
)