commit 60fbff76df1c1c8bf8d76f722a343bc0b42e4e3d
Author: Siwat Sirichai <siwat@siwatinc.com>
Date:   Wed Jan 8 23:08:57 2025 +0700

    fastapi yolo detector

diff --git a/app.py b/app.py
new file mode 100644
index 0000000..f4ecaa1
--- /dev/null
+++ b/app.py
@@ -0,0 +1,36 @@
+from fastapi import FastAPI, WebSocket
+from ultralytics import YOLO
+import torch
+import cv2
+import base64
+import numpy as np
+
+app = FastAPI()
+
+model = YOLO("yolov8n.pt")
+if torch.cuda.is_available():
+    model.to('cuda')
+
+@app.websocket("/detect")
+async def detect(websocket: WebSocket):
+    await websocket.accept()
+    try:
+        while True:
+            data = await websocket.receive_text()
+            # Decode base64 image bytes
+            img_data = base64.b64decode(data)
+            np_arr = np.frombuffer(img_data, np.uint8)
+            frame = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
+
+            results = model(frame, stream=False)
+            boxes = []
+            for r in results:
+                for box in r.boxes:
+                    boxes.append({
+                        "class": int(box.cls[0]),
+                        "confidence": float(box.conf[0]),
+                    })
+
+            await websocket.send_json({"detections": boxes})
+    except:
+        pass
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..2122d86
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,6 @@
+fastapi
+uvicorn
+torch
+torchvision
+ultralytics
+opencv-python