File size: 5,009 Bytes
6bbbfda
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
import os
import sys
import logging
import cv2
import numpy as np
from pathlib import Path

# 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, detect_ships, check_for_plastic_bottle,
    check_for_ship, check_for_plastic_waste
)

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
    if is_beach:
        logger.info("Using beach-specific bottle detection")
        bottle_detections = detect_plastic_bottles_in_beach(img, hsv)
    else:
        logger.info("Using standard bottle detection")
        bottle_detections = detect_plastic_bottles(img, hsv)
    
    logger.info(f"Detected {len(bottle_detections)} potential plastic bottles")
    
    # Draw bottle detections
    for det in bottle_detections:
        x1, y1, x2, y2 = det["bbox"]
        conf = det["confidence"]
        
        # Draw red rectangle for bottles
        cv2.rectangle(img_result, (x1, y1), (x2, y2), (0, 0, 255), 2)
        cv2.putText(img_result, f"Bottle: {conf:.2f}", (x1, y1-10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
    
    # Detect ships if in water scene
    ship_detections = []
    if is_water:
        logger.info("Detecting ships in water scene")
        ship_detections = detect_ships(img, hsv)
        
        logger.info(f"Detected {len(ship_detections)} potential ships")
        
        # Draw ship detections
        for det in ship_detections:
            x1, y1, x2, y2 = det["bbox"]
            conf = det["confidence"]
            
            # Draw blue rectangle for ships
            cv2.rectangle(img_result, (x1, y1), (x2, y2), (255, 0, 0), 2)
            cv2.putText(img_result, f"Ship: {conf:.2f}", (x1, y1-10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
    
    # Save the result
    output_dir = Path("test_output/enhanced_detection")
    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,
        "bottle_detections": len(bottle_detections),
        "ship_detections": len(ship_detections),
        "output_path": str(output_path)
    }

def main():
    """Main function to test enhanced detection on sample images"""
    # Test directory
    test_dir = "test_files"
    
    # Check if test directory exists
    if not os.path.isdir(test_dir):
        logger.error(f"Test directory not found: {test_dir}")
        return
    
    # Get all image files in the test directory
    image_files = [f for f in os.listdir(test_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
    
    if not image_files:
        logger.error(f"No image files found in {test_dir}")
        return
    
    results = {}
    
    # Process each image
    for img_file in image_files:
        img_path = os.path.join(test_dir, img_file)
        results[img_file] = test_on_image(img_path)
    
    # Print summary
    logger.info("\n\n--- 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"  Plastic bottles: {result['bottle_detections']}")
            logger.info(f"  Ships: {result['ship_detections']}")
            logger.info(f"  Output: {result['output_path']}")
            logger.info("---")

if __name__ == "__main__":
    main()