Spaces:
Paused
Paused
File size: 2,072 Bytes
9c0eec0 | 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 | import asyncio
import json
from pathlib import Path
import cv2
import numpy as np
import torch
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
app = FastAPI(title="Reading Trace L2CS")
pipeline = None
def get_pipeline():
global pipeline
if pipeline is None:
from l2cs import Pipeline
pipeline = Pipeline(
weights=Path("models/L2CSNet_gaze360.pkl"),
arch="ResNet50",
device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
)
return pipeline
@app.get("/health")
def health():
get_pipeline()
return {"status": "ok", "device": "cuda" if torch.cuda.is_available() else "cpu"}
def infer(data: bytes):
frame = cv2.imdecode(np.frombuffer(data, np.uint8), cv2.IMREAD_COLOR)
if frame is None:
return {"error": "invalid frame"}
height, width = frame.shape[:2]
try:
result = get_pipeline().step(frame)
except ValueError:
return {"faces": 0}
yaw = np.asarray(result.yaw).reshape(-1)
pitch = np.asarray(result.pitch).reshape(-1)
if not yaw.size:
return {"faces": 0}
boxes = np.asarray(getattr(result, "bboxes", []))
index, fx, fy = 0, .5, .5
if boxes.ndim == 2 and boxes.shape[0] == yaw.size and boxes.shape[1] >= 4:
areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
index = int(np.argmax(areas))
x0, y0, x1, y1 = boxes[index, :4]
fx, fy = float((x0 + x1) / 2 / width), float((y0 + y1) / 2 / height)
return {"yaw": float(yaw[index]), "pitch": float(pitch[index]), "fx": fx, "fy": fy, "faces": int(yaw.size)}
@app.websocket("/ws")
async def websocket_gaze(socket: WebSocket):
await socket.accept()
try:
while True:
data = await socket.receive_bytes()
if len(data) > 2_000_000:
await socket.close(code=1009)
return
result = await asyncio.to_thread(infer, data)
await socket.send_text(json.dumps(result))
except WebSocketDisconnect:
pass
|