98 lines
3.0 KiB
Python
98 lines
3.0 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Create a comparison image showing detection results from all models
|
|
Run this after the tests have completed to generate a visual comparison
|
|
"""
|
|
from pathlib import Path
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
import sys
|
|
|
|
|
|
def create_comparison_grid():
|
|
"""Create a grid comparison of all detection results"""
|
|
test_outputs = Path(__file__).parent / "test_outputs"
|
|
|
|
# Define the images we want to compare
|
|
images_to_compare = [
|
|
("Florence-2", "florence2_eager_detection.png"),
|
|
("GroundingDINO", "grounding_dino_detection.png"),
|
|
("OWLv2", "owlv2_detection.png"),
|
|
("YOLO-World", "yolo_world_detection.png"),
|
|
]
|
|
|
|
# Load all images
|
|
loaded_images = []
|
|
for title, filename in images_to_compare:
|
|
img_path = test_outputs / filename
|
|
if img_path.exists():
|
|
img = Image.open(img_path)
|
|
loaded_images.append((title, img))
|
|
print(f"✓ Loaded {filename}")
|
|
else:
|
|
print(f"✗ Missing {filename}")
|
|
|
|
if not loaded_images:
|
|
print("No images found in test_outputs directory!")
|
|
return False
|
|
|
|
# Calculate grid dimensions
|
|
# Use 2 columns
|
|
cols = 2
|
|
rows = (len(loaded_images) + cols - 1) // cols
|
|
|
|
# Get image dimensions (assuming all are same size)
|
|
img_width, img_height = loaded_images[0][1].size
|
|
|
|
# Add space for titles
|
|
title_height = 40
|
|
padding = 20
|
|
|
|
# Create output image
|
|
grid_width = cols * img_width + (cols + 1) * padding
|
|
grid_height = rows * (img_height + title_height) + (rows + 1) * padding
|
|
|
|
comparison = Image.new('RGB', (grid_width, grid_height), color='white')
|
|
draw = ImageDraw.Draw(comparison)
|
|
|
|
# Try to use a nice font, fall back to default if not available
|
|
try:
|
|
font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 24)
|
|
except:
|
|
font = ImageFont.load_default()
|
|
|
|
# Place images in grid
|
|
for idx, (title, img) in enumerate(loaded_images):
|
|
row = idx // cols
|
|
col = idx % cols
|
|
|
|
x = padding + col * (img_width + padding)
|
|
y = padding + row * (img_height + title_height + padding)
|
|
|
|
# Draw title
|
|
title_y = y
|
|
# Center the title
|
|
bbox = draw.textbbox((0, 0), title, font=font)
|
|
title_width = bbox[2] - bbox[0]
|
|
title_x = x + (img_width - title_width) // 2
|
|
draw.text((title_x, title_y), title, fill='black', font=font)
|
|
|
|
# Paste image
|
|
comparison.paste(img, (x, y + title_height))
|
|
|
|
# Save comparison
|
|
output_path = test_outputs / "comparison_all_models.png"
|
|
comparison.save(output_path)
|
|
print(f"\n✅ Saved comparison to {output_path}")
|
|
|
|
# Also save as high quality JPG
|
|
jpg_path = test_outputs / "comparison_all_models.jpg"
|
|
comparison.save(jpg_path, quality=95)
|
|
print(f"✅ Saved comparison to {jpg_path}")
|
|
|
|
return True
|
|
|
|
|
|
if __name__ == "__main__":
|
|
success = create_comparison_grid()
|
|
sys.exit(0 if success else 1)
|