import logging
from ultralytics import YOLO

# Suppress Ultralytics logging
logging.getLogger('ultralytics').setLevel(logging.WARNING)

# Define models and the test image
models = {
    "yolov8n": "yolov8n.pt",  # Pretrained model
    "fine_tuned": "yolov8n_rubberducks.pt"   # Fine-tuned model
}
image_path = "test_image.jpg"

# Initialize a dictionary to store results
performance = {}

# Run inference for each model
for model_name, model_path in models.items():
    # Load the model
    model = YOLO(model_path)
    
    # Run inference on the test image
    results = model(image_path)
    first_result = results[0]  # Extract the first result
    
    # Count the number of detections (boxes)
    num_detections = len(first_result.boxes) if hasattr(first_result, 'boxes') and first_result.boxes is not None else 0
    
    # Calculate total confidence score of detections
    if num_detections > 0:
        total_confidence = sum(float(box.conf) for box in first_result.boxes)  # Convert tensor to float
    else:
        total_confidence = 0.0  # No detections
    
    # Store performance data
    performance[model_name] = {
        "detections": num_detections,
        "confidence": total_confidence
    }

# Extract results for comparison
yolo_detections = performance['yolov8n']['detections']
yolo_confidence = performance['yolov8n']['confidence']
fine_tuned_detections = performance['fine_tuned']['detections']
fine_tuned_confidence = performance['fine_tuned']['confidence']

# Calculate the difference
diff_detections = fine_tuned_detections - yolo_detections
diff_confidence = fine_tuned_confidence - yolo_confidence
detection_diff_word = "more" if diff_detections > 0 else "less"
confidence_diff_word = "more" if diff_confidence > 0 else "less"

# Print streamlined results
print()
print(f"    YOLOv8n detected {yolo_detections} ducks with a total confidence of {yolo_confidence:.2f}")
print(f"    The fine-tuned model detected {fine_tuned_detections} ducks with a total confidence of {fine_tuned_confidence:.2f}")
print(f"    The fine-tuned model detects ducks with {abs(diff_confidence * 100):.0f}% {confidence_diff_word} confidence.")
print()