File size: 2,163 Bytes
d9c6ad3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
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()
|