""" Image preprocessing utilities. Handles image loading, resizing, normalization for model inference. """ import torch import torchvision.transforms as transforms from PIL import Image import io import logging from typing import Union, Tuple import numpy as np logger = logging.getLogger(__name__) # Tile extraction size (matches Kaggle SimpleSlideDataset standardize_transform) TARGET_SIZE = 256 # ViT backbone input size (matches Kaggle test_transform: transforms.Resize(224)) # CRITICAL: The model was trained with 224x224 input to the ViT backbone, # NOT 256x256. Using 256 here causes feature space mismatch and random predictions. MODEL_INPUT_SIZE = 224 # ImageNet normalization statistics (used during model training) IMAGENET_MEAN = [0.485, 0.456, 0.406] IMAGENET_STD = [0.229, 0.224, 0.225] # Preprocessing pipeline - matches Kaggle test_transform exactly: # transforms.Resize(224), transforms.ToTensor(), transforms.Normalize(...) preprocess_transform = transforms.Compose([ transforms.Resize((MODEL_INPUT_SIZE, MODEL_INPUT_SIZE)), transforms.ToTensor(), transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD) ]) # For standardization before model (matches training) resize_transform = transforms.Resize((TARGET_SIZE, TARGET_SIZE)) def load_image_from_bytes(image_bytes: bytes) -> Image.Image: """ Load image from bytes. Args: image_bytes: Image data as bytes Returns: PIL Image in RGB format Raises: ValueError: If image cannot be loaded """ try: image = Image.open(io.BytesIO(image_bytes)).convert('RGB') return image except Exception as e: logger.error(f"Failed to load image from bytes: {e}") raise ValueError(f"Invalid image data: {e}") def load_image_from_path(image_path: str) -> Image.Image: """ Load image from file path. Args: image_path: Path to image file Returns: PIL Image in RGB format Raises: FileNotFoundError: If file doesn't exist ValueError: If image cannot be loaded """ try: image = Image.open(image_path).convert('RGB') return image except FileNotFoundError: logger.error(f"Image file not found: {image_path}") raise except Exception as e: logger.error(f"Failed to load image from {image_path}: {e}") raise ValueError(f"Invalid image file: {e}") def preprocess_image(image: Image.Image) -> torch.Tensor: """ Preprocess single image for model inference. Args: image: PIL Image in RGB format Returns: Preprocessed tensor of shape (3, 224, 224) Process: 1. Resize to 224x224 (matches Kaggle test_transform) 2. Convert to tensor 3. Normalize with ImageNet statistics """ try: tensor = preprocess_transform(image) return tensor except Exception as e: logger.error(f"Failed to preprocess image: {e}") raise def preprocess_images_batch(images: list) -> torch.Tensor: """ Preprocess batch of images (tiles). Args: images: List of PIL Images Returns: Batch tensor of shape (num_images, 3, 256, 256) """ tensors = [] for img in images: try: tensor = preprocess_image(img) tensors.append(tensor) except Exception as e: logger.warning(f"Failed to preprocess image in batch: {e}") continue if not tensors: raise ValueError("No valid images in batch") return torch.stack(tensors) def extract_tiles_from_image( image: Image.Image, tile_size: int = 256, max_tiles: int = 1000 ) -> list: """ Extract tile patches from a large image. Useful for whole slide images (WSI) or large medical images. Args: image: PIL Image tile_size: Size of each tile patch (256x256) max_tiles: Maximum number of tiles to extract Returns: List of PIL Images (tiles) """ width, height = image.size tiles = [] try: # Extract non-overlapping tiles for y in range(0, height, tile_size): for x in range(0, width, tile_size): if len(tiles) >= max_tiles: break # Extract tile with padding if at edges right = min(x + tile_size, width) bottom = min(y + tile_size, height) tile = image.crop((x, y, right, bottom)) # Pad if necessary to maintain tile_size if tile.size != (tile_size, tile_size): padded_tile = Image.new('RGB', (tile_size, tile_size), color=(0, 0, 0)) padded_tile.paste(tile, (0, 0)) tile = padded_tile tiles.append(tile) if len(tiles) >= max_tiles: break logger.info(f"Extracted {len(tiles)} tiles from image ({width}x{height})") return tiles except Exception as e: logger.error(f"Failed to extract tiles: {e}") raise def get_image_info(image: Image.Image) -> dict: """Get metadata about an image.""" return { 'size': image.size, 'width': image.width, 'height': image.height, 'mode': image.mode, 'format': image.format } class ImagePreprocessor: """Image preprocessing pipeline.""" def __init__(self, target_size: int = MODEL_INPUT_SIZE, mean: list = IMAGENET_MEAN, std: list = IMAGENET_STD): """ Initialize preprocessor. Args: target_size: Target image size for ViT backbone input (224 to match Kaggle test_transform) mean: Normalization mean values std: Normalization std values """ self.target_size = target_size self.mean = mean self.std = std self.transform = transforms.Compose([ transforms.Resize((target_size, target_size)), transforms.ToTensor(), transforms.Normalize(mean=mean, std=std) ]) def process(self, image_input: Union[str, bytes, Image.Image]) -> torch.Tensor: """ Process image from various input formats. Args: image_input: Image path (str), image bytes, or PIL Image Returns: Preprocessed tensor (3, 256, 256) """ # Load image if needed if isinstance(image_input, str): image = load_image_from_path(image_input) elif isinstance(image_input, bytes): image = load_image_from_bytes(image_input) elif isinstance(image_input, Image.Image): image = image_input else: raise TypeError(f"Unsupported input type: {type(image_input)}") # Preprocess return preprocess_image(image) def process_batch(self, images: list) -> torch.Tensor: """ Process batch of images. Args: images: List of image inputs (paths, bytes, or PIL Images) Returns: Batch tensor (batch_size, 3, 256, 256) """ tensors = [] for img_input in images: try: tensor = self.process(img_input) tensors.append(tensor) except Exception as e: logger.warning(f"Failed to process image: {e}") continue if not tensors: raise ValueError("No images could be processed") return torch.stack(tensors)