add rknn conversion - install and use scripts, auto convert to rknn after AI training is done and put pt and rknn models to /azaion/models directory

This commit is contained in:
zxsanny
2025-02-21 14:26:54 +02:00
parent 6e4b0f40ef
commit ceb50bf48a
11 changed files with 205 additions and 46 deletions
+20
View File
@@ -0,0 +1,20 @@
import cv2
import numpy as np
import yaml
from predictor import Predictor
from ultralytics import YOLO
class YOLOPredictor(Predictor):
def __init__(self):
self.model = YOLO('/azaion/models/azaion.onnx')
self.model.task = 'detect'
with open('data.yaml', 'r') as f:
data_yaml = yaml.safe_load(f)
class_names = data_yaml['names']
names = self.model.names
def predict(self, frame):
results = self.model.track(frame, persist=True, tracker='bytetrack.yaml')
return results[0].plot()