include-metrics-from-training (#6)
Reviewed-on: #6
This commit was merged in pull request #6.
This commit is contained in:
@@ -12,6 +12,7 @@ from typing import Any
|
||||
|
||||
import yaml
|
||||
from sanitize_onnx import sanitize_onnx
|
||||
from training_metrics import write_training_metrics
|
||||
from ultralytics import YOLO # type: ignore[reportMissingImports]
|
||||
|
||||
|
||||
@@ -101,6 +102,7 @@ def main() -> None:
|
||||
if not trained_weights.exists():
|
||||
raise FileNotFoundError(f"Could not find trained weights in {save_dir / 'weights'}")
|
||||
|
||||
write_training_metrics(save_dir / "results.csv", model_dir / "training_metrics.json")
|
||||
copy_if_exists(trained_weights, model_dir / "best.pt")
|
||||
trained_model = YOLO(str(trained_weights))
|
||||
onnx_path = Path(trained_model.export(format="onnx", imgsz=args.imgsz))
|
||||
|
||||
Reference in New Issue
Block a user