Spaces:
Sleeping
Sleeping
| import numpy as np | |
| from PIL import Image | |
| import io | |
| from typing import Tuple, Dict, Union | |
| import asyncio | |
| from concurrent.futures import ThreadPoolExecutor | |
| # Optimized thread pool for CPU-intensive image operations | |
| image_executor = ThreadPoolExecutor(max_workers=2, thread_name_prefix="ImageProcessor") | |
| # Common model input sizes | |
| MODEL_SIZES = { | |
| 'xceptionnet': (299, 299), | |
| 'mesonet': (256, 256) | |
| } | |
| def _preprocess_image_core(image_data: bytes, target_size: Tuple[int, int]) -> np.ndarray: | |
| """ | |
| Core image preprocessing function (synchronous). | |
| Args: | |
| image_data: Raw image bytes | |
| target_size: Target size as (width, height) | |
| Returns: | |
| Preprocessed numpy array ready for model prediction | |
| """ | |
| # Read and validate image | |
| image = Image.open(io.BytesIO(image_data)) | |
| # Convert to RGB if needed | |
| if image.mode != 'RGB': | |
| image = image.convert('RGB') | |
| # Resize and normalize | |
| image = image.resize(target_size, Image.Resampling.LANCZOS) | |
| img_array = np.array(image, dtype=np.float32) / 255.0 | |
| # Add batch dimension | |
| return np.expand_dims(img_array, axis=0) | |
| async def preprocess_image(image_data: bytes, target_size: Tuple[int, int]) -> np.ndarray: | |
| """ | |
| Async image preprocessing for FastAPI. | |
| Args: | |
| image_data: Raw image bytes from uploaded file | |
| target_size: Target size as (width, height) | |
| Returns: | |
| Preprocessed numpy array ready for model prediction | |
| """ | |
| loop = asyncio.get_event_loop() | |
| return await loop.run_in_executor(image_executor, _preprocess_image_core, image_data, target_size) | |
| async def preprocess_for_models(image_data: bytes) -> Dict[str, np.ndarray]: | |
| """ | |
| Preprocess image for all models simultaneously (optimized). | |
| Args: | |
| image_data: Raw image bytes from uploaded file | |
| Returns: | |
| Dictionary with preprocessed arrays for each model | |
| """ | |
| def _process_all_sizes(data: bytes) -> Dict[str, np.ndarray]: | |
| # Load image once | |
| image = Image.open(io.BytesIO(data)) | |
| if image.mode != 'RGB': | |
| image = image.convert('RGB') | |
| results = {} | |
| for model_name, size in MODEL_SIZES.items(): | |
| # Resize and normalize | |
| resized = image.resize(size, Image.Resampling.LANCZOS) | |
| img_array = np.array(resized, dtype=np.float32) / 255.0 | |
| results[model_name] = np.expand_dims(img_array, axis=0) | |
| return results | |
| loop = asyncio.get_event_loop() | |
| return await loop.run_in_executor(image_executor, _process_all_sizes, image_data) | |
| def validate_image(image_data: bytes) -> bool: | |
| """ | |
| Validate if the provided bytes represent a valid image. | |
| Args: | |
| image_data: Raw image bytes | |
| Returns: | |
| True if valid image, False otherwise | |
| """ | |
| try: | |
| with Image.open(io.BytesIO(image_data)) as img: | |
| img.verify() | |
| return True | |
| except Exception: | |
| return False | |
| # Legacy compatibility function (for any remaining Flask code) | |
| def preprocess_image_sync(file_or_bytes: Union[bytes, object], target_size: Tuple[int, int]) -> np.ndarray: | |
| """Legacy synchronous preprocessing function for backward compatibility.""" | |
| if isinstance(file_or_bytes, bytes): | |
| return _preprocess_image_core(file_or_bytes, target_size) | |
| else: | |
| # Assume it's a file-like object | |
| image = Image.open(file_or_bytes) | |
| if image.mode != 'RGB': | |
| image = image.convert('RGB') | |
| image = image.resize(target_size, Image.Resampling.LANCZOS) | |
| img_array = np.array(image, dtype=np.float32) / 255.0 | |
| return np.expand_dims(img_array, axis=0) |