from ultralytics import YOLO
import os

def test():
    # Load the custom trained model
    model_path = "best.pt"
    if not os.path.exists(model_path):
        print(f"Error: {model_path} not found.")
        return

    model = YOLO(model_path)

    # Path to data.yaml
    data_yaml = "unified_dataset/data.yaml"
    
    # Run validation
    print(f"Running validation on {data_yaml}...")
    metrics = model.val(data=data_yaml, split='val')
    
    # Print results
    print("\n--- Validation Results ---")
    print(f"mAP50: {metrics.results_dict['metrics/mAP50(B)']:.4f}")
    print(f"mAP50-95: {metrics.results_dict['metrics/mAP50-95(B)']:.4f}")
    print(f"Fitness: {metrics.fitness:.4f}")
    
    # Check if mAP50 > 0.90
    if metrics.results_dict['metrics/mAP50(B)'] > 0.90:
        print("\nSUCCESS: mAP50 is greater than 0.90")
    else:
        print("\nWARNING: mAP50 is below 0.90")

if __name__ == "__main__":
    test()