test / app.py
Piarasingh85's picture
Update app.py
864dc88 verified
raw
history blame contribute delete
899 Bytes
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