src/core/preprocessing.py

"""Image preprocessing for MelanoScope AI."""
import logging
from typing import Union, Optional
import numpy as np
from PIL import Image
from torchvision import transforms
from ..config.settings import ModelConfig

logger = logging.getLogger(__name__)

class ImagePreprocessor:
    """Handles image preprocessing for model inference."""
    
    def __init__(self):
        self.transforms = self._create_transform_pipeline()
        logger.info("ImagePreprocessor initialized")
    
    def _create_transform_pipeline(self) -> transforms.Compose:
        """Create image transformation pipeline."""
        return transforms.Compose([
            transforms.Resize(ModelConfig.IMAGE_SIZE),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=ModelConfig.NORMALIZATION_MEAN,
                std=ModelConfig.NORMALIZATION_STD
            ),
        ])
    
    def preprocess(self, image_input: Union[Image.Image, np.ndarray]) -> Optional[np.ndarray]:
        """Preprocess image for model inference."""
        try:
            pil_image = self._convert_to_pil(image_input)
            if pil_image is None:
                return None
            
            tensor = self.transforms(pil_image).unsqueeze(0).numpy()
            return tensor
            
        except Exception as e:
            logger.error(f"Error preprocessing image: {e}")
            return None
    
    def _convert_to_pil(self, image_input: Union[Image.Image, np.ndarray]) -> Optional[Image.Image]:
        """Convert image input to PIL Image in RGB mode."""
        try:
            if isinstance(image_input, Image.Image):
                return image_input.convert("RGB")
            else:
                return Image.fromarray(image_input).convert("RGB")
        except Exception as e:
            logger.error(f"Error converting to PIL: {e}")
            return None