From d862ebb1c22a45051c8c10ad33671b13c35aa0ae Mon Sep 17 00:00:00 2001 From: Peik Etzell Date: Sat, 11 Nov 2023 15:35:06 +0200 Subject: [PATCH] Add simple yolo function --- backend/yolo_model/yolo.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) create mode 100644 backend/yolo_model/yolo.py diff --git a/backend/yolo_model/yolo.py b/backend/yolo_model/yolo.py new file mode 100644 index 0000000..fa30169 --- /dev/null +++ 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