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

train is ready

Browse Source 
manual_run reuses train's export
add current_model to constants
This commit is contained in:
Alex Bezdieniezhnykh
2025-05-31 18:41:10 +03:00
parent 0bd5bd6d27
commit 538dc8efa9
4 changed files with 18 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
train.py
+11 -16
View File
@@ -158,25 +158,20 @@ def train_dataset():
model_dir = path.join(models_dir, today_folder)
shutil.copytree(results.save_dir, model_dir)
model_path = path.join(models_dir, f'{prefix[:-1]}.pt')
shutil.copy(path.join(model_dir, 'weights', 'best.pt'), model_path)
return model_path
shutil.copy(path.join(model_dir, 'weights', 'best.pt'), constants.CURRENT_PT_MODEL)
def validate(model_path):
model = YOLO(model_path)
print(model.val())
if __name__ == '__main__':
model_path = train_dataset()
validate(path.join('runs', 'detect', 'train7', 'weights', 'best.pt'))
onnx_path = export_onnx(model_path)
def export_current_model():
export_onnx(constants.CURRENT_PT_MODEL)
api_client = ApiClient()
with open(onnx_path, 'rb') as binary_file:
with open(constants.CURRENT_ONNX_MODEL, 'rb') as binary_file:
onnx_bytes = binary_file.read()
key = Security.get_model_encryption_key()
api_client.upload_big_small_resource(onnx_bytes, onnx_path, constants.MODELS_FOLDER, key)
api_client.upload_big_small_resource(onnx_bytes, 'azaion.onnx', constants.MODELS_FOLDER, key)
if __name__ == '__main__':
train_dataset()
export_current_model()
print('success!')