test_beach_plastic_detection.py

import os
import sys
import logging
import cv2
import numpy as np
from pathlib import Path
import urllib.request

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Add the app directory to the path so we can import modules
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from app.services.image_processing import (
    detect_beach_scene, detect_water_scene, detect_plastic_bottles, 
    detect_plastic_bottles_in_beach, check_for_plastic_bottle
)

# Sample beach plastic images (public domain or creative commons)
SAMPLE_IMAGES = [
    # Beach with plastic bottles
    "https://cdn.pixabay.com/photo/2019/07/30/11/13/plastic-waste-4372436_1280.jpg",
    # Beach with plastic bottles
    "https://live.staticflickr.com/4499/37193114384_25b662f3b3_b.jpg",
    # Plastic bottle on beach
    "https://cdn.pixabay.com/photo/2019/06/15/16/28/plastic-4275696_1280.jpg"
]

def download_sample_images():
    """Download sample images for testing"""
    output_dir = Path("test_files/beach_plastic")
    output_dir.mkdir(parents=True, exist_ok=True)
    
    downloaded_files = []
    
    for i, url in enumerate(SAMPLE_IMAGES):
        try:
            output_path = output_dir / f"beach_plastic_{i+1}.jpg"
            
            # Skip if already downloaded
            if output_path.exists():
                logger.info(f"File already exists: {output_path}")
                downloaded_files.append(str(output_path))
                continue
                
            # Download the image
            logger.info(f"Downloading: {url}")
            urllib.request.urlretrieve(url, output_path)
            downloaded_files.append(str(output_path))
            logger.info(f"Downloaded to: {output_path}")
        except Exception as e:
            logger.error(f"Error downloading {url}: {e}")
    
    return downloaded_files

def test_on_image(image_path):
    """Test all detection functions on a single image"""
    logger.info(f"Testing detection on: {image_path}")
    
    # Read the image
    img = cv2.imread(image_path)
    if img is None:
        logger.error(f"Could not read image: {image_path}")
        return False
    
    # Get image dimensions
    height, width = img.shape[:2]
    logger.info(f"Image dimensions: {width}x{height}")
    
    # Create a copy for drawing results
    img_result = img.copy()
    
    # Convert to HSV for color-based detection
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    
    # Detect scene type
    is_beach = detect_beach_scene(img, hsv)
    is_water = detect_water_scene(img, hsv)
    
    scene_type = "unknown"
    if is_beach and is_water:
        scene_type = "coastal"
    elif is_beach:
        scene_type = "beach"
    elif is_water:
        scene_type = "water"
    
    logger.info(f"Scene type: {scene_type}")
    
    # Add scene type text to image
    cv2.putText(img_result, f"Scene: {scene_type}", (10, 30), 
                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 2)
    
    # Detect plastic bottles with both methods for comparison
    
    # Standard detection
    standard_bottles = detect_plastic_bottles(img, hsv)
    logger.info(f"Standard detection found {len(standard_bottles)} bottles")
    
    # Beach-specific detection
    beach_bottles = detect_plastic_bottles_in_beach(img, hsv)
    logger.info(f"Beach-specific detection found {len(beach_bottles)} bottles")
    
    # Use the appropriate detection based on scene type
    if is_beach:
        bottle_detections = beach_bottles
        logger.info("Using beach-specific bottle detection")
    else:
        bottle_detections = standard_bottles
        logger.info("Using standard bottle detection")
    
    # Draw standard detection in green
    for det in standard_bottles:
        x1, y1, x2, y2 = det["bbox"]
        conf = det["confidence"]
        
        # Draw green rectangle for standard detection
        cv2.rectangle(img_result, (x1, y1), (x2, y2), (0, 255, 0), 1)
        cv2.putText(img_result, f"Std: {conf:.2f}", (x1, y1-10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
    
    # Draw beach-specific detection in red (thicker line)
    for det in beach_bottles:
        x1, y1, x2, y2 = det["bbox"]
        conf = det["confidence"]
        
        # Draw red rectangle for beach-specific detection
        cv2.rectangle(img_result, (x1, y1), (x2, y2), (0, 0, 255), 2)
        cv2.putText(img_result, f"Beach: {conf:.2f}", (x1, y1-25),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
    
    # Save the result
    output_dir = Path("test_output/beach_plastic")
    output_dir.mkdir(parents=True, exist_ok=True)
    
    base_name = os.path.basename(image_path)
    output_path = output_dir / f"result_{base_name}"
    
    cv2.imwrite(str(output_path), img_result)
    logger.info(f"Result saved to: {output_path}")
    
    return {
        "scene_type": scene_type,
        "standard_bottles": len(standard_bottles),
        "beach_bottles": len(beach_bottles),
        "output_path": str(output_path)
    }

def main():
    """Main function to test beach plastic bottle detection"""
    # Download sample images
    image_paths = download_sample_images()
    
    if not image_paths:
        logger.error("No images to test")
        return
    
    results = {}
    
    # Process each image
    for img_path in image_paths:
        results[os.path.basename(img_path)] = test_on_image(img_path)
    
    # Print summary
    logger.info("\n\n--- Beach Plastic Detection Results Summary ---")
    for img_file, result in results.items():
        if result:
            logger.info(f"{img_file}:")
            logger.info(f"  Scene type: {result['scene_type']}")
            logger.info(f"  Standard detection: {result['standard_bottles']} bottles")
            logger.info(f"  Beach-specific detection: {result['beach_bottles']} bottles")
            logger.info(f"  Output: {result['output_path']}")
            logger.info("---")

if __name__ == "__main__":
    main()