JSWOOK/gazeTracker

1import asyncio2import json3from pathlib import Path4 5import cv26import numpy as np7import torch8from fastapi import FastAPI, WebSocket, WebSocketDisconnect9 10app = FastAPI(title="Reading Trace L2CS")11pipeline = None12 13 14def get_pipeline():15    global pipeline16    if pipeline is None:17        from l2cs import Pipeline18        pipeline = Pipeline(19            weights=Path("models/L2CSNet_gaze360.pkl"),20            arch="ResNet50",21            device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),22        )23    return pipeline24 25 26@app.get("/health")27def health():28    get_pipeline()29    return {"status": "ok", "device": "cuda" if torch.cuda.is_available() else "cpu"}30 31 32def infer(data: bytes):33    frame = cv2.imdecode(np.frombuffer(data, np.uint8), cv2.IMREAD_COLOR)34    if frame is None:35        return {"error": "invalid frame"}36    height, width = frame.shape[:2]37    try:38        result = get_pipeline().step(frame)39    except ValueError:40        return {"faces": 0}41    yaw = np.asarray(result.yaw).reshape(-1)42    pitch = np.asarray(result.pitch).reshape(-1)43    if not yaw.size:44        return {"faces": 0}45    boxes = np.asarray(getattr(result, "bboxes", []))46    index, fx, fy = 0, .5, .547    if boxes.ndim == 2 and boxes.shape[0] == yaw.size and boxes.shape[1] >= 4:48        areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])49        index = int(np.argmax(areas))50        x0, y0, x1, y1 = boxes[index, :4]51        fx, fy = float((x0 + x1) / 2 / width), float((y0 + y1) / 2 / height)52    return {"yaw": float(yaw[index]), "pitch": float(pitch[index]), "fx": fx, "fy": fy, "faces": int(yaw.size)}53 54 55@app.websocket("/ws")56async def websocket_gaze(socket: WebSocket):57    await socket.accept()58    try:59        while True:60            data = await socket.receive_bytes()61            if len(data) > 2_000_000:62                await socket.close(code=1009)63                return64            result = await asyncio.to_thread(infer, data)65            await socket.send_text(json.dumps(result))66    except WebSocketDisconnect:67        pass68