azaion/ai-training
1
0
Fork 0
mirror of https://github.com/azaion/ai-training.git synced 2026-04-22 08:56:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

add onnx inference

Browse Source
This commit is contained in:
Alex Bezdieniezhnykh
2025-03-24 21:48:22 +02:00
parent 3113d59a3a
commit 7d99e377f1
7 changed files with 242 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files
train.py
+10 -9
View File
@@ -237,9 +237,9 @@ def validate(model_path):
pass
def upload_model(model_path: str):
model = YOLO(model_path)
model.export(format="onnx", imgsz=1280, nms=True, batch=4)
def upload_model(model_path: str, size_small_in_kb: int=3):
# model = YOLO(model_path)
# model.export(format="onnx", imgsz=1280, nms=True, batch=4)
onnx_model = path.dirname(model_path) + Path(model_path).stem + '.onnx'
with open(onnx_model, 'rb') as f_in:
@@ -248,7 +248,7 @@ def upload_model(model_path: str):
key = Security.get_model_encryption_key()
onnx_encrypted = Security.encrypt_to(onnx_bytes, key)
part1_size = min(10 * 1024, int(0.9 * len(onnx_encrypted)))
part1_size = min(size_small_in_kb * 1024, int(0.9 * len(onnx_encrypted)))
onnx_part_small = onnx_encrypted[:part1_size] # slice bytes for part1
onnx_part_big = onnx_encrypted[part1_size:]
@@ -264,8 +264,9 @@ def upload_model(model_path: str):
api.upload_file('azaion.onnx.small', onnx_part_small)
if __name__ == '__main__':
model_path = train_dataset(from_scratch=True)
validate(path.join('runs', 'detect', 'train7', 'weights', 'best.pt'))
form_data_sample(500)
convert2rknn()
upload_model('azaion.pt')
# model_path = train_dataset(from_scratch=True)
# validate(path.join('runs', 'detect', 'train7', 'weights', 'best.pt'))
# form_data_sample(500)
# convert2rknn()
upload_model('azaion-2024-10-26.onnx')