mirror of
https://github.com/azaion/ai-training.git
synced 2026-04-22 10:46:35 +00:00
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:
@@ -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()
|
||||
Reference in New Issue
Block a user