Spaces:
Sleeping
Sleeping
| 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() |