Add simple yolo function

LinuzJ · Nov 11, 2023 · d862ebb · d862ebb
1 parent 0617472
commit d862ebb
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/backend/yolo_model/yolo.py b/backend/yolo_model/yolo.py
@@ -0,0 +1,12 @@
+from ultralytics import YOLO
+
+model = YOLO("yolov8n-pose.pt")
+
+
+def inference(input):
+    results_generator = model(source=input, conf=0.5, stream=True)
+    res = []
+    for frame in results_generator:
+        frame.orig_img = None  # don't send back whole image
+        res.append(frame.numpy())
+    return res