test_model.py 942 B

123456789101112131415161718192021222324252627282930313233
  1. from ultralytics import YOLO
  2. import os
  3. def test():
  4. # Load the custom trained model
  5. model_path = "best.pt"
  6. if not os.path.exists(model_path):
  7. print(f"Error: {model_path} not found.")
  8. return
  9. model = YOLO(model_path)
  10. # Path to data.yaml
  11. data_yaml = "unified_dataset/data.yaml"
  12. # Run validation
  13. print(f"Running validation on {data_yaml}...")
  14. metrics = model.val(data=data_yaml, split='val')
  15. # Print results
  16. print("\n--- Validation Results ---")
  17. print(f"mAP50: {metrics.results_dict['metrics/mAP50(B)']:.4f}")
  18. print(f"mAP50-95: {metrics.results_dict['metrics/mAP50-95(B)']:.4f}")
  19. print(f"Fitness: {metrics.fitness:.4f}")
  20. # Check if mAP50 > 0.90
  21. if metrics.results_dict['metrics/mAP50(B)'] > 0.90:
  22. print("\nSUCCESS: mAP50 is greater than 0.90")
  23. else:
  24. print("\nWARNING: mAP50 is below 0.90")
  25. if __name__ == "__main__":
  26. test()