test_detection.py

"""
Test script to verify the object detection functionality with a sample image.
"""

import sys
from PIL import Image, ImageDraw
import pandas as pd

# Import our detection functions
from app import initialize_model, detect_objects

def create_test_image():
    """Create a simple test image with some basic shapes."""
    # Create a white background
    img = Image.new('RGB', (640, 480), 'white')
    draw = ImageDraw.Draw(img)
    
    # Draw some simple shapes that might be detected as objects
    # Rectangle (might be detected as a book or similar)
    draw.rectangle([100, 100, 300, 200], fill='blue', outline='black', width=3)
    
    # Circle (might be detected as a ball or similar)
    draw.ellipse([400, 150, 500, 250], fill='red', outline='black', width=3)
    
    # Another rectangle
    draw.rectangle([200, 300, 400, 400], fill='green', outline='black', width=3)
    
    return img

def main():
    print("Testing YOLOv8 object detection...")
    
    # Initialize model
    if not initialize_model():
        print("Failed to initialize model")
        return
    
    # Create test image
    test_img = create_test_image()
    test_img.save("test_image.jpg")
    print("Created test image: test_image.jpg")
    
    # Run detection
    annotated_img, detection_df = detect_objects(test_img)
    
    # Save results
    annotated_img.save("annotated_result.jpg")
    print("Saved annotated result: annotated_result.jpg")
    
    # Print detection results
    print(f"\nDetection Results:")
    print(f"Number of detections: {len(detection_df)}")
    
    if len(detection_df) > 0:
        print("\nDetailed results:")
        print(detection_df.to_string(index=False))
        
        # Verify data formatting
        print("\nData formatting verification:")
        for _, row in detection_df.iterrows():
            print(f"Class: {row['class']}")
            print(f"Confidence: {row['confidence']} (type: {type(row['confidence'])})")
            print(f"Coordinates: ({row['x1']}, {row['y1']}, {row['x2']}, {row['y2']})")
            print(f"Coordinate types: {type(row['x1'])}, {type(row['y1'])}, {type(row['x2'])}, {type(row['y2'])}")
            print("---")
    else:
        print("No objects detected")
    
    print("\nTest completed successfully!")

if __name__ == "__main__":
    main()