Spaces:
Sleeping
Sleeping
File size: 3,681 Bytes
02d085a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | import json
import logging
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
import cv2
import numpy as np
import uvicorn
import config
from pipeline import load_sign_pipeline, make_translator_for_connection
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("sign-to-text-server")
app = FastAPI(title="Sign-to-Text Streaming API")
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=False,
allow_methods=["*"],
allow_headers=["*"],
)
pipeline_state = {"ready": False, "pipeline": None}
@app.on_event("startup")
def on_startup():
logger.info("Loading sign pipeline...")
pipeline_state["pipeline"] = load_sign_pipeline()
pipeline_state["ready"] = True
logger.info("Sign pipeline ready. Server accepting connections.")
@app.get("/health")
def health():
return {"status": "ok" if pipeline_state["ready"] else "loading"}
@app.get("/")
def root():
return {
"name": "Sign-to-Text Streaming API",
"status": "ok" if pipeline_state["ready"] else "loading",
"websocket_endpoint": "/ws",
"protocol": {
"client_sends": "binary WebSocket frames, each one JPEG-encoded RGB video frame, at config.TARGET_FPS",
"server_sends": [
{"type": "segment", "start_seconds": "float", "end_seconds": "float",
"num_frames": "int", "text": "string",
"dino_batch_ms": "float", "translation_ms": "float"},
{"type": "error", "message": "string"},
],
},
}
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
if not pipeline_state["ready"]:
await websocket.close(code=1013)
return
await websocket.accept()
logger.info("Client connected -- new translator created")
translator = make_translator_for_connection(pipeline_state["pipeline"])
frame_index = 0
try:
while True:
raw_bytes = await websocket.receive_bytes()
frame_bgr = cv2.imdecode(np.frombuffer(raw_bytes, dtype=np.uint8), cv2.IMREAD_COLOR)
if frame_bgr is None:
await websocket.send_text(json.dumps({"type": "error", "message": f"Could not decode frame {frame_index}"}))
frame_index += 1
continue
frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
timestamp_seconds = frame_index / config.TARGET_FPS
try:
completed_segments = translator.push_frame(frame_rgb, frame_index, timestamp_seconds)
except Exception:
logger.exception(f"Error while processing frame {frame_index}")
await websocket.send_text(json.dumps({"type": "error", "message": f"Error while processing frame {frame_index}"}))
completed_segments = []
for segment_result in completed_segments:
await websocket.send_text(json.dumps({"type": "segment", **segment_result}))
frame_index += 1
except WebSocketDisconnect:
logger.info("Client disconnected")
try:
trailing_segments = translator.flush()
for segment_result in trailing_segments:
logger.info(f"Trailing segment on disconnect: {segment_result['text']}")
except Exception:
logger.exception("Error flushing trailing segment on disconnect")
except Exception:
logger.exception("Unexpected error in websocket loop")
finally:
translator.close()
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=7860) |