Spaces:
Sleeping
Sleeping
| """ | |
| Process test images using the EnhancedImageDryer. | |
| This script takes images from the test_images directory, | |
| processes them with the EnhancedImageDryer, and saves the results | |
| to the test_results directory. | |
| """ | |
| import os | |
| import sys | |
| import datetime | |
| from PIL import Image | |
| from pathlib import Path | |
| from dotenv import load_dotenv | |
| # Import the enhanced image dryer | |
| from src.enhanced_image_dryer import EnhancedImageDryer | |
| # Load environment variables | |
| load_dotenv() | |
| def process_image(image_path, use_fallback=False): | |
| """Process an image using the EnhancedImageDryer.""" | |
| # Ensure test_results directory exists | |
| os.makedirs("test_results", exist_ok=True) | |
| # Load the image | |
| try: | |
| image = Image.open(image_path) | |
| print(f"Loaded image: {image_path}") | |
| print(f"Image size: {image.size}, Mode: {image.mode}") | |
| except Exception as e: | |
| print(f"Error loading image: {str(e)}") | |
| return False | |
| # Process the image with the EnhancedImageDryer | |
| print("Initializing EnhancedImageDryer...") | |
| dryer = EnhancedImageDryer() | |
| # Check if API key is set | |
| if not dryer.api_key: | |
| print("Error: STABILITY_API_KEY not found in environment variables.") | |
| print("Please set your Stability API key in the .env file.") | |
| return False | |
| print(f"API Key found: {dryer.api_key[:5]}...{dryer.api_key[-5:] if len(dryer.api_key) > 10 else ''}") | |
| if use_fallback: | |
| # Use the fallback method directly | |
| print("Using fallback drying method as requested...") | |
| processed_image = dryer.apply_fallback_drying_effect(image) | |
| else: | |
| # Try the API first | |
| print("Attempting to process with Stability AI API...") | |
| processed_image = dryer.process_image(image) | |
| # If API fails, use fallback method | |
| if processed_image is None: | |
| print("API processing failed. Using fallback drying method...") | |
| processed_image = dryer.apply_fallback_drying_effect(image) | |
| if processed_image: | |
| # Save the processed image | |
| timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") | |
| output_filename = f"test_results/enhanced_dried_{Path(image_path).stem}_{timestamp}.png" | |
| processed_image.save(output_filename) | |
| print(f"Successfully processed image and saved to: {output_filename}") | |
| return True | |
| else: | |
| print(f"Failed to process image: {image_path}") | |
| return False | |
| def main(): | |
| """Main function to process images from the test_images directory.""" | |
| print("\n=== Processing Test Images with EnhancedImageDryer ===\n") | |
| # Parse command line arguments | |
| use_fallback = "--fallback" in sys.argv | |
| if use_fallback: | |
| print("Fallback mode enabled: Will use local processing instead of API") | |
| # Check if test_images directory exists | |
| if not os.path.exists("test_images"): | |
| print("Error: test_images directory not found.") | |
| return | |
| # Get list of image files in the test_images directory | |
| image_files = [f for f in os.listdir("test_images") | |
| if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif'))] | |
| if not image_files: | |
| print("No image files found in the test_images directory.") | |
| print("Please place your test images in the test_images directory.") | |
| return | |
| print(f"Found {len(image_files)} image(s) in the test_images directory:") | |
| for i, file in enumerate(image_files, 1): | |
| print(f"{i}. {file}") | |
| # Process each image | |
| print("\nProcessing images...") | |
| for file in image_files: | |
| image_path = os.path.join("test_images", file) | |
| print(f"\nProcessing: {file}") | |
| success = process_image(image_path, use_fallback) | |
| if success: | |
| print(f"[SUCCESS] Successfully processed {file}") | |
| else: | |
| print(f"[FAILED] Failed to process {file}") | |
| print("\nAll processing complete!") | |
| print(f"Processed images saved to the test_results directory.") | |
| if __name__ == "__main__": | |
| main() | |