from fastapi import FastAPI, WebSocket
import cv2
from transformers import pipeline

app = FastAPI()

# Load the model from Hugging Face
model = pipeline("object-detection", model="facebook/detr-resnet-50")

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    cap = cv2.VideoCapture(0)  # Open the first webcam device
    
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        
        # Convert the frame to a format suitable for the model
        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        
        # Use the model to make predictions
        predictions = model(rgb_frame)
        
        # Send the predictions back to the client
        await websocket.send_json(predictions)

    cap.release()
    await websocket.close()

# To run the app, use: uvicorn app:app --reload