344 lines
10 KiB
Python
344 lines
10 KiB
Python
"""
|
|
FastAPI Web Application for Beacon Tracker Video Streaming
|
|
|
|
This application reads JPEG frames from shared memory and streams them
|
|
to web browsers via Server-Sent Events (SSE) or WebSocket (H.264 stream).
|
|
"""
|
|
|
|
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request
|
|
from fastapi.responses import HTMLResponse, StreamingResponse, Response
|
|
from fastapi.staticfiles import StaticFiles
|
|
from fastapi.templating import Jinja2Templates
|
|
from shared_memory_reader import SharedMemoryFrameReader
|
|
import asyncio
|
|
import time
|
|
import base64
|
|
import json
|
|
import os
|
|
from typing import Set
|
|
from contextlib import asynccontextmanager
|
|
|
|
# Global state
|
|
reader = None
|
|
reader_lock = asyncio.Lock()
|
|
last_frame_data = None
|
|
last_frame_header = None
|
|
frame_lock = asyncio.Lock()
|
|
|
|
# Video streaming state
|
|
active_websocket_clients: Set[WebSocket] = set()
|
|
streaming_task = None
|
|
streaming_active = False
|
|
TCP_HOST = '127.0.0.1'
|
|
TCP_PORT = 8888
|
|
|
|
|
|
async def init_reader():
|
|
"""Initialize the shared memory reader"""
|
|
global reader
|
|
async with reader_lock:
|
|
if reader is None:
|
|
try:
|
|
reader = SharedMemoryFrameReader()
|
|
print("Connected to shared memory successfully")
|
|
except Exception as e:
|
|
print(f"Failed to connect to shared memory: {e}")
|
|
reader = None
|
|
|
|
|
|
async def cleanup_reader():
|
|
"""Clean up shared memory reader"""
|
|
global reader
|
|
async with reader_lock:
|
|
if reader:
|
|
reader.close()
|
|
reader = None
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI):
|
|
"""Lifespan context manager for startup and shutdown events"""
|
|
# Startup
|
|
await init_reader()
|
|
print("FastAPI application started")
|
|
yield
|
|
# Shutdown
|
|
await cleanup_reader()
|
|
print("FastAPI application shutdown")
|
|
|
|
|
|
# Create FastAPI app
|
|
app = FastAPI(
|
|
title="Beacon Tracker Video Stream",
|
|
description="High-performance video streaming with WebSocket and SSE support",
|
|
version="2.0",
|
|
lifespan=lifespan
|
|
)
|
|
|
|
# Mount static files if directory exists
|
|
if os.path.exists("static"):
|
|
app.mount("/static", StaticFiles(directory="static"), name="static")
|
|
else:
|
|
print("Warning: 'static' directory not found, static files will not be served")
|
|
|
|
# Templates
|
|
templates = Jinja2Templates(directory="templates")
|
|
|
|
|
|
@app.get("/", response_class=HTMLResponse)
|
|
async def index(request: Request):
|
|
"""Render the main page"""
|
|
return templates.TemplateResponse("index.html", {"request": request})
|
|
|
|
|
|
@app.get("/stream")
|
|
async def stream_sse():
|
|
"""
|
|
Server-Sent Events stream endpoint
|
|
|
|
Continuously sends JPEG frames to the client as they become available.
|
|
"""
|
|
async def generate():
|
|
global reader
|
|
|
|
# Initialize reader for this stream
|
|
await init_reader()
|
|
|
|
# Check again after init
|
|
async with reader_lock:
|
|
local_reader = reader
|
|
|
|
if local_reader is None:
|
|
yield f"data: {json.dumps({'error': 'Failed to connect to shared memory'})}\n\n"
|
|
return
|
|
|
|
consecutive_failures = 0
|
|
max_failures = 100
|
|
|
|
while True:
|
|
try:
|
|
async with reader_lock:
|
|
if reader is None:
|
|
yield f"data: {json.dumps({'error': 'Reader is None'})}\n\n"
|
|
return
|
|
|
|
# Run blocking read in thread pool
|
|
result = await asyncio.to_thread(reader.read_frame)
|
|
|
|
if result:
|
|
header, jpeg_data = result
|
|
consecutive_failures = 0
|
|
|
|
# Encode JPEG as base64 for transmission
|
|
jpeg_base64 = base64.b64encode(jpeg_data).decode('utf-8')
|
|
|
|
# Create event data
|
|
event_data = {
|
|
'frame_number': header.frame_number,
|
|
'timestamp_us': header.timestamp_us,
|
|
'width': header.width,
|
|
'height': header.height,
|
|
'data_size': header.data_size,
|
|
'jpeg': jpeg_base64
|
|
}
|
|
|
|
# Send as SSE event
|
|
yield f"data: {json.dumps(event_data)}\n\n"
|
|
|
|
# Update global state
|
|
async with frame_lock:
|
|
global last_frame_data, last_frame_header
|
|
last_frame_data = jpeg_data
|
|
last_frame_header = header
|
|
|
|
else:
|
|
consecutive_failures += 1
|
|
if consecutive_failures >= max_failures:
|
|
yield f"data: {json.dumps({'error': 'No frames available'})}\n\n"
|
|
consecutive_failures = 0
|
|
|
|
# Small delay to prevent busy waiting
|
|
await asyncio.sleep(0.001) # 1ms - fast updates
|
|
|
|
except Exception as e:
|
|
print(f"Error in stream: {e}")
|
|
yield f"data: {json.dumps({'error': str(e)})}\n\n"
|
|
await asyncio.sleep(1)
|
|
|
|
return StreamingResponse(
|
|
generate(),
|
|
media_type="text/event-stream",
|
|
headers={
|
|
"Cache-Control": "no-cache",
|
|
"Connection": "keep-alive",
|
|
"X-Accel-Buffering": "no"
|
|
}
|
|
)
|
|
|
|
|
|
@app.get("/status")
|
|
async def status():
|
|
"""Get the current status of the stream"""
|
|
await init_reader()
|
|
|
|
async with frame_lock:
|
|
if last_frame_header:
|
|
return {
|
|
'connected': True,
|
|
'last_frame': last_frame_header.frame_number,
|
|
'timestamp_us': last_frame_header.timestamp_us,
|
|
'resolution': f"{last_frame_header.width}x{last_frame_header.height}"
|
|
}
|
|
else:
|
|
return {
|
|
'connected': reader is not None,
|
|
'last_frame': None,
|
|
'message': 'Waiting for frames...'
|
|
}
|
|
|
|
|
|
@app.get("/latest_frame")
|
|
async def latest_frame():
|
|
"""Get the latest frame as a JPEG image"""
|
|
async with frame_lock:
|
|
if last_frame_data:
|
|
return Response(content=last_frame_data, media_type="image/jpeg")
|
|
else:
|
|
return Response(content="No frame available", status_code=404)
|
|
|
|
|
|
async def stream_video_from_tcp():
|
|
"""
|
|
Read MPEG-TS video stream from TCP socket and broadcast via WebSocket.
|
|
This runs as a background task.
|
|
"""
|
|
global streaming_active
|
|
|
|
print(f"Starting video stream reader from TCP: {TCP_HOST}:{TCP_PORT}")
|
|
|
|
# Wait a bit for C++ application to start TCP server
|
|
await asyncio.sleep(2)
|
|
|
|
try:
|
|
# Connect to TCP server
|
|
print(f"Connecting to TCP server {TCP_HOST}:{TCP_PORT}...")
|
|
|
|
reader, writer = await asyncio.open_connection(TCP_HOST, TCP_PORT)
|
|
|
|
print("Connected to TCP server, starting stream...")
|
|
|
|
streaming_active = True
|
|
chunk_size = 32768 # 32KB chunks for MPEG-TS
|
|
|
|
while streaming_active and active_websocket_clients:
|
|
try:
|
|
# Read chunk from TCP
|
|
data = await asyncio.wait_for(reader.read(chunk_size), timeout=5.0)
|
|
|
|
if not data:
|
|
print("No data from TCP, stream may have ended")
|
|
break
|
|
|
|
# Broadcast binary data to all connected WebSocket clients
|
|
disconnected_clients = set()
|
|
for client in active_websocket_clients.copy():
|
|
try:
|
|
await client.send_bytes(data)
|
|
except Exception as e:
|
|
print(f"Error sending to client: {e}")
|
|
disconnected_clients.add(client)
|
|
|
|
# Remove disconnected clients
|
|
active_websocket_clients.difference_update(disconnected_clients)
|
|
|
|
except asyncio.TimeoutError:
|
|
# No data in 5 seconds, continue waiting
|
|
continue
|
|
except Exception as e:
|
|
print(f"Error reading from TCP: {e}")
|
|
break
|
|
|
|
writer.close()
|
|
await writer.wait_closed()
|
|
|
|
except ConnectionRefusedError:
|
|
print(f"ERROR: Could not connect to TCP server {TCP_HOST}:{TCP_PORT}")
|
|
print("Make sure C++ application is running first")
|
|
except Exception as e:
|
|
print(f"ERROR in video streaming: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
finally:
|
|
streaming_active = False
|
|
print("Video streaming stopped")
|
|
|
|
|
|
@app.websocket("/ws/video")
|
|
async def websocket_video_endpoint(websocket: WebSocket):
|
|
"""WebSocket endpoint for video streaming"""
|
|
global streaming_task, streaming_active
|
|
|
|
await websocket.accept()
|
|
print(f"Client connected to video stream: {id(websocket)}")
|
|
|
|
# Add client to active set
|
|
active_websocket_clients.add(websocket)
|
|
|
|
# Start streaming task if not already running
|
|
if not streaming_active:
|
|
streaming_task = asyncio.create_task(stream_video_from_tcp())
|
|
print("Started video streaming task")
|
|
|
|
try:
|
|
# Keep connection alive and handle client messages
|
|
while True:
|
|
# Wait for client messages (like ping/pong or control messages)
|
|
try:
|
|
data = await asyncio.wait_for(websocket.receive_text(), timeout=1.0)
|
|
# Handle client messages if needed
|
|
if data == "ping":
|
|
await websocket.send_text("pong")
|
|
except asyncio.TimeoutError:
|
|
# No message received, continue
|
|
continue
|
|
|
|
except WebSocketDisconnect:
|
|
print(f"Client disconnected from video stream: {id(websocket)}")
|
|
except Exception as e:
|
|
print(f"WebSocket error: {e}")
|
|
finally:
|
|
# Remove client from active set
|
|
active_websocket_clients.discard(websocket)
|
|
|
|
# Stop streaming if no more clients
|
|
if not active_websocket_clients:
|
|
streaming_active = False
|
|
print("No more clients, stopping stream")
|
|
|
|
|
|
@app.get("/health")
|
|
async def health():
|
|
"""Health check endpoint"""
|
|
return {
|
|
"status": "healthy",
|
|
"active_websocket_clients": len(active_websocket_clients),
|
|
"streaming_active": streaming_active,
|
|
"tcp_endpoint": f"{TCP_HOST}:{TCP_PORT}"
|
|
}
|
|
|
|
|
|
if __name__ == '__main__':
|
|
import uvicorn
|
|
|
|
print("Starting FastAPI server on http://0.0.0.0:5000")
|
|
print(f"Video stream will be available at ws://0.0.0.0:5000/ws/video")
|
|
print(f"Expecting C++ TCP stream at {TCP_HOST}:{TCP_PORT}")
|
|
|
|
uvicorn.run(
|
|
app,
|
|
host='0.0.0.0',
|
|
port=5000,
|
|
log_level='info',
|
|
access_log=True
|
|
)
|